{ "cells": [ { "cell_type": "markdown", "id": "3e0daf24", "metadata": {}, "source": [ "# Supervised fine-tuning for LFM2.5-8B-A1B with Halo\n", "\n", "This notebook fine-tunes [LFM2.5-8B-A1B](https://huggingface.co/LiquidAI/LFM2.5-8B-A1B) on [UltraChat 200K](https://huggingface.co/datasets/HuggingFaceH4/ultrachat_200k) using [Halo](https://whitecircle.com/research/halo), which is a distributed-training layer for Hugging Face models. The [Halo repository](https://github.com/whitecircle/halo) contains the source code.\n", "\n", "Expert parallelism assigns different mixture-of-experts (MoE) experts to different GPUs. Halo sends each token to the GPU that owns its selected expert, so each GPU stores only part of the expert weights.\n", "\n", "This two-way expert-parallel recipe distributes the model's 32 routed experts across two GPUs. Each process owns 16 experts. Halo gathers the weights during save and writes a standard Hugging Face checkpoint.\n" ] }, { "cell_type": "markdown", "id": "bca6dee8", "metadata": {}, "source": [ "## Prerequisites\n", "\n", "- Two matching GPUs that matches the published Halo images, which are built for Hopper (H100, H200) and Blackwell (B200, B300).\n", "- An NVIDIA driver that provides CUDA 13.0 or higher\n", "- Docker with the NVIDIA Container Toolkit\n", "- A large writable volume for model, dataset, and checkpoint files\n", "\n", "This notebook does not run in Google Colab. Use a compatible two-GPU machine from your cloud provider or your own cluster. Allow several hours for a full UltraChat epoch on two B300 GPUs. This is a planning estimate, not a benchmark. For scale, Halo reports 31,262 tokens per second per GPU for LFM2.5-8B-A1B at sequence 8,192 on one B300. After the first 20 steps, use the logged step time and total step count to refine the estimate.\n", "\n", "## Recipe at a Glance\n", "\n", "This recipe uses full-parameter BF16 training with two-way expert parallelism.\n", "\n", "| Setting | Value |\n", "| --- | --- |\n", "| GPUs | 2 |\n", "| Expert parallel size | 2 |\n", "| Dataset | UltraChat 200K, supervised split |\n", "| Sequence length | 8,192 |\n", "| Effective batch size | 16 |\n", "| Output | Gathered Hugging Face checkpoint |\n", "\n", "Do not enable context parallelism. LFM2 short-convolution layers operate across the sequence axis." ] }, { "cell_type": "markdown", "id": "62342004", "metadata": {}, "source": [ "## Installation & Setup\n", "\n", "This notebook uses the public Halo image that matches the detected GPUs. Define these host paths before you start the container:\n", "\n", "- `HALO_DIR`: an optional Halo clone for source reference\n", "- `COOKBOOK_DIR`: a clone of this cookbook repository\n", "- `DATA_DIR`: a large persistent volume for downloads and run files\n", "\n", "~~~bash\n", "export HALO_DIR=/path/to/halo\n", "export COOKBOOK_DIR=/path/to/cookbook\n", "export DATA_DIR=/path/to/large/volume\n", "\n", "git clone --recurse-submodules https://github.com/whitecircle/halo.git \"$HALO_DIR\"\n", "git clone https://github.com/Liquid4All/cookbook.git \"$COOKBOOK_DIR\"\n", "~~~\n", "\n", "Pull the Docker image for the detected GPUs.\n", "~~~bash\n", "GPU_NAME=\"$(nvidia-smi --query-gpu=name --format=csv,noheader | head -n 1)\"\n", "case \"$GPU_NAME\" in\n", " *H100*|*H200*) HALO_IMAGE=public.ecr.aws/whitecircle/halo:hopper ;;\n", " *B200*|*B300*|*GB200*|*GB300*) HALO_IMAGE=public.ecr.aws/whitecircle/halo:blackwell ;;\n", " *) echo \"Unsupported GPU: $GPU_NAME\"; exit 1 ;;\n", "esac\n", "export HALO_IMAGE\n", "\n", "docker pull \"$HALO_IMAGE\"\n", "~~~\n", "\n", "Launch the Docker image:\n", "~~~bash\n", "docker run --rm -it --gpus '\"device=0,1\"' \\\n", " --ipc=host --shm-size=128g \\\n", " --ulimit memlock=-1 --ulimit stack=67108864 \\\n", " -p 8888:8888 \\\n", " -e HF_HOME=/mnt/hf \\\n", " -e HF_DATASETS_CACHE=/mnt/hf/datasets \\\n", " -e TMPDIR=/mnt/tmp \\\n", " -e HALO_DATA_ROOT=/mnt \\\n", " -e CUDA_DEVICE_MAX_CONNECTIONS=1 \\\n", " -v \"$COOKBOOK_DIR\":/cookbook \\\n", " -v \"$DATA_DIR\":/mnt \\\n", " -w /workspace \\\n", " \"$HALO_IMAGE\" \\\n", " jupyter lab --ip=0.0.0.0 --port=8888 --no-browser \\\n", " --allow-root --notebook-dir=/cookbook/finetuning/notebooks\n", "~~~\n", "\n", "The container runs the Halo version packaged in the image.\n", "The `HALO_DIR` clone is for source reference and is not mounted over that package.\n", "`-p 8888:8888` publishes only the Jupyter port.\n", "`CUDA_DEVICE_MAX_CONNECTIONS=1` supports the DeepEP communication schedule used by this two-GPU MoE run.\n", "\n", "Open the URL printed by Jupyter. Then open this notebook.\n" ] }, { "cell_type": "markdown", "id": "6622f4f4", "metadata": {}, "source": [ "## Check the Runtime\n", "\n", "This cell sets the run paths, confirms that two GPUs are visible, and creates the checkpoint directory. It stops if a different number of GPUs is visible.\n" ] }, { "cell_type": "code", "execution_count": 1, "id": "986c880e", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "['NVIDIA B300 SXM6 AC', 'NVIDIA B300 SXM6 AC']\n" ] } ], "source": [ "import subprocess\n", "from pathlib import Path\n", "\n", "RUN_ROOT = Path(\"/mnt/checkpoints\")\n", "CONFIG_PATH = RUN_ROOT / \"lfm25-8b-a1b-ultrachat-ep2.yaml\"\n", "OUTPUT_PATH = RUN_ROOT / \"lfm25-8b-a1b-ultrachat-ep2\"\n", "\n", "gpu_names = subprocess.check_output(\n", " [\n", " \"nvidia-smi\",\n", " \"--query-gpu=name\",\n", " \"--format=csv,noheader\",\n", " ],\n", " text=True,\n", ").strip().splitlines()\n", "assert len(gpu_names) == 2, f\"Expected 2 visible GPUs, found {len(gpu_names)}\"\n", "\n", "RUN_ROOT.mkdir(parents=True, exist_ok=True)\n", "print(gpu_names)\n" ] }, { "cell_type": "markdown", "id": "18f0f4f8", "metadata": {}, "source": [ "## Data Prep\n", "\n", "The [UltraChat 200K](https://huggingface.co/datasets/HuggingFaceH4/ultrachat_200k) `train_sft` split already stores each conversation in a `messages` column. Halo downloads the split and applies the model's chat template. No separate conversion step is necessary. The configuration reserves 1% of the split for evaluation.\n" ] }, { "cell_type": "markdown", "id": "0f023667", "metadata": {}, "source": [ "## Create the Training Configuration\n", "\n", "This configuration follows Halo's LFM2 MoE recipe. The MoE settings have these purposes:\n", "\n", "- `moe_balancing: bias_update` adjusts expert-selection biases during training because LFM2 does not use a router auxiliary loss.\n", "- `save_sharded_ep: false` gathers expert weights and saves one standard Hugging Face checkpoint.\n", "- `use_grouped_gemm: true` combines expert projections to reduce small kernel launches.\n", "- `fp32_router: true` keeps routing calculations in FP32 for stable expert selection.\n", "- `fp32_experts: false` keeps expert computation in BF16 to reduce memory use and training time.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "22f827f5", "metadata": {}, "outputs": [], "source": [ "CONFIG_PATH.write_text(\n", " r\"\"\"\n", "# --- Model ---\n", "model_name_or_path: LiquidAI/LFM2.5-8B-A1B\n", "\n", "# --- Dataset ---\n", "dataset:\n", "- HuggingFaceH4/ultrachat_200k@train_sft\n", "conversation_field: messages\n", "test_size: 0.01\n", "train_on_completions_only: true\n", "\n", "# --- Tokens ---\n", "assistant_message_template: \"<|im_start|>assistant\\n\"\n", "pad_token: \"<|pad|>\"\n", "eos_token: \"<|im_end|>\"\n", "\n", "# --- LoRA ---\n", "use_peft: false\n", "\n", "# --- Parallelism & MoE precision ---\n", "expert_parallel_size: 2\n", "moe_balancing: bias_update\n", "save_sharded_ep: false\n", "use_grouped_gemm: true\n", "fp32_router: true\n", "fp32_experts: false\n", "\n", "# --- Sequence & packing ---\n", "max_length: 8192\n", "packing: true\n", "use_liger_kernel: false\n", "bf16: true\n", "\n", "# --- Training ---\n", "attn_implementation: flash_attention_2\n", "per_device_train_batch_size: 1\n", "per_device_eval_batch_size: 1\n", "gradient_accumulation_steps: 8\n", "num_train_epochs: 1.0\n", "gradient_checkpointing: true\n", "gradient_checkpointing_kwargs:\n", " use_reentrant: false\n", "\n", "# --- Optimizer & LR schedule ---\n", "optim: adamw_torch_fused\n", "learning_rate: 5.0e-06\n", "lr_scheduler_type: cosine\n", "warmup_steps: 32\n", "max_grad_norm: 1.0\n", "\n", "# --- Output ---\n", "output_dir: /mnt/checkpoints/lfm25-8b-a1b-ultrachat-ep2\n", "save_strategy: steps\n", "save_steps: 1000\n", "eval_strategy: steps\n", "eval_steps: 300\n", "save_total_limit: 1\n", "save_only_model: true\n", "\n", "# --- Logging ---\n", "logging_steps: 1\n", "logging_first_step: true\n", "report_to: none\n", "\n", "# --- Data ---\n", "dataloader_num_workers: 2\n", "remove_unused_columns: false\n", "\"\"\".lstrip()\n", ")\n", "print(CONFIG_PATH.read_text())" ] }, { "cell_type": "markdown", "id": "4752f992", "metadata": {}, "source": [ "## Check the Launch Command\n", "\n", "The dry run prints the two-process command without loading the model.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "bc56a0bc", "metadata": {}, "outputs": [], "source": [ "!halo launch sft {CONFIG_PATH} -n 2 --dry-run\n" ] }, { "cell_type": "markdown", "id": "ac5a9617", "metadata": {}, "source": [ "## Launch Training\n", "\n", "This cell starts the full UltraChat run on both visible GPUs.\n", "\n", "Any configuration field can be overridden on the command line, and arguments after `--` take precedence over the file. To exercise the whole pipeline on a short run before committing to the epoch:\n", "\n", "~~~bash\n", "halo launch sft /mnt/checkpoints/lfm25-8b-a1b-ultrachat-ep2.yaml -n 2 -- --max_steps=20 --save_steps=20 --eval_steps=10\n", "~~~\n", "\n", "Halo resumes a run only on the same number of GPUs it started with, and this configuration sets `save_only_model: true`, which writes no optimizer or scheduler state. A run interrupted before the first checkpoint at step 1,000 therefore restarts from the beginning. Lower `save_steps` on a preemptible machine.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "97f81f59", "metadata": {}, "outputs": [], "source": [ "!halo launch sft {CONFIG_PATH} -n 2\n" ] }, { "cell_type": "markdown", "id": "bfc8e559", "metadata": {}, "source": [ "## Inference\n", "\n", "The saved checkpoint uses the native Transformers format. This prompt checks that the checkpoint loads and follows a specific instruction. It is not a quality benchmark. To measure training gain, compare the base and fine-tuned checkpoints on the same held-out UltraChat examples with the same decoding settings.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "07ed3615", "metadata": {}, "outputs": [], "source": [ "import torch\n", "from transformers import AutoModelForCausalLM, AutoTokenizer\n", "\n", "CHECKPOINT = max(\n", " OUTPUT_PATH.glob(\"checkpoint-*\"),\n", " key=lambda p: int(p.name.removeprefix(\"checkpoint-\")),\n", ")\n", "print(CHECKPOINT)\n", "\n", "tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)\n", "model = AutoModelForCausalLM.from_pretrained(\n", " CHECKPOINT,\n", " dtype=torch.bfloat16,\n", " device_map=\"auto\",\n", ")\n", "\n", "messages = [\n", " {\n", " \"role\": \"user\",\n", " \"content\": (\n", " \"Explain how a heat pump warms a home in winter. \"\n", " \"Use three concise bullet points and avoid jargon.\"\n", " ),\n", " }\n", "]\n", "\n", "inputs = tokenizer.apply_chat_template(\n", " messages,\n", " add_generation_prompt=True,\n", " return_tensors=\"pt\",\n", " return_dict=True,\n", ").to(model.device)\n", "\n", "output = model.generate(\n", " **inputs,\n", " max_new_tokens=256,\n", " do_sample=False,\n", ")\n", "reply = tokenizer.decode(\n", " output[0, inputs[\"input_ids\"].shape[-1] :],\n", " skip_special_tokens=True,\n", ")\n", "print(reply)\n" ] }, { "cell_type": "markdown", "id": "8c1ca927", "metadata": {}, "source": [ "## Cleanup\n", "\n", "Shut down Jupyter, then press `Ctrl+C` in the terminal that runs the container. The `--rm` option removes the stopped container. The files under `DATA_DIR` remain available. Stop or terminate the cloud GPU machine when you finish to avoid further charges.\n" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 5 }