|
| 1 | +{ |
| 2 | + "cells": [ |
| 3 | + { |
| 4 | + "cell_type": "markdown", |
| 5 | + "id": "intro", |
| 6 | + "metadata": {}, |
| 7 | + "source": "# Fine-tune LFM2.5-8B-A1B with Halo\n\nThis notebook fine-tunes [LFM2.5-8B-A1B](https://huggingface.co/LiquidAI/LFM2.5-8B-A1B) on UltraChat 200K.\n\nHalo distributes the model's 32 routed experts across two GPUs. Each process owns 16 experts during training.\n\nHalo gathers the expert weights when it saves. The result is a standard Hugging Face checkpoint.\n" |
| 8 | + }, |
| 9 | + { |
| 10 | + "cell_type": "markdown", |
| 11 | + "id": "support", |
| 12 | + "metadata": {}, |
| 13 | + "source": "## Configuration\n\nThis 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\nDo not enable context parallelism. LFM2 short-convolution layers operate across the sequence axis.\n" |
| 14 | + }, |
| 15 | + { |
| 16 | + "cell_type": "markdown", |
| 17 | + "id": "runtime", |
| 18 | + "metadata": {}, |
| 19 | + "source": "## Start the Halo notebook server\n\nThis notebook uses the public Halo image that matches the detected GPUs. It needs two matching NVIDIA GPUs and a large data volume.\n\nRun these commands on the host. Replace the three paths before you start the container.\n\n~~~bash\ngit clone --recurse-submodules https://github.com/whitecircle/halo.git\ngit clone https://github.com/Liquid4All/cookbook.git\n\nexport HALO_DIR=/path/to/halo\nexport COOKBOOK_DIR=/path/to/cookbook\nexport DATA_DIR=/path/to/large/volume\nGPU_NAME=\"$(nvidia-smi --query-gpu=name --format=csv,noheader | head -n 1)\"\ncase \"$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 ;;\nesac\nexport HALO_IMAGE\n\ndocker pull \"$HALO_IMAGE\"\ndocker run --rm -it --gpus '\"device=0,1\"' \\\n --network host \\\n --ipc=host --shm-size=128g \\\n --ulimit memlock=-1 --ulimit stack=67108864 \\\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 PYTHONPATH=/workspace \\\n -e CUDA_DEVICE_MAX_CONNECTIONS=1 \\\n -v \"$HALO_DIR\":/workspace \\\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\nOpen the URL printed by Jupyter. Then open this notebook.\n" |
| 20 | + }, |
| 21 | + { |
| 22 | + "cell_type": "markdown", |
| 23 | + "id": "preflight-title", |
| 24 | + "metadata": {}, |
| 25 | + "source": "## Check the runtime\n\nThe notebook does not create a CUDA context before the distributed launch.\n" |
| 26 | + }, |
| 27 | + { |
| 28 | + "cell_type": "code", |
| 29 | + "execution_count": null, |
| 30 | + "id": "preflight", |
| 31 | + "metadata": {}, |
| 32 | + "outputs": [], |
| 33 | + "source": "import subprocess\nfrom pathlib import Path\n\nRUN_ROOT = Path(\"/mnt/checkpoints\")\nCONFIG_PATH = RUN_ROOT / \"lfm25-8b-a1b-ultrachat-ep2.yaml\"\nOUTPUT_PATH = RUN_ROOT / \"lfm25-8b-a1b-ultrachat-ep2\"\n\ngpu_names = subprocess.check_output(\n [\n \"nvidia-smi\",\n \"--query-gpu=name\",\n \"--format=csv,noheader\",\n ],\n text=True,\n).strip().splitlines()\nassert len(gpu_names) == 2, f\"Expected 2 visible GPUs, found {len(gpu_names)}\"\n\nRUN_ROOT.mkdir(parents=True, exist_ok=True)\nprint(gpu_names)\n" |
| 34 | + }, |
| 35 | + { |
| 36 | + "cell_type": "markdown", |
| 37 | + "id": "config-title", |
| 38 | + "metadata": {}, |
| 39 | + "source": "## Create the training configuration\n\nThis configuration matches Halo's validated LFM2 MoE recipe.\n" |
| 40 | + }, |
| 41 | + { |
| 42 | + "cell_type": "code", |
| 43 | + "execution_count": null, |
| 44 | + "id": "config", |
| 45 | + "metadata": {}, |
| 46 | + "outputs": [], |
| 47 | + "source": "CONFIG_PATH.write_text(\n r\"\"\"\nmodel_name_or_path: LiquidAI/LFM2.5-8B-A1B\nmoe_balancing: bias_update\n\ndataset:\n- HuggingFaceH4/ultrachat_200k@train_sft\nconversation_field: messages\ntest_size: 0.01\ntrain_only_on_completions: true\nassistant_message_template: \"<|im_start|>assistant\\n\"\npad_token: \"<|pad|>\"\neos_token: \"<|im_end|>\"\n\nexpert_parallel_size: 2\nsave_sharded_ep: false\nuse_grouped_gemm: true\nfp32_router: true\nfp32_experts: false\n\nattn_implementation: flash_attention_2\nuse_liger_kernel: false\npacking: true\nmax_length: 8192\nbf16: true\n\nper_device_train_batch_size: 1\nper_device_eval_batch_size: 1\ngradient_accumulation_steps: 8\nnum_train_epochs: 1.0\ngradient_checkpointing: true\ngradient_checkpointing_kwargs:\n use_reentrant: false\n\noptim: adamw_torch_fused\nlearning_rate: 5.0e-06\nlr_scheduler_type: cosine\nwarmup_steps: 32\nmax_grad_norm: 1.0\n\nsave_strategy: steps\nsave_steps: 1000\neval_strategy: steps\neval_steps: 300\nsave_total_limit: 1\nsave_only_model: true\noutput_dir: /mnt/checkpoints/lfm25-8b-a1b-ultrachat-ep2\n\nlogging_steps: 1\nlogging_first_step: true\nreport_to: none\nremove_unused_columns: false\ndataloader_num_workers: 2\n\nuse_peft: false\n\"\"\".lstrip()\n)\nprint(CONFIG_PATH.read_text())\n" |
| 48 | + }, |
| 49 | + { |
| 50 | + "cell_type": "markdown", |
| 51 | + "id": "fields", |
| 52 | + "metadata": {}, |
| 53 | + "source": "Halo uses DeepEP for token dispatch and grouped GEMM for expert computation.\n\nThe router runs in FP32. The experts run in BF16.\n\nThe bias update balances expert selection without a router auxiliary loss.\n" |
| 54 | + }, |
| 55 | + { |
| 56 | + "cell_type": "markdown", |
| 57 | + "id": "dry-run-title", |
| 58 | + "metadata": {}, |
| 59 | + "source": "## Check the launch command\n\nThe dry run prints the two-process command without loading the model.\n" |
| 60 | + }, |
| 61 | + { |
| 62 | + "cell_type": "code", |
| 63 | + "execution_count": null, |
| 64 | + "id": "dry-run", |
| 65 | + "metadata": {}, |
| 66 | + "outputs": [], |
| 67 | + "source": "!halo launch sft {CONFIG_PATH} -n 2 --dry-run\n" |
| 68 | + }, |
| 69 | + { |
| 70 | + "cell_type": "markdown", |
| 71 | + "id": "train-title", |
| 72 | + "metadata": {}, |
| 73 | + "source": "## Start training\n\nThis cell starts the full UltraChat run on both visible GPUs.\n" |
| 74 | + }, |
| 75 | + { |
| 76 | + "cell_type": "code", |
| 77 | + "execution_count": null, |
| 78 | + "id": "train", |
| 79 | + "metadata": {}, |
| 80 | + "outputs": [], |
| 81 | + "source": "!halo launch sft {CONFIG_PATH} -n 2\n" |
| 82 | + }, |
| 83 | + { |
| 84 | + "cell_type": "markdown", |
| 85 | + "id": "load-title", |
| 86 | + "metadata": {}, |
| 87 | + "source": "## Load the saved checkpoint\n\nThe saved checkpoint uses the native Transformers format.\n" |
| 88 | + }, |
| 89 | + { |
| 90 | + "cell_type": "code", |
| 91 | + "execution_count": null, |
| 92 | + "id": "load", |
| 93 | + "metadata": {}, |
| 94 | + "outputs": [], |
| 95 | + "source": "import torch\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\ntokenizer = AutoTokenizer.from_pretrained(OUTPUT_PATH)\nmodel = AutoModelForCausalLM.from_pretrained(\n OUTPUT_PATH,\n dtype=torch.bfloat16,\n device_map=\"auto\",\n)\n\nmessages = [\n {\n \"role\": \"user\",\n \"content\": \"Summarize this support ticket and list the next actions.\",\n }\n]\ninputs = tokenizer.apply_chat_template(\n messages,\n add_generation_prompt=True,\n return_tensors=\"pt\",\n).to(model.device)\noutput = model.generate(\n inputs,\n max_new_tokens=256,\n do_sample=False,\n)\nreply = tokenizer.decode(\n output[0, inputs.shape[-1] :],\n skip_special_tokens=True,\n)\nprint(reply)\n" |
| 96 | + } |
| 97 | + ], |
| 98 | + "metadata": { |
| 99 | + "kernelspec": { |
| 100 | + "display_name": "Python 3", |
| 101 | + "language": "python", |
| 102 | + "name": "python3" |
| 103 | + }, |
| 104 | + "language_info": { |
| 105 | + "name": "python", |
| 106 | + "version": "3.12" |
| 107 | + } |
| 108 | + }, |
| 109 | + "nbformat": 4, |
| 110 | + "nbformat_minor": 5 |
| 111 | +} |
0 commit comments