[Feature] Add recompute_cfg and per-model recompute declarations (3/3)

autotest/config/npu_qwen3_recompute.py

@@ -10,20 +10,20 @@
 from xtuner.v1.loss.ce_loss import CELossConfig
 from xtuner.v1.model.moe.qwen3 import Qwen3MoE30BA3Config
 from xtuner.v1.train import TrainerConfig
+from xtuner.v1.utils import RecomputeConfig
 
 
 QWEN3_MOE_PATH = os.environ["QWEN3_MOE_PATH"]
 ALPACA_PATH = os.environ["ALPACA_PATH"]
 
 
-moe_cfg = Qwen3MoE30BA3Config(compile_cfg=False)
+moe_cfg = Qwen3MoE30BA3Config(recompute_cfg=RecomputeConfig(ratio=0.25), compile_cfg=False)
 optim_cfg = AdamWConfig(lr=6e-05)
 lr_cfg = LRConfig(lr_type="cosine", lr_min=1e-6)
 fsdp_cfg = FSDPConfig(
     cpu_offload=False,
     ep_size=moe_cfg.ep_size,
     tp_size=4,
-    recompute_ratio=0.25,
 )
 
 dataset_config = [

autotest/config/qwen3_5_recompute.py

@@ -10,18 +10,18 @@
 from xtuner.v1.loss.ce_loss import CELossConfig
 from xtuner.v1.model import Qwen3_5_VLMoE35BA3Config
 from xtuner.v1.train import TrainerConfig
+from xtuner.v1.utils import RecomputeConfig
 
 
 QWEN3_MOE_PATH = os.environ["QWEN3_MOE_PATH"]
 ALPACA_PATH = os.environ["ALPACA_PATH"]
 
 
-moe_cfg = Qwen3_5_VLMoE35BA3Config(compile_cfg=False)
+moe_cfg = Qwen3_5_VLMoE35BA3Config(recompute_cfg=RecomputeConfig(ratio=0.25), compile_cfg=False)
 optim_cfg = AdamWConfig(lr=6e-05)
 lr_cfg = LRConfig(lr_type="cosine", lr_min=1e-6)
 fsdp_cfg = FSDPConfig(
     cpu_offload=False,  # qwen3.5 don't support True
-    recompute_ratio=0.25,
 )
 
 dataset_config = [

autotest/config/qwen3_recompute.py

@@ -10,19 +10,19 @@
 from xtuner.v1.loss.ce_loss import CELossConfig
 from xtuner.v1.model.moe.qwen3 import Qwen3MoE30BA3Config
 from xtuner.v1.train import TrainerConfig
+from xtuner.v1.utils import RecomputeConfig
 
 
 QWEN3_MOE_PATH = os.environ["QWEN3_MOE_PATH"]
 ALPACA_PATH = os.environ["ALPACA_PATH"]
 
 
-moe_cfg = Qwen3MoE30BA3Config(compile_cfg=True)
+moe_cfg = Qwen3MoE30BA3Config(recompute_cfg=RecomputeConfig(ratio=0.25), compile_cfg=True)
 optim_cfg = AdamWConfig(lr=6e-05)
 lr_cfg = LRConfig(lr_type="cosine", lr_min=1e-6)
 fsdp_cfg = FSDPConfig(
     cpu_offload=True,
     ep_size=moe_cfg.ep_size,
-    recompute_ratio=0.25,
 )
 
 dataset_config = [

ci/config/qwen3_moe_30BA3_sac.py

@@ -0,0 +1,82 @@
+"""Qwen3-MoE-30BA3 with region-level selective checkpointing switched on.
+
+Copied from `qwen3_moe_30BA3.py`; the only substantive difference is `recompute_cfg` on the model
+config. Everything below `recompute_ratio` (1.0 by default) is still wrapped in a checkpoint -- a
+unit named here is kept resident *inside* those layers instead of being recomputed.
+
+To sweep arms, copy this file and change RECOMPUTE_CFG. Do not pass it on the command line.
+
+    RECOMPUTE_CFG = None                      # recompute everything (baseline)
+    RECOMPUTE_CFG = True                      # keep every unit the model supports
+    RECOMPUTE_CFG = [RecomputeUnit.SAVE_ATTN] # keep exactly these
+
+Available units for an MoE stack (it supports all of them):
+    SAVE_ATTN           attention call and its projections
+    SAVE_MOE_GATE       router: gating projection, top-k, routing weights
+    SAVE_MOE_DISPATCH   permutation/padding around the dispatch and combine all-to-alls
+    SAVE_MLP            shared-expert MLP, plus the dense MLP of the first_k_dense_replace layers
+
+Needs torch >= 2.10: regions are carried as fx annotations, which older torch cannot record. With
+RECOMPUTE_CFG = None this config runs on 2.9 as well.
+
+Run (8xH200):
+    export QWEN3_MOE_PATH=... ALPACA_PATH=...   # see ci/scripts/CI_ENV.sh
+    torchrun --nproc-per-node 8 -m xtuner.v1.train.cli.sft --config ci/config/qwen3_moe_30BA3_sac.py
+"""
+
+import os
+
+from xtuner.v1.config import AdamWConfig, FSDPConfig, LRConfig
+from xtuner.v1.datasets import FTDPTokenizeFnConfig
+from xtuner.v1.datasets.config import DataloaderConfig, DatasetConfig
+from xtuner.v1.loss.ce_loss import CELossConfig
+from xtuner.v1.model.moe.qwen3 import Qwen3MoE30BA3Config
+from xtuner.v1.train import TrainerConfig
+from xtuner.v1.utils.selective_checkpointing import RecomputeUnit
+
+
+QWEN3_MOE_PATH = os.environ["QWEN3_MOE_PATH"]
+ALPACA_PATH = os.environ["ALPACA_PATH"]
+
+RECOMPUTE_CFG = [RecomputeUnit.SAVE_ATTN]
+
+moe_cfg = Qwen3MoE30BA3Config(recompute_cfg=RECOMPUTE_CFG)
+# moe_cfg = Qwen3MoE30BA3Config(recompute_cfg=RECOMPUTE_CFG)
+optim_cfg = AdamWConfig(lr=6e-05)
+lr_cfg = LRConfig(lr_type="cosine", lr_min=1e-6)
+fsdp_cfg = FSDPConfig(
+    torch_compile=True,
+    cpu_offload=False,
+    ep_size=moe_cfg.ep_size,
+)
+
+dataset_config = [
+    {
+        "dataset": DatasetConfig(name="alpaca", anno_path=ALPACA_PATH, sample_ratio=1.0),
+        "tokenize_fn": FTDPTokenizeFnConfig(max_length=16386),
+    },
+]
+
+dataloader_config = DataloaderConfig(pack_max_length=16384)
+
+loss_cfg = CELossConfig()
+
+
+trainer = TrainerConfig(
+    load_from=QWEN3_MOE_PATH,
+    model_cfg=moe_cfg,
+    optim_cfg=optim_cfg,
+    fsdp_cfg=fsdp_cfg,
+    dataset_cfg=dataset_config,
+    dataloader_cfg=dataloader_config,
+    lr_cfg=lr_cfg,
+    loss_cfg=loss_cfg,
+    tokenizer_path=QWEN3_MOE_PATH,
+    global_batch_size=16,
+    # total_epoch=1,
+    total_step=15,
+    # A 30B HF checkpoint per arm fills the disk fast, and nothing here needs the weights back.
+    debug_skip_save=True,
+    work_dir="/tmp/qwen3_moe_30BA3_sac",
+    seed=0,
+)

ci/config/sac_verify_attn.py

@@ -0,0 +1,39 @@
+"""SAC 验证配置（attn）。配套 tools/verify_sac.sh，勿手动单独运行。
+
+eager + all2all：关掉 compile 才能把 loss 差异归因到 SAC 本身而不是编译核。
+"""
+
+import os
+
+from xtuner.v1.config import AdamWConfig, FSDPConfig, LRConfig
+from xtuner.v1.datasets import FTDPTokenizeFnConfig
+from xtuner.v1.datasets.config import DataloaderConfig, DatasetConfig
+from xtuner.v1.loss.ce_loss import CELossConfig
+from xtuner.v1.model.moe.qwen3 import Qwen3MoE30BA3Config
+from xtuner.v1.train import TrainerConfig
+
+
+moe_cfg = Qwen3MoE30BA3Config(ep_size=4, dispatcher="all2all", recompute_cfg=['save_attn'])
+fsdp_cfg = FSDPConfig(torch_compile=False, cpu_offload=False, ep_size=moe_cfg.ep_size)
+
+trainer = TrainerConfig(
+    load_from=os.environ["QWEN3_MOE_PATH"],
+    model_cfg=moe_cfg,
+    optim_cfg=AdamWConfig(lr=6e-05),
+    fsdp_cfg=fsdp_cfg,
+    dataset_cfg=[
+        {
+            "dataset": DatasetConfig(name="alpaca", anno_path=os.environ["ALPACA_PATH"], sample_ratio=1.0),
+            "tokenize_fn": FTDPTokenizeFnConfig(max_length=8194),
+        },
+    ],
+    dataloader_cfg=DataloaderConfig(pack_max_length=8192),
+    lr_cfg=LRConfig(lr_type="cosine", lr_min=1e-6),
+    loss_cfg=CELossConfig(),
+    tokenizer_path=os.environ["QWEN3_MOE_PATH"],
+    global_batch_size=16,
+    total_step=3,
+    debug_skip_save=True,
+    work_dir="/tmp/sac_verify_attn",
+    seed=0,
+)

ci/config/sac_verify_none.py

@@ -0,0 +1,39 @@
+"""SAC 验证配置（none）。配套 tools/verify_sac.sh，勿手动单独运行。
+
+eager + all2all：关掉 compile 才能把 loss 差异归因到 SAC 本身而不是编译核。
+"""
+
+import os
+
+from xtuner.v1.config import AdamWConfig, FSDPConfig, LRConfig
+from xtuner.v1.datasets import FTDPTokenizeFnConfig
+from xtuner.v1.datasets.config import DataloaderConfig, DatasetConfig
+from xtuner.v1.loss.ce_loss import CELossConfig
+from xtuner.v1.model.moe.qwen3 import Qwen3MoE30BA3Config
+from xtuner.v1.train import TrainerConfig
+
+
+moe_cfg = Qwen3MoE30BA3Config(ep_size=4, dispatcher="all2all")
+fsdp_cfg = FSDPConfig(torch_compile=False, cpu_offload=False, ep_size=moe_cfg.ep_size)
+
+trainer = TrainerConfig(
+    load_from=os.environ["QWEN3_MOE_PATH"],
+    model_cfg=moe_cfg,
+    optim_cfg=AdamWConfig(lr=6e-05),
+    fsdp_cfg=fsdp_cfg,
+    dataset_cfg=[
+        {
+            "dataset": DatasetConfig(name="alpaca", anno_path=os.environ["ALPACA_PATH"], sample_ratio=1.0),
+            "tokenize_fn": FTDPTokenizeFnConfig(max_length=8194),
+        },
+    ],
+    dataloader_cfg=DataloaderConfig(pack_max_length=8192),
+    lr_cfg=LRConfig(lr_type="cosine", lr_min=1e-6),
+    loss_cfg=CELossConfig(),
+    tokenizer_path=os.environ["QWEN3_MOE_PATH"],
+    global_batch_size=16,
+    total_step=3,
+    debug_skip_save=True,
+    work_dir="/tmp/sac_verify_none",
+    seed=0,
+)