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

examples/qwen3-next-80b-megatron/README.md

@@ -0,0 +1,31 @@
+## Baseten Qwen 80b 2.5 Training
+
+### Prerequisites
+https://github.com/basetenlabs/ml-cookbook/tree/main?tab=readme-ov-file#prerequisites
+
+This example uses megatron training
+- `run_megatron.sh`
+
+Current files:
+- `config.py`: Baseten training job config
+- `run_megatron.sh`: Megatron SFT launcher.
+
+### Required Baseten secrets
+- `hf_access_token`
+
+### Runtime knobs
+Edit variables directly in the script you use:
+- `run_megatron.sh`:
+  - `MODEL_ID`, `DATASET_ID`, `DATASET_SPLIT`
+  - `LORA_RANK`, `LORA_ALPHA`
+  - `TENSOR_PARALLEL_SIZE`, `PIPELINE_PARALLEL_SIZE`, `CONTEXT_PARALLEL_SIZE`, `EXPERT_PARALLEL_SIZE`
+  - `MICRO_BATCH_SIZE`, `GLOBAL_BATCH_SIZE`, `MAX_LENGTH`
+  - `LR_DECAY_STYLE`, `MIN_LR`, `SAVE_INTERVAL`, `LOG_INTERVAL`, `SAVE_FULL_MODEL`
+
+### Launch
+`config.py` runs whichever script is in `training_runtime.start_commands`.
+
+Run:
+```bash
+truss train push config.py
+```

examples/qwen3-next-80b-megatron/config.py

@@ -0,0 +1,38 @@
+from truss_train import definitions
+from truss.base import truss_config
+
+project_name = "qwen80b-instruct-megatron-lora"
+
+# Prebuilt image with Megatron + ms-swift.
+BASE_IMAGE = "baseten/megatron:py3.11.11-cuda12.8.1-torch2.8.0-fa2.8.1-megatron0.14.1-msswift3.10.3"
+
+training_runtime = definitions.Runtime(
+    start_commands=[
+        "/bin/bash -c 'chmod +x ./run_megatron.sh && ./run_megatron.sh'",
+    ],
+    environment_variables={
+        "HF_TOKEN": definitions.SecretReference(name="hf_access_token"),
+    },
+    cache_config=definitions.CacheConfig(
+        enabled=True,
+        enable_legacy_hf_mount=True,
+    ),
+    checkpointing_config=definitions.CheckpointingConfig(enabled=True),
+)
+
+# Multi-node distributed job for large Qwen80B tuning.
+training_compute = definitions.Compute(
+    node_count=1,
+    accelerator=truss_config.AcceleratorSpec(
+        accelerator=truss_config.Accelerator.H200,
+        count=8,
+    ),
+)
+
+my_training_job = definitions.TrainingJob(
+    image=definitions.Image(base_image=BASE_IMAGE),
+    compute=training_compute,
+    runtime=training_runtime,
+)
+
+first_project_with_job = definitions.TrainingProject(name=project_name, job=my_training_job)

examples/qwen3-next-80b-megatron/run_megatron.sh

@@ -0,0 +1,111 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Validate HF token
+[[ -n "${HF_TOKEN:-}${HUGGING_FACE_HUB_TOKEN:-}${HUGGINGFACE_HUB_TOKEN:-}" ]] || {
+  echo "ERROR: HF token required for checkpoint upload"; exit 1
+}
+
+# Environment
+export HF_HOME="/tmp/huggingface"
+export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True"
+export TORCH_NCCL_ASYNC_ERROR_HANDLING="1"
+export NCCL_SOCKET_IFNAME="^docker0,lo"
+export MASTER_PORT="29500"
+mkdir -p "$HF_HOME"
+
+# Install dependencies
+pip install -q --upgrade pip
+pip install -q "ms-swift[llm]==3.12.5" datasets huggingface_hub "transformers==4.57.1"
+
+# Checkpoint directory
+checkpoint_dir="${BT_CHECKPOINT_DIR:-/mnt/ckpts}/qwen80b-instruct-megatron-lora"
+mkdir -p "$checkpoint_dir"
+printf '{}' > "$checkpoint_dir/args.json"
+
+# Workaround: sync args.json to timestamped subdirs created by ms-swift
+(while true; do
+  for d in "$checkpoint_dir"/v*-*; do
+    [[ -d "$d" && ! -f "$d/args.json" ]] && cp "$checkpoint_dir/args.json" "$d/args.json" 2>/dev/null || true
+  done
+  sleep 1
+done) &
+trap "kill $! 2>/dev/null" EXIT
+
+# Run training
+echo "Starting training: model=Qwen/Qwen3-Next-80B-A3B-Instruct nodes=${BT_GROUP_SIZE}x${BT_NUM_GPUS}gpu"
+train_exit=0
+NPROC_PER_NODE="$BT_NUM_GPUS" \
+NNODES="$BT_GROUP_SIZE" \
+NODE_RANK="$BT_NODE_RANK" \
+MASTER_ADDR="$BT_LEADER_ADDR" \
+MASTER_PORT="$MASTER_PORT" \
+megatron sft \
+  --model Qwen/Qwen3-Next-80B-A3B-Instruct \
+  --model_type qwen3_next \
+  --save "$checkpoint_dir" \
+  --dataset winglian/pirate-ultrachat-10k \
+  --template minimax_m2 \
+  --check_model false \
+  --load_safetensors true \
+  --train_type lora \
+  --lora_rank 8 \
+  --lora_alpha 16 \
+  --merge_lora false \
+  --target_modules all-linear \
+  --max_epochs 1 \
+  --lr_decay_style constant \
+  --clip_grad 1.0 \
+  --split_dataset_ratio 0.01 \
+  --tensor_model_parallel_size 1 \
+  --pipeline_model_parallel_size 1 \
+  --context_parallel_size 1 \
+  --expert_model_parallel_size 8 \
+  --bf16 true \
+  --loss_scale default \
+  --micro_batch_size 1 \
+  --global_batch_size 8 \
+  --packing false \
+  --cross_entropy_loss_fusion true \
+  --recompute_granularity selective \
+  --recompute_modules core_attn moe \
+  --lr 2e-4 \
+  --lr_warmup_fraction 0.05 \
+  --min_lr 1e-5 \
+  --max_length 16384 \
+  --save_interval 5 \
+  --log_interval 1 \
+  --num_workers 8 \
+  --dataset_num_proc 8 \
+  --lazy_tokenize true \
+  --load_from_cache_file true \
+  --no_save_optim true \
+  --no_save_rng true \
+  --sequence_parallel true \
+  --attention_backend flash \
+  --overlap_grad_reduce false \
+  --overlap_param_gather false \
+  --use_distributed_optimizer false \
+  --use_hf 1 || train_exit=$?
+
+# Upload checkpoint (single-node: node 0 uploads; multi-node: node 1 uploads, node 0 waits)
+hub_repo="baseten-admin/qwen80b-instruct-megatron-lora"
+upload_marker="$checkpoint_dir/.upload_done"
+script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+node_rank="${BT_NODE_RANK:-0}"
+num_nodes="${BT_GROUP_SIZE:-1}"
+
+if [[ "$num_nodes" == "1" || "$node_rank" == "1" ]]; then
+  python "$script_dir/upload_checkpoint.py" "$checkpoint_dir" "$hub_repo"
+  touch "$upload_marker"
+elif [[ "$node_rank" == "0" ]]; then
+  waited=0
+  while [[ ! -f "$upload_marker" && $waited -lt 3600 ]]; do sleep 5; ((waited+=5)); done
+  [[ -f "$upload_marker" ]] || { echo "Upload timeout"; exit 1; }
+fi
+
+# Exit with training status (success if checkpoints exist despite non-zero exit)
+if [[ $train_exit -ne 0 ]]; then
+  find "$checkpoint_dir" -name "*.safetensors" -type f | grep -q . && exit 0
+  exit $train_exit
+fi

examples/qwen3-next-80b-megatron/upload_checkpoint.py

@@ -0,0 +1,25 @@
+#!/usr/bin/env python3
+"""Upload checkpoint to Hugging Face Hub."""
+import argparse
+import datetime
+from huggingface_hub import HfApi
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("checkpoint_dir", help="Path to checkpoint directory")
+    parser.add_argument("repo_id", help="HF Hub repo ID (e.g., user/model-name)")
+    args = parser.parse_args()
+
+    api = HfApi()
+    api.create_repo(args.repo_id, repo_type="model", private=True, exist_ok=True)
+    api.upload_folder(
+        repo_id=args.repo_id,
+        folder_path=args.checkpoint_dir,
+        commit_message=f"Checkpoint {datetime.datetime.utcnow().isoformat()}Z",
+    )
+    print(f"Uploaded to {args.repo_id}")
+
+
+if __name__ == "__main__":
+    main()