[Feature] Support GLM-5.2 source-layer DSA indexer training

.gitignore

@@ -118,6 +118,11 @@ data
 *.pkl.json
 *.log.json
 work_dirs/
+
+# Local training-analysis artifacts generated by plot_xtuner_losses.py.
+/loss_comparison*/
+/examples/v1/scripts/plot_xtuner_losses.py
+/tests/scripts/test_plot_xtuner_losses.py
 work_dir/
 
 # Pytorch

examples/v1/config/sft_glm5p2.py

@@ -6,6 +6,7 @@
 from xtuner.v1.float8.config import Float8Config, ScalingGranularity
 from xtuner.v1.loss import CELossConfig
 from xtuner.v1.model import get_model_config_from_hf
+from xtuner.v1.module.attention import DSAIndexerTrainingConfig
 from xtuner.v1.train import TrainerConfig
 from xtuner.v1.train.trainer import LoadCheckpointConfig
 
@@ -36,6 +37,7 @@ def _get_float8_config() -> Float8Config | None:
 # On single-node 8-GPU SFT, EP=8 leaves FSDP size at 1 and replicates non-expert params.
 ep_size = int(os.environ.get("EP_SIZE", "1"))
 intra_layer_micro_batch = int(os.environ.get("INTRA_LAYER_MICRO_BATCH", "1"))
+sp_size = int(os.environ.get("SP_SIZE", "1"))
 global_batch_size = int(os.environ.get("GLOBAL_BATCH_SIZE", os.environ.get("WORLD_SIZE", "8")))
 sample_max_length = int(os.environ.get("SAMPLE_MAX_LENGTH", "4096"))
 pack_max_length = int(os.environ.get("PACK_MAX_LENGTH", "16384"))
@@ -54,6 +56,15 @@ def _get_float8_config() -> Float8Config | None:
 model_cfg.lm_loss_cfg = loss_cfg
 if hasattr(model_cfg.attention, "sparse_mla_backend"):
     model_cfg.attention.sparse_mla_backend = os.environ.get("SPARSE_MLA_BACKEND", "tilelang")
+train_dsa_indexer = _get_bool_env("TRAIN_DSA_INDEXER", False)
+if train_dsa_indexer:
+    if model_cfg.attention.sparse_mla_backend != "cudnn_dsa":
+        raise ValueError("DSA indexer training requires SPARSE_MLA_BACKEND=cudnn_dsa.")
+    model_cfg.attention.indexer_training = DSAIndexerTrainingConfig(
+        loss_coeff=float(os.environ.get("INDEXER_LOSS_COEFF", "1.0")),
+        indexer_only=_get_bool_env("INDEXER_ONLY", False),
+        debug_interval=int(os.environ.get("INDEXER_DEBUG_INTERVAL", "0")),
+    )
 
 cache_dir = os.path.join(work_dir, "jsonl_cache")
 cache_tag = os.environ.get("CACHE_TAG", f"glm52_{sample_max_length}")
@@ -98,16 +109,30 @@ def _get_float8_config() -> Float8Config | None:
 elif optimizer == "adamw":
     optim_cfg = AdamWConfig(
         lr=lr,
+        weight_decay=float(os.environ.get("WEIGHT_DECAY", "0.01")),
         foreach=_get_bool_env("ADAMW_FOREACH", False),
         swap_optimizer=_get_bool_env("SWAP_OPTIMIZER", False),
     )
 else:
     raise ValueError(f"Unsupported OPTIMIZER={optimizer!r}. Use adamw or muon.")
 lr_cfg = LRConfig(lr_type=os.environ.get("LR_TYPE", "cosine"), warmup_ratio=float(os.environ.get("WARMUP_RATIO", "0")))
+recompute_ratio = float(os.environ.get("RECOMPUTE_RATIO", "1.0"))
+torch_compile = _get_bool_env("TORCH_COMPILE", False)
+if train_dsa_indexer:
+    if sp_size != 1:
+        raise ValueError("DSA indexer training requires SP_SIZE=1.")
+    if intra_layer_micro_batch != 1:
+        raise ValueError("DSA indexer training requires INTRA_LAYER_MICRO_BATCH=1.")
+    if model_cfg.compile_cfg or torch_compile:
+        raise ValueError("DSA indexer training requires MODEL_COMPILE=0 and TORCH_COMPILE=0.")
+    if recompute_ratio != 0:
+        raise ValueError("DSA indexer training requires RECOMPUTE_RATIO=0 (no activation checkpointing).")
+
 fsdp_cfg = FSDPConfig(
     cpu_offload=_get_bool_env("CPU_OFFLOAD", False),
     ep_size=ep_size,
-    torch_compile=_get_bool_env("TORCH_COMPILE", False),
+    torch_compile=torch_compile,
+    recompute_ratio=recompute_ratio,
 )
 
 trainer = TrainerConfig(
@@ -123,7 +148,7 @@ def _get_float8_config() -> Float8Config | None:
     global_batch_size=global_batch_size,
     total_step=total_step,
     intra_layer_micro_batch=intra_layer_micro_batch,
-    sp_size=int(os.environ.get("SP_SIZE", "1")),
+    sp_size=sp_size,
     load_checkpoint_cfg=LoadCheckpointConfig(checkpoint_path=os.environ.get("LOAD_CHECKPOINT_PATH")),
     checkpoint_interval=int(os.environ.get("CHECKPOINT_INTERVAL", "200")),
     checkpoint_maxkeep=int(os.environ.get("CHECKPOINT_MAX_KEEP", "3")),
@@ -134,4 +159,5 @@ def _get_float8_config() -> Float8Config | None:
     profile_time=_get_bool_env("PROFILE_TIME", False),
     profile_step=[int(x) for x in os.environ.get("PROFILE_STEP", "2,3").split(",") if x],
     debug_skip_save=_get_bool_env("DEBUG_SKIP_SAVE", False),
+    do_clip=_get_bool_env("DO_CLIP", True),
 )

examples/v1/scripts/train_glm52_indexer.sh

@@ -0,0 +1,91 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Train GLM-5.2 and its main-stack source indexers jointly. Activate the
+# intended Python environment before invoking this script.
+: "${GLM5_2_MODEL_PATH:?GLM5_2_MODEL_PATH is required}"
+
+export DATASET_TYPE="${DATASET_TYPE:-alpaca}"
+case "${DATASET_TYPE}" in
+  alpaca)
+    : "${ALPACA_PATH:?ALPACA_PATH is required when DATASET_TYPE=alpaca}"
+    ;;
+  alpaca_long)
+    : "${ALPACA_LONG_PATH:?ALPACA_LONG_PATH is required when DATASET_TYPE=alpaca_long}"
+    ;;
+  *)
+    echo "Unsupported DATASET_TYPE=${DATASET_TYPE}; use alpaca or alpaca_long." >&2
+    exit 2
+    ;;
+esac
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
+CONFIG_PATH="${1:-${REPO_ROOT}/examples/v1/config/sft_glm5p2.py}"
+export WORK_DIR="${2:-${WORK_DIR:-work_dirs/glm52_indexer_sft}}"
+export PYTHONPATH="${REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}"
+
+export TRAIN_DSA_INDEXER="${TRAIN_DSA_INDEXER:-1}"
+export INDEXER_LOSS_COEFF="${INDEXER_LOSS_COEFF:-1.0}"
+export INDEXER_ONLY="${INDEXER_ONLY:-0}"
+export INDEXER_DEBUG_INTERVAL="${INDEXER_DEBUG_INTERVAL:-0}"
+export SPARSE_MLA_BACKEND="${SPARSE_MLA_BACKEND:-cudnn_dsa}"
+
+# These values reflect the constraints enforced by sft_glm5p2.py while
+# source-indexer training is enabled.
+export SP_SIZE="${SP_SIZE:-1}"
+export INTRA_LAYER_MICRO_BATCH="${INTRA_LAYER_MICRO_BATCH:-1}"
+export RECOMPUTE_RATIO="${RECOMPUTE_RATIO:-0}"
+export MODEL_COMPILE="${MODEL_COMPILE:-0}"
+export TORCH_COMPILE="${TORCH_COMPILE:-0}"
+
+export EP_SIZE="${EP_SIZE:-4}"
+export TOTAL_STEP="${TOTAL_STEP:-300}"
+export GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-8}"
+export LR="${LR:-1e-6}"
+export WEIGHT_DECAY="${WEIGHT_DECAY:-0.01}"
+export DO_CLIP="${DO_CLIP:-1}"
+
+export DATASET_SAMPLE_RATIO="${DATASET_SAMPLE_RATIO:-1.0}"
+export SAMPLE_MAX_LENGTH="${SAMPLE_MAX_LENGTH:-4096}"
+export PACK_MAX_LENGTH="${PACK_MAX_LENGTH:-4096}"
+export CACHE_TAG="${CACHE_TAG:-glm52_indexer_4096}"
+
+export FP8="${FP8:-1}"
+export DEBUG_SKIP_SAVE="${DEBUG_SKIP_SAVE:-0}"
+export CHECKPOINT_INTERVAL="${CHECKPOINT_INTERVAL:-200}"
+export HF_INTERVAL="${HF_INTERVAL:-${TOTAL_STEP}}"
+export HF_MAX_KEEP="${HF_MAX_KEEP:-1}"
+export PROFILE_TIME="${PROFILE_TIME:-0}"
+export PROFILE_MEMORY="${PROFILE_MEMORY:-0}"
+
+NNODES="${NNODES:-${NODE_COUNT:-1}}"
+NODE_RANK="${NODE_RANK:-0}"
+MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}"
+MASTER_PORT="${MASTER_PORT:-6000}"
+NPROC_PER_NODE="${NPROC_PER_NODE:-8}"
+
+cd "${REPO_ROOT}"
+test -f "${CONFIG_PATH}"
+mkdir -p "${WORK_DIR}"
+ulimit -n 65536
+
+command=(
+  torchrun
+  "--nproc-per-node=${NPROC_PER_NODE}"
+  "--master-addr=${MASTER_ADDR}"
+  "--master-port=${MASTER_PORT}"
+  "--nnodes=${NNODES}"
+  "--node-rank=${NODE_RANK}"
+  --tee 3
+  -m xtuner.v1.train.cli.sft
+  --config "${CONFIG_PATH}"
+)
+
+if [[ "${DRY_RUN:-0}" != "0" ]]; then
+  printf '%q ' "${command[@]}"
+  printf '\n'
+  exit 0
+fi
+
+"${command[@]}" 2>&1 | tee -a "${WORK_DIR}/node_${NODE_RANK}.txt"

tests/model/test_glm52_moe.py

@@ -29,7 +29,7 @@
 from xtuner.v1.data_proto import SequenceContext
 from xtuner.v1.loss.ce_loss import CELossConfig
 from xtuner.v1.model import Glm52MoEConfig, get_model_config, get_model_config_from_hf
-from xtuner.v1.module.attention import DSAMLAConfig
+from xtuner.v1.module.attention import DSAIndexerTrainingConfig, DSAMLAConfig
 from xtuner.v1.module.mtp import MTPConfig
 from xtuner.v1.module.router.noaux_router import NoAuxRouterConfig
 from xtuner.v1.utils.test_utils import init_data_mesh
@@ -182,6 +182,36 @@ def test_rejects_shared_physical_mtp_indexer(self):
         with pytest.raises(ValueError, match="physical MTP indexer_types"):
             config.build()
 
+    def test_indexer_training_keeps_physical_mtp_indexer_frozen_by_default(self):
+        # 主干 source indexer 解冻时，physical MTP indexer 仍保持 frozen。
+        config = _tiny_glm52_config()
+        config.attention.indexer_training = DSAIndexerTrainingConfig(loss_coeff=1.0)
+        config.mtp_config = MTPConfig(num_layers=1, share_weights=True)
+
+        with mock.patch("torch.cuda.Stream"):
+            model = config.build()
+
+        main_attention = model.layers["0"].self_attn
+        mtp_attention = model.mtp_block.layers[0].decoder_layer.self_attn  # type: ignore[union-attr]
+        assert all(parameter.requires_grad for parameter in main_attention.indexer.parameters())
+        assert all(not parameter.requires_grad for parameter in mtp_attention.indexer.parameters())
+        assert mtp_attention.indexer_training is None
+
+    def test_indexer_only_trains_main_source_indexers_exclusively(self):
+        # 严格过拟合模式固定 attention teacher，只训练主干 source indexer。
+        config = _tiny_glm52_config()
+        config.attention.indexer_training = DSAIndexerTrainingConfig(loss_coeff=1.0, indexer_only=True)
+        config.mtp_config = MTPConfig(num_layers=1, share_weights=True)
+
+        with mock.patch("torch.cuda.Stream"):
+            model = config.build()
+
+        trainable_names = [name for name, parameter in model.named_parameters() if parameter.requires_grad]
+        assert trainable_names
+        assert all(name.startswith("layers.0.self_attn.indexer.") for name in trainable_names)
+        assert all(not parameter.requires_grad for parameter in model.layers["1"].parameters())
+        assert all(not parameter.requires_grad for parameter in model.mtp_block.parameters())  # type: ignore[union-attr]
+
 
 @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
 class TestGlm52CheckpointConversion(DeterministicDDPTestCase):