Qwen3-Coder-Next long-context LoRA fine-tuning example by ervinwang-jpg · Pull Request #86 · basetenlabs/ml-cookbook

ervinwang-jpg · 2026-03-11T17:59:10Z

Summary

Replaces qwen3-next-80b-megatron with a tested qwen3-80b-msswift example for Qwen3-Coder-Next (80B MoE) LoRA fine-tuning
Default config: 1 node, 8x H200, 32K seq length, ep=8, LoRA rank 8
README includes tested configurations table covering 16K–64K sequence lengths across 1–4 node setups
Uses MS-Swift + MegatronLM with flash-linear-attention, optimizer CPU offload, and activation recomputation

Tested Configurations

Nodes	GPUs	Seq Length	TP	PP	EP	Peak Memory	Status
1	8	16K	—	—	8	105 GiB	✅
1	8	32K	—	—	8	121 GiB	✅
4	32	48K	2	—	16	98 GiB	✅
4	32	64K	2	2	4	118 GiB	✅

Test plan

1-node 32K config verified end-to-end (training + checkpoint save)
4-node 48K and 64K configs verified end-to-end
Verify truss train init --examples qwen3-80b-msswift works after merge

🤖 Generated with Claude Code

Copilot

Pull request overview

Adds a new Baseten training example for fine-tuning Qwen3-Next-80B with Megatron (ms-swift), targeting 16k context length, and includes a helper to upload resulting checkpoints to the Hugging Face Hub.

Changes:

Introduces a Megatron SFT launcher script for Qwen3-Next-80B with 16k max length and LoRA settings.
Adds a Baseten truss_train job config to run the launcher in a prebuilt Megatron image.
Adds a small Python utility to upload the checkpoint directory to the Hugging Face Hub, plus README instructions.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 6 comments.

File	Description
examples/qwen3-next-80b-megatron/upload_checkpoint.py	New HF Hub upload helper for checkpoint folders.
examples/qwen3-next-80b-megatron/run_megatron.sh	New Megatron/ms-swift training launcher + post-training upload coordination.
examples/qwen3-next-80b-megatron/config.py	New Baseten training job configuration targeting H200s.
examples/qwen3-next-80b-megatron/README.md	New README describing how to run the example.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

You can also share your feedback on Copilot code review. Take the survey.

Copilot · 2026-03-11T18:07:10Z

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


training_compute is configured for node_count=1, but run_megatron.sh assumes multi-node (BT_NODE_RANK==1 performs the upload and BT_NODE_RANK==0 waits). With a single node, rank 0 will wait for an upload marker that can never be created and the job will timeout/fail. Align the config with the script by either setting node_count >= 2 or updating the upload logic to handle the single-node case (e.g., rank 0 uploads when BT_GROUP_SIZE==1).

Copilot · 2026-03-11T18:07:11Z

+# Upload checkpoint (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)"
+
+if [[ "$BT_NODE_RANK" == "1" ]]; then
+  python "$script_dir/upload_checkpoint.py" "$checkpoint_dir" "$hub_repo"
+  touch "$upload_marker"
+elif [[ "$BT_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


The upload/wait logic hard-codes uploads to BT_NODE_RANK==1. If this script is run with a single node (or any setup without a rank 1), rank 0 will always hit the wait loop and eventually error with "Upload timeout". Consider making the uploader conditional on world size (e.g., upload on rank 0 when BT_GROUP_SIZE==1, otherwise upload on rank 1) or selecting the uploader as BT_NODE_RANK==$((BT_GROUP_SIZE-1)).

Copilot · 2026-03-11T18:07:11Z

+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" \


With set -u, referencing ${BT_GROUP_SIZE} / ${BT_NUM_GPUS} / ${BT_NODE_RANK} / ${BT_LEADER_ADDR} will immediately error if the script is run outside the Baseten runtime (or if any of these env vars are missing). Other example scripts in this repo use default fallbacks (e.g., ${BT_GROUP_SIZE:-1}, ${BT_NODE_RANK:-0}, ${BT_LEADER_ADDR:-localhost}) to avoid unbound-variable exits while still working on-platform.

Suggested change

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" \

echo "Starting training: model=Qwen/Qwen3-Next-80B-A3B-Instruct nodes=${BT_GROUP_SIZE:-1}x${BT_NUM_GPUS:-1}gpu"

train_exit=0

NPROC_PER_NODE="${BT_NUM_GPUS:-1}" \

NNODES="${BT_GROUP_SIZE:-1}" \

NODE_RANK="${BT_NODE_RANK:-0}" \

MASTER_ADDR="${BT_LEADER_ADDR:-localhost}" \

Copilot · 2026-03-11T18:07:11Z

+  done
+  sleep 1
+done) &
+trap "kill $! 2>/dev/null" EXIT


The background args.json sync loop is terminated via trap "kill $!" EXIT, which relies on $! still referring to the intended background PID at exit time. This is brittle if another background job is added later. Capture the PID into a dedicated variable right after starting the loop and kill that variable in the trap (similar to other scripts in this repo).

Suggested change

trap "kill $! 2>/dev/null" EXIT

args_sync_pid=$!

trap 'kill "$args_sync_pid" 2>/dev/null' EXIT

Copilot · 2026-03-11T18:07:12Z

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


upload_folder is called without repo_type. Elsewhere in the repo, uploads explicitly pass repo_type="model" to avoid ambiguity. Consider also explicitly passing the token to HfApi(token=...) (from HF_TOKEN / HUGGING_FACE_HUB_TOKEN / HUGGINGFACE_HUB_TOKEN) so the script is deterministic even if a different cached token exists on disk.

@copilot open a new pull request to apply changes based on this feedback

Copilot · 2026-03-11T18:07:12Z

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


The README's title and "Runtime knobs" list appear copied from the MiniMax example and do not match this example: the script hard-codes the model/dataset/parallelism flags and does not define variables like MODEL_ID, DATASET_SPLIT, etc. Update the README to reflect the actual knobs that can be edited (or refactor run_megatron.sh to use the documented variables).

Suggested change

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`

The training configuration is currently hard-coded inside `run_megatron.sh`.

To change how training runs, edit that script directly. Common things to adjust include:

- model and checkpoint paths

- dataset path / loader configuration

- tensor / pipeline / context / expert parallelism settings

- batch sizes and sequence length

- learning rate schedule and optimizer settings

- logging, checkpointing, and save-format options

Copilot · 2026-03-11T18:16:38Z

@ervinwang-jpg I've opened a new pull request, #87, to work on those changes. Once the pull request is ready, I'll request review from you.

rcano-baseten · 2026-03-11T18:26:55Z

+printf '{}' > "$checkpoint_dir/args.json"
+
+# Workaround: sync args.json to timestamped subdirs created by ms-swift
+(while true; do


why do we need this?

rcano-baseten · 2026-03-11T18:27:25Z

+  --model_type qwen3_next \
+  --save "$checkpoint_dir" \
+  --dataset winglian/pirate-ultrachat-10k \
+  --template minimax_m2 \


is this the right template?

rcano-baseten · 2026-03-11T18:27:59Z

+  --use_distributed_optimizer false \
+  --use_hf 1 || train_exit=$?
+
+# Upload checkpoint (node 1 uploads, node 0 waits)


let's avoid uploading checkpoints to HF - i think we can just write this to the baseten checkpoint dir

rcano-baseten · 2026-03-11T18:30:32Z

+  --use_hf 1 || train_exit=$?
+
+# Upload checkpoint (node 1 uploads, node 0 waits)
+hub_repo="baseten-admin/qwen80b-instruct-megatron-lora"


If we just have megatron write the checkpoint to our shared checkpoint dir, none of this is necessary

rcano-baseten · 2026-03-11T18:31:53Z

+# Install dependencies
+pip install -q --upgrade pip
+pip install -q "ms-swift[llm]==3.12.5" datasets huggingface_hub "transformers==4.57.1"
+


I think we want to install flash-linear-attention

FLA_PKG_DIR=$BT_RW_CACHE_DIR/fla_packages export PYTHONPATH=$FLA_PKG_DIR:$PYTHONPATH # Remove old polluted cache if it exists rm -rf $BT_RW_CACHE_DIR/pip_packages 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

rcano-baseten · 2026-03-11T18:34:00Z

+  --tensor_model_parallel_size 1 \
+  --pipeline_model_parallel_size 1 \
+  --context_parallel_size 1 \
+  --expert_model_parallel_size 8 \


if we set --recompute_num_layers 2 we can get FLA + 32k Seq Len

Remove the previous qwen3-next-80b-megatron example and replace with qwen3-80b-msswift, a tested Qwen3-Coder-Next LoRA fine-tuning example using MS-Swift + MegatronLM. Default config runs on 1x H200 node at 32K sequence length. README includes a tested configurations table showing working setups from 16K (1 node) to 64K (4 nodes). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Verified tp=2, pp=2, ep=4 works at 64K sequence length on just 2 H200 nodes (106 GiB peak memory). Updated README table and scaling example to show 2-node as the primary 64K option. Restored default code to 1-node 32K config. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Verified tp=2, pp=4, ep=4 works at 128K sequence length on 4 H200 nodes (134 GiB peak memory, 7 GiB headroom). tp=4 is not viable due to Qwen3-Coder-Next having only 2 KV heads (GQA). Code restored to 1-node 32K defaults. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

ervinwang-jpg added 6 commits March 11, 2026 10:57

qwen 80b

b742ddc

Update README.md

a8192a4

clean up

ed733dc

cleanup

be83e5c

cleanup

30657e6

name folder

96d95be

ervinwang-jpg requested review from Copilot and rcano-baseten March 11, 2026 18:03

Copilot started reviewing on behalf of ervinwang-jpg March 11, 2026 18:03 View session

Copilot AI reviewed Mar 11, 2026

View reviewed changes

Copilot AI mentioned this pull request Mar 11, 2026

Fix upload_checkpoint.py: explicit HF token and repo_type for upload_folder #87

Draft

name folder

ab80a18

rcano-baseten reviewed Mar 11, 2026

View reviewed changes

ervinwang-jpg force-pushed the qwen80b branch from 28b2aa3 to ab80a18 Compare March 11, 2026 18:29

rcano-baseten reviewed Mar 11, 2026

View reviewed changes

ervinwang-jpg and others added 2 commits March 11, 2026 14:38

Update run_megatron.sh

a67c3b0

rcano-baseten changed the title ~~Qwen80b~~ Qwen3-Coder-Next long-context LoRA fine-tuning example Mar 12, 2026

rcano-baseten and others added 2 commits March 12, 2026 12:09

rcano-baseten mentioned this pull request Mar 12, 2026

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

Merged

2 tasks

rcano-baseten closed this Mar 12, 2026

	trap "kill $! 2>/dev/null" EXIT
	args_sync_pid=$!
	trap 'kill "$args_sync_pid" 2>/dev/null' EXIT

-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`
+The training configuration is currently hard-coded inside `run_megatron.sh`.
+To change how training runs, edit that script directly. Common things to adjust include:
+- model and checkpoint paths
+- dataset path / loader configuration
+- tensor / pipeline / context / expert parallelism settings
+- batch sizes and sequence length
+- learning rate schedule and optimizer settings
+- logging, checkpointing, and save-format options

Uh oh!

Conversation

ervinwang-jpg commented Mar 11, 2026 • edited by rcano-baseten Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Tested Configurations

Test plan

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Mar 11, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 11, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 11, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 11, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 11, 2026

Choose a reason for hiding this comment

Uh oh!

ervinwang-jpg Mar 11, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 11, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI commented Mar 11, 2026

Uh oh!

rcano-baseten Mar 11, 2026

Choose a reason for hiding this comment

Uh oh!

rcano-baseten Mar 11, 2026

Choose a reason for hiding this comment

Uh oh!

rcano-baseten Mar 11, 2026

Choose a reason for hiding this comment

Uh oh!

rcano-baseten Mar 11, 2026

Choose a reason for hiding this comment

Uh oh!

rcano-baseten Mar 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rcano-baseten Mar 11, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

ervinwang-jpg commented Mar 11, 2026 •

edited by rcano-baseten

Loading

rcano-baseten Mar 11, 2026 •

edited

Loading