From b742ddc61dd13498d40058686797c384049a8117 Mon Sep 17 00:00:00 2001
From: Ervin Wang <ervin.wang@baseten.co>
Date: Wed, 11 Mar 2026 10:57:32 -0400
Subject: [PATCH 01/11] qwen 80b

---
 examples/qwen80b/README.md       |  31 ++++
 examples/qwen80b/config.py       |  38 ++++
 examples/qwen80b/run_megatron.sh | 302 +++++++++++++++++++++++++++++++
 3 files changed, 371 insertions(+)
 create mode 100644 examples/qwen80b/README.md
 create mode 100644 examples/qwen80b/config.py
 create mode 100644 examples/qwen80b/run_megatron.sh

diff --git a/examples/qwen80b/README.md b/examples/qwen80b/README.md
new file mode 100644
index 0000000..3f4a5e5
--- /dev/null
+++ b/examples/qwen80b/README.md
@@ -0,0 +1,31 @@
+## Baseten MiniMax 2.5 Training (Multi-Node, Megatron)
+
+### Prerequisites
+https://github.com/basetenlabs/ml-cookbook/tree/main?tab=readme-ov-file#prerequisites
+
+This example uses megatron training
+- `run_megatron.sh` (Megatron path with TP/EP and FP8 training flags).
+
+Current files:
+- `config.py`: Baseten training job config (2 nodes x 8 GPUs, H200 in current 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
+```
diff --git a/examples/qwen80b/config.py b/examples/qwen80b/config.py
new file mode 100644
index 0000000..5a65f61
--- /dev/null
+++ b/examples/qwen80b/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)
diff --git a/examples/qwen80b/run_megatron.sh b/examples/qwen80b/run_megatron.sh
new file mode 100644
index 0000000..68db0be
--- /dev/null
+++ b/examples/qwen80b/run_megatron.sh
@@ -0,0 +1,302 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+LOG_BASE_DIR="${BT_CHECKPOINT_DIR:-/mnt/ckpts}/debug_logs"
+mkdir -p "${LOG_BASE_DIR}"
+LOG_FILE="${LOG_BASE_DIR}/run-megatron-node-${BT_NODE_RANK:-0}-$(date +%Y%m%d-%H%M%S).log"
+exec > >(tee -a "${LOG_FILE}") 2>&1
+
+echo "==== run_megatron.sh start node=${BT_NODE_RANK:-0} $(date -Is) ===="
+echo "log_file=${LOG_FILE}"
+
+if [[ -z "${HF_TOKEN:-}" && -z "${HUGGING_FACE_HUB_TOKEN:-}" && -z "${HUGGINGFACE_HUB_TOKEN:-}" ]]; then
+  echo "ERROR: HF token is not set. An HF token is required to upload final checkpoints to Hugging Face. Configure Baseten secret 'hf_access_token' and map it to HF_TOKEN in config.py."
+  exit 1
+fi
+
+CACHE_ROOT="/tmp"
+export HF_HOME="${CACHE_ROOT}/huggingface"
+export HUGGINGFACE_HUB_CACHE="${HF_HOME}/hub"
+export TRANSFORMERS_CACHE="${HF_HOME}/transformers"
+export HF_DATASETS_CACHE="${HF_HOME}/datasets"
+export PIP_CACHE_DIR="${CACHE_ROOT}/pip"
+export TRITON_CACHE_DIR="${CACHE_ROOT}/triton-cache"
+export TORCH_EXTENSIONS_DIR="${CACHE_ROOT}/torch-extensions"
+export NCCL_DEBUG="NONE"
+export TORCH_DISTRIBUTED_DEBUG="OFF"
+export TORCH_NCCL_ASYNC_ERROR_HANDLING="1"
+export NCCL_SOCKET_IFNAME="^docker0,lo"
+export OMP_NUM_THREADS="4"
+export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True"
+export TORCH_COMPILE_DISABLE="1"
+export TORCHINDUCTOR_COMPILE_THREADS="1"
+export TORCH_DISABLE_ADDR2LINE=1
+
+export TORCH_DISTRIBUTED_DEBUG=NONE
+export NCCL_DEBUG=INFO
+
+export MASTER_PORT="29500"
+
+mkdir -p "${HF_HOME}" "${PIP_CACHE_DIR}" "${TRITON_CACHE_DIR}" "${TORCH_EXTENSIONS_DIR}"
+export CUDA_LAUNCH_BLOCKING=1
+
+
+PY_BIN="$(command -v python3 || command -v python)"
+if [[ -z "${PY_BIN}" ]]; then
+  echo "No Python interpreter found in PATH." >&2
+  exit 1
+fi
+
+SWIFT_VERSION="3.12.5"
+"${PY_BIN}" -m pip install --upgrade pip
+"${PY_BIN}" -m pip install "ms-swift[llm]==${SWIFT_VERSION}" datasets huggingface_hub
+"${PY_BIN}" -m pip install "transformers==4.57.1" -U
+
+# Training variables (edit these directly; all set to fixed defaults).
+MODEL_ID="Qwen/Qwen3-Next-80B-A3B-Instruct"
+MODEL_TYPE="qwen3_next"
+DATASET_ID="winglian/pirate-ultrachat-10k"
+DATASET_SPLIT="train"
+CHECKPOINT_NAME="qwen80b-instruct-megatron-lora"
+RUN_NAME="qwen80b-instruct-megatron-lora"
+MODEL_ARG="${MODEL_ID}"
+
+LORA_RANK=8
+LORA_ALPHA=16
+SPLIT_DATASET_RATIO=0.01
+TENSOR_PARALLEL_SIZE=1
+PIPELINE_PARALLEL_SIZE=1
+CONTEXT_PARALLEL_SIZE=1
+EXPERT_PARALLEL_SIZE=8
+DESIRED_DATA_PARALLEL_SIZE=1
+MICRO_BATCH_SIZE=1
+GLOBAL_BATCH_SIZE=8
+MAX_EPOCHS=1
+LR=2e-4
+LR_DECAY_STYLE="constant"
+LR_WARMUP_FRACTION=0.05
+MIN_LR=1e-5
+MAX_LENGTH=2048
+NUM_WORKERS=8
+DATASET_NUM_PROC=8
+SAVE_FULL_MODEL="false"
+SAVE_INTERVAL=5
+LOG_INTERVAL=1
+REPORT_TO="none"
+MSSWIFT_COMPAT_MODE="true"
+
+
+if [[ ! -d "${MODEL_ID}" ]]; then
+  LOCAL_MODEL_DIR="${CACHE_ROOT}/model-snapshots/${MODEL_ID//\//__}"
+  mkdir -p "${LOCAL_MODEL_DIR}"
+  echo "Pre-downloading model snapshot to cache: ${LOCAL_MODEL_DIR}"
+  MODEL_ID_ENV="${MODEL_ID}" LOCAL_MODEL_DIR_ENV="${LOCAL_MODEL_DIR}" "${PY_BIN}" - <<'PY'
+import os
+from huggingface_hub import snapshot_download
+
+snapshot_download(
+    repo_id=os.environ["MODEL_ID_ENV"],
+    local_dir=os.environ["LOCAL_MODEL_DIR_ENV"],
+    local_dir_use_symlinks=False,
+    resume_download=True,
+)
+print(f"snapshot_ready={os.environ['LOCAL_MODEL_DIR_ENV']}")
+PY
+  MODEL_ARG="${LOCAL_MODEL_DIR}"
+fi
+
+checkpoint_dir="${BT_CHECKPOINT_DIR:-/mnt/ckpts}/${CHECKPOINT_NAME}"
+mkdir -p "${checkpoint_dir}"
+printf '{}' > "${checkpoint_dir}/args.json"
+
+# Megatron constraint:
+# global_batch_size % (micro_batch_size * data_parallel_size) == 0
+WORLD_SIZE=$(( ${BT_GROUP_SIZE:-1} * ${BT_NUM_GPUS:-1} ))
+MODEL_PARALLEL_SIZE=$(( TENSOR_PARALLEL_SIZE * PIPELINE_PARALLEL_SIZE * CONTEXT_PARALLEL_SIZE * EXPERT_PARALLEL_SIZE ))
+if (( MODEL_PARALLEL_SIZE <= 0 )); then
+  MODEL_PARALLEL_SIZE=1
+fi
+DATA_PARALLEL_SIZE=$(( WORLD_SIZE / MODEL_PARALLEL_SIZE ))
+if (( DATA_PARALLEL_SIZE <= 0 )); then
+  DATA_PARALLEL_SIZE=1
+fi
+if (( WORLD_SIZE != MODEL_PARALLEL_SIZE * DESIRED_DATA_PARALLEL_SIZE )); then
+  echo "ERROR: WORLD_SIZE=${WORLD_SIZE} does not match requested parallelism:"
+  echo "  tp=${TENSOR_PARALLEL_SIZE} pp=${PIPELINE_PARALLEL_SIZE} cp=${CONTEXT_PARALLEL_SIZE} ep=${EXPERT_PARALLEL_SIZE} dp=${DESIRED_DATA_PARALLEL_SIZE}"
+  echo "Expected WORLD_SIZE=$(( MODEL_PARALLEL_SIZE * DESIRED_DATA_PARALLEL_SIZE ))."
+  exit 1
+fi
+if (( DATA_PARALLEL_SIZE != DESIRED_DATA_PARALLEL_SIZE )); then
+  echo "ERROR: computed DATA_PARALLEL_SIZE=${DATA_PARALLEL_SIZE}, expected ${DESIRED_DATA_PARALLEL_SIZE}."
+  exit 1
+fi
+MIN_DIVISOR=$(( MICRO_BATCH_SIZE * DATA_PARALLEL_SIZE ))
+if (( MIN_DIVISOR <= 0 )); then
+  MIN_DIVISOR=1
+fi
+if (( GLOBAL_BATCH_SIZE % MIN_DIVISOR != 0 )); then
+  echo "Adjusting GLOBAL_BATCH_SIZE from ${GLOBAL_BATCH_SIZE} to ${MIN_DIVISOR} to satisfy Megatron divisibility."
+  GLOBAL_BATCH_SIZE="${MIN_DIVISOR}"
+fi
+
+echo "Launching ms-swift Megatron SFT for ${MODEL_ID}"
+echo "model_path=${MODEL_ARG}"
+echo "checkpoint_dir=${checkpoint_dir}"
+echo "BT_GROUP_SIZE=${BT_GROUP_SIZE} BT_NUM_GPUS=${BT_NUM_GPUS} BT_NODE_RANK=${BT_NODE_RANK}"
+echo "MASTER_ADDR=${BT_LEADER_ADDR} MASTER_PORT=${MASTER_PORT}"
+echo "ms_swift_version=${SWIFT_VERSION}"
+echo "profile=working_config TP=${TENSOR_PARALLEL_SIZE} EP=${EXPERT_PARALLEL_SIZE} PP=${PIPELINE_PARALLEL_SIZE} CP=${CONTEXT_PARALLEL_SIZE} LORA_RANK=${LORA_RANK} MAX_LENGTH=${MAX_LENGTH} MICRO_BATCH_SIZE=${MICRO_BATCH_SIZE} GLOBAL_BATCH_SIZE=${GLOBAL_BATCH_SIZE}"
+echo "model_type=${MODEL_TYPE}"
+echo "MSSWIFT_COMPAT_MODE=${MSSWIFT_COMPAT_MODE}"
+
+# ms-swift may write checkpoints into timestamped subdirs (e.g., v0-YYYYMMDD-HHMMSS).
+# Keep args.json mirrored there to avoid checkpoint save failures.
+ARGS_SYNC_PID=""
+sync_args_json_loop() {
+  while true; do
+    for run_dir in "${checkpoint_dir}"/v*-*; do
+      if [[ -d "${run_dir}" ]] && [[ ! -f "${run_dir}/args.json" ]]; then
+        cp "${checkpoint_dir}/args.json" "${run_dir}/args.json" || true
+      fi
+    done
+    sleep 1
+  done
+}
+
+sync_args_json_loop &
+ARGS_SYNC_PID=$!
+cleanup_args_sync() {
+  if [[ -n "${ARGS_SYNC_PID}" ]]; then
+    kill "${ARGS_SYNC_PID}" 2>/dev/null || true
+  fi
+}
+trap cleanup_args_sync EXIT
+
+set +e
+SPLIT_ARGS=()
+if [[ "${DATASET_SPLIT}" != "train" ]]; then
+  SPLIT_JSON="[{\"hf_dataset_id\": \"${DATASET_ID}\", \"split\": [\"${DATASET_SPLIT}\"]}]"
+  SPLIT_ARGS=(--custom_dataset_info "${SPLIT_JSON}")
+fi
+
+COMPAT_ARGS=()
+if [[ "${MSSWIFT_COMPAT_MODE}" == "true" ]]; then
+  COMPAT_ARGS=(
+    --overlap_grad_reduce false
+    --overlap_param_gather false
+    --use_distributed_optimizer false
+  )
+fi
+
+PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF}" \
+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 "${MODEL_ARG}" \
+  --model_type "${MODEL_TYPE}" \
+  --save "${checkpoint_dir}" \
+  --dataset "${DATASET_ID}" \
+  --template minimax_m2 \
+  --check_model false \
+  --load_safetensors true \
+  --train_type lora \
+  --lora_rank "${LORA_RANK}" \
+  --lora_alpha "${LORA_ALPHA}" \
+  --merge_lora "${SAVE_FULL_MODEL}" \
+  --target_modules all-linear \
+  --max_epochs "${MAX_EPOCHS}" \
+  --lr_decay_style "${LR_DECAY_STYLE}" \
+  --clip_grad 1.0 \
+  --split_dataset_ratio "${SPLIT_DATASET_RATIO}" \
+  --tensor_model_parallel_size "${TENSOR_PARALLEL_SIZE}" \
+  --pipeline_model_parallel_size "${PIPELINE_PARALLEL_SIZE}" \
+  --context_parallel_size "${CONTEXT_PARALLEL_SIZE}" \
+  --expert_model_parallel_size "${EXPERT_PARALLEL_SIZE}" \
+  --bf16 true \
+  --loss_scale default \
+  --micro_batch_size "${MICRO_BATCH_SIZE}" \
+  --global_batch_size "${GLOBAL_BATCH_SIZE}" \
+  --packing false \
+  --cross_entropy_loss_fusion true \
+  --recompute_granularity selective \
+  --recompute_modules core_attn moe \
+  --lr "${LR}" \
+  --lr_warmup_fraction "${LR_WARMUP_FRACTION}" \
+  --min_lr "${MIN_LR}" \
+  --max_length "${MAX_LENGTH}" \
+  --save_interval "${SAVE_INTERVAL}" \
+  --log_interval "${LOG_INTERVAL}" \
+  --num_workers "${NUM_WORKERS}" \
+  --dataset_num_proc "${DATASET_NUM_PROC}" \
+  --lazy_tokenize true \
+  --load_from_cache_file true \
+  --no_save_optim true \
+  --no_save_rng true \
+  --sequence_parallel true \
+  --attention_backend flash \
+  "${COMPAT_ARGS[@]}" \
+  "${SPLIT_ARGS[@]}" \
+  --use_hf 1
+TRAIN_EXIT_CODE=$?
+set -e
+
+
+
+export CHECKPOINT_DIR="${checkpoint_dir}"
+export HUB_MODEL_ID="baseten-admin/qwen80b-instruct-megatron-lora"
+NODE1_UPLOAD_DONE_MARKER="${CHECKPOINT_DIR}/.node1_upload_done"
+# Final checkpoint upload to Hugging Face Hub from rank 0.
+if [[ "${BT_NODE_RANK}" == "1" ]]; then
+  echo "Starting final HF Hub upload from ${CHECKPOINT_DIR} to ${HUB_MODEL_ID}..."
+  "${PY_BIN}" - <<'PY'
+import datetime
+import os
+from huggingface_hub import HfApi
+
+checkpoint_dir = os.environ["CHECKPOINT_DIR"]
+repo_id = os.environ["HUB_MODEL_ID"]
+token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")
+
+api = HfApi(token=token)
+api.create_repo(repo_id=repo_id, repo_type="model", private=True, exist_ok=True)
+api.upload_folder(
+    repo_id=repo_id,
+    repo_type="model",
+    folder_path=checkpoint_dir,
+    commit_message=f"Final checkpoint upload {datetime.datetime.utcnow().isoformat()}Z",
+)
+print(f"hf_upload_complete repo={repo_id} folder={checkpoint_dir}")
+PY
+  touch "${NODE1_UPLOAD_DONE_MARKER}"
+fi
+
+if [[ "${BT_NODE_RANK}" == "0" ]]; then
+  NODE0_WAIT_TIMEOUT_SECONDS="${NODE0_WAIT_TIMEOUT_SECONDS:-3600}"
+  NODE0_WAIT_POLL_SECONDS="${NODE0_WAIT_POLL_SECONDS:-5}"
+  echo "Node 0 waiting for node 1 upload marker: ${NODE1_UPLOAD_DONE_MARKER} (timeout ${NODE0_WAIT_TIMEOUT_SECONDS}s)"
+  waited=0
+  while [[ ! -f "${NODE1_UPLOAD_DONE_MARKER}" ]]; do
+    sleep "${NODE0_WAIT_POLL_SECONDS}"
+    waited=$(( waited + NODE0_WAIT_POLL_SECONDS ))
+    if (( waited >= NODE0_WAIT_TIMEOUT_SECONDS )); then
+      echo "Timed out waiting for node 1 upload completion marker."
+      exit 1
+    fi
+  done
+  echo "Node 1 upload marker found."
+fi
+
+
+
+
+echo "[rank ${BT_NODE_RANK}] megatron exit code=${TRAIN_EXIT_CODE}"
+if [[ "${TRAIN_EXIT_CODE}" -ne 0 ]]; then
+  if [[ -d "${checkpoint_dir}" ]] && find "${checkpoint_dir}" -name "*.safetensors" -type f | grep -q .; then
+    echo "Training exited ${TRAIN_EXIT_CODE}, but safetensors exist in ${checkpoint_dir}; treating as success."
+    exit 0
+  fi
+  exit "${TRAIN_EXIT_CODE}"
+fi

From a8192a41f96b5ca9710f35fdee347ab91bd7f3ce Mon Sep 17 00:00:00 2001
From: ervinwang-jpg <ervin.wang@baseten.co>
Date: Wed, 11 Mar 2026 11:00:36 -0400
Subject: [PATCH 02/11] Update README.md

---
 examples/qwen80b/README.md | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/examples/qwen80b/README.md b/examples/qwen80b/README.md
index 3f4a5e5..7646f6b 100644
--- a/examples/qwen80b/README.md
+++ b/examples/qwen80b/README.md
@@ -1,13 +1,13 @@
-## Baseten MiniMax 2.5 Training (Multi-Node, Megatron)
+## 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` (Megatron path with TP/EP and FP8 training flags).
+- `run_megatron.sh`
 
 Current files:
-- `config.py`: Baseten training job config (2 nodes x 8 GPUs, H200 in current config).
+- `config.py`: Baseten training job config
 - `run_megatron.sh`: Megatron SFT launcher.
 
 ### Required Baseten secrets

From ed733dcb1cb6d26a9f8c6201084977c796dfac45 Mon Sep 17 00:00:00 2001
From: Ervin Wang <ervin.wang@baseten.co>
Date: Wed, 11 Mar 2026 13:06:14 -0400
Subject: [PATCH 03/11] clean up

---
 examples/qwen80b/run_megatron.sh | 185 +++++--------------------------
 1 file changed, 29 insertions(+), 156 deletions(-)

diff --git a/examples/qwen80b/run_megatron.sh b/examples/qwen80b/run_megatron.sh
index 68db0be..917f93b 100644
--- a/examples/qwen80b/run_megatron.sh
+++ b/examples/qwen80b/run_megatron.sh
@@ -1,14 +1,6 @@
 #!/usr/bin/env bash
 set -euo pipefail
 
-LOG_BASE_DIR="${BT_CHECKPOINT_DIR:-/mnt/ckpts}/debug_logs"
-mkdir -p "${LOG_BASE_DIR}"
-LOG_FILE="${LOG_BASE_DIR}/run-megatron-node-${BT_NODE_RANK:-0}-$(date +%Y%m%d-%H%M%S).log"
-exec > >(tee -a "${LOG_FILE}") 2>&1
-
-echo "==== run_megatron.sh start node=${BT_NODE_RANK:-0} $(date -Is) ===="
-echo "log_file=${LOG_FILE}"
-
 if [[ -z "${HF_TOKEN:-}" && -z "${HUGGING_FACE_HUB_TOKEN:-}" && -z "${HUGGINGFACE_HUB_TOKEN:-}" ]]; then
   echo "ERROR: HF token is not set. An HF token is required to upload final checkpoints to Hugging Face. Configure Baseten secret 'hf_access_token' and map it to HF_TOKEN in config.py."
   exit 1
@@ -22,132 +14,32 @@ export HF_DATASETS_CACHE="${HF_HOME}/datasets"
 export PIP_CACHE_DIR="${CACHE_ROOT}/pip"
 export TRITON_CACHE_DIR="${CACHE_ROOT}/triton-cache"
 export TORCH_EXTENSIONS_DIR="${CACHE_ROOT}/torch-extensions"
-export NCCL_DEBUG="NONE"
-export TORCH_DISTRIBUTED_DEBUG="OFF"
 export TORCH_NCCL_ASYNC_ERROR_HANDLING="1"
 export NCCL_SOCKET_IFNAME="^docker0,lo"
 export OMP_NUM_THREADS="4"
 export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True"
-export TORCH_COMPILE_DISABLE="1"
 export TORCHINDUCTOR_COMPILE_THREADS="1"
 export TORCH_DISABLE_ADDR2LINE=1
-
-export TORCH_DISTRIBUTED_DEBUG=NONE
-export NCCL_DEBUG=INFO
-
 export MASTER_PORT="29500"
 
 mkdir -p "${HF_HOME}" "${PIP_CACHE_DIR}" "${TRITON_CACHE_DIR}" "${TORCH_EXTENSIONS_DIR}"
-export CUDA_LAUNCH_BLOCKING=1
-
-
-PY_BIN="$(command -v python3 || command -v python)"
-if [[ -z "${PY_BIN}" ]]; then
-  echo "No Python interpreter found in PATH." >&2
-  exit 1
-fi
 
 SWIFT_VERSION="3.12.5"
-"${PY_BIN}" -m pip install --upgrade pip
-"${PY_BIN}" -m pip install "ms-swift[llm]==${SWIFT_VERSION}" datasets huggingface_hub
-"${PY_BIN}" -m pip install "transformers==4.57.1" -U
+python -m pip install --upgrade pip
+python -m pip install "ms-swift[llm]==${SWIFT_VERSION}" datasets huggingface_hub
+python -m pip install "transformers==4.57.1" -U
 
-# Training variables (edit these directly; all set to fixed defaults).
 MODEL_ID="Qwen/Qwen3-Next-80B-A3B-Instruct"
-MODEL_TYPE="qwen3_next"
-DATASET_ID="winglian/pirate-ultrachat-10k"
-DATASET_SPLIT="train"
 CHECKPOINT_NAME="qwen80b-instruct-megatron-lora"
-RUN_NAME="qwen80b-instruct-megatron-lora"
-MODEL_ARG="${MODEL_ID}"
-
-LORA_RANK=8
-LORA_ALPHA=16
-SPLIT_DATASET_RATIO=0.01
-TENSOR_PARALLEL_SIZE=1
-PIPELINE_PARALLEL_SIZE=1
-CONTEXT_PARALLEL_SIZE=1
-EXPERT_PARALLEL_SIZE=8
-DESIRED_DATA_PARALLEL_SIZE=1
-MICRO_BATCH_SIZE=1
-GLOBAL_BATCH_SIZE=8
-MAX_EPOCHS=1
-LR=2e-4
-LR_DECAY_STYLE="constant"
-LR_WARMUP_FRACTION=0.05
-MIN_LR=1e-5
-MAX_LENGTH=2048
-NUM_WORKERS=8
-DATASET_NUM_PROC=8
-SAVE_FULL_MODEL="false"
-SAVE_INTERVAL=5
-LOG_INTERVAL=1
-REPORT_TO="none"
-MSSWIFT_COMPAT_MODE="true"
-
-
-if [[ ! -d "${MODEL_ID}" ]]; then
-  LOCAL_MODEL_DIR="${CACHE_ROOT}/model-snapshots/${MODEL_ID//\//__}"
-  mkdir -p "${LOCAL_MODEL_DIR}"
-  echo "Pre-downloading model snapshot to cache: ${LOCAL_MODEL_DIR}"
-  MODEL_ID_ENV="${MODEL_ID}" LOCAL_MODEL_DIR_ENV="${LOCAL_MODEL_DIR}" "${PY_BIN}" - <<'PY'
-import os
-from huggingface_hub import snapshot_download
-
-snapshot_download(
-    repo_id=os.environ["MODEL_ID_ENV"],
-    local_dir=os.environ["LOCAL_MODEL_DIR_ENV"],
-    local_dir_use_symlinks=False,
-    resume_download=True,
-)
-print(f"snapshot_ready={os.environ['LOCAL_MODEL_DIR_ENV']}")
-PY
-  MODEL_ARG="${LOCAL_MODEL_DIR}"
-fi
 
 checkpoint_dir="${BT_CHECKPOINT_DIR:-/mnt/ckpts}/${CHECKPOINT_NAME}"
 mkdir -p "${checkpoint_dir}"
 printf '{}' > "${checkpoint_dir}/args.json"
 
-# Megatron constraint:
-# global_batch_size % (micro_batch_size * data_parallel_size) == 0
-WORLD_SIZE=$(( ${BT_GROUP_SIZE:-1} * ${BT_NUM_GPUS:-1} ))
-MODEL_PARALLEL_SIZE=$(( TENSOR_PARALLEL_SIZE * PIPELINE_PARALLEL_SIZE * CONTEXT_PARALLEL_SIZE * EXPERT_PARALLEL_SIZE ))
-if (( MODEL_PARALLEL_SIZE <= 0 )); then
-  MODEL_PARALLEL_SIZE=1
-fi
-DATA_PARALLEL_SIZE=$(( WORLD_SIZE / MODEL_PARALLEL_SIZE ))
-if (( DATA_PARALLEL_SIZE <= 0 )); then
-  DATA_PARALLEL_SIZE=1
-fi
-if (( WORLD_SIZE != MODEL_PARALLEL_SIZE * DESIRED_DATA_PARALLEL_SIZE )); then
-  echo "ERROR: WORLD_SIZE=${WORLD_SIZE} does not match requested parallelism:"
-  echo "  tp=${TENSOR_PARALLEL_SIZE} pp=${PIPELINE_PARALLEL_SIZE} cp=${CONTEXT_PARALLEL_SIZE} ep=${EXPERT_PARALLEL_SIZE} dp=${DESIRED_DATA_PARALLEL_SIZE}"
-  echo "Expected WORLD_SIZE=$(( MODEL_PARALLEL_SIZE * DESIRED_DATA_PARALLEL_SIZE ))."
-  exit 1
-fi
-if (( DATA_PARALLEL_SIZE != DESIRED_DATA_PARALLEL_SIZE )); then
-  echo "ERROR: computed DATA_PARALLEL_SIZE=${DATA_PARALLEL_SIZE}, expected ${DESIRED_DATA_PARALLEL_SIZE}."
-  exit 1
-fi
-MIN_DIVISOR=$(( MICRO_BATCH_SIZE * DATA_PARALLEL_SIZE ))
-if (( MIN_DIVISOR <= 0 )); then
-  MIN_DIVISOR=1
-fi
-if (( GLOBAL_BATCH_SIZE % MIN_DIVISOR != 0 )); then
-  echo "Adjusting GLOBAL_BATCH_SIZE from ${GLOBAL_BATCH_SIZE} to ${MIN_DIVISOR} to satisfy Megatron divisibility."
-  GLOBAL_BATCH_SIZE="${MIN_DIVISOR}"
-fi
-
 echo "Launching ms-swift Megatron SFT for ${MODEL_ID}"
-echo "model_path=${MODEL_ARG}"
 echo "checkpoint_dir=${checkpoint_dir}"
 echo "BT_GROUP_SIZE=${BT_GROUP_SIZE} BT_NUM_GPUS=${BT_NUM_GPUS} BT_NODE_RANK=${BT_NODE_RANK}"
 echo "MASTER_ADDR=${BT_LEADER_ADDR} MASTER_PORT=${MASTER_PORT}"
-echo "ms_swift_version=${SWIFT_VERSION}"
-echo "profile=working_config TP=${TENSOR_PARALLEL_SIZE} EP=${EXPERT_PARALLEL_SIZE} PP=${PIPELINE_PARALLEL_SIZE} CP=${CONTEXT_PARALLEL_SIZE} LORA_RANK=${LORA_RANK} MAX_LENGTH=${MAX_LENGTH} MICRO_BATCH_SIZE=${MICRO_BATCH_SIZE} GLOBAL_BATCH_SIZE=${GLOBAL_BATCH_SIZE}"
-echo "model_type=${MODEL_TYPE}"
-echo "MSSWIFT_COMPAT_MODE=${MSSWIFT_COMPAT_MODE}"
 
 # ms-swift may write checkpoints into timestamped subdirs (e.g., v0-YYYYMMDD-HHMMSS).
 # Keep args.json mirrored there to avoid checkpoint save failures.
@@ -173,21 +65,6 @@ cleanup_args_sync() {
 trap cleanup_args_sync EXIT
 
 set +e
-SPLIT_ARGS=()
-if [[ "${DATASET_SPLIT}" != "train" ]]; then
-  SPLIT_JSON="[{\"hf_dataset_id\": \"${DATASET_ID}\", \"split\": [\"${DATASET_SPLIT}\"]}]"
-  SPLIT_ARGS=(--custom_dataset_info "${SPLIT_JSON}")
-fi
-
-COMPAT_ARGS=()
-if [[ "${MSSWIFT_COMPAT_MODE}" == "true" ]]; then
-  COMPAT_ARGS=(
-    --overlap_grad_reduce false
-    --overlap_param_gather false
-    --use_distributed_optimizer false
-  )
-fi
-
 PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF}" \
 NPROC_PER_NODE="${BT_NUM_GPUS}" \
 NNODES="${BT_GROUP_SIZE}" \
@@ -195,56 +72,55 @@ NODE_RANK="${BT_NODE_RANK}" \
 MASTER_ADDR="${BT_LEADER_ADDR}" \
 MASTER_PORT="${MASTER_PORT}" \
 megatron sft \
-  --model "${MODEL_ARG}" \
-  --model_type "${MODEL_TYPE}" \
+  --model Qwen/Qwen3-Next-80B-A3B-Instruct \
+  --model_type qwen3_next \
   --save "${checkpoint_dir}" \
-  --dataset "${DATASET_ID}" \
+  --dataset winglian/pirate-ultrachat-10k \
   --template minimax_m2 \
   --check_model false \
   --load_safetensors true \
   --train_type lora \
-  --lora_rank "${LORA_RANK}" \
-  --lora_alpha "${LORA_ALPHA}" \
-  --merge_lora "${SAVE_FULL_MODEL}" \
+  --lora_rank 8 \
+  --lora_alpha 16 \
+  --merge_lora false \
   --target_modules all-linear \
-  --max_epochs "${MAX_EPOCHS}" \
-  --lr_decay_style "${LR_DECAY_STYLE}" \
+  --max_epochs 1 \
+  --lr_decay_style constant \
   --clip_grad 1.0 \
-  --split_dataset_ratio "${SPLIT_DATASET_RATIO}" \
-  --tensor_model_parallel_size "${TENSOR_PARALLEL_SIZE}" \
-  --pipeline_model_parallel_size "${PIPELINE_PARALLEL_SIZE}" \
-  --context_parallel_size "${CONTEXT_PARALLEL_SIZE}" \
-  --expert_model_parallel_size "${EXPERT_PARALLEL_SIZE}" \
+  --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 "${MICRO_BATCH_SIZE}" \
-  --global_batch_size "${GLOBAL_BATCH_SIZE}" \
+  --micro_batch_size 1 \
+  --global_batch_size 8 \
   --packing false \
   --cross_entropy_loss_fusion true \
   --recompute_granularity selective \
   --recompute_modules core_attn moe \
-  --lr "${LR}" \
-  --lr_warmup_fraction "${LR_WARMUP_FRACTION}" \
-  --min_lr "${MIN_LR}" \
-  --max_length "${MAX_LENGTH}" \
-  --save_interval "${SAVE_INTERVAL}" \
-  --log_interval "${LOG_INTERVAL}" \
-  --num_workers "${NUM_WORKERS}" \
-  --dataset_num_proc "${DATASET_NUM_PROC}" \
+  --lr 2e-4 \
+  --lr_warmup_fraction 0.05 \
+  --min_lr 1e-5 \
+  --max_length 2048 \
+  --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 \
-  "${COMPAT_ARGS[@]}" \
-  "${SPLIT_ARGS[@]}" \
+  --overlap_grad_reduce false \
+  --overlap_param_gather false \
+  --use_distributed_optimizer false \
   --use_hf 1
 TRAIN_EXIT_CODE=$?
 set -e
 
-
-
 export CHECKPOINT_DIR="${checkpoint_dir}"
 export HUB_MODEL_ID="baseten-admin/qwen80b-instruct-megatron-lora"
 NODE1_UPLOAD_DONE_MARKER="${CHECKPOINT_DIR}/.node1_upload_done"
@@ -289,9 +165,6 @@ if [[ "${BT_NODE_RANK}" == "0" ]]; then
   echo "Node 1 upload marker found."
 fi
 
-
-
-
 echo "[rank ${BT_NODE_RANK}] megatron exit code=${TRAIN_EXIT_CODE}"
 if [[ "${TRAIN_EXIT_CODE}" -ne 0 ]]; then
   if [[ -d "${checkpoint_dir}" ]] && find "${checkpoint_dir}" -name "*.safetensors" -type f | grep -q .; then

From be83e5ce776bcedc38e534af1558c6d2455ffb72 Mon Sep 17 00:00:00 2001
From: Ervin Wang <ervin.wang@baseten.co>
Date: Wed, 11 Mar 2026 13:30:35 -0400
Subject: [PATCH 04/11] cleanup

---
 examples/qwen80b/run_megatron.sh | 158 +++++++++----------------------
 1 file changed, 46 insertions(+), 112 deletions(-)

diff --git a/examples/qwen80b/run_megatron.sh b/examples/qwen80b/run_megatron.sh
index 917f93b..dba982b 100644
--- a/examples/qwen80b/run_megatron.sh
+++ b/examples/qwen80b/run_megatron.sh
@@ -1,80 +1,49 @@
 #!/usr/bin/env bash
 set -euo pipefail
 
-if [[ -z "${HF_TOKEN:-}" && -z "${HUGGING_FACE_HUB_TOKEN:-}" && -z "${HUGGINGFACE_HUB_TOKEN:-}" ]]; then
-  echo "ERROR: HF token is not set. An HF token is required to upload final checkpoints to Hugging Face. Configure Baseten secret 'hf_access_token' and map it to HF_TOKEN in config.py."
-  exit 1
-fi
+# Validate HF token
+[[ -n "${HF_TOKEN:-}${HUGGING_FACE_HUB_TOKEN:-}${HUGGINGFACE_HUB_TOKEN:-}" ]] || {
+  echo "ERROR: HF token required for checkpoint upload"; exit 1
+}
 
-CACHE_ROOT="/tmp"
-export HF_HOME="${CACHE_ROOT}/huggingface"
-export HUGGINGFACE_HUB_CACHE="${HF_HOME}/hub"
-export TRANSFORMERS_CACHE="${HF_HOME}/transformers"
-export HF_DATASETS_CACHE="${HF_HOME}/datasets"
-export PIP_CACHE_DIR="${CACHE_ROOT}/pip"
-export TRITON_CACHE_DIR="${CACHE_ROOT}/triton-cache"
-export TORCH_EXTENSIONS_DIR="${CACHE_ROOT}/torch-extensions"
+# 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 OMP_NUM_THREADS="4"
-export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True"
-export TORCHINDUCTOR_COMPILE_THREADS="1"
-export TORCH_DISABLE_ADDR2LINE=1
 export MASTER_PORT="29500"
+mkdir -p "$HF_HOME"
 
-mkdir -p "${HF_HOME}" "${PIP_CACHE_DIR}" "${TRITON_CACHE_DIR}" "${TORCH_EXTENSIONS_DIR}"
-
-SWIFT_VERSION="3.12.5"
-python -m pip install --upgrade pip
-python -m pip install "ms-swift[llm]==${SWIFT_VERSION}" datasets huggingface_hub
-python -m pip install "transformers==4.57.1" -U
-
-MODEL_ID="Qwen/Qwen3-Next-80B-A3B-Instruct"
-CHECKPOINT_NAME="qwen80b-instruct-megatron-lora"
-
-checkpoint_dir="${BT_CHECKPOINT_DIR:-/mnt/ckpts}/${CHECKPOINT_NAME}"
-mkdir -p "${checkpoint_dir}"
-printf '{}' > "${checkpoint_dir}/args.json"
+# Install dependencies
+pip install -q --upgrade pip
+pip install -q "ms-swift[llm]==3.12.5" datasets huggingface_hub "transformers==4.57.1"
 
-echo "Launching ms-swift Megatron SFT for ${MODEL_ID}"
-echo "checkpoint_dir=${checkpoint_dir}"
-echo "BT_GROUP_SIZE=${BT_GROUP_SIZE} BT_NUM_GPUS=${BT_NUM_GPUS} BT_NODE_RANK=${BT_NODE_RANK}"
-echo "MASTER_ADDR=${BT_LEADER_ADDR} MASTER_PORT=${MASTER_PORT}"
+# Checkpoint directory
+checkpoint_dir="${BT_CHECKPOINT_DIR:-/mnt/ckpts}/qwen80b-instruct-megatron-lora"
+mkdir -p "$checkpoint_dir"
+printf '{}' > "$checkpoint_dir/args.json"
 
-# ms-swift may write checkpoints into timestamped subdirs (e.g., v0-YYYYMMDD-HHMMSS).
-# Keep args.json mirrored there to avoid checkpoint save failures.
-ARGS_SYNC_PID=""
-sync_args_json_loop() {
-  while true; do
-    for run_dir in "${checkpoint_dir}"/v*-*; do
-      if [[ -d "${run_dir}" ]] && [[ ! -f "${run_dir}/args.json" ]]; then
-        cp "${checkpoint_dir}/args.json" "${run_dir}/args.json" || true
-      fi
-    done
-    sleep 1
+# 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
-}
-
-sync_args_json_loop &
-ARGS_SYNC_PID=$!
-cleanup_args_sync() {
-  if [[ -n "${ARGS_SYNC_PID}" ]]; then
-    kill "${ARGS_SYNC_PID}" 2>/dev/null || true
-  fi
-}
-trap cleanup_args_sync EXIT
+  sleep 1
+done) &
+trap "kill $! 2>/dev/null" EXIT
 
-set +e
-PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF}" \
-NPROC_PER_NODE="${BT_NUM_GPUS}" \
-NNODES="${BT_GROUP_SIZE}" \
-NODE_RANK="${BT_NODE_RANK}" \
-MASTER_ADDR="${BT_LEADER_ADDR}" \
-MASTER_PORT="${MASTER_PORT}" \
+# 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}" \
+  --save "$checkpoint_dir" \
   --dataset winglian/pirate-ultrachat-10k \
   --template minimax_m2 \
   --check_model false \
@@ -117,59 +86,24 @@ megatron sft \
   --overlap_grad_reduce false \
   --overlap_param_gather false \
   --use_distributed_optimizer false \
-  --use_hf 1
-TRAIN_EXIT_CODE=$?
-set -e
-
-export CHECKPOINT_DIR="${checkpoint_dir}"
-export HUB_MODEL_ID="baseten-admin/qwen80b-instruct-megatron-lora"
-NODE1_UPLOAD_DONE_MARKER="${CHECKPOINT_DIR}/.node1_upload_done"
-# Final checkpoint upload to Hugging Face Hub from rank 0.
-if [[ "${BT_NODE_RANK}" == "1" ]]; then
-  echo "Starting final HF Hub upload from ${CHECKPOINT_DIR} to ${HUB_MODEL_ID}..."
-  "${PY_BIN}" - <<'PY'
-import datetime
-import os
-from huggingface_hub import HfApi
+  --use_hf 1 || train_exit=$?
 
-checkpoint_dir = os.environ["CHECKPOINT_DIR"]
-repo_id = os.environ["HUB_MODEL_ID"]
-token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")
+# 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)"
 
-api = HfApi(token=token)
-api.create_repo(repo_id=repo_id, repo_type="model", private=True, exist_ok=True)
-api.upload_folder(
-    repo_id=repo_id,
-    repo_type="model",
-    folder_path=checkpoint_dir,
-    commit_message=f"Final checkpoint upload {datetime.datetime.utcnow().isoformat()}Z",
-)
-print(f"hf_upload_complete repo={repo_id} folder={checkpoint_dir}")
-PY
-  touch "${NODE1_UPLOAD_DONE_MARKER}"
-fi
-
-if [[ "${BT_NODE_RANK}" == "0" ]]; then
-  NODE0_WAIT_TIMEOUT_SECONDS="${NODE0_WAIT_TIMEOUT_SECONDS:-3600}"
-  NODE0_WAIT_POLL_SECONDS="${NODE0_WAIT_POLL_SECONDS:-5}"
-  echo "Node 0 waiting for node 1 upload marker: ${NODE1_UPLOAD_DONE_MARKER} (timeout ${NODE0_WAIT_TIMEOUT_SECONDS}s)"
+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 "${NODE1_UPLOAD_DONE_MARKER}" ]]; do
-    sleep "${NODE0_WAIT_POLL_SECONDS}"
-    waited=$(( waited + NODE0_WAIT_POLL_SECONDS ))
-    if (( waited >= NODE0_WAIT_TIMEOUT_SECONDS )); then
-      echo "Timed out waiting for node 1 upload completion marker."
-      exit 1
-    fi
-  done
-  echo "Node 1 upload marker found."
+  while [[ ! -f "$upload_marker" && $waited -lt 3600 ]]; do sleep 5; ((waited+=5)); done
+  [[ -f "$upload_marker" ]] || { echo "Upload timeout"; exit 1; }
 fi
 
-echo "[rank ${BT_NODE_RANK}] megatron exit code=${TRAIN_EXIT_CODE}"
-if [[ "${TRAIN_EXIT_CODE}" -ne 0 ]]; then
-  if [[ -d "${checkpoint_dir}" ]] && find "${checkpoint_dir}" -name "*.safetensors" -type f | grep -q .; then
-    echo "Training exited ${TRAIN_EXIT_CODE}, but safetensors exist in ${checkpoint_dir}; treating as success."
-    exit 0
-  fi
-  exit "${TRAIN_EXIT_CODE}"
+# 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

From 30657e6888db5bf36f294297ff4a05efa4157639 Mon Sep 17 00:00:00 2001
From: Ervin Wang <ervin.wang@baseten.co>
Date: Wed, 11 Mar 2026 13:36:54 -0400
Subject: [PATCH 05/11] cleanup

---
 examples/qwen80b/run_megatron.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/qwen80b/run_megatron.sh b/examples/qwen80b/run_megatron.sh
index dba982b..f9fb013 100644
--- a/examples/qwen80b/run_megatron.sh
+++ b/examples/qwen80b/run_megatron.sh
@@ -72,7 +72,7 @@ megatron sft \
   --lr 2e-4 \
   --lr_warmup_fraction 0.05 \
   --min_lr 1e-5 \
-  --max_length 2048 \
+  --max_length 16384 \
   --save_interval 5 \
   --log_interval 1 \
   --num_workers 8 \

From 96d95be3b460d2629ab1db633211bc5c533cebae Mon Sep 17 00:00:00 2001
From: Ervin Wang <ervin.wang@baseten.co>
Date: Wed, 11 Mar 2026 13:58:47 -0400
Subject: [PATCH 06/11] name folder

---
 .../README.md                                 |  0
 .../config.py                                 |  0
 .../run_megatron.sh                           |  0
 .../upload_checkpoint.py                      | 25 +++++++++++++++++++
 4 files changed, 25 insertions(+)
 rename examples/{qwen80b => qwen3-next-80b-megatron}/README.md (100%)
 rename examples/{qwen80b => qwen3-next-80b-megatron}/config.py (100%)
 rename examples/{qwen80b => qwen3-next-80b-megatron}/run_megatron.sh (100%)
 create mode 100644 examples/qwen3-next-80b-megatron/upload_checkpoint.py

diff --git a/examples/qwen80b/README.md b/examples/qwen3-next-80b-megatron/README.md
similarity index 100%
rename from examples/qwen80b/README.md
rename to examples/qwen3-next-80b-megatron/README.md
diff --git a/examples/qwen80b/config.py b/examples/qwen3-next-80b-megatron/config.py
similarity index 100%
rename from examples/qwen80b/config.py
rename to examples/qwen3-next-80b-megatron/config.py
diff --git a/examples/qwen80b/run_megatron.sh b/examples/qwen3-next-80b-megatron/run_megatron.sh
similarity index 100%
rename from examples/qwen80b/run_megatron.sh
rename to examples/qwen3-next-80b-megatron/run_megatron.sh
diff --git a/examples/qwen3-next-80b-megatron/upload_checkpoint.py b/examples/qwen3-next-80b-megatron/upload_checkpoint.py
new file mode 100644
index 0000000..cf37fa2
--- /dev/null
+++ b/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()

From ab80a1800439b580809dbe49756570bfec45fe64 Mon Sep 17 00:00:00 2001
From: Ervin Wang <ervin.wang@baseten.co>
Date: Wed, 11 Mar 2026 14:18:13 -0400
Subject: [PATCH 07/11] name folder

---
 examples/qwen3-next-80b-megatron/run_megatron.sh | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/examples/qwen3-next-80b-megatron/run_megatron.sh b/examples/qwen3-next-80b-megatron/run_megatron.sh
index f9fb013..28de674 100644
--- a/examples/qwen3-next-80b-megatron/run_megatron.sh
+++ b/examples/qwen3-next-80b-megatron/run_megatron.sh
@@ -88,15 +88,17 @@ megatron sft \
   --use_distributed_optimizer false \
   --use_hf 1 || train_exit=$?
 
-# Upload checkpoint (node 1 uploads, node 0 waits)
+# 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 [[ "$BT_NODE_RANK" == "1" ]]; then
+if [[ "$num_nodes" == "1" || "$node_rank" == "1" ]]; then
   python "$script_dir/upload_checkpoint.py" "$checkpoint_dir" "$hub_repo"
   touch "$upload_marker"
-elif [[ "$BT_NODE_RANK" == "0" ]]; then
+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; }

From a67c3b0436a3a6e1d13e44b928239458b8ee5c40 Mon Sep 17 00:00:00 2001
From: ervinwang-jpg <ervin.wang@baseten.co>
Date: Wed, 11 Mar 2026 14:38:31 -0400
Subject: [PATCH 08/11] Update run_megatron.sh

---
 examples/qwen3-next-80b-megatron/run_megatron.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/qwen3-next-80b-megatron/run_megatron.sh b/examples/qwen3-next-80b-megatron/run_megatron.sh
index 28de674..03fbce4 100644
--- a/examples/qwen3-next-80b-megatron/run_megatron.sh
+++ b/examples/qwen3-next-80b-megatron/run_megatron.sh
@@ -45,7 +45,7 @@ megatron sft \
   --model_type qwen3_next \
   --save "$checkpoint_dir" \
   --dataset winglian/pirate-ultrachat-10k \
-  --template minimax_m2 \
+  --template qwen3_nothinking \
   --check_model false \
   --load_safetensors true \
   --train_type lora \

From 9e7f78584d94e9b46c0ecff725aa4e65b391399d Mon Sep 17 00:00:00 2001
From: rcano-baseten <raymond.cano@baseten.co>
Date: Thu, 12 Mar 2026 11:40:08 -0400
Subject: [PATCH 09/11] Replace Qwen3-80B example with tested long-context LoRA
 config

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>
---
 examples/qwen3-80b-msswift/training/README.md |  56 +++++++++
 examples/qwen3-80b-msswift/training/config.py |  49 ++++++++
 examples/qwen3-80b-msswift/training/run.sh    |  77 ++++++++++++
 examples/qwen3-next-80b-megatron/README.md    |  31 -----
 examples/qwen3-next-80b-megatron/config.py    |  38 ------
 .../qwen3-next-80b-megatron/run_megatron.sh   | 111 ------------------
 .../upload_checkpoint.py                      |  25 ----
 7 files changed, 182 insertions(+), 205 deletions(-)
 create mode 100644 examples/qwen3-80b-msswift/training/README.md
 create mode 100644 examples/qwen3-80b-msswift/training/config.py
 create mode 100644 examples/qwen3-80b-msswift/training/run.sh
 delete mode 100644 examples/qwen3-next-80b-megatron/README.md
 delete mode 100644 examples/qwen3-next-80b-megatron/config.py
 delete mode 100644 examples/qwen3-next-80b-megatron/run_megatron.sh
 delete mode 100644 examples/qwen3-next-80b-megatron/upload_checkpoint.py

diff --git a/examples/qwen3-80b-msswift/training/README.md b/examples/qwen3-80b-msswift/training/README.md
new file mode 100644
index 0000000..f0c0620
--- /dev/null
+++ b/examples/qwen3-80b-msswift/training/README.md
@@ -0,0 +1,56 @@
+# 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      |
+| 4     | 32   | 48K        | 2  | —  | 16 | 8         | 1                | 98                 | ~270s     |
+| 4     | 32   | 64K        | 2  | 2  | 4  | 8         | 1                | 118                | ~455s     |
+
+## 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 4 nodes:
+
+```python
+# config.py
+node_count=4
+```
+
+```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 32
+--recompute_num_layers 1
+--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
diff --git a/examples/qwen3-next-80b-megatron/README.md b/examples/qwen3-next-80b-megatron/README.md
deleted file mode 100644
index 7646f6b..0000000
--- a/examples/qwen3-next-80b-megatron/README.md
+++ /dev/null
@@ -1,31 +0,0 @@
-## 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
-```
diff --git a/examples/qwen3-next-80b-megatron/config.py b/examples/qwen3-next-80b-megatron/config.py
deleted file mode 100644
index 5a65f61..0000000
--- a/examples/qwen3-next-80b-megatron/config.py
+++ /dev/null
@@ -1,38 +0,0 @@
-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)
diff --git a/examples/qwen3-next-80b-megatron/run_megatron.sh b/examples/qwen3-next-80b-megatron/run_megatron.sh
deleted file mode 100644
index 03fbce4..0000000
--- a/examples/qwen3-next-80b-megatron/run_megatron.sh
+++ /dev/null
@@ -1,111 +0,0 @@
-#!/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 qwen3_nothinking \
-  --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
diff --git a/examples/qwen3-next-80b-megatron/upload_checkpoint.py b/examples/qwen3-next-80b-megatron/upload_checkpoint.py
deleted file mode 100644
index cf37fa2..0000000
--- a/examples/qwen3-next-80b-megatron/upload_checkpoint.py
+++ /dev/null
@@ -1,25 +0,0 @@
-#!/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()

From c37d8b36b0cd72f53d0c70772b0de921046b351e Mon Sep 17 00:00:00 2001
From: rcano-baseten <raymond.cano@baseten.co>
Date: Thu, 12 Mar 2026 12:09:54 -0400
Subject: [PATCH 10/11] Add 2-node 64K config to tested configurations

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>
---
 examples/qwen3-80b-msswift/training/README.md | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/examples/qwen3-80b-msswift/training/README.md b/examples/qwen3-80b-msswift/training/README.md
index f0c0620..86dfbe4 100644
--- a/examples/qwen3-80b-msswift/training/README.md
+++ b/examples/qwen3-80b-msswift/training/README.md
@@ -10,6 +10,7 @@ This example fine-tunes the [Qwen3-Coder-Next](https://huggingface.co/Qwen/Qwen3
 |-------|------|------------|----|----|----|-----------|------------------|--------------------|-----------|
 | 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     |
 
@@ -36,11 +37,11 @@ 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 4 nodes:
+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=4
+node_count=2
 ```
 
 ```bash
@@ -48,8 +49,8 @@ node_count=4
 --tensor_model_parallel_size 2
 --pipeline_model_parallel_size 2
 --expert_model_parallel_size 4
---global_batch_size 32
---recompute_num_layers 1
+--global_batch_size 4
+--recompute_num_layers 2
 --max_length 64000
 ```
 

From d1368a68b451d9ae2b72bd6b9c072347fb99f144 Mon Sep 17 00:00:00 2001
From: rcano-baseten <raymond.cano@baseten.co>
Date: Thu, 12 Mar 2026 13:01:19 -0400
Subject: [PATCH 11/11] Add 128K sequence length config to tested
 configurations

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>
---
 examples/qwen3-80b-msswift/training/README.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/examples/qwen3-80b-msswift/training/README.md b/examples/qwen3-80b-msswift/training/README.md
index 86dfbe4..89e990c 100644
--- a/examples/qwen3-80b-msswift/training/README.md
+++ b/examples/qwen3-80b-msswift/training/README.md
@@ -13,6 +13,7 @@ This example fine-tunes the [Qwen3-Coder-Next](https://huggingface.co/Qwen/Qwen3
 | 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