Add Halo LFM2.5 Expert Parallel and vision training notebooks #120
There are no files selected for viewing
|
Contributor
There was a problem hiding this comment. can we please rename this to "sft_moe_with_halo.ipynb" to match naming convention of our fine-tuning notebooks? |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,111 @@ | ||
| { | ||
| "cells": [ | ||
| { | ||
| "cell_type": "markdown", | ||
| "id": "intro", | ||
| "metadata": {}, | ||
| "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" | ||
|
Contributor
There was a problem hiding this comment. As we're not able to run this notebook in a Google Colab like our other notebooks, I think it would be helpful to developers if we could get a short "Prerequisites" section for what kind of GPU is needed (you mention two matching NVIDIA GPUs and large data volume later, but I think having it as overview would be helpful) Maybe a note like this could be helpful as well.
I think it could also be helpful to give the readers a rough estimate of how much training time to anticipate for at least one GPU type (e.g. X hours for two H100s)
Contributor
There was a problem hiding this comment. mention "[...on UltraChat 200K] using Halo" and link to your docs please
Contributor
There was a problem hiding this comment. please add a link to the Hugging face dataset. |
||
| }, | ||
| { | ||
| "cell_type": "markdown", | ||
| "id": "support", | ||
| "metadata": {}, | ||
| "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" | ||
|
Contributor
There was a problem hiding this comment. I think readers would benefit from a short explanation of what "Expert parallelism" is (can be one or two sentences) Generally, I think the most valuable aspect of this notebook is that Halo allows us to fine-tune MoEs, so anything you might want to explain here for context might be really useful to developers. |
||
| }, | ||
| { | ||
| "cell_type": "markdown", | ||
| "id": "runtime", | ||
| "metadata": {}, | ||
|
Contributor
There was a problem hiding this comment. can we get some explanations here for a few of the settings: for example:
|
||
| "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" | ||
|
Contributor
There was a problem hiding this comment. "Replace the three paths" -> please name them explicitly (HALO_DIr, COOKBOOK_DIR, DATA_DIR) |
||
| }, | ||
| { | ||
| "cell_type": "markdown", | ||
| "id": "preflight-title", | ||
| "metadata": {}, | ||
| "source": "## Check the runtime\n\nThe notebook does not create a CUDA context before the distributed launch.\n" | ||
| }, | ||
| { | ||
| "cell_type": "code", | ||
| "execution_count": null, | ||
| "id": "preflight", | ||
| "metadata": {}, | ||
| "outputs": [], | ||
| "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" | ||
|
Contributor
There was a problem hiding this comment. would be great if we could have the cell outputs here |
||
| }, | ||
| { | ||
|
Contributor
There was a problem hiding this comment. even if no data processing is needed here, it could be useful to developers to have a "Data Prep" section here, that just briefly explains what data format is expected and why we dont need any data prep here |
||
| "cell_type": "markdown", | ||
| "id": "config-title", | ||
| "metadata": {}, | ||
| "source": "## Create the training configuration\n\nThis configuration matches Halo's validated LFM2 MoE recipe.\n" | ||
| }, | ||
| { | ||
| "cell_type": "code", | ||
| "execution_count": null, | ||
| "id": "config", | ||
| "metadata": {}, | ||
| "outputs": [], | ||
| "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" | ||
|
Contributor
There was a problem hiding this comment. It would be helpful to developers new to fine-tuning MoEs with Halo explain a few of the most important training configurations here? especially the halo specific ones:
|
||
| }, | ||
| { | ||
| "cell_type": "markdown", | ||
| "id": "fields", | ||
| "metadata": {}, | ||
| "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" | ||
| }, | ||
| { | ||
| "cell_type": "markdown", | ||
| "id": "dry-run-title", | ||
| "metadata": {}, | ||
| "source": "## Check the launch command\n\nThe dry run prints the two-process command without loading the model.\n" | ||
|
Contributor
There was a problem hiding this comment. Should we have this dry run check in the VL notebook as well? |
||
| }, | ||
| { | ||
| "cell_type": "code", | ||
| "execution_count": null, | ||
| "id": "dry-run", | ||
| "metadata": {}, | ||
| "outputs": [], | ||
| "source": "!halo launch sft {CONFIG_PATH} -n 2 --dry-run\n" | ||
| }, | ||
| { | ||
| "cell_type": "markdown", | ||
| "id": "train-title", | ||
| "metadata": {}, | ||
| "source": "## Start training\n\nThis cell starts the full UltraChat run on both visible GPUs.\n" | ||
|
Contributor
There was a problem hiding this comment. Can we streamline naming conventions with our toher notebooks to call this section "Launch training" |
||
| }, | ||
| { | ||
| "cell_type": "code", | ||
| "execution_count": null, | ||
| "id": "train", | ||
| "metadata": {}, | ||
| "outputs": [], | ||
| "source": "!halo launch sft {CONFIG_PATH} -n 2\n" | ||
| }, | ||
| { | ||
| "cell_type": "markdown", | ||
| "id": "load-title", | ||
| "metadata": {}, | ||
| "source": "## Load the saved checkpoint\n\nThe saved checkpoint uses the native Transformers format.\n" | ||
|
Contributor
There was a problem hiding this comment. can we call this section "Inference" please to match our other notebooks |
||
| }, | ||
| { | ||
| "cell_type": "code", | ||
| "execution_count": null, | ||
| "id": "load", | ||
| "metadata": {}, | ||
| "outputs": [], | ||
| "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" | ||
|
Contributor
There was a problem hiding this comment. Is "Summarize this support ticket and list the next actions." the best prompt to test if the fine-tuning run did anything here? |
||
| } | ||
| ], | ||
| "metadata": { | ||
|
Contributor
There was a problem hiding this comment. Should we add a clean up section with stopping container, etc. so develoerps wont run into unexpected costs? |
||
| "kernelspec": { | ||
| "display_name": "Python 3", | ||
| "language": "python", | ||
| "name": "python3" | ||
| }, | ||
| "language_info": { | ||
| "name": "python", | ||
| "version": "3.12" | ||
| } | ||
| }, | ||
| "nbformat": 4, | ||
| "nbformat_minor": 5 | ||
| } | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
can you please move this below the cell for SFT with TRL with the following edits: