Qwen3-Coder-Next long-context LoRA fine-tuning example

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.

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
+)

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