diff --git a/examples/qwen3-80b-msswift/training/README.md b/examples/qwen3-80b-msswift/training/README.md new file mode 100644 index 0000000..89e990c --- /dev/null +++ b/examples/qwen3-80b-msswift/training/README.md @@ -0,0 +1,58 @@ +# Qwen3-Coder-Next Long Context Fine-Tuning with MS-Swift + +This example fine-tunes the [Qwen3-Coder-Next](https://huggingface.co/Qwen/Qwen3-Coder-Next) (80B MoE) model using LoRA with the MS-Swift framework and MegatronLM on Baseten. It uses the [LongAlign-10k](https://huggingface.co/datasets/zai-org/LongAlign-10k) dataset for long-context supervised fine-tuning. + +**Default configuration:** 1 node, 8x H200 GPUs, 32K sequence length + +## Tested Configurations + +| Nodes | GPUs | Seq Length | TP | PP | EP | LoRA Rank | Recompute Layers | Peak Memory (GiB) | Time/Iter | +|-------|------|------------|----|----|----|-----------|------------------|--------------------|-----------| +| 1 | 8 | 16K | — | — | 8 | 64 | 4 | 105 | ~35s | +| 1 | 8 | 32K | — | — | 8 | 8 | 2 | 121 | ~40s | +| 2 | 16 | 64K | 2 | 2 | 4 | 8 | 2 | 106 | ~338s | +| 4 | 32 | 48K | 2 | — | 16 | 8 | 1 | 98 | ~270s | +| 4 | 32 | 64K | 2 | 2 | 4 | 8 | 1 | 118 | ~455s | +| 4 | 32 | 128K | 2 | 4 | 4 | 8 | 2 | 134 | ~663s | + +## Prerequisites + +1. [Create a Baseten account](https://baseten.co/signup) if you don't already have one. +2. Install the Truss CLI: + ```bash + # pip + pip install -U truss + # or uv + uv add truss + ``` + +## Getting Started + +Initialize the example, navigate into the directory, and push the training job: + +```bash +truss train init --examples qwen3-80b-msswift +cd qwen3-80b-msswift +truss train push training/config.py +``` + +### Scaling to longer sequences + +To train at longer sequence lengths, increase `node_count` in `config.py` and adjust parallelism flags in `run.sh`. For example, for 64K on 2 nodes: + +```python +# config.py +node_count=2 +``` + +```bash +# run.sh — key flags to change +--tensor_model_parallel_size 2 +--pipeline_model_parallel_size 2 +--expert_model_parallel_size 4 +--global_batch_size 4 +--recompute_num_layers 2 +--max_length 64000 +``` + +> **Note:** This example requires H200 GPUs. You may need to [contact Baseten](https://www.baseten.co/contact) to get approval for this instance type before running the job. diff --git a/examples/qwen3-80b-msswift/training/config.py b/examples/qwen3-80b-msswift/training/config.py new file mode 100644 index 0000000..c3437ca --- /dev/null +++ b/examples/qwen3-80b-msswift/training/config.py @@ -0,0 +1,49 @@ +# Import necessary classes from the Baseten Training SDK +from truss_train import definitions +from truss.base import truss_config + +project_name = "LoRA Qwen3-Coder-Next Long Context - ML Cookbook" + +# 1. Define a base image for your training job +BASE_IMAGE = "baseten/megatron:py3.11.11-cuda12.8.1-torch2.8.0-fa2.8.1-megatron0.14.1-msswift3.10.3" + +# 2. Define the Runtime Environment for the Training Job +# This includes start commands and environment variables. +# Secrets from the baseten workspace like API keys are referenced using +# `SecretReference`. + +training_runtime = definitions.Runtime( + start_commands=[ # Example: list of commands to run your training script + "chmod +x ./run.sh && ./run.sh" + ], + checkpointing_config=definitions.CheckpointingConfig( + enabled=True, + ), + cache_config=definitions.CacheConfig( + enabled=True, + ), +) + +# 3. Define the Compute Resources for the Training Job +training_compute = definitions.Compute( + node_count=1, + accelerator=truss_config.AcceleratorSpec( + accelerator=truss_config.Accelerator.H200, + count=8, + ), +) + +# 4. Define the Training Job +# This brings together the image, compute, and runtime configurations. +my_training_job = definitions.TrainingJob( + image=definitions.Image(base_image=BASE_IMAGE), + compute=training_compute, + runtime=training_runtime, +) + + +# This config will be pushed using the Truss CLI. +# The association of the job to the project happens at the time of push. +first_project_with_job = definitions.TrainingProject( + name=project_name, job=my_training_job +) diff --git a/examples/qwen3-80b-msswift/training/run.sh b/examples/qwen3-80b-msswift/training/run.sh new file mode 100644 index 0000000..09946ca --- /dev/null +++ b/examples/qwen3-80b-msswift/training/run.sh @@ -0,0 +1,77 @@ +#!/bin/bash + +FLA_PKG_DIR=$BT_PROJECT_CACHE_DIR/fla_packages +export PYTHONPATH=$FLA_PKG_DIR:$PYTHONPATH +if python -c "import fla" 2>/dev/null; then + echo "flash-linear-attention already installed in cache, skipping" +else + echo "Installing flash-linear-attention to cache" + pip install --target=$FLA_PKG_DIR --no-deps flash-linear-attention fla-core +fi + +SAVE_FULL_MODEL=false +checkpoint_dir="$BT_CHECKPOINT_DIR/qwen3-coder-next-lora-8-16" + +PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True NPROC_PER_NODE=$BT_NUM_GPUS NNODES=$BT_GROUP_SIZE NODE_RANK=$BT_NODE_RANK MASTER_ADDR=$BT_LEADER_ADDR megatron sft \ + --model Qwen/Qwen3-Coder-Next \ + --model_type qwen3_next \ + --save $checkpoint_dir \ + --dataset 'zai-org/LongAlign-10k' \ + --load_safetensors true \ + --save_safetensors true \ + --train_type lora \ + --lora_rank 8 \ + --lora_alpha 16 \ + --target_modules all-linear \ + --no_initialization false \ + --split_dataset_ratio 0.01 \ + --expert_model_parallel_size 8 \ + --moe_permute_fusion true \ + --moe_grouped_gemm true \ + --moe_shared_expert_overlap true \ + --moe_aux_loss_coeff 1e-3 \ + --micro_batch_size 1 \ + --global_batch_size 8 \ + --packing true \ + --recompute_granularity full \ + --recompute_method uniform \ + --recompute_num_layers 2 \ + --train_iters 100 \ + --eval_iters 10 \ + --finetune true \ + --cross_entropy_loss_fusion true \ + --lr 1e-4 \ + --lr_warmup_fraction 0.05 \ + --min_lr 1e-5 \ + --eval_interval 2 \ + --max_length 32000 \ + --num_workers 8 \ + --dataset_num_proc 8 \ + --no_save_optim true \ + --no_save_rng true \ + --sequence_parallel true \ + --attention_backend flash \ + --optimizer_cpu_offload true \ + --use_precision_aware_optimizer true \ + --merge_lora $SAVE_FULL_MODEL \ + --use_hf 1 + +# Only check for safetensors on the last node +if [ $BT_NODE_RANK -ne $(($BT_GROUP_SIZE - 1)) ]; then + # Non-master nodes spin forever; master node sends the exit code + sleep infinity +fi + +# Capture the exit code +MEGATRON_EXIT_CODE=$? + +# If the command failed, check if safetensors exist in checkpoint_dir +if [ $MEGATRON_EXIT_CODE -ne 0 ]; then + if [ -d "$checkpoint_dir" ] && [ -n "$(find "$checkpoint_dir" -name "*.safetensors" -type f 2>/dev/null)" ]; then + echo "Safetensors found in $checkpoint_dir. Exiting successfully." + exit 0 + else + echo "Megatron command failed and no safetensors found in $checkpoint_dir. Exiting with error code." + exit $MEGATRON_EXIT_CODE + fi +fi