InternLM · albertimff · Jul 31, 2026 · Aug 3, 2026
diff --git a/example_run/sft_qwen35_35b_256k_32gpu.py b/example_run/sft_qwen35_35b_256k_32gpu.py
@@ -0,0 +1,198 @@
+import json
+import os
+import shutil
+from pathlib import Path
+from typing import Any
+
+from xtuner.v1.config import FSDPConfig, LRConfig, MuonConfig
+from xtuner.v1.datasets import (
+    PretrainTokenizeFunctionConfig,
+    Qwen3VLTokenizeFnConfig,
+)
+from xtuner.v1.datasets.config import DataloaderConfig, DatasetConfig
+from xtuner.v1.datasets.mllm_tokenize_fn import OSSLoaderConfig
+from xtuner.v1.loss import CELossConfig
+from xtuner.v1.model import Qwen3_5_VLMoE35BA3Config
+from xtuner.v1.model.compose.qwen3_vl.modeling_qwen3_vl import (
+    QWEN3VL_COMPILE_CFG,
+)
+# from xtuner.v1.model.moe.moe import MTPConfig
+from xtuner.v1.train import ResumeConfig, TrainerConfig
+
+
+# This vision-layer compile rule is incompatible with Qwen3.5-35B-A3B.
+QWEN3VL_COMPILE_CFG.pop(
+    "xtuner.v1.model.compose.qwen3_vl.modeling_vision."
+    "Qwen3VLVisionLayer.forward",
+    None,
+)
+
+
+def _get_int_env(name: str, default: int) -> int:
+    return int(os.getenv(name, str(default)))
+
+
+def _get_float_env(name: str, default: float) -> float:
+    return float(os.getenv(name, str(default)))
+
+
+def _get_bool_env(name: str, default: bool) -> bool:
+    value = os.getenv(name)
+    if value is None:
+        return default
+    return value.lower() in {"1", "true", "yes", "y", "on"}
+
+
+# Paths are provided by one of the submit scripts in this directory.
+ceph_config = os.getenv("CEPH_CONFIG", "")
+meta_data_path = Path(os.environ["META_DATA_PATH"])
+model_path = Path(os.environ["MODEL_PATH"])
+work_dir = Path(os.environ["WORK_DIR"])
+tokenizer_cache_dir = os.environ["TOKENIZER_CACHE_DIR"]
+chat_template_name = os.getenv("CHAT_TEMPLATE_NAME", "qwen3.5-vl")
+
+work_dir.mkdir(parents=True, exist_ok=True)
+shutil.copy2(__file__, work_dir)
+
+# 256k-context SFT defaults for 32 GPUs (4 nodes x 8 GPUs).
+sample_max_length = _get_int_env("SAMPLE_MAX_LENGTH", 256 * 1024)
+pack_max_length = _get_int_env("PACK_MAX_LENGTH", 256 * 1024)
+rand_video_max_frames = _get_int_env("RAND_VIDEO_MAX_FRAMES", 24)
+num_workers = _get_int_env("NUM_WORKERS", 4)
+global_batch_size = _get_int_env("GLOBAL_BATCH_SIZE", 8)
+total_epoch = _get_int_env("TOTAL_EPOCH", 1)
+hf_interval = _get_int_env("HF_INTERVAL", 500)
+hf_max_keep = _get_int_env("HF_MAX_KEEP", 2)
+checkpoint_interval = _get_int_env("CHECKPOINT_INTERVAL", 500)
+checkpoint_maxkeep = _get_int_env("CHECKPOINT_MAXKEEP", 2)
+
+lr = _get_float_env("LR", 2e-5)
+lr_min = _get_float_env("LR_MIN", 1e-6)
+weight_decay = _get_float_env("WEIGHT_DECAY", 0.05)
+warmup_ratio = _get_float_env("WARMUP_RATIO", 0.1)
+recompute_ratio = _get_float_env("RECOMPUTE_RATIO", 1.0)
+loss_reduction = os.getenv("LOSS_REDUCTION", "square")
+max_pixels = _get_int_env("MAX_PIXELS", 16_777_216)
+
+sp_size = _get_int_env("SP_SIZE", 4)
+ep_size = _get_int_env("EP_SIZE", 1)
+tp_size = _get_int_env("TP_SIZE", 1)
+torch_compile = _get_bool_env("TORCH_COMPILE", True)
+
+# Qwen3.5-35B-A3B model settings.
+model_cfg = Qwen3_5_VLMoE35BA3Config()
+
+with (model_path / "config.json").open("r", encoding="utf-8") as file:
+    model_hf_config: dict[str, Any] = json.load(file)
+
+model_cfg.text_config.vocab_size = model_hf_config["text_config"]["vocab_size"]
+# model_cfg.text_config.mtp_config = [
+#     MTPConfig(
+#         name="normal",
+#         mask_type=None,
+#         num_layers=4,
+#         share_weights=True,
+#         loss_scaling_factor=1.0,
+#     ),
+# ]
+
+if ep_size > 1:
+    model_cfg.text_config.ep_size = ep_size
+    model_cfg.text_config.dispatcher = "deepep"
+
+# Dataset recipe: META_DATA_PATH points to a metadata JSON file.
+oss_loader_cfg = (
+    OSSLoaderConfig(backend_kwargs={"conf_path": ceph_config})
+    if ceph_config
+    else None
+)
+
+ds_collections: dict[str, Any] = json.loads(
+    meta_data_path.read_text(encoding="utf-8")
+)
+has_pretrain = any(
+    data.get("text_pretrain", False) for data in ds_collections.values()
+)
+dataset_config: list[dict[str, Any]] = []
+
+for name, data in ds_collections.items():
+    is_pretrain = data.get("text_pretrain", False)
+    if is_pretrain:
+        tokenize_fn = PretrainTokenizeFunctionConfig(hash=data.get("hash"))
+    else:
+        tokenize_fn = Qwen3VLTokenizeFnConfig(
+            chat_template=chat_template_name,
+            llm_pack_weight=-3.2,
+            visual_pack_weight=5.0,
+            max_length=sample_max_length,
+            processor_path=str(model_path),
+            rand_video_max_frames=rand_video_max_frames,
+            oss_loader_cfg=oss_loader_cfg,
+            max_pixels=max_pixels,
+            debug=True,
+        )
+
+    dataset_config.append(
+        {
+            "dataset": DatasetConfig(
+                name=name,
+                anno_path=data["annotation"],
+                media_root=data.get("media_root") or "",
+                sample_ratio=data.get("sample_ratio", 1.0),
+                class_name="JsonlDataset" if is_pretrain else "VLMJsonlDataset",
+                enable_sequential_sampler=True,
+                cache_tag="xtuner_train_v2",
+                cache_dir=tokenizer_cache_dir,
+            ),
+            "tokenize_fn": tokenize_fn,
+        }
+    )
+
+dataloader_config = DataloaderConfig(
+    dataset_config_list=dataset_config,
+    pack_max_length=pack_max_length,
+    pack_level="mllm_hybrid" if has_pretrain else "soft",
+    pack_to_max_length=True,
+    collator="qwen3_vl_sft_collator",
+    num_workers=num_workers,
+    pack_extra_buffer_size=_get_int_env("PACK_EXTRA_BUFFER_SIZE", 20),
+)
+
+optim_cfg = MuonConfig(lr=lr, weight_decay=weight_decay)
+lr_cfg = LRConfig(
+    lr_type="cosine",
+    warmup_ratio=warmup_ratio,
+    lr_min=lr_min,
+)
+fsdp_cfg = FSDPConfig(
+    tp_size=tp_size,
+    ep_size=ep_size,
+    recompute_ratio=recompute_ratio,
+    torch_compile=torch_compile,
+    checkpoint_preserve_rng_state=False,
+)
+
+trainer = TrainerConfig(
+    sp_size=sp_size,
+    load_from=str(model_path),
+    resume_cfg=ResumeConfig(auto_resume=True),
+    tokenizer_path=str(model_path),
+    fsdp_cfg=fsdp_cfg,
+    exp_tracker="tensorboard",
+    model_cfg=model_cfg,
+    optim_cfg=optim_cfg,
+    dataloader_cfg=dataloader_config,
+    lr_cfg=lr_cfg,
+    loss_cfg=CELossConfig(
+        mode="chunk",
+        chunk_size=1024,
+        loss_reduction=loss_reduction,
+    ),
+    global_batch_size=global_batch_size,
+    total_epoch=total_epoch,
+    hf_interval=hf_interval,
+    checkpoint_interval=checkpoint_interval,
+    checkpoint_maxkeep=checkpoint_maxkeep,
+    hf_max_keep=hf_max_keep,
+    work_dir=work_dir,
+)
diff --git a/example_run/submit_qwen35_sft_16gpu_sp8.sh b/example_run/submit_qwen35_sft_16gpu_sp8.sh
@@ -0,0 +1,153 @@
+#!/usr/bin/env bash
+set -euo pipefail
+set -x
+
+# 16 GPUs = 2 nodes x 8 GPUs per node.
+#
+# Muon currently does not support EP > 1 in this XTuner version, so use
+# EP=1 and shard the language model over all 16 ranks with FSDP. SP=8 splits
+# every packed 256K sequence over 8 ranks (32K tokens per rank).
+gpu_group="${GPU_GROUP:?Set GPU_GROUP to the rjob charged group}"
+namespace="${NAMESPACE:?Set NAMESPACE to the rjob namespace}"
+gpus_per_node=8
+num_nodes=2
+
+script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+xtuner_path="$(cd -- "${script_dir}/.." && pwd)"
+# The Python config reads all world-size-dependent values from the environment,
+# so the same config can be used for both 16-GPU and 32-GPU submissions.
+config_file="${script_dir}/sft_qwen35_35b_256k_32gpu.py"
+meta_data_path="${META_DATA_PATH:-${script_dir}/meta.json}"
+
+model_path="${MODEL_PATH:?Set MODEL_PATH to the Qwen3.5 model snapshot directory}"
+output_root="${OUTPUT_ROOT:?Set OUTPUT_ROOT to the training output directory}"
+tokenizer_cache_dir="${TOKENIZER_CACHE_DIR:?Set TOKENIZER_CACHE_DIR to the tokenizer cache directory}"
+log_dir="${LOG_DIR:-${output_root}/logs}"
+ceph_config="${CEPH_CONFIG:-}"
+
+image="${IMAGE:?Set IMAGE to the training image}"
+meta_name=$(basename "${meta_data_path}")
+meta_name=${meta_name%.*}
+run_tag="${meta_name}-16gpu-ep1-sp8-muon-$(date +%Y%m%d-%H%M%S)"
+job_name="qwen35-sft-${run_tag}"
+work_dir="${output_root}/${run_tag}"
+
+required_files=(
+    "${config_file}"
+    "${meta_data_path}"
+    "${model_path}/config.json"
+    "${model_path}/tokenizer_config.json"
+    "${model_path}/model.safetensors.index.json"
+)
+for required_file in "${required_files[@]}"; do
+    if [[ ! -f "${required_file}" ]]; then
+        echo "ERROR: required file does not exist: ${required_file}" >&2
+        exit 2
+    fi
+done
+
+submit_mode_args=()
+if [[ "${PREDICT_ONLY:-false}" == "true" ]]; then
+    submit_mode_args+=(--predict-only=true)
+fi
+
+mount_args=()
+if [[ -n "${RJOB_MOUNTS:-}" ]]; then
+    IFS=',' read -r -a mount_specs <<< "${RJOB_MOUNTS}"
+    for mount_spec in "${mount_specs[@]}"; do
+        [[ -n "${mount_spec}" ]] && mount_args+=(--mount="${mount_spec}")
+    done
+fi
+
+rjob submit \
+    "${submit_mode_args[@]}" \
+    --name="${job_name}" \
+    --task_name t0 \
+    --gpu="${gpus_per_node}" \
+    --memory=1500000 \
+    --cpu=50 \
+    --charged-group="${gpu_group}" \
+    --namespace="${namespace}" \
+    --private-machine=group \
+    -P "${num_nodes}" \
+    --image="${image}" \
+    "${mount_args[@]}" \
+    --host-network=true \
+    --gang-start=true \
+    --custom-resources=rdma/mlnx_shared=8 \
+    --custom-resources=mellanox.com/mlnx_rdma=1 \
+    -e DISTRIBUTED_JOB=true \
+    -e XTUNER_PATH="${xtuner_path}" \
+    -e CONFIG_FILE="${config_file}" \
+    -e MODEL_PATH="${model_path}" \
+    -e META_DATA_PATH="${meta_data_path}" \
+    -e CEPH_CONFIG="${ceph_config}" \
+    -e WORK_DIR="${work_dir}" \
+    -e TOKENIZER_CACHE_DIR="${tokenizer_cache_dir}" \
+    -e XTUNER_TOKENIZE_DEBUG_SAMPLES="${XTUNER_TOKENIZE_DEBUG_SAMPLES:-0}" \
+    -e LOG_DIR="${log_dir}" \
+    -e GPUS_PER_NODE="${gpus_per_node}" \
+    -e TORCHRUN_NNODES="${num_nodes}" \
+    -e SAMPLE_MAX_LENGTH=262144 \
+    -e PACK_MAX_LENGTH=262144 \
+    -e GLOBAL_BATCH_SIZE=8 \
+    -e SP_SIZE=8 \
+    -e TP_SIZE=1 \
+    -e EP_SIZE=1 \
+    -e NUM_WORKERS=4 \
+    -e PACK_EXTRA_BUFFER_SIZE=20 \
+    -e RAND_VIDEO_MAX_FRAMES=24 \
+    -e MAX_PIXELS=16777216 \
+    -e LR=2e-5 \
+    -e LR_MIN=1e-6 \
+    -e WEIGHT_DECAY=0.05 \
+    -e WARMUP_RATIO=0.1 \
+    -e RECOMPUTE_RATIO=1.0 \
+    -e LOSS_REDUCTION=square \
+    -e TORCH_COMPILE=true \
+    -e TOTAL_EPOCH=1 \
+    -e HF_INTERVAL=500 \
+    -e HF_MAX_KEEP=2 \
+    -e CHECKPOINT_INTERVAL=500 \
+    -e CHECKPOINT_MAXKEEP=2 \
+    -- bash -lc '
+        set -euo pipefail
+        set -x
+
+        export PYTHONPATH="${XTUNER_PATH}:${PYTHONPATH:-}"
+        export TORCHRUN_NODE_RANK="${NODE_RANK:-${RANK:-}}"
+        export MASTER_PORT="${MASTER_PORT:-29500}"
+        export LOG_FILE="${LOG_DIR}/qwen35-sft-16gpu-ep1-sp8-muon-node${TORCHRUN_NODE_RANK:-unknown}.log"
+
+        if [ -z "${TORCHRUN_NODE_RANK}" ] || [ -z "${MASTER_ADDR:-}" ]; then
+            echo "ERROR: NODE_RANK/RANK and MASTER_ADDR are required."
+            env | sort | grep -E \
+                "^(NODE_RANK|RANK|NODE_COUNT|WORLD_SIZE|MASTER_ADDR|MASTER_PORT|HOSTNAME)=" \
+                || true
+            exit 2
+        fi
+
+        if [ "${TORCHRUN_NNODES}" != "1" ] && {
+            [ "${MASTER_ADDR}" = "127.0.0.1" ] ||
+                [ "${MASTER_ADDR}" = "localhost" ]
+        }; then
+            echo "ERROR: MASTER_ADDR=${MASTER_ADDR} is invalid for a multi-node job."
+            exit 2
+        fi
+
+        mkdir -p "${LOG_DIR}"
+        exec > >(tee "${LOG_FILE}") 2>&1
+
+        cd "${XTUNER_PATH}"
+        ls -l "${CONFIG_FILE}" "${META_DATA_PATH}"
+        python -c "import xtuner; print(\"xtuner import:\", xtuner.__file__)"
+
+        torchrun \
+            --nproc-per-node="${GPUS_PER_NODE}" \
+            --nnodes="${TORCHRUN_NNODES}" \
+            --node_rank="${TORCHRUN_NODE_RANK}" \
+            --master_addr="${MASTER_ADDR}" \
+            --master_port="${MASTER_PORT}" \
+            xtuner/v1/train/cli/sft.py \
+            --config="${CONFIG_FILE}"
+    '