diff --git a/autotest/config/npu_qwen3_recompute.py b/autotest/config/npu_qwen3_recompute.py index 197ea1b0ac..3f8211afc4 100644 --- a/autotest/config/npu_qwen3_recompute.py +++ b/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 = [ diff --git a/autotest/config/qwen3_5_recompute.py b/autotest/config/qwen3_5_recompute.py index ada5dfd2e3..2338392275 100644 --- a/autotest/config/qwen3_5_recompute.py +++ b/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 = [ diff --git a/autotest/config/qwen3_recompute.py b/autotest/config/qwen3_recompute.py index 5d7ba68a2c..467bcffd35 100644 --- a/autotest/config/qwen3_recompute.py +++ b/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 = [ diff --git a/ci/config/qwen3_moe_30BA3_sac.py b/ci/config/qwen3_moe_30BA3_sac.py new file mode 100644 index 0000000000..b632a87d5f --- /dev/null +++ b/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, +) diff --git a/ci/config/sac_verify_attn.py b/ci/config/sac_verify_attn.py new file mode 100644 index 0000000000..889f529e1b --- /dev/null +++ b/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, +) diff --git a/ci/config/sac_verify_none.py b/ci/config/sac_verify_none.py new file mode 100644 index 0000000000..35c7204e42 --- /dev/null +++ b/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, +) diff --git a/docs/design/selective_checkpointing.md b/docs/design/selective_checkpointing.md new file mode 100644 index 0000000000..536f5076f5 --- /dev/null +++ b/docs/design/selective_checkpointing.md @@ -0,0 +1,240 @@ +# Selective Checkpointing & Activation Offload + +## 背景 + +目前 XTuner 的 recompute（gradient checkpointing）和 activation offload 都只能做到 +`DecoderLayer` 级别：在 FSDP sharding 阶段，对整层 `layer` 套 `checkpoint_wrapper` +（见 `xtuner/v1/model/moe/moe.py:1039`），并由 `fsdp_config.recompute_ratio` +（`xtuner/v1/config/fsdp.py:18`）控制重算前若干层。 + +这种粒度太粗：一层里 attention / MoE / MLP 的显存和重算代价差异很大，无法只重算其中一部分。 +本设计的目标是把 recompute 和 activation offload 下沉到 **region（埋点区间）级别**：模型作者在 +forward 里用 `checkpoint_record` 埋点，每个模型在 `model_cfg` 里以 marker 区间声明「哪些段要 +重算 / offload」，再由 `recompute_ratio` 控制作用到哪几层。区间可**跨模块**，不受架构边界约束。 + +## Goals / Non-Goals + +**Goals**（按优先级排序） + +- **[最高优先] 在 torch 2.10 上用非 legacy（`use_reentrant=False`）的 checkpoint 替换现有 legacy 接口。** + 这是本设计的前置项，需落实两件事：①domino EP（`intra_layer_micro_batch > 1`）下非 reentrant 能正确 + 重算；②非 legacy 走 `saved_tensors_hooks`，`DecoderLayer` 可返回 TypedDict、`seq_ctx` 等可 dict 直传， + `_check_signature_of_forward` 随之可删、model / mtp 简化。（后续 region 级 SAC recompute 也依赖此升级， + 见「SAC：region 级 recompute」。）**先做这一步，再据其结论落地后续 schema 与实现。** +- 支持 region（埋点区间）级别的 recompute，用户通过 `recompute_cfg`（语义枚举）声明要留驻哪些子结构，粒度可跨模块。 +- `recompute_cfg` 是每个模型的 per-layer 默认方案，`recompute_ratio` 决定这套方案作用到哪些 layer。 + +**Non-Goals（本 PR 不做，但需保证不冲突）** + +- **activation offload 的实现本身**。本 PR 只做 recompute，`activation_offload_cfg` 的 schema + 与 recompute 对齐、预留位置，保证后续开发不会与本次改动冲突，具体 offload 机制（落在 + `async_save_on_cpu` 上、以及与环境变量 `XTUNER_ACTIVATION_OFFLOAD` 的 BC）留到后续设计。 +- 脱离代码埋点的全局 monkeypatch。recompute 区间只通过模型 forward 里显式的 `checkpoint_record` + 埋点界定，不做对任意自由函数的运行时替换。 + +## 第一优先：调研升级到非 legacy checkpoint wrapper + +> 这是本设计的前置调研项，先于下面的 schema / 实现完成。本章只谈 wrapper 升级本身，就两件事；SAC +> 对非 legacy 的依赖与相关约束放在「SAC：region 级 recompute」里讲。 + +**现状**：`checkpoint_wrapper`（`xtuner/v1/model/utils/checkpointing.py`）封装 legacy 的 +`torch.distributed.algorithms._checkpoint.checkpoint_wrapper` + `CheckpointImpl.REENTRANT`。 +reentrant 版是一个 `autograd.Function`，要求需要求梯度的 tensor **以位置参数直传直返**；`dict` 嵌套 +tensor 会被当成单个非 tensor 参数，wrapper 判定「没有输入需要梯度」→ 断梯度 / 报错。 +`_check_signature_of_forward`（`checkpointing.py:28`）那套严格签名检查正是为在这种情况下早失败而存在。 + +升级到 `use_reentrant=False` 要落实两件事： + +**① domino EP 正确性。** 现在用 `REENTRANT` 的主要原因是 domino EP（`trainer_cfg` 里 +`intra_layer_micro_batch > 1`）下重算要能正常进行。先复现旧接口为何需要 reentrant / 非 reentrant +会失败的现象与根因，再验证 `use_reentrant=False` 在 `intra_layer_micro_batch > 1` 下能否正确重算。 + +**② TypedDict I/O + 删签名检查。** `use_reentrant=False` 用 `saved_tensors_hooks` 实现，对嵌套输入 / +返回值天然友好。升级后 `DecoderLayer` 可以**返回 TypedDict 而不是位置约定的 tuple**、`seq_ctx` / +`position_embeddings` 等可作为 dict 直传，mtp 那类第二路输出不必再挤进 tuple 某个约定位置、由下游按 +index 解包。因此本步的**直接产出**包括：decoder layer 返回值 tuple → TypedDict，删除 +`_check_signature_of_forward`，model / mtp 层代码随之简化。机制已用 POC（玩具模型）验证：非 legacy 下 +「嵌套 dict 输入 + TypedDict 输出」的梯度与无 checkpoint **精确一致**、到达 dict 内每个 tensor；同一个 +dict 进 / 出的 layer 用 legacy REENTRANT 则直接报 `element 0 of tensors does not require grad`。落到 +真实 layer 需覆盖 `seq_ctx` / `position_embeddings` 等实际结构。 + +**验收**：两件事都到位后，domino EP 下做 grad-norm parity + 显存对比（升级前后），确认数值一致、显存 +收益符合预期。 + +## SAC：region 级 recompute + +region 级 recompute 的主体是 SAC（selective activation checkpointing）：整层套一个 checkpoint、由 +policy 逐 op 裁决存或重算。用户侧的 `recompute_cfg`（语义枚举）只是这套机制最终的**表达形式**，先讲机制、再落到配置。 + +### 机制 + +recompute 落在 `torch.utils.checkpoint.checkpoint(..., use_reentrant=False, context_fn=...)` +的 SAC 路径上，而**不是**对 submodule 逐个套 wrapper。被选中的整个 `DecoderLayer` 套**一个** +checkpoint，`context_fn` 由 torch 的 `create_selective_checkpoint_contexts(policy_fn)` 构造， +`policy_fn` 逐 op 裁决：默认 `MUST_RECOMPUTE`（全重算），被点名的区间返回 `MUST_SAVE`（该段激活留驻、 +不重算）。用户看到的是 region 粒度（marker），底层是 op 粒度的 policy 在执行。 + +```python +from torch.utils.checkpoint import checkpoint, create_selective_checkpoint_contexts, CheckpointPolicy + +def policy_fn(ctx, op, *args, **kwargs) -> CheckpointPolicy: + # marker 状态机决定当前是否处于某个 SAVE 区间 + return CheckpointPolicy.MUST_SAVE if _in_save_interval() else CheckpointPolicy.MUST_RECOMPUTE + +ctx_fn = lambda: create_selective_checkpoint_contexts(policy_fn) +out = checkpoint(layer_forward, *args, use_reentrant=False, context_fn=ctx_fn) +``` + +这条路径**必须** `use_reentrant=False`——`create_selective_checkpoint_contexts` / `context_fn` +只有非 legacy 才支持，因此它与「第一优先」的非 legacy wrapper 调研是同一件事：调研结论为真是本 schema +成立的前提。 + +**默认口径的选择**：policy 对每个 op 只回答「SAVE（存下来、不重算）还是 RECOMPUTE（丢掉、backward +重算）」。于是「默认 + 点名例外」有两种反向配法： + +- **A**：默认 `MUST_RECOMPUTE`，点名的区间设成 `MUST_SAVE` —— 默认全重算，点名少数**不重算**。 +- **B**：默认 `MUST_SAVE`，点名的区间设成 `MUST_RECOMPUTE` —— 默认全存，点名少数**重算**。 + +两种 SAC 都能做、也都能跨模块，差别只是 `policy_fn` 默认返回值那一行。**本设计只用 A**：recompute 的 +目的是省显存，常态一定是「重算绝大多数、只把少数重算代价高的段（如 flash-attn）留驻」，所以默认全重算、 +点名 SAVE 才符合直觉，`recompute_cfg` 的语义也单一。不为对称性把 B 也塞进 schema。 + +### 埋点：`checkpoint_record` + +模型作者在 forward 里埋**点位**（不是区间），标注可寻址的语义边界： + +```python +def forward(self, x): + checkpoint_record("attn.core") + ... + checkpoint_record("mlp.up") + ... +``` + +`checkpoint_record(name)` 是命令式点位，不在 SAC 会话内时是 no-op。它翻转一个 context-local 的 +「当前活跃区间」状态机，policy 据此决定当前 op 存还是重算。 + +**为什么是点位 marker 而不是 `with` 上下文**：重算区间可能**跨模块**——例如 `A.a2 + B.b1` +（A 尾部 + B 头部），二者分处 `A.forward` 与 `B.forward`，中间隔着上层的 `A()→B()` 调用。`with` +块受词法作用域限制，无法横跨两个兄弟模块的 forward；命令式 marker 依赖**运行时程序顺序 + 外部状态**， +天然可跨。 + +### 用户配置：语义枚举，不是 marker 字符串 + +用户**不直接填 marker 字符串**（那反人类，且要求用户了解模型内部埋点）。面向用户的是一组**语义枚举**， +每个枚举表示「某类子结构不重算（留驻）」，与 `recompute_ratio`（选层）配合： + +```python +class RecomputeUnit(Enum): + SAVE_ATTN = auto() # attention 段不重算、留驻 + SAVE_MOE_GATE = auto() + ... + +# model config +recompute_cfg: list[RecomputeUnit] | bool | None +``` + +枚举 → 具体 marker 区间的映射，仿照 `default_compile_cfg`，定义在每个 model 的默认 cfg 里（各模型的 +attention / MoE 结构不同，同一个 `SAVE_ATTN` 落到的 marker 区间也不同）： + +```python +@property +def default_recompute_cfg(self) -> dict[RecomputeUnit, list[tuple[str, str]]]: + return { + RecomputeUnit.SAVE_ATTN: [("attn.core", "attn.out")], + RecomputeUnit.SAVE_MOE_GATE: [("moe.gate", "moe.dispatch")], + } +``` + +三层职责因此各自稳定： + +- **用户**：config 里挑语义枚举（`SAVE_ATTN`），不碰字符串、不需了解内部结构。 +- **model 作者**：定义枚举 → marker 区间，每个模型一次、随架构而定（与 `checkpoint_record` 埋点同处一层）。 +- **内部**：把选中的枚举解析成 marker 区间，喂给 policy 的状态机（下节）。 + +> 与 `compile_cfg` 的 `bool | None | dict` 三态一致：`None` / `True` 用 model 默认方案，`False` 关闭 +> （全重算、不留驻）。是否再开放「直传区间」的高级口径（调试用）留待实现时定，默认路径只有枚举。 + +### 内部表示：marker 区间 → policy + +枚举解析出的每个 `(start, end)` 是一段 `[start, end)` 半开区间（`end` 取「下一步起点 marker」）， +区间内的 op 被标 `MUST_SAVE`。区间重叠 / 嵌套用「活跃区间集合非空即 SAVE」定义；policy 只做「命中 +start 置位、命中 end 复位」的状态机，不做配对断言。 + +这带来一个稳健性红利：marker 区间配不平（如 `end` 落在没走到的分支）只会**多留驻一段显存，绝不影响 +梯度**——SAVE 与 RECOMPUTE 两条路数值都精确正确。因此 model 作者写 `default_recompute_cfg` 时无需 +担心边界配对的正确性。 + +### 状态生命周期（实现关键） + +会话状态必须在**被 checkpoint 的 callable 顶部**重新绑定、每趟重置。因为 `use_reentrant=False` +的 backward 会**重跑整个 forward** 来重算，marker 随之重新命中；若状态挂在 module 全局而非随 callable +重置，forward 与 recompute 两趟的 SAVE 集合会错位。用 `contextvars` 在 callable 入口 set、 +`finally` reset。 + +### 待验证的 SAC 风险点 + +玩具模型 POC 已验证：梯度精确、跨模块区间成立、SAVE 段确实不重算。但下面两点玩具模型覆盖不到，需在 +真实 layer 上验证： + +- **op 序确定性**：SAC 的 save-list 靠 forward 与 recompute 两趟 **op 顺序一一对齐**。 + `intra_layer_micro_batch > 1`（domino）下 micro-batch 交织会不会改变 op 顺序，是本方案能否落地的 + 头号拦路虎。 +- **torch.compile 交互**：SAC 的 `context_fn` + `saved_tensors_hooks` 在 compile 下的行为需单独确认。 + +## 寻址语义：与 `recompute_ratio` 的两步分工 + +生效范围两步式，两者正交、不是覆盖关系： + +1. **选层**：`recompute_ratio` 决定哪些 layer 套上 SAC checkpoint（沿用现有 `_should_recompute` 语义）。 +2. **选区间**：在被选中的 layer 内部，`recompute_cfg` 选中的枚举解析成 marker 区间，抠出**留驻不重算**的段，其余全重算。 + +即 `recompute_ratio` 定义「模板铺到前多少层」，`recompute_cfg` 定义「每层内部哪些段不重算」。 + +> 旧设计的「按 module 类匹配 + 最外层去重」不再适用：SAC 下整层只有**一个** checkpoint，不存在 +> 嵌套 wrapper 的双重 checkpoint 问题，SAVE / RECOMPUTE 由 policy 逐 op 裁决，没有「谁套谁」。 + +## recompute 与 activation offload 的关系 + +两者不互斥，但作用对象正好互补：SAC 下 offload 的对象是被**留驻的那部分激活**（`MUST_SAVE` 段） +——被重算的段本就不落地、无需 offload，能省显存的只有留驻下来的 SAVE 张量。 + +本 PR **只实现 recompute**。对 offload 只有一个约束：`activation_offload_cfg` 复用 `recompute_cfg` +同一套语义枚举面向用户、同一套「枚举 → marker 区间」的 model 默认映射机制，且 SAC 的 saved-tensor +注入点要为后续 offload 的 hook 机制（预期落在 `async_save_on_cpu` 上、按留驻 tensor 建 saved-tensor +hook）留出干净入口，不写死实现细节。offload 的具体机制与环境变量 BC 留到后续设计文档。 + +## FSDP 包裹顺序 + +FSDP（`fully_shard`）永远是**最外层**：先对 `DecoderLayer` 套上 SAC checkpoint，再对该 layer +做 `fully_shard`。由于 `fully_shard` 以整层为单位、SAC checkpoint 也以整层为边界，backward 重算 +发生时该层参数已由 FSDP all-gather 就位，重算路径无需额外的 gather 时机处理。 + +## 环境 + +torch 2.10 环境路径：`/mnt/shared-storage-user/yehaochen/miniconda3/envs/py312-pt210/bin/torchrun` + +## 开发计划：分工与 stacked PR + +**硬约束**：非 legacy 升级（第一优先）是**门禁**，不能与后续并行——domino EP 下 +`use_reentrant=False` 能否正确重算的结论，决定整套 SAC schema 是否成立。因此是「门禁串行 → 通过后扇出」， +不是一次并行 N 块。 + +切成 **3 个 owning agent / 3 层 stacked PR**： + +| 层 | agent | 内容 | 主要文件面（尽量不重叠） | +| --- | --- | --- | --- | +| **A（base）** | 非 legacy 基座 | wrapper → 非 legacy + domino EP 验证；`DecoderLayer` 返回 tuple → TypedDict、`seq_ctx` dict 直传、删 `_check_signature_of_forward`、改 model / mtp 调用点。**并定义共享契约**：`checkpoint_record`（先落 no-op 桩）、`RecomputeUnit` 枚举、枚举解析出的区间类型 | `checkpointing.py`、decoder layers、model / mtp、sharding 路径 | +| **B（stack 在 A 上）** | SAC 引擎 | `checkpoint_record` 背后的 `contextvars` session + `policy_fn` + `create_selective_checkpoint_contexts` 接线；FSDP 路径换成整层 SAC；op 序 / `torch.compile` 在真实 layer 上验证 | 新 SAC 引擎模块 + 包裹调用点 | +| **C（stack 在 A 上，与 B 并行）** | 配置 + 逐模型 | `recompute_cfg` 三态解析（`bool \| None \| list[RecomputeUnit]`）+ 与 `recompute_ratio` 打通；各模型 `default_recompute_cfg`；在模型 forward 里埋 `checkpoint_record` | config 模块 + 各 model 文件 | + +**为什么合并原「非 legacy 升级」与「TypedDict 重构」为一个 A**：两者都要改 `checkpointing.py`、decoder +layer、sharding 路径，拆两个 agent 会在同一批文件上冲突；同一 owner 一次到位。 + +**为什么 B、C 能真并行**：B 只碰引擎与包裹点、不碰模型 forward；C 只碰 config 与模型文件、不碰引擎内部。 +C 埋的 `checkpoint_record` 在 B 合入前就是 no-op，无害；两者只在 A 的契约上相遇。 + +**唯一 seam**：FSDP 包裹处「把解析出的区间应用到选中层」这一行归 **B** 拥有（它是机制）；C 只负责把 +config 解析成区间、通过契约函数交给 B，不直接改包裹点，避免 B / C 同时改 sharding 路径。 + +**stack 顺序**：`A → B`、`A → C`，B / C 互为 sibling；A 合入后各自 rebase 到 A，平行推进。 diff --git a/docs/design/selective_checkpointing_and_compile.md b/docs/design/selective_checkpointing_and_compile.md new file mode 100644 index 0000000000..5e74d0ae67 --- /dev/null +++ b/docs/design/selective_checkpointing_and_compile.md @@ -0,0 +1,266 @@ +# Selective checkpointing under `torch.compile` + +Why unit-level SAC needs the compiler's cooperation, which routes exist, and what each one +measured. All numbers are Qwen3-MoE-30BA3, `ep_size=4` on 8 GPUs, `dispatcher="deepep"`, +torch 2.10, reported as the mean of steps 5-8 of an 8-step run. Two shapes recur: + +- **8k** — `pack_max_length=8192`, no domino. Baseline 8904.4 tgs, 84.32 GB peak. +- **4k** — `pack_max_length=4096` with domino (`intra_layer_micro_batch=2`). Baseline 7812.4 tgs. + +The 4k shape is CPU-bound (GPU busy 55-60%), so every CPU-side cost lands fully in wall time there +and is largely hidden at 8k. The same change reads as −0.9% at 8k and −13.9% at 4k; neither number +is wrong, they measure different regimes. + +## The problem + +The checkpoint policy is asked about every op that reaches the dispatcher, and it has to answer one +question: does this op belong to a unit the user asked to keep. There are exactly two ways to know, +and which one applies is a property of the unit, not a style choice. + +**Naming the op.** Works when the op is specific enough to identify the unit on its own — +`flash_attn::_flash_attn_varlen_forward` appears nowhere else in the model. Costs nothing: the +compile set is untouched, no graph breaks, and it works identically inside and outside compiled +code, because the ops worth keeping are exactly the ones inductor cannot fuse and those are the ones +still reaching the dispatcher. + +**Naming the callable.** Works for anything, at a price. The marker is ordinary python state — a +`ContextVar` set around the callable — and compiled code neither writes nor reads it: reading one +inside a `fullgraph=True` region is a hard compile error, and a value written during tracing is +traced away. So the callable has to leave the compiled set, via `torch._dynamo.disable`. + +This is not a visibility problem. With the whole layer compiled the policy is still consulted +**117597 times per run**, because ops that inductor cannot fuse still go through the dispatcher. +Measured on this model, 26 distinct ops reach it, among them `flash_attn::_flash_attn_varlen_forward_v2`, +`moe::m_grouped_gemm`, `moe::permute`/`unpermute` and `aten::mm.out`; 61% of the calls are +`aten::record_stream`. What is missing for a callable-scoped unit is the region identity, not the op. + +## The route taken: op identity where possible, withdrawal where not + +`RecomputeTargetMap` binds each unit a model supports to one of the two resolutions. MoE declares +`SAVE_ATTN` as `KeptOps` and the gate and dispatch stages as `KeptCallables`. + +Two mechanics of the withdrawal are easy to get wrong: + +- Removing an entry from `compile_cfg` does **not** keep it out of the compiled set. Dynamo inlines + it into whichever compiled caller reaches it, and the marker is traced away exactly as before. + `torch._dynamo.disable` is what makes it run in python. +- A disabled callee inside a `fullgraph=True` region is a hard error + (`Skip inlining torch.compiler.disable()d function`), not a split. The surviving entries are + therefore relaxed to `fullgraph=False`. + +Measured, 8k, `attn`, when it was still resolved by withdrawing the callable that encloses it: + +| | tgs | peak | compiled graphs / captured calls | kept | +|---|---|---|---|---| +| no unit | 8904.4 | 84.32 GB | 6 / 154 | — | +| withdrawal, marker never fired | 8883.1 | 84.32 GB | 6 / 154 | nothing | +| withdrawal, marker firing | 8021.7 | 112.14 GB | 4 / 26 | 42112 tensors | + +−9.9% for the same unit that op identity delivers for free. The cost is concentrated, not diffuse: +`_pre_moe_forward` accounts for 128 of the 154 captured calls, because it exists to give +`torch.compile` the ops on either side of attention. Withdrawing it is most of what an MoE layer +compiles. That is why attention resolves by op identity and why `KeptCallables` should name the +*smallest* callable that covers the unit: the compilation given up is the whole callable's. + +## What the shipped units cost + +Measured on the shipped resolutions, 16k domino (`pack_max_length=8192`, +`intra_layer_micro_batch=2`), `ep_size=4`, deepep. **Three independent runs per setting**, because +one run per setting is not enough to say anything about throughput here: + +| unit | tgs (3 runs) | mean | peak allocated | peak reserved | +|---|---|---|---|---| +| none | 10206.0 / 9990.6 / 10205.9 | 10134.2 | 84.4-85.3 GB | 108.5-109.3 GB | +| `attn` | 10165.6 / 10204.6 / 10071.2 | 10147.2 | 90.4-91.2 GB | 114.4-115.3 GB | +| `moe_gate` | 10095.2 / 10179.5 / 10144.4 | 10139.7 | 93.7-94.5 GB | 117.7-118.7 GB | +| `moe_dispatch` | OOM | — | — | — | + +**Throughput is unchanged.** The three means are within 0.1% of each other, while the baseline's own +three runs span 2.1%. Any single-run comparison of these units reads as ±2% in whichever direction +the run-to-run variance happened to fall, and means nothing. + +**Memory is the reproducible effect**: `attn` costs +6.0 GB allocated, `moe_gate` +9.2 GB, +each within ±0.4 GB across runs. + +So a unit here buys nothing on this model and costs memory. That is not a statement about the +mechanism -- it is a statement about which activations an MoE layer's recompute is actually spent +on, and the census below says why. + +`moe_dispatch` is not usable at this shape at all. It keeps the permutation and padding buffers +on both sides of the all-to-all, the widest tensors in the layer, and domino already runs at 109 GB +reserved of 140 GB. It OOMs during the first step -- and also OOMs without domino, and at +`pack_max_length=4096`, where it peaks at 111.9 GB against the baseline's 85.9 GB. It is verified +numerically for one step (see below) and declared for smaller models, not as a default here. + +## Numerical correctness + +Checkpointing changes only the backward pass, so **step-1 loss must be bit-identical** whatever is +kept. It is, across every setting and both shapes -- `2.46262765` at 8k, `2.38410378` at 4k, +including the `moe_dispatch` run that OOMs immediately afterwards. + +Gradients need a noise floor to interpret, which is why the baseline was run twice under identical +settings: + +| | step-1 loss | step-1 grad_norm | vs baseline | +|---|---|---|---| +| baseline | 2.46262765 | 24.39401245 | — | +| baseline, second run | 2.46262765 | 24.39329147 | 3.0e-5 | +| `attn` | 2.46262765 | 24.39325905 | 3.1e-5 | +| `moe_gate` | 2.46262765 | 24.39400864 | 1.6e-8 | + +Every unit sits at or below the floor two identical runs produce. The floor itself comes from +reduction ordering in the grouped GEMM and the all-to-all, and exists with no unit selected. + +Measured eager (`torch_compile=False`, all2all, no domino) so that compilation cannot be the source +of a difference; the compiled path is covered by the throughput runs above. + +## Whether a unit repays its cost + +On this model: no. Every unit raises peak reserved -- keeping activations is the trade -- and none +of them buys back measurable throughput. The census below says why. + +A resident-activation census under the shipped configuration (domino, 8k, whole layer checkpointed) +found **245 tensors totalling 15.27 GiB**: + +| producer | count | bytes | +|---|---|---| +| `LogSoftmaxBackward0 (16384, 151936) float32` | 1 | 9.27 GiB | +| layer-boundary `(1, 8192, 2048)` bf16 | 94 | 2.94 GiB | +| `TBackward0 (2048, 151936)` bf16 | 1 | 0.58 GiB | +| everything else (60 producers) | 149 | ~2.5 GiB | + +The loss logits alone are 61% of the resident set, and no recompute unit addresses them. Nothing +inside the MoE layers survives — the whole-layer checkpoint is doing its job. That bounds how much +any unit-level policy can win here. + +## Routes measured + +### Keeping the expert GEMM by op identity + +The same resolution `SAVE_ATTN` uses, pointed at the ops that dominate an MoE layer's recompute: + +| kept | tgs (8k) | peak | captured | +|---|---|---|---| +| baseline | 8904.4 | 84.32 GB | 154 | +| `flash_attn::_flash_attn_varlen_forward_v2` | 8828.2 | 87.31 GB | 154 | +| `moe::m_grouped_gemm` | **9477.3 (+6.4%)** | 105.25 GB | 154 | +| `moe::m_grouped_gemm` + `permute` + `unpermute` | **9616.7 (+8.0%)** | 122.61 GB | 154 | + +Not declared as a unit here because +20 to +38 GB is not a trade a user can currently tune — the +selection is per-model, not per-layer. It is the obvious next unit once the selection can say +"in the first N layers". The limit of op identity is expressiveness: it cannot reach anything fused +into a generated kernel, and it cannot scope a unit to part of the model. It is what torchtitan does +(`torchtitan/distributed/activation_checkpoint.py`), with a preset op list and no region concept. + +### Withdrawing a smaller callable + +Making the excluded callable smaller does not make the exclusion cheaper in proportion, because +`torch._dynamo.disable` breaks the graph once **per level of the inline stack it has to unwind**, +not once per call: + +| cut point | inline depth | breaks | captured | tgs (8k) | +|---|---|---|---|---| +| `MultiHeadAttention.forward` | 1 | 2 | 60 | 8101.7 | +| `attn_imp.flash_attention` | 2 | — | 96 | 8664.5 | +| `flash_attn_varlen_func` | 3 | 4 | 104 | 8660.3 | + +Cutting deeper leaves more code compiled but unwinds more levels; the two effects cancel. A +withdrawal is cheapest when its boundary is *shallow*, near the compile entry point — the opposite of intuition. + +The two resolutions also cannot be made equivalent. Even at the tightest cut, the withdrawal keeps +`aten::alias` alongside the attention kernel, because a `dynamo.disable`d segment runs eagerly and +eager execution dispatches bookkeeping ops that do not exist in the compiled path. + +### Letting the marker run while compiling, and taking the graph break + +Instead of withdrawing the callable, let the marker execute during tracing and accept the graph +break Dynamo takes at it. The marker becomes live and the policy sees the unit — but only for ops +that reach the dispatcher, so it keeps 752 tensors where withdrawing the callable keeps 42112. + +Measured, 8k, `attn`: 8578.4 tgs, 86.03 GB, 11 graphs / 90 captured, 5 breaks. Cheaper than +withdrawing the callable and more expensive than op identity, with a third semantics again. Not +adopted because "keep this unit" should not silently mean "keep the handful of ops in it that +inductor happened not to fuse". + +### `fx_traceback.annotate` plus a joint-graph pass + +The only marker that survives compilation intact. Dynamo executes it during tracing +(`_dynamo/variables/ctx_manager.py`), stamps `node.meta["custom"]`, and AOT carries it through +decomposition and into the backward nodes. Verified on GPU under `fullgraph=True`: 6 graphs, 154 +captured calls, **0 graph breaks** — identical to the baseline — with the joint pass seeing 410 +annotated nodes and pinning 386. + +And it changes nothing. Peak memory stayed at 84.32 GB, byte for byte. + +The reason is a level mismatch. `torch.utils.checkpoint(use_reentrant=False)` is implemented with +saved-tensor hooks: its `pack_hook` replaces every tensor autograd saves inside the region with a +holder and frees it. A compiled region is an ordinary `autograd.Function`, so the tensors the +partitioner chose to save are handed to `ctx.save_for_backward` and discarded like everything else. +The partitioner's decision is made and honoured — its *product* is thrown away. + +So the tag route requires the partitioner to be the only decider, which means removing the outer +checkpoint. + +### Removing the outer checkpoint + +Then the partitioner governs, and the tags work. It also costs 61.8 GiB. + +| | resident tensors | resident bytes | +|---|---|---| +| whole layer checkpointed | 245 | 15.27 GiB | +| no checkpoint, everything forced to `MUST_RECOMPUTE` | 1706 | 77.04 GiB | + +`MUST_RECOMPUTE` applies *within* one joint graph. It cannot recompute across a graph boundary, +because the backward of each graph needs that graph's inputs. Domino + EP splits a layer into many +graphs, and the boundaries land on the widest tensors in the model: one `(49152, 2048)` per layer +(9.00 GiB total), one `(65536, 768)` per layer (4.50 GiB), 96 layer-boundary `(1, 8192, 2048)` +(3.00 GiB). Of the 77.04 GiB, 18.33 GiB has no `grad_fn` at all — pure graph-boundary input. + +Forcing every node to recompute changes nothing versus a `budget=0.0`: 131.14 GB against 131.34 GB. +The memory that cannot be reclaimed is not the partitioner's to reclaim. + +`torch._functorch.config.activation_memory_budget` on the same shape, with no checkpoint: + +| budget | tgs | peak | +|---|---|---| +| baseline (whole layer checkpointed) | 8904.4 | 84.32 GB | +| 0.0 | 10032.5 (+12.7%) | 107.71 GB | +| 0.25 | **10282.0 (+15.5%)** | 111.68 GB | +| 0.5 | 9750.1 (+9.5%) | 113.25 GB | + +The fastest result measured anywhere, for +23 GB and no way back down: `budget=0.0` is the floor and +it is 23.4 GB above the baseline. The curve is not monotone (0.25 beats both neighbours) and that is +unexplained; it is a single point per budget, not a characterised curve. + +**Under domino this route is unusable.** Domino already needs ~19 GB more reserved, and removing the +checkpoint adds ~24 GB more, which crosses the 140 GB device limit: `budget=0.0` runs at 1925 tgs +(5.4× slower than the checkpointed baseline) with 134.60 GB reserved, and 0.25 and 0.5 both OOM. + +## The route that would be complete, and what blocks it + +If the checkpoint became a `tag_activation_checkpoint` HOP — that is, if `torch.compile` wrapped the +checkpoint call rather than sitting inside it — the two mechanisms collapse into one. The policy +result is written to `node.meta["recompute"]` during tracing (`torch/utils/checkpoint.py`) and the +partitioner consumes it; there is no second decider to overrule it, no dispatch mode at runtime, and +no callable to withdraw from the compile set. + +The HOP only forms if the checkpointed region speculates into a single subgraph. On the domino path +it never does: `torch_all2all.py` reads `permuted_hidden_states.grad_fn` to register a backward +pre-hook for the CUDA-event choreography, and Dynamo cannot trace a tensor's `grad_fn`. Speculation +fails, everything traced is discarded, and the graph before the checkpoint call is empty — measured +at 30B as `unique_graphs=0, calls_captured=0`, a silent degradation to eager rather than an error. + +torchtitan reaches this topology because its MoE dispatch is a custom op (`deepep.dispatch.default`), +which enters the graph whole. Making xtuner's dispatcher expressible the same way — replacing the +`grad_fn.register_prehook` choreography — is the one change that unlocks it. + +## Memory wall, for context + +Independent of checkpointing: domino at 16384 is 4.44× *slower* than not using it (14.20 s vs +3.199 s per step) purely because reserved memory reaches 134.43 GB of 140.06 GB and the caching +allocator thrashes. At 8192 the same code is **16.1% faster** with domino than without (1.583 s vs +1.838 s). `expandable_segments` does not help, so this is total demand rather than fragmentation, and +switching between the all2all and deepep backends does not help either (deepep is 7.1% faster only +when domino is off). + +Any activation-memory work on the domino path should check headroom before it profiles kernels. diff --git a/tests/model/test_recompute.py b/tests/model/test_recompute.py index 1d323cd926..07160861ad 100644 --- a/tests/model/test_recompute.py +++ b/tests/model/test_recompute.py @@ -5,16 +5,23 @@ test_non_tensor_signature_preserves_gradients: 关键字参数 + dict 返回值下梯度与不重算一致。 test_checkpointing_keeps_the_module_itself: 换类而非套壳，isinstance/属性/容器协议原生可用。 test_module_without_a_protocol_does_not_gain_one: 被包裹模块没有的协议不会凭空出现。 +TestRecomputeCfgResolution test_unset_cfg_keeps_full_recompute: `None` 不改变显存行为，解析为不留驻。 test_true_selects_every_supported_unit: `True` 选中模型声明的全部 unit。 test_explicit_units_select_only_themselves: 显式 list 只选中对应 unit。 - test_string_units_are_accepted: 配置文件里的字符串能解析成 RecomputeUnit。 + test_string_units_are_accepted: 配置文件里的字符串能解析成 SaveUnit。 test_unsupported_unit_is_rejected: 模型不支持的 unit 在构造时报错并列出支持项。 test_disable_propagates_into_nested_configs: `False` 递归关闭嵌套子模型配置。 test_disable_reaches_every_sub_model_of_a_real_compose_config: 真实 compose 配置的三个子配置都被关闭。 test_units_round_trip_through_json: enum 序列化成可读字符串并能读回。 +TestRecomputeRatioMigration + test_old_field_still_takes_effect: 旧的 fsdp_cfg.recompute_ratio 仍然生效。 + test_unset_old_field_leaves_the_new_one_alone: 没设旧字段时不覆盖新位置的值。 + test_setting_both_is_an_error: 两处都设时报错而不是静默择一。 +TestDeclaredTargets test_declared_targets_resolve: 声明表里的 op 名与 callable 名都能解析到真实对象。 test_no_unit_names_the_method_that_holds_most_compilation: 没有 unit 点名承载最多编译的那个方法。 +TestUnitCostIsProportionate test_an_op_identity_unit_costs_no_compilation: KeptOps 不改动编译集合。 test_a_callable_unit_keeps_its_callers_compiled: KeptCallables 只退出自身，调用者仍编译。 test_no_unit_withdraws_the_method_that_holds_most_compilation: 没有 unit 撤出编译占比最大的方法。 @@ -23,13 +30,28 @@ """ +import pydoc import pytest import torch from torch import nn from torch.autograd.graph import saved_tensors_hooks -from xtuner.v1.model.utils import apply_gradient_checkpointing +from xtuner.v1.config import FSDPConfig +from xtuner.v1.model.base import BaseModel, TorchCompileOption, XTunerBaseModelConfig, _disable_nested_switch +from xtuner.v1.model.compose.qwen3_vl import Qwen3VLMoE30BA3Config +from xtuner.v1.model.dense.dense import DENSE_RECOMPUTE_CFG +from xtuner.v1.model.moe.moe import MOE_RECOMPUTE_CFG, MoE, MoEConfig +from xtuner.v1.model.utils import ( + KeptCallables, + KeptOps, + RecomputeConfig, + SaveUnit, + apply_gradient_checkpointing, + resolve_kept_ops, +) +from xtuner.v1.module.attention import MHAConfig +from xtuner.v1.module.router import NoAuxRouterConfig class _KeywordOnlyBlock(nn.Module): @@ -167,3 +189,213 @@ def __init__(self) -> None: make_call(wrapped, x)["out"].square().sum().backward() assert x.data_ptr() in packed + + +def _build_tiny_moe_config(**overrides) -> MoEConfig: + """A MoE small enough to instantiate on the meta device in a config-only test.""" + router_config = NoAuxRouterConfig( + scoring_func="sigmoid", + router_scaling_factor=1.0, + n_group=1, + topk_group=1, + norm_topk_prob=True, + ) + return MoEConfig( + vocab_size=256, + max_position_embeddings=128, + pad_token_id=0, + eos_token_id=0, + num_hidden_layers=2, + hidden_size=64, + intermediate_size=128, + rms_norm_eps=1e-6, + rope_theta=1e6, + hidden_act="silu", + attention=MHAConfig(num_attention_heads=4, num_key_value_heads=4, head_dim=16), + tie_word_embeddings=False, + n_routed_experts=4, + n_shared_experts=1, + num_experts_per_tok=2, + first_k_dense_replace=1, + hidden_factor=1.0, + moe_intermediate_size=64, + router=router_config, + compile_cfg=False, + **overrides, + ) + + +def _resolve_units(**overrides) -> set: + with torch.device("meta"): + return set(MoE(config=_build_tiny_moe_config(**overrides))._selected_recompute_units) + + +class _NestedProbeConfig(XTunerBaseModelConfig): + """Stand-in for a sub-model config, as a compose model nests one.""" + + +class _ProbeConfig(XTunerBaseModelConfig): + text_config: XTunerBaseModelConfig + + +class _ProbeModel(BaseModel): + """A model that contributes nothing but ``BaseModel.__init__``'s config resolution.""" + + config: _ProbeConfig + + +class TestRecomputeCfgResolution: + def test_unset_cfg_keeps_full_recompute(self): + # `None` must not change the memory profile of an existing training run: unlike `compile_cfg`, + # it resolves to "retain nothing" rather than to the model's declared units. + assert _resolve_units(recompute_cfg=RecomputeConfig(save=None)) == set() + + def test_true_selects_every_supported_unit(self): + assert _resolve_units(recompute_cfg=RecomputeConfig(save=True)) == set(MOE_RECOMPUTE_CFG) + + def test_explicit_units_select_only_themselves(self): + assert _resolve_units(recompute_cfg=RecomputeConfig(save=[SaveUnit.MOE_GATE])) == {SaveUnit.MOE_GATE} + + def test_string_units_are_accepted(self): + # Configs arrive as JSON/py files where units are written as plain strings. + assert _resolve_units(recompute_cfg=RecomputeConfig(save=["attn"])) == {SaveUnit.ATTN} + + def test_unsupported_unit_is_rejected(self): + # A model declaring no units cannot honour any selection, so this is a user configuration + # error rather than something to silently drop. It surfaces at construction, before the run + # spends anything on materializing and sharding weights. + with pytest.raises(ValueError, match="does not support"): + _ProbeModel(_ProbeConfig(text_config=_NestedProbeConfig(), recompute_cfg=RecomputeConfig(save=[SaveUnit.ATTN]))) + + def test_disable_propagates_into_nested_configs(self): + # A sub-model resolves its own switch, so `False` on the outer config only means something + # if it reaches the nested ones. `compile_cfg` must stay untouched: the walk is per switch. + config = _ProbeConfig(text_config=_NestedProbeConfig(), recompute_cfg=RecomputeConfig(save=False)) + + model = _ProbeModel(config) + + assert model._selected_recompute_units == set() + assert config.text_config.recompute_cfg.save is False + assert config.text_config.compile_cfg is None + + def test_disable_reaches_every_sub_model_of_a_real_compose_config(self): + # The probe above has one nested config; a shipped compose config has three, one of them a + # further-derived MoE config. Exercised on the config walk rather than through the model, + # because constructing a 30B compose model is the expensive part and contributes nothing: + # what can regress here is which nested configs the walk reaches. + config = Qwen3VLMoE30BA3Config(recompute_cfg=RecomputeConfig(save=False)) + + _disable_nested_switch(config, "recompute_cfg", subfield="save") + + for sub_config in (config.vision_config, config.projector_config, config.text_config): + assert sub_config.recompute_cfg.save is False + + def test_units_round_trip_through_json(self): + # Trainer resume reads the config back, and serialized runs are read by humans, so units + # must survive as their readable names. + config = _build_tiny_moe_config(recompute_cfg=RecomputeConfig(save=[SaveUnit.ATTN, SaveUnit.MOE_GATE])) + + dumped = config.model_dump(mode="json")["recompute_cfg"] + assert dumped["save"] == ["attn", "moe_gate"] + + restored = _build_tiny_moe_config(recompute_cfg=RecomputeConfig(**dumped)) + assert restored.recompute_cfg.save == [SaveUnit.ATTN, SaveUnit.MOE_GATE] + + +class TestDeclaredTargets: + @pytest.mark.parametrize("target_map", [MOE_RECOMPUTE_CFG, DENSE_RECOMPUTE_CFG], ids=["moe", "dense"]) + def test_declared_targets_resolve(self, target_map): + # A renamed method or op would not fail anywhere at runtime on its own: the unit would + # simply keep nothing and the region would stay recomputed, silently costing the memory the + # user asked to keep. Resolve every name a model declares so a rename fails here instead. + for unit, target in target_map.items(): + if isinstance(target, KeptOps): + # A build registers only one flash-attention version, so *some* name must resolve + # rather than all of them. + assert resolve_kept_ops(target.names), f"{unit} names no op that resolves: {target.names}" + else: + for name in target.names: + assert pydoc.locate(name) is not None, f"{unit} names {name}, which does not resolve" + + def test_no_unit_names_the_method_that_holds_most_compilation(self): + # `_pre_moe_forward` exists to give `torch.compile` the ops on either side of attention. + # Naming it would withdraw most of what an MoE layer compiles, which is the cost this + # design exists to avoid -- attention is kept by op identity and the gate by its own + # callable instead. + for unit, target in MOE_RECOMPUTE_CFG.items(): + if isinstance(target, KeptCallables): + assert _PRE_MOE_FORWARD not in target.names, f"{unit} withdraws {_PRE_MOE_FORWARD}" + + +_PRE_MOE_FORWARD = "xtuner.v1.module.decoder_layer.moe_decoder_layer.MoEDecoderLayer._pre_moe_forward" + + +class TestUnitCostIsProportionate: + """What a unit costs follows from how it is resolved, and nothing costs a withdrawn method. + + An op-identity unit changes nothing about compilation. A callable unit is excluded from the + compiled set, which a compiled caller sees as a graph break -- so the callers are relaxed to + `fullgraph=False`, but they stay compiled. No unit withdraws `_pre_moe_forward`, where most of + an MoE layer's compilation lives. + """ + + @staticmethod + def _compile_cfg(**overrides) -> dict[str, TorchCompileOption]: + # The tiny config disables compilation; this asks what would be compiled if it did not. + config = _build_tiny_moe_config(**overrides).model_copy(update={"compile_cfg": None}) + with torch.device("meta"): + return MoE(config=config).compile_cfg + + def test_an_op_identity_unit_costs_no_compilation(self): + assert self._compile_cfg(recompute_cfg=RecomputeConfig(save=[SaveUnit.ATTN])) == self._compile_cfg() + + def test_a_callable_unit_keeps_its_callers_compiled(self): + unit = SaveUnit.MOE_GATE + # Relaxed, not removed: the caller still compiles, it just splits at the excluded callee. + # Being absent from `compile_cfg` is not enough to be outside the compiled set, because + # Dynamo inlines a callee into whichever compiled caller reaches it. + baseline = self._compile_cfg() + relaxed = self._compile_cfg(recompute_cfg=RecomputeConfig(save=[unit])) + + assert set(relaxed) == set(baseline), "a unit must not remove a caller from the compiled set" + assert not any(option.get("fullgraph") for option in relaxed.values()) + + def test_no_unit_withdraws_the_method_that_holds_most_compilation(self): + for unit in MOE_RECOMPUTE_CFG: + assert _PRE_MOE_FORWARD in self._compile_cfg(recompute_cfg=RecomputeConfig(save=[unit])), f"{unit} withdrew {_PRE_MOE_FORWARD}" + + def test_attention_is_kept_by_op_identity(self): + # The attention kernel is a custom op, so it reaches the policy from inside a compiled + # region -- which is why this unit costs nothing. + config = _build_tiny_moe_config(recompute_cfg=RecomputeConfig(save=[SaveUnit.ATTN])) + with torch.device("meta"): + model = MoE(config=config) + assert model.kept_ops + assert model.keeps_any_recompute_unit + + +class TestRecomputeRatioMigration: + """`fsdp_cfg.recompute_ratio` 已迁到 `recompute_cfg.ratio`，旧字段留作过渡。""" + + def _model(self, fsdp_kwargs, **config_kwargs): + model = _ProbeModel(_ProbeConfig(text_config=_NestedProbeConfig(), **config_kwargs)) + model._migrate_recompute_ratio(FSDPConfig(**fsdp_kwargs)) + return model + + def test_old_field_still_takes_effect(self): + # 过渡期内旧配置不能突然失效，值要被搬到新位置。 + model = self._model({"recompute_ratio": 0.25}) + + assert model.config.recompute_cfg.ratio == 0.25 + + def test_unset_old_field_leaves_the_new_one_alone(self): + # 旧字段的哨兵是 None 而不是 1.0，否则「没设过」会被当成「显式设成 1.0」， + # 把用户在新位置写的值覆盖掉。 + model = self._model({}, recompute_cfg=RecomputeConfig(ratio=0.5)) + + assert model.config.recompute_cfg.ratio == 0.5 + + def test_setting_both_is_an_error(self): + # 同一个设置有两个写法时，静默挑一个是最糟的处理。 + with pytest.raises(ValueError, match="both set"): + self._model({"recompute_ratio": 0.25}, recompute_cfg=RecomputeConfig(ratio=0.5)) diff --git a/tests/model/test_selective_checkpointing.py b/tests/model/test_selective_checkpointing.py index b1738f9366..2ac6be43da 100644 --- a/tests/model/test_selective_checkpointing.py +++ b/tests/model/test_selective_checkpointing.py @@ -6,6 +6,8 @@ TestKeptCallables test_kept_callable_reproduces_full_recompute: 按 callable 留驻与全重算逐位相同。 test_unit_marker_outside_a_region_is_noop: 未被包装时不产生任何可观察行为。 +TestContractLayering + test_module_layer_imports_the_contract_without_the_model_layer: 契约模块必须能被 module/ 层单独导入。 TestUnsupportedUnits test_in_place_op_in_kept_unit_is_recomputed_not_refused: 只动自己缓冲区的 in-place 写不该被拒。 test_mutating_a_kept_tensor_is_caught_by_torch: 写到被留驻张量上时由 torch 精确拦下。 @@ -19,6 +21,8 @@ test_kept_unit_matches_full_recompute_under_compile: compile 下同上，且不触发 cached-tensor-mutated。 """ +import subprocess +import sys from collections import Counter import pytest @@ -27,7 +31,8 @@ from torch.utils._python_dispatch import TorchDispatchMode from xtuner.v1.model.utils import ( - RecomputeUnit, + KeptOps, + SaveUnit, apply_selective_checkpointing, in_recompute_unit, resolve_kept_ops, @@ -52,7 +57,7 @@ class _UnitBlock(_Block): def __init__(self) -> None: super().__init__() - self.second_stage = in_recompute_unit(RecomputeUnit.SAVE_ATTN, self._second_stage) + self.second_stage = in_recompute_unit(SaveUnit.ATTN, self._second_stage) def _second_stage(self, hidden: torch.Tensor) -> torch.Tensor: return torch.tanh(self.second(hidden)) @@ -67,7 +72,7 @@ class _InPlaceUnitBlock(nn.Module): def __init__(self) -> None: super().__init__() self.linear = nn.Linear(4, 4) - self.unit = in_recompute_unit(RecomputeUnit.SAVE_ATTN, self._unit) + self.unit = in_recompute_unit(SaveUnit.ATTN, self._unit) def _unit(self, x: torch.Tensor) -> torch.Tensor: hidden = self.linear(x) @@ -85,7 +90,7 @@ class _MutatesKeptTensorBlock(nn.Module): def __init__(self) -> None: super().__init__() self.linear = nn.Linear(4, 4) - self.unit = in_recompute_unit(RecomputeUnit.SAVE_ATTN, self._unit) + self.unit = in_recompute_unit(SaveUnit.ATTN, self._unit) def _unit(self, x: torch.Tensor) -> torch.Tensor: kept = torch.tanh(self.linear(x)) @@ -186,6 +191,19 @@ def test_unit_marker_outside_a_region_is_noop(self): assert inputs.grad is not None +class TestContractLayering: + def test_module_layer_imports_the_contract_without_the_model_layer(self): + # 契约之所以在 xtuner/v1/utils 而不是挨着 engine，是因为它命名的 callable 在 + # xtuner/v1/module 里：一旦契约里出现指向 model/ 的 import，这条独立导入就会变成 + # 循环导入而失败。必须用干净的解释器：同进程里 xtuner.v1.model 早就被导入了。 + result = subprocess.run( + [sys.executable, "-c", "import xtuner.v1.module.decoder_layer.moe_decoder_layer"], + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stderr + + class TestUnsupportedUnits: def test_in_place_op_in_kept_unit_is_recomputed_not_refused(self): # in-place op 本身不危险，危险的是它写到了被 unit 留驻的张量上——那由 torch 的 @@ -202,6 +220,19 @@ def test_in_place_op_outside_kept_unit_is_fine(self): _run(_InPlaceOutsideBlock(), keeps_any_unit=True) +class TestDeclarations: + def test_kept_ops_and_callables_are_distinct_targets(self): + # 两种解析方式的代价完全不同（一个零编译代价，一个要退出编译集合），所以类型必须能区分。 + from xtuner.v1.model.moe.moe import MOE_RECOMPUTE_CFG + + assert isinstance(MOE_RECOMPUTE_CFG[SaveUnit.ATTN], KeptOps) + assert set(MOE_RECOMPUTE_CFG) == { + SaveUnit.ATTN, + SaveUnit.MOE_GATE, + SaveUnit.MOE_DISPATCH, + } + + class TestRecomputeIsObservable: """按 reviewer 的三个目标直接观察：精度对齐、重算生效、重算+SAC 生效。""" diff --git a/tools/verify_sac.sh b/tools/verify_sac.sh new file mode 100755 index 0000000000..6a7de75f65 --- /dev/null +++ b/tools/verify_sac.sh @@ -0,0 +1,57 @@ +#!/bin/bash +# 验证 recompute_cfg：跑两次同一份训练，一次不选 unit、一次选 save_attn，比对结果。 +# +# 判据（两条都必须成立）： +# 1. step-1 loss 逐位相同 —— checkpoint 只改 backward，前向不能被动到； +# 2. 峰值显存显著上升 —— unit 真的把激活留驻了，而不是配置被静默忽略。 +# +# 用 eager（torch_compile=False）跑：编译会改变 kernel，从而改变 loss 的低位， +# 那样就无法把差异归因到 SAC 本身。 +set -u +REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +OUT="${SAC_VERIFY_OUT:-/tmp/sac_verify}" +PY="${PYTHON:-/mnt/shared-storage-user/yehaochen/miniconda3/envs/py312-pt210/bin/torchrun}" +export QWEN3_MOE_PATH="${QWEN3_MOE_PATH:-/mnt/shared-storage-user/llmrazor-share/model/Qwen3-30B-A3B}" +export ALPACA_PATH="${ALPACA_PATH:-/mnt/shared-storage-user/llmrazor-share/data/alpaca}" +export EP_DISABLE_GIN=1 EP_REUSE_NCCL_COMM=0 +mkdir -p "$OUT" + +for unit in none attn; do + echo "==> 跑 recompute_cfg=$unit ..." + (cd /tmp && PYTHONPATH="$REPO" timeout 3600 "$PY" --master-port "${MASTER_PORT:-21971}" \ + --nproc-per-node "${NPROC:-8}" -m xtuner.v1.train.cli.sft \ + --config "$REPO/ci/config/sac_verify_$unit.py") > "$OUT/$unit.log" 2>&1 + [ $? -eq 0 ] || { echo "失败，日志见 $OUT/$unit.log"; grep -m3 -iE "error|Traceback" "$OUT/$unit.log"; exit 1; } +done + +"${PYTHON_BIN:-python}" - "$OUT" <<'PY' +import re, sys, pathlib + +out = pathlib.Path(sys.argv[1]) +res = {} +for unit in ("none", "attn"): + log = out / f"{unit}.log" + text = log.read_text(errors="ignore") + # 只看 rank 0：各 rank 峰值本就不同，跨 rank 比较得到的是 rank 差异而非配置差异。 + rows = re.findall(r"RANK 0\].*?Step (\d+)/\d+.*?reduced_llm_loss: ([0-9.]+).*?max_memory: ([0-9.]+) GB", text) + by_step = {int(s): (loss, float(mem)) for s, loss, mem in rows} + if 1 not in by_step or len(by_step) < 2: + print("没读到 %s 至少两步的 rank-0 指标，日志：%s" % (unit, log)) + sys.exit(1) + # loss 取第一步（纯前向）；显存取最后一步：第一步的峰值被权重加载/优化器状态初始化这类 + # 一次性分配盖住，激活的稳态峰值要到第二步才显出来。 + res[unit] = (by_step[1][0], by_step[max(by_step)][1]) + +l0, m0 = res["none"] +l1, m1 = res["attn"] +print("") +print(" 无 unit step1_loss=%s 稳态峰值=%.2f GB" % (l0, m0)) +print(" save_attn step1_loss=%s 稳态峰值=%.2f GB (+%.2f GB)" % (l1, m1, m1 - m0)) +print("") + +ok_loss = l0 == l1 +ok_mem = (m1 - m0) > 1.0 +print(" [%s] 前向未受影响：step-1 loss 逐位相同" % ("PASS" if ok_loss else "FAIL")) +print(" [%s] unit 确实生效：稳态峰值显存上升 > 1 GB" % ("PASS" if ok_mem else "FAIL")) +sys.exit(0 if (ok_loss and ok_mem) else 1) +PY diff --git a/xtuner/v1/config/fsdp.py b/xtuner/v1/config/fsdp.py index 7335d6d7a5..0d0b5f070b 100644 --- a/xtuner/v1/config/fsdp.py +++ b/xtuner/v1/config/fsdp.py @@ -15,8 +15,15 @@ class FSDPConfig(BaseModel): tp_size: Annotated[int, Parameter(help="Tensor parallel size")] = 1 ep_size: Annotated[int, Parameter(help="Expert parallel size")] = 1 reshard_after_forward: Annotated[bool, Parameter(help="Reshard model parameters after forward pass")] = True - recompute_ratio: Annotated[float, Parameter(help="Gradient checkpointing ratio for memory optimization")] = 1.0 - vision_recompute_ratio: Annotated[float, Parameter(help="Recompute ratio for vision modules")] = 1.0 + # Deprecated: moved to `model_cfg.recompute_cfg.ratio` / `.vision_ratio`, so that "which layers + # are recomputed" and "what stays saved inside them" live together -- neither means much without + # the other. `None` marks "not set by the user", which is what lets the migration tell a real + # request from a default. Remove both fields, and `BaseModel._migrate_recompute_ratio`, once + # downstream configs have moved. + recompute_ratio: Annotated[float | None, Parameter(help="Deprecated, use `model_cfg.recompute_cfg.ratio`")] = None + vision_recompute_ratio: Annotated[ + float | None, Parameter(help="Deprecated, use `model_cfg.recompute_cfg.vision_ratio`") + ] = None checkpoint_preserve_rng_state: Annotated[bool, Parameter(help="Preserve RNG state during checkpointing")] = True # Training-time FSDP CPU offload is version-sensitive for XTuner model configs # that keep selected fp32 trainable parameters outside FSDP via diff --git a/xtuner/v1/model/base.py b/xtuner/v1/model/base.py index 697e5f9933..655f140af9 100644 --- a/xtuner/v1/model/base.py +++ b/xtuner/v1/model/base.py @@ -19,7 +19,6 @@ import torch.nn as nn import torch.nn.functional as F from cyclopts import Parameter -from more_itertools import consume from pydantic import BaseModel as PydanticBaseModel from pydantic import ConfigDict, Field, computed_field, model_validator from pydantic.fields import FieldInfo @@ -62,7 +61,16 @@ set_async_save_process_qos, ) -from .utils import ModelForwardExtraLogInfo +from .utils import ( + KeptCallables, + KeptOps, + ModelForwardExtraLogInfo, + RecomputeConfig, + RecomputeTargetMap, + SaveUnit, + in_recompute_unit, + resolve_kept_ops, +) logger = get_logger() @@ -143,6 +151,14 @@ class XTunerBaseModelConfig(PydanticBaseModel): "`dict[str, TorchCompileOption]`: Customize the compile option", ), ] = None + # `activation_offload_cfg` belongs here, next to `recompute_cfg` and sharing its `SaveUnit` + # vocabulary and per-model interval declarations: offloading applies to the regions SAC keeps + # resident, since recomputed regions never land in memory to begin with. Not declared until it is + # implemented -- a config field nothing reads is a switch that silently does nothing. + recompute_cfg: Annotated[ + RecomputeConfig, + Parameter(group="model", help="Which layers are recomputed, and what stays saved inside them"), + ] = RecomputeConfig() hf_key_mapping: Annotated[dict[str, str] | None, "Remapping hf key based on the `to_hf_key_list`"] = None dcp_ignore_frozen_params: bool = True lm_loss_cfg: BaseLossConfig = CELossConfig() @@ -538,6 +554,42 @@ def _save_file( save_file(tensors, filename, metadata=metadata) +def _disable_nested_switch(obj: Any, field_name: str, *, subfield: str | None = None) -> None: + """Turn a tri-state feature switch off on every config reachable from + ``obj``. + + Sub-model configs carry their own copy of switches such as ``compile_cfg`` and ``recompute_cfg``, and each + sub-model resolves its own. Turning a feature off on the outer config therefore only means something if the + ``False`` reaches them, which is what this walk does. + + Traversal stops at configs that do not declare ``field_name``: a config that does not take part in the feature + cannot hide a participant behind it, and stopping keeps unrelated sub-configs out of the walk. + + Args: + obj (Any): Config, container, or leaf value to walk. + field_name (str): Name of the field holding the switch. + subfield (str | None): When the switch is one field *inside* ``field_name`` -- as + ``recompute_cfg.save`` is inside ``RecomputeConfig`` -- the name of that inner field. + Defaults to None, i.e. ``field_name`` is itself the switch. + """ + if isinstance(obj, PydanticBaseModel): + if not hasattr(obj, field_name): + return + if subfield is None: + setattr(obj, field_name, False) + else: + setattr(getattr(obj, field_name), subfield, False) + for nested_field in type(obj).model_fields: + _disable_nested_switch(getattr(obj, nested_field), field_name, subfield=subfield) + elif isinstance(obj, Mapping): + for value in obj.values(): + _disable_nested_switch(value, field_name, subfield=subfield) + # str & bytes are Iterables of themselves and would recurse forever. + elif isinstance(obj, Iterable) and not isinstance(obj, (str, bytes)): + for value in obj: + _disable_nested_switch(value, field_name, subfield=subfield) + + class BaseModel(nn.Module): load_spec_mapping: dict[str, LoadSpec] = {} fsdp_mesh: DeviceMesh | None = None @@ -556,6 +608,10 @@ def __init__(self, config: XTunerBaseModelConfig): self._pending_async_hf: AsyncHFSaveHandle | None = None self._async_hf_resources: AsyncHFResources | None = None + # Recompute resolves first: selecting a unit withdraws the callables enclosing its region + # from the compiled set, so the compile config depends on the answer. + self._selected_recompute_units: set[SaveUnit] = set() + self._kept_ops, self._kept_callables = self._resolve_recompute_cfg(self.config) self._compile_cfg = self._resolve_compile_cfg(self.config) self._float8_handler: Float8Handler | None = None @@ -592,6 +648,7 @@ def fully_shard( ) -> Self: """Fully shard the model parameters.""" self.fsdp_config = fsdp_config + self._migrate_recompute_ratio(fsdp_config) self.fsdp_mesh = self._init_world_mesh() self._world_mesh = self.fsdp_mesh @@ -984,6 +1041,22 @@ def compile_cfg(self) -> dict[str, TorchCompileOption]: return _compile_cfg + @property + def default_recompute_cfg(self) -> RecomputeTargetMap: + """What each recompute unit this architecture supports resolves to. + + This is the model author's vocabulary: it declares which :class:`SaveUnit` s the architecture supports and + what each one is made of, not which of them are worth enabling. An architecture that supports none declares + nothing. + + Like ``default_compile_cfg``, an override must be answerable from ``self.config`` alone: it is read while + ``BaseModel.__init__`` resolves the user's selection, which is before the subclass has built its layers. + + Returns: + RecomputeTargetMap: Supported units mapped to what implements them for this architecture. + """ + return {} + @property def kept_ops(self) -> frozenset: """Op overloads kept resident wherever they run, as selected by @@ -992,20 +1065,16 @@ def kept_ops(self) -> frozenset: Returns: frozenset: Overloads to keep. Empty means no op-identity unit was selected. """ - return frozenset() + return self._kept_ops @property def keeps_any_recompute_unit(self) -> bool: """Whether any unit at all was selected. - This is what tells the sharding paths whether to install the per-op policy at all: with - nothing selected, torch's own default already recomputes everything, so a policy that - reaches the same answer would only put a dispatch mode in the way of every op. - Returns: bool: True when ``config.recompute_cfg`` selected something, by either resolution. """ - return False + return bool(self._selected_recompute_units) @property def float8_handler(self): @@ -2571,14 +2640,14 @@ def _resolve_compile_cfg( custom_cfg = config.compile_cfg if custom_cfg is False: - self._disable_compile_cfg(self.config) + _disable_nested_switch(self.config, "compile_cfg") return {} # torch.compile is not supported on NPU if DEVICE == "npu": if custom_cfg is not False: log_rank0.warning("torch.compile is not supported on NPU, disabling torch.compile.") - self._disable_compile_cfg(self.config) + _disable_nested_switch(self.config, "compile_cfg") return {} if custom_cfg is True or custom_cfg is None: @@ -2586,27 +2655,172 @@ def _resolve_compile_cfg( else: compile_cfg = custom_cfg - return compile_cfg - - def _disable_compile_cfg(self, obj): - if isinstance(obj, PydanticBaseModel) and hasattr(obj, "compile_cfg"): - obj.compile_cfg = False - consume(self._disable_compile_cfg(getattr(obj, x)) for x in obj.__class__.model_fields) - elif isinstance(obj, Mapping): - consume(map(self._disable_compile_cfg, obj.values())) - # str&bytes are special Iterable, need to exclude it, otherwise it will infinite loop - elif isinstance(obj, Iterable) and not isinstance(obj, (str, bytes)): - consume(map(self._disable_compile_cfg, obj)) - else: - return + return self._without_compiled_selected_regions(compile_cfg) + + def _migrate_recompute_ratio(self, fsdp_config: FSDPConfig) -> None: + # Compatibility shim for the two ratios that used to live on `FSDPConfig`. Remove together + # with those fields once downstream configs have moved to `recompute_cfg`. + moves = ( + ("recompute_ratio", "ratio", "recompute_cfg.ratio"), + ("vision_recompute_ratio", "vision_ratio", "recompute_cfg.vision_ratio"), + ) + for old_name, new_name, new_path in moves: + old_value = getattr(fsdp_config, old_name, None) + if old_value is None: + continue + if new_name in self.config.recompute_cfg.model_fields_set: + raise ValueError( + f"`fsdp_cfg.{old_name}` and `model_cfg.{new_path}` are both set. They are the same setting; " + f"keep only `model_cfg.{new_path}`." + ) + log_rank0.warning( + f"`fsdp_cfg.{old_name}` is deprecated and will be removed; use `model_cfg.{new_path}` instead." + ) + setattr(self.config.recompute_cfg, new_name, old_value) + + def _resolve_recompute_cfg(self, config: XTunerBaseModelConfig) -> tuple[frozenset, set[str]]: + """Turn the user's ``recompute_cfg.save`` selection into what the + policy needs. + + Args: + config (XTunerBaseModelConfig): The model config carrying the selection. + + Returns: + tuple[frozenset, set[str]]: The op overloads to keep by identity, and the qualified + names of the callables to keep by marker. + """ + selected = config.recompute_cfg.save + + # `False` has to reach the nested sub-model configs, which a compose model builds after this + # returns and which resolve their own switch against them. + if selected is False: + _disable_nested_switch(self.config, "recompute_cfg", subfield="save") + return frozenset(), set() + + # `None` means "keep the memory profile of plain full recompute", not "use the model default" as it does + # for `compile_cfg`. `default_recompute_cfg` is the vocabulary of what an architecture *can* keep resident, + # not a recommendation: keeping a unit trades memory for speed, and retaining everything can even exceed + # the peak of not checkpointing at all, because SAC storage duplicates what autograd already holds. Only the + # user knows which side of that trade they want, so an unset config changes nothing. + if selected is None: + return frozenset(), set() + + supported = self.default_recompute_cfg + + # `True` asks for whatever the model offers, so an empty vocabulary is not a user error -- + # but it does mean the request is a no-op, which is worth saying out loud rather than + # letting the run look configured when it is not. + if selected is True and not supported: + log_rank0.warning( + f"`recompute_cfg.save=True` has no effect: {type(self).__name__} declares no recompute units, so every " + "region is recomputed." + ) + units = list(supported) if selected is True else selected + + op_names: list[str] = [] + callables: set[str] = set() + for unit in units: + if unit not in supported: + supported_desc = ", ".join(sorted(supported)) if supported else "none" + raise ValueError( + f"`recompute_cfg.save` selects {unit!r}, which {type(self).__name__} does not support. " + f"Units supported by this model: {supported_desc}. Note that a compose model declares no units " + "of its own -- set `recompute_cfg.save` on the sub-model config that owns the layers instead." + ) + target = supported[unit] + if isinstance(target, KeptOps): + op_names.extend(target.names) + else: + callables.update(target.names) + self._selected_recompute_units.add(unit) + return resolve_kept_ops(tuple(op_names)), callables + + def _without_compiled_selected_regions( + self, compile_cfg: dict[str, TorchCompileOption] + ) -> dict[str, TorchCompileOption]: + """Let the compiled callables tolerate the gap a selected unit leaves + in them. + + A unit resolved by callable is excluded from compilation, and a compiled caller that reaches it sees that as + a graph break -- which ``fullgraph=True`` reports as an error rather than a split. Which callers those are is + not knowable from the config: Dynamo inlines across call sites, so the break can surface in any compiled + method upstream of the unit. Relaxing all of them is the only rule that is correct without tracing. + + Units resolved by op identity cost nothing here; they never leave a gap. + """ + if not self._kept_callables: + return compile_cfg + + log_rank0.info( + f"Keeping {sorted(self._selected_recompute_units)} resident excludes {sorted(self._kept_callables)} " + f"from torch.compile, so the callers that reach them are compiled with `fullgraph=False`." + ) + return { + target: TorchCompileOption(**{**option, "fullgraph": False}) + for target, option in compile_cfg.items() + if target not in self._kept_callables + } def _maybe_enable_compile(self, compile_cfg: dict[str, TorchCompileOption]): if compile_cfg: torch._dynamo.config.cache_size_limit = 256 + # Before compiling anything: install the unit wrappers, each excluded from compilation. + # Being absent from `compile_cfg` is not enough to be outside the compiled set -- Dynamo + # inlines a callee into whichever compiled caller reaches it, and the wrapper's ContextVar + # is then a hard error rather than a graph break. On a callable no compiled caller reaches, + # the exclusion costs nothing. + self._install_recompute_units() + for target, option in compile_cfg.items(): self._compile_overwrite(target, option) + def _install_recompute_units(self) -> None: + """Wrap each selected unit's callables so the policy can recognise + their ops.""" + for unit in sorted(self._selected_recompute_units): + target = self.default_recompute_cfg.get(unit) + if not isinstance(target, KeptCallables): + continue + for name in sorted(target.names): + self._recompute_unit_overwrite(unit, name) + + def _recompute_unit_overwrite(self, unit: SaveUnit, func_name: str) -> None: + """Install the unit wrapper on one callable, and keep it out of the + compiled set. + + Args: + unit (SaveUnit): The unit this callable implements. + func_name (str): Qualified name of the callable, as written in ``compile_cfg``. + """ + function = cast(FunctionType | MaybeCompile, pydoc.locate(func_name)) + if function is None: + raise AttributeError(f"Recompute config error! Cannot locate the callable: {func_name}") + + def install(fn: Any) -> Any: + return torch._dynamo.disable(in_recompute_unit(unit, fn)) + + if isinstance(function, MaybeCompile): + function.disable_compile() + function.func = install(function.func) + return + + function = cast(FunctionType, function) + if get_function_type(function) is not FunctionEnum.CLASS_LEVEL_FUNCTION: + raise ValueError( + f"Recompute config error! {func_name} must be a method or a `@maybe_compile` function to implement " + f"a recompute unit." + ) + qualname_split = function.__qualname__.split(".") + assert len(qualname_split) == 2, ( + f"XTuner Internal Error! the name of {function} should be recognized as " + f"., but got {qualname_split}" + ) + class_name, method_name = qualname_split + cls = getattr(import_module(function.__module__), class_name) + setattr(cls, method_name, install(function)) + logger.debug(f"{func_name} now implements {unit} and is excluded from compilation.") + def _mark_dynamic(self, seq_ctx: SequenceContext, dim=0): """`cu_seq_lens_q` and `cu_seq_lens_k` are dynamic shapes in each fwd/bwd pass. diff --git a/xtuner/v1/model/compose/intern_s1/modeling_vision.py b/xtuner/v1/model/compose/intern_s1/modeling_vision.py index b0407e38b4..0d4a0e486d 100644 --- a/xtuner/v1/model/compose/intern_s1/modeling_vision.py +++ b/xtuner/v1/model/compose/intern_s1/modeling_vision.py @@ -396,7 +396,7 @@ def fully_shard( for param in self.parameters(): param.requires_grad = False - recompute_ratio = fsdp_config.vision_recompute_ratio + recompute_ratio = self.config.recompute_cfg.vision_ratio num_recompute_layers = int(len(self.encoder.layer) * recompute_ratio) generator = torch.Generator() @@ -412,6 +412,7 @@ def fully_shard( self.kept_ops, keeps_any_unit=self.keeps_any_recompute_unit, preserve_rng_state=checkpoint_preserve_rng_state, + # The layer's own forward is compiled just below, making it one opaque region. ) if self.config.drop_path_rate == 0.0 and self.compile_cfg: # Compile the class function, then restore descriptor binding on this instance. diff --git a/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py b/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py index f2758b0d5e..85e8a5d49c 100644 --- a/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py +++ b/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py @@ -333,7 +333,7 @@ def fully_shard( param.requires_grad = False checkpoint_preserve_rng_state = fsdp_config.checkpoint_preserve_rng_state - num_recompute_layers = int(len(self.blocks) * fsdp_config.vision_recompute_ratio) + num_recompute_layers = int(len(self.blocks) * self.config.recompute_cfg.vision_ratio) for layer_idx in tqdm(list(range(len(self.blocks))), desc="[Vision Fully Shard]"): layer = self.blocks[layer_idx] @@ -343,6 +343,7 @@ def fully_shard( self.kept_ops, keeps_any_unit=self.keeps_any_recompute_unit, preserve_rng_state=checkpoint_preserve_rng_state, + # The layer's own forward is compiled just below, making it one opaque region. ) if self.compile_cfg: # Compile the class function, then restore descriptor binding on this instance. diff --git a/xtuner/v1/model/dense/dense.py b/xtuner/v1/model/dense/dense.py index 274aebd00b..bf4457d852 100644 --- a/xtuner/v1/model/dense/dense.py +++ b/xtuner/v1/model/dense/dense.py @@ -26,7 +26,7 @@ TorchCompileOption, TransformerConfig, ) -from xtuner.v1.model.utils import apply_selective_checkpointing +from xtuner.v1.model.utils import KeptOps, RecomputeTargetMap, SaveUnit, apply_selective_checkpointing from xtuner.v1.module import ( GatedDeltaNetConfig, LMHead, @@ -51,6 +51,15 @@ **DEFAULT_FLOAT8_CFG, } +DENSE_RECOMPUTE_CFG: RecomputeTargetMap = { + # A dense stack has no router and no expert dispatch, so attention is the only unit it can keep + # -- and it keeps it by op identity, which costs no compilation. + SaveUnit.ATTN: KeptOps( + "flash_attn::_flash_attn_varlen_forward_v3", + "flash_attn::_flash_attn_varlen_forward_v2", + ), +} + class Dense(BaseModel): config: TransformerConfig @@ -164,6 +173,26 @@ def build_layers(self, config: TransformerConfig) -> nn.ModuleDict: def default_compile_cfg(self) -> dict[str, TorchCompileOption]: return DENSE_COMPILE_CFG + @property + @override + def default_recompute_cfg(self) -> RecomputeTargetMap: + """Marker intervals this architecture can keep resident, keyed by + semantic unit. + + Both units are **eager-only**: ``DenseDecoderLayer.forward`` is compiled as one fullgraph region, so the + markers that delimit them are folded away and every region falls back to being recomputed. They are declared + because eager training addresses them normally, and because the regions become effective as soon as the layer + is compiled at a finer granularity. + + Returns: + RecomputeTargetMap: Supported units mapped to the marker intervals that implement them. + """ + # Linear-attention layers carry in-place convolution state across the forward, so replaying them under a + # selective checkpoint is untested. Hybrid models declare no units until it is. + if "linear_attention" in self.config.layers_type: + return {} + return DENSE_RECOMPUTE_CFG + # NOTE: Add this overload for inferring the return type for easier type checking and using @overload # type: ignore def __call__( # type: ignore @@ -225,7 +254,7 @@ def fully_shard( lm_head_mp_policy = MixedPrecisionPolicy(param_dtype=torch.float32, reduce_dtype=torch.float32) else: lm_head_mp_policy = mp_policy - num_recompute_layers = int(self.config.num_hidden_layers * self.fsdp_config.recompute_ratio) + num_recompute_layers = int(self.config.num_hidden_layers * self.config.recompute_cfg.ratio) generator = torch.Generator() generator.manual_seed(dist.get_rank()) @@ -240,6 +269,7 @@ def fully_shard( self.kept_ops, keeps_any_unit=self.keeps_any_recompute_unit, preserve_rng_state=checkpoint_preserve_rng_state, + # The layer's own forward is compiled just below, making it one opaque region. ) # Linear-attention (GatedDeltaNet) layers write ``seq_ctx.seq_idx`` inside the diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py index c90033870e..e11937b4d7 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/xtuner/v1/model/moe/moe.py @@ -45,7 +45,11 @@ TransformerConfig, ) from xtuner.v1.model.utils import ( + KeptCallables, + KeptOps, ModelForwardExtraLogInfo, + RecomputeTargetMap, + SaveUnit, apply_selective_checkpointing, module_dict_repr, ) @@ -114,6 +118,44 @@ MOE_EP_COMPILE_CFG = MOE_NON_EP_COMPILE_CFG.copy() MOE_EP_COMPILE_CFG.pop(MOE_DECODER_LAYER_FORWARD) +_MOE_GATE = "xtuner.v1.module.decoder_layer.moe_decoder_layer.MoEGate" +_DISPATCHERS = ( + "xtuner.v1.module.dispatcher.torch_all2all.TorchAll2AllDispatcher", + "xtuner.v1.module.dispatcher.deepep.DeepEPDispatcher", + "xtuner.v1.module.dispatcher.agrs.MoEAGRSDispatcher", + "xtuner.v1.module.dispatcher.base.NaiveDispatcher", +) +_DISPATCH_STAGES = ( + "dispatch_preprocess", + "dispatch", + "dispatch_postprocess", + "combine_preprocess", + "combine", + "combine_postprocess", +) + +MOE_RECOMPUTE_CFG: RecomputeTargetMap = { + # The attention kernel identifies itself, so this unit needs nothing taken out of the compiled + # set: a custom op is never fused, so it reaches the checkpoint policy compiled or not. Both + # flash-attention spellings are listed because only one is registered in a given build. + SaveUnit.ATTN: KeptOps( + "flash_attn::_flash_attn_varlen_forward_v3", + "flash_attn::_flash_attn_varlen_forward_v2", + ), + # The router's ops are ordinary `addmm`/`topk` that appear all over the layer, so it has to be + # named by callable -- and the callable is the projection alone, not `MoEGate.forward`. The + # projection holds all the activation memory the router has (the logits are + # `tokens x n_routed_experts` while the router's own tensors are top-k sized), and unlike the + # router it contains no in-place write, which the policy refuses to keep. + SaveUnit.MOE_GATE: KeptCallables(f"{_MOE_GATE}.project"), + # The dispatcher stages are already the finest callables there are, and none of them is + # compiled, so this unit costs no compilation either. Its stages call into the backend library, + # so ops the library performs on its own buffers -- deep_ep swaps storage with `aten.set_` -- + # fall inside the unit and are reported once; they are recomputed rather than kept, and torch + # catches the case where such a write lands on a tensor the unit did keep. + SaveUnit.MOE_DISPATCH: KeptCallables(*(f"{cls}.{stage}" for cls in _DISPATCHERS for stage in _DISPATCH_STAGES)), +} + class MoEModelOutputs(ModelOutputs): router_logits: dict[str, torch.Tensor] | None = None @@ -585,8 +627,9 @@ def _micro_batch_forward( d2h_stream=self.offload_stream, block_idx=layer_idx - self.config.first_k_dense_replace, group="text", - custom_check_fn=lambda x: x.data_ptr() - in [hidden_states.data_ptr() for hidden_states in hidden_states_list], + custom_check_fn=lambda x: ( + x.data_ptr() in [hidden_states.data_ptr() for hidden_states in hidden_states_list] + ), prefetch=True, reserve_pin_memory=True, ): @@ -1243,8 +1286,8 @@ def fully_shard( if self._should_recompute(None, mtp_idx=mtp_idx) or ( self.config.mtp_config is not None and self.config.mtp_config.share_weights ): # share mtp head must recompute - # Marker intervals are declared against the decoder layer, so an MTP layer - # gets the same mechanism with nothing kept: recomputed whole, as before. + # Recompute units are declared against the decoder layer, so an MTP layer gets + # the same mechanism with nothing kept: recomputed whole, as before. mtp_layer = apply_selective_checkpointing(mtp_layer) self.mtp_block.layers[mtp_idx] = mtp_layer @@ -1294,6 +1337,31 @@ def default_compile_cfg(self) -> dict[str, TorchCompileOption]: else: return MOE_NON_EP_COMPILE_CFG + @property + @override + def default_recompute_cfg(self) -> RecomputeTargetMap: + """What each supported unit resolves to for this architecture. + + Both units are fine-grained: neither names ``_pre_moe_forward``, the method that holds attention, the + layernorms and the gate together so that ``torch.compile`` can cover the ops on either side of attention. + Withdrawing that would cost most of what an MoE layer compiles. + + - ``SAVE_ATTN`` is resolved by op identity and costs no compilation at all. + - ``SAVE_MOE_GATE`` names ``MoEGate.project``, which holds the router's activation memory. Its callers are + compiled with ``fullgraph=False`` so they can split at it. + + ``SAVE_MOE_DISPATCH`` is not declared here; see the note next to :data:`MOE_RECOMPUTE_CFG`. + + Returns: + RecomputeTargetMap: Supported units mapped to the marker intervals that implement them. + """ + # Linear-attention layers carry in-place convolution state across the forward, so replaying them under a + # selective checkpoint is untested. Hybrid models declare no units until it is. + if "linear_attention" in self.config.layers_type: + return {} + + return MOE_RECOMPUTE_CFG + @property def need_update_bias(self) -> bool: router_config = self.config.router @@ -1456,13 +1524,6 @@ def patched_emb_forward(self, input): self.sparse, ) - def _compiles_whole_decoder_layer(self) -> bool: - # Without EP the decoder layer's own forward is compiled, so the layer is one opaque region - # and no marker inside it can delimit anything. With EP that entry is dropped and only the - # methods below it are compiled, which leaves the dispatcher-call boundaries in eager python - # for markers to land on. - return MOE_DECODER_LAYER_FORWARD in self.compile_cfg - def _should_recompute( self, layer_idx: int | None, @@ -1501,7 +1562,7 @@ def _should_recompute( mtp_layers = 1 if self.config.mtp_config.share_weights else self.config.mtp_config.num_layers else: mtp_layers = 0 - recompute_ratio = self.fsdp_config.recompute_ratio if self.fsdp_config is not None else 0.0 + recompute_ratio = self.config.recompute_cfg.ratio total_layers = num_layers + mtp_layers num_recompute_layers = int(total_layers * recompute_ratio) diff --git a/xtuner/v1/model/utils/__init__.py b/xtuner/v1/model/utils/__init__.py index 09cd02bff7..948e45653d 100644 --- a/xtuner/v1/model/utils/__init__.py +++ b/xtuner/v1/model/utils/__init__.py @@ -1,31 +1,26 @@ -from .checkpointing import apply_gradient_checkpointing -from .misc import ModelForwardExtraLogInfo, module_dict_repr -from .selective_checkpointing import ( +from xtuner.v1.utils.selective_checkpointing import ( KeptCallables, KeptOps, - RecomputeTarget, + RecomputeConfig, RecomputeTargetMap, - RecomputeUnit, - active_recompute_unit, - apply_selective_checkpointing, - in_recompute_unit, - recompute_unit, - resolve_kept_ops, + SaveUnit, ) +from .checkpointing import apply_gradient_checkpointing +from .misc import ModelForwardExtraLogInfo, module_dict_repr +from .selective_checkpointing import apply_selective_checkpointing, in_recompute_unit, resolve_kept_ops + __all__ = [ "apply_gradient_checkpointing", "apply_selective_checkpointing", - "in_recompute_unit", - "resolve_kept_ops", "module_dict_repr", "ModelForwardExtraLogInfo", "KeptCallables", "KeptOps", - "RecomputeTarget", "RecomputeTargetMap", - "RecomputeUnit", - "active_recompute_unit", - "recompute_unit", + "RecomputeConfig", + "SaveUnit", + "in_recompute_unit", + "resolve_kept_ops", ] diff --git a/xtuner/v1/model/utils/selective_checkpointing.py b/xtuner/v1/model/utils/selective_checkpointing.py index 8891064953..3b39a00915 100644 --- a/xtuner/v1/model/utils/selective_checkpointing.py +++ b/xtuner/v1/model/utils/selective_checkpointing.py @@ -1,181 +1,34 @@ -"""Shared contract for region-level selective activation checkpointing (SAC). - -This module holds the vocabulary that the SAC layers agree on and nothing else: - -- Users select :class:`RecomputeUnit` members in the model config. -- Model authors declare a :data:`RecomputeTargetMap` saying what each unit they support resolves to - for their architecture. -- The SAC engine turns the selection into a per-op checkpoint policy. - -It lives under ``xtuner.v1.utils`` rather than next to the models because the targets name callables -in ``xtuner.v1.module`` while the unit vocabulary is consumed by ``xtuner.v1.model`` configs; a home -inside either package would make the two import each other. - -**A unit resolves to one of two things, and which one is not a style choice.** The checkpoint policy -is asked about every op that reaches the dispatcher, and it has to answer "is this op part of a kept -unit". There are exactly two ways to know: - -- :class:`KeptOps` names the ops directly. This is the cheapest possible form -- it changes nothing - about compilation -- but it only works when the op is *specific enough to identify the unit on its - own*. Attention qualifies: `flash_attn::_flash_attn_varlen_forward` appears nowhere else. A gate's - ``addmm`` does not. -- :class:`KeptCallables` names callables whose whole body belongs to the unit. Everything they - dispatch is kept. This can express any region, at the price of taking those callables out of the - compiled set -- a marker is ordinary python state and the ops it would cover run as fused kernels, - so a region inside compiled code is invisible either way. - -Prefer :class:`KeptOps`. Reach for :class:`KeptCallables` when no op identifies the unit, and then -name the *smallest* callable that covers it: the compilation given up is the whole callable's, not -the region's. +"""The selective activation checkpointing engine. + +Drives a per-op checkpoint policy for the units a user selected. The vocabulary those units are +written in lives in :mod:`xtuner.v1.utils.selective_checkpointing`; the config resolution that turns +a user's ``recompute_cfg`` into resolved targets lives with the model configs. + +A unit reaches the policy by one of two routes, and the policy answers the same question either way +-- "does this op belong to a kept unit": + +- :class:`~xtuner.v1.utils.selective_checkpointing.KeptOps` targets are recognised from the op + itself, so they need nothing installed and work identically inside and outside compiled code. +- :class:`~xtuner.v1.utils.selective_checkpointing.KeptCallables` targets are recognised from a + ``ContextVar`` the engine's wrapper sets around the callable. That only works because the model + layer has also taken those callables out of the compiled set: a ``ContextVar`` set inside compiled + code is neither written nor readable. """ -from collections.abc import Iterator -from contextlib import contextmanager -from contextvars import ContextVar -from dataclasses import dataclass from functools import partial -from typing import Any, TypeAlias +from typing import Any import torch import torch.nn as nn from torch.utils.checkpoint import CheckpointPolicy, create_selective_checkpoint_contexts from xtuner.v1.utils import log_rank0 -from xtuner.v1.utils.enum_helper import StrEnum +from xtuner.v1.utils.selective_checkpointing import SaveUnit, active_recompute_unit, recompute_unit from .checkpointing import apply_gradient_checkpointing, checkpoint_flattened, install_checkpointing -__all__ = [ - "RecomputeUnit", - "KeptOps", - "KeptCallables", - "RecomputeTarget", - "RecomputeTargetMap", - "active_recompute_unit", - "recompute_unit", - "apply_selective_checkpointing", - "in_recompute_unit", - "resolve_kept_ops", -] - - -class RecomputeUnit(StrEnum): - """Semantic units of activation that may be kept resident instead of - recomputed. - - Each member names a *class of sub-structure* whose activations are worth keeping in memory because recomputing it - is expensive relative to what it costs to store. A unit selected by the user means "do not recompute this part"; - everything not selected is recomputed. What a unit resolves to is architecture-specific and is declared per model, - so the same unit can cover different code in different models. - - Members are strings so pydantic round-trips them to readable names in serialized configs. - """ - - SAVE_ATTN = "save_attn" - """Keep the attention kernel's output -- the flash-attention call itself, not the projections - around it. - - This is the narrowest unit and the only one that costs no compilation, because the attention - kernel is a custom op: inductor cannot fuse it, so it is always called as a fallback kernel and - is always visible to the checkpoint policy, compiled or not. - """ - - SAVE_MOE_GATE = "save_moe_gate" - """Keep the MoE router: gating projection, top-k selection, and routing weights.""" - - SAVE_MOE_DISPATCH = "save_moe_dispatch" - """Keep the tensors produced around expert dispatch and combine: the permutation, padding and - unpermutation buffers on either side of the all-to-all. - - The collective itself is always recomputed, never kept. Keeping a collective would elide it from - the recompute pass, which is only sound if nothing it communicates with is replayed; the safe - rule is to replay it. So this unit trades memory for the surrounding tensor work, not for the - communication. - """ - - -@dataclass(frozen=True) -class KeptOps: - """Resolve a unit by op identity: keep these ops wherever they run. - - Costs nothing -- compilation is untouched and no graph breaks are introduced -- because the ops - worth naming this way are exactly the ones inductor cannot fuse, which are the ones that still - reach the dispatcher from inside a compiled region. - - Args: - names (tuple[str, ...]): Qualified op names, e.g. ``"flash_attn::_flash_attn_varlen_forward_v3"``. - A name that is not registered in this build is skipped, so a model may list several - backends' spellings of the same kernel. - """ - - names: tuple[str, ...] - - def __init__(self, *names: str) -> None: - object.__setattr__(self, "names", tuple(names)) - - -@dataclass(frozen=True) -class KeptCallables: - """Resolve a unit by callable: keep everything these callables dispatch. - - The callables are taken out of the compiled set, because a region inside compiled code cannot be - addressed at all -- so the compilation given up is the whole callable's. Name the smallest - callable that covers the unit. - - Args: - names (tuple[str, ...]): Qualified callable names, in the same form ``compile_cfg`` uses. - """ - - names: tuple[str, ...] - - def __init__(self, *names: str) -> None: - object.__setattr__(self, "names", tuple(names)) - - -RecomputeTarget: TypeAlias = KeptOps | KeptCallables -"""What a single :class:`RecomputeUnit` resolves to for one architecture.""" - -RecomputeTargetMap: TypeAlias = dict[RecomputeUnit, RecomputeTarget] -"""Per-model declaration of what each supported :class:`RecomputeUnit` resolves -to. - -Models expose this as a ``default_recompute_cfg`` property, mirroring ``default_compile_cfg``. A unit absent from the -mapping is not supported by that architecture. -""" - - -def active_recompute_unit() -> RecomputeUnit | None: - """Return the unit whose callable is currently executing, if any. - - Meaningful only while a checkpoint policy is running, which is the only caller. Units resolved by - :class:`KeptOps` never set this -- they are recognised from the op itself. - - Returns: - RecomputeUnit | None: The innermost open unit, or None outside one. - """ - return _ACTIVE_UNIT.get() - - -@contextmanager -def recompute_unit(unit: RecomputeUnit) -> Iterator[None]: - """Mark the enclosed call as belonging to ``unit``. - - Entered by the wrapper the engine installs on a :class:`KeptCallables` target, never by model - code. It is a plain ``ContextVar``, which is only readable because the wrapped callable has been - taken out of the compiled set; inside compiled code it would neither be set nor read. - - Args: - unit (RecomputeUnit): The unit the enclosed call belongs to. - """ - token = _ACTIVE_UNIT.set(unit) - try: - yield - finally: - _ACTIVE_UNIT.reset(token) - - -_ACTIVE_UNIT: ContextVar[RecomputeUnit | None] = ContextVar("xtuner_recompute_unit", default=None) +__all__ = ["apply_selective_checkpointing", "in_recompute_unit", "resolve_kept_ops"] def apply_selective_checkpointing( @@ -213,15 +66,15 @@ def apply_selective_checkpointing( return install_checkpointing(module, partial(_run_checkpointed_region, kept_ops, preserve_rng_state)) -def in_recompute_unit(unit: RecomputeUnit, func: Any) -> Any: +def in_recompute_unit(unit: SaveUnit, func: Any) -> Any: """Wrap ``func`` so everything it dispatches belongs to ``unit``. Installed by the model layer on each callable a - :class:`KeptCallables` target names, together with the - exclusion from compilation that makes the marker readable at all. + :class:`~xtuner.v1.utils.selective_checkpointing.KeptCallables` target names, together with + the exclusion from compilation that makes the marker readable at all. Args: - unit (RecomputeUnit): The unit this callable implements. + unit (SaveUnit): The unit this callable implements. func (Any): The callable to wrap. Returns: diff --git a/xtuner/v1/module/decoder_layer/moe_decoder_layer.py b/xtuner/v1/module/decoder_layer/moe_decoder_layer.py index 37ec22bb12..edc8af0fe6 100644 --- a/xtuner/v1/module/decoder_layer/moe_decoder_layer.py +++ b/xtuner/v1/module/decoder_layer/moe_decoder_layer.py @@ -139,11 +139,21 @@ def __init__( if self.gate_bias: self.bias = nn.Parameter(torch.zeros(self.n_routed_experts)) - def forward( - self, hidden_states: torch.Tensor, rollout_routed_experts: torch.Tensor | None = None - ) -> RouterResults: + def project(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Compute the routing logits. + + Split from :meth:`forward` so that the projection can be addressed on its own: it holds all the activation + memory the router has -- the logits are ``tokens x n_routed_experts`` while everything the router itself + produces is top-k sized -- and unlike the router it contains no in-place write, which selective checkpointing + cannot keep. + + Args: + hidden_states (torch.Tensor): Layer input, of any leading shape. + + Returns: + torch.Tensor: Routing logits, flattened to ``(tokens, n_routed_experts)``. + """ _, _, h = hidden_states.shape - ### compute gating score hidden_states = hidden_states.view(-1, h) if isinstance(self.weight, DTensor): @@ -156,11 +166,14 @@ def forward( bias = self.bias.to_local() if isinstance(self.bias, DTensor) else self.bias if self.router_compute_dtype == "native": - logits = F.linear(hidden_states, weight, bias) - else: - bias = bias.float() if bias is not None else None - logits = F.linear(hidden_states.float(), weight.float(), bias) - return self.router(logits, rollout_routed_experts) + return F.linear(hidden_states, weight, bias) + bias = bias.float() if bias is not None else None + return F.linear(hidden_states.float(), weight.float(), bias) + + def forward( + self, hidden_states: torch.Tensor, rollout_routed_experts: torch.Tensor | None = None + ) -> RouterResults: + return self.router(self.project(hidden_states), rollout_routed_experts) # Debug for aligning with hf implementation. # logits = F.linear(hidden_states, weight, bias) @@ -306,6 +319,7 @@ def __init__( self.dispatcher = build_dispatcher( dispatcher=dispatcher, n_routed_experts=n_routed_experts, + hidden_size=hidden_size, ep_group=process_group, training_dtype="fp8" if float8_cfg is not None else "bf16", generate_dtype=generate_config.dtype if generate_config is not None else "bf16", @@ -417,11 +431,14 @@ def _forward( origin_shape = hidden_states.shape # reshape hidden_states to (batch_size * seq_len, hidden_size) + # Flattened before the marker so that the dispatch region covers the same operations here as + # it does in `_micro_batch_forward`, which reshapes outside the region too. + flat_hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) # ProberList.before_dispatch( # self.layer_idx, hidden_states, router_results["topk_ids"], router_results["topk_weights"] # ) pre_dispatched = self.dispatcher.dispatch_preprocess( - hidden_states=hidden_states.view(-1, hidden_states.shape[-1]), + hidden_states=flat_hidden_states, topk_ids=router_results["topk_ids"], ) dispatched = self.dispatcher.dispatch( @@ -473,14 +490,15 @@ def _forward( pre_combined=pre_combined, combined=combined, ) - combined_hidden_states = post_combined["hidden_states"] - combined_hidden_states = combined_hidden_states.view(*origin_shape) + combined_hidden_states = post_combined["hidden_states"].view(*origin_shape) # debug for aligning with hf implementation. # combined_hidden_states = self._hf_expert_forward_for_debug(hidden_states, router_results, origin_shape) # ProberList.after_combine(self.layer_idx, combined_hidden_states) + # Recorded outside the branch so the region never depends on a configuration-dependent path: + # without shared experts it is simply empty rather than left open until the next marker. if self.n_shared_experts > 0: shared_experts_out = self._shared_experts_forward(hidden_states=hidden_states) else: @@ -600,6 +618,11 @@ def _micro_batch_forward( shared_experts_out_list: list[torch.Tensor | None] + # Recorded outside the branch so the region never depends on a configuration-dependent path: + # without shared experts it is simply empty rather than left open until the next marker. It + # also sits outside the loop, unlike the single-batch path's per-call region: this stage runs + # every micro-batch back to back with nothing else between them, so one region spanning the + # whole stage covers exactly the same operations as one region per micro-batch would. if self.n_shared_experts > 0: shared_experts_out_list = [] for pre_moe_forward_out in pre_moe_forward_out_list: diff --git a/xtuner/v1/module/dispatcher/__init__.py b/xtuner/v1/module/dispatcher/__init__.py index e4981392d0..4017e340d8 100644 --- a/xtuner/v1/module/dispatcher/__init__.py +++ b/xtuner/v1/module/dispatcher/__init__.py @@ -30,6 +30,7 @@ def build_dispatcher( dispatcher: Literal["deepep", "all2all", "agrs"] | None, n_routed_experts: int, + hidden_size: int, ep_group: dist.ProcessGroup | None = None, training_dtype: Literal["bf16", "fp8"] = "bf16", generate_dtype: Literal["bf16", "fp8"] = "bf16", @@ -55,6 +56,7 @@ def build_dispatcher( # TODO: remove type ignore here return DeepEPDispatcher( n_routed_experts=n_routed_experts, + hidden_size=hidden_size, process_group=ep_group, training_dtype=training_dtype, generate_dtype=generate_dtype, diff --git a/xtuner/v1/module/dispatcher/deepep.py b/xtuner/v1/module/dispatcher/deepep.py index 679253e2ba..10d5f38f20 100644 --- a/xtuner/v1/module/dispatcher/deepep.py +++ b/xtuner/v1/module/dispatcher/deepep.py @@ -13,6 +13,7 @@ combine_forward, dispatch_backward, dispatch_forward, + get_low_latency_buffer, ) from xtuner.v1.utils import copy_method_signature, get_device, get_logger @@ -257,6 +258,7 @@ def __init__( self, *, n_routed_experts: int, + hidden_size: int, process_group: torch.distributed.ProcessGroup, training_dtype: Literal["fp8", "bf16"] = "bf16", generate_dtype: Literal["fp8", "bf16"] = "bf16", @@ -273,6 +275,13 @@ def __init__( "Process group must be provided for `DeepEPDispatcher`. " "If you are training a MoE model, it means that `expert parallel` is not enabled in the config." ) + # Built here rather than on first dispatch. The buffer is a process-wide singleton whose + # constructor all-gathers device ids and IPC handles, and `all_gather_object` swaps a + # tensor's storage with `aten.set_`. Left lazy, that lands inside whichever forward happens + # to run first -- which under selective checkpointing is a checkpointed one, where the write + # hits a tensor the policy kept and torch rejects the step with "Tensor cached during + # selective activation checkpoint has been mutated". + get_low_latency_buffer(self._process_group, hidden=hidden_size, num_experts=n_routed_experts) @override def dispatch_preprocess( diff --git a/xtuner/v1/train/trainer.py b/xtuner/v1/train/trainer.py index bd780a7deb..68a26d6a6d 100644 --- a/xtuner/v1/train/trainer.py +++ b/xtuner/v1/train/trainer.py @@ -1177,6 +1177,12 @@ def build_engine( if engine.model.compile_cfg is not None: log_rank0.info(f"The `compile_cfg` of model is {json.dumps(engine.model.compile_cfg, indent=4)}") + # Only reported when something is actually kept resident: this is the memory knob to look at + # first when a run's peak does not match the one it was tuned for. + if engine.model.keeps_any_recompute_unit: + log_rank0.info( + f"The `recompute_cfg` of model keeps these units resident: {sorted(engine.model._selected_recompute_units)}" + ) return engine def build_lr_scheduler(self, lr_cfg: LRConfig, scheduler_step: int) -> torch.optim.lr_scheduler.LRScheduler: diff --git a/xtuner/v1/utils/__init__.py b/xtuner/v1/utils/__init__.py index 915e2a0c79..638d67a640 100644 --- a/xtuner/v1/utils/__init__.py +++ b/xtuner/v1/utils/__init__.py @@ -22,6 +22,14 @@ ) from .pad import pad_to_max_length, pad_to_multiple_of from .profile import profile_time, profile_time_and_memory, timer, timer_logger +from .selective_checkpointing import ( + KeptCallables, + KeptOps, + RecomputeConfig, + RecomputeTarget, + RecomputeTargetMap, + SaveUnit, +) from .state import ForwardState from .type_helper import copy_method_signature, copy_signature, ray_method from .update_weights_utils import monkey_unpatch_torch_reductions @@ -68,4 +76,10 @@ "trim_memory", "group_tensors_by_device_mesh_and_placements", "cal_total_norm", + "KeptCallables", + "KeptOps", + "RecomputeTarget", + "RecomputeTargetMap", + "RecomputeConfig", + "SaveUnit", ] diff --git a/xtuner/v1/utils/selective_checkpointing.py b/xtuner/v1/utils/selective_checkpointing.py new file mode 100644 index 0000000000..50d7a41e49 --- /dev/null +++ b/xtuner/v1/utils/selective_checkpointing.py @@ -0,0 +1,208 @@ +"""Shared contract for region-level selective activation checkpointing (SAC). + +This module holds the vocabulary that the SAC layers agree on and nothing else: + +- Users select :class:`SaveUnit` members in the model config. +- Model authors declare a :data:`RecomputeTargetMap` saying what each unit they support resolves to + for their architecture. +- The SAC engine turns the selection into a per-op checkpoint policy. + +It lives under ``xtuner.v1.utils`` rather than next to the models because the targets name callables +in ``xtuner.v1.module`` while the unit vocabulary is consumed by ``xtuner.v1.model`` configs; a home +inside either package would make the two import each other. + +**A unit resolves to one of two things, and which one is not a style choice.** The checkpoint policy +is asked about every op that reaches the dispatcher, and it has to answer "is this op part of a kept +unit". There are exactly two ways to know: + +- :class:`KeptOps` names the ops directly. This is the cheapest possible form -- it changes nothing + about compilation -- but it only works when the op is *specific enough to identify the unit on its + own*. Attention qualifies: `flash_attn::_flash_attn_varlen_forward` appears nowhere else. A gate's + ``addmm`` does not. +- :class:`KeptCallables` names callables whose whole body belongs to the unit. Everything they + dispatch is kept. This can express any region, at the price of taking those callables out of the + compiled set -- a marker is ordinary python state and the ops it would cover run as fused kernels, + so a region inside compiled code is invisible either way. + +Prefer :class:`KeptOps`. Reach for :class:`KeptCallables` when no op identifies the unit, and then +name the *smallest* callable that covers it: the compilation given up is the whole callable's, not +the region's. +""" + +from collections.abc import Iterator +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass +from typing import TypeAlias + +from cyclopts import Parameter +from pydantic import BaseModel as PydanticBaseModel +from pydantic import ConfigDict +from typing_extensions import Annotated + +from .enum_helper import StrEnum + + +__all__ = [ + "SaveUnit", + "RecomputeConfig", + "KeptOps", + "KeptCallables", + "RecomputeTarget", + "RecomputeTargetMap", + "active_recompute_unit", + "recompute_unit", +] + + +class SaveUnit(StrEnum): + """Semantic units of activation that are *saved* -- kept resident instead + of recomputed. + + Naming is deliberate: selecting a member means "save this", not "recompute this". Everything not + selected is recomputed, which is the default and costs no memory. + + Each member names a *class of sub-structure* whose activations are worth keeping because + recomputing it is expensive relative to what it costs to store. What a unit resolves to is + architecture-specific and is declared per model, so the same unit can cover different code in + different models. + + Members are strings so pydantic round-trips them to readable names in serialized configs. + """ + + ATTN = "attn" + """Keep the attention kernel's output -- the flash-attention call itself, + not the projections around it. + + This is the narrowest unit and the only one that costs no compilation, because the attention + kernel is a custom op: inductor cannot fuse it, so it is always called as a fallback kernel and + is always visible to the checkpoint policy, compiled or not. + """ + + MOE_GATE = "moe_gate" + """Keep the MoE router: gating projection, top-k selection, and routing weights.""" + + MOE_DISPATCH = "moe_dispatch" + """Keep the tensors produced around expert dispatch and combine: the permutation, padding and + unpermutation buffers on either side of the all-to-all. + + The collective itself is always recomputed, never kept. Keeping a collective would elide it from + the recompute pass, which is only sound if nothing it communicates with is replayed; the safe + rule is to replay it. So this unit trades memory for the surrounding tensor work, not for the + communication. + """ + + +@dataclass(frozen=True) +class KeptOps: + """Resolve a unit by op identity: keep these ops wherever they run. + + Costs nothing -- compilation is untouched and no graph breaks are introduced -- because the ops + worth naming this way are exactly the ones inductor cannot fuse, which are the ones that still + reach the dispatcher from inside a compiled region. + + Args: + names (tuple[str, ...]): Qualified op names, e.g. ``"flash_attn::_flash_attn_varlen_forward_v3"``. + A name that is not registered in this build is skipped, so a model may list several + backends' spellings of the same kernel. + """ + + names: tuple[str, ...] + + def __init__(self, *names: str) -> None: + object.__setattr__(self, "names", tuple(names)) + + +@dataclass(frozen=True) +class KeptCallables: + """Resolve a unit by callable: keep everything these callables dispatch. + + The callables are taken out of the compiled set, because a region inside compiled code cannot be + addressed at all -- so the compilation given up is the whole callable's. Name the smallest + callable that covers the unit. + + Args: + names (tuple[str, ...]): Qualified callable names, in the same form ``compile_cfg`` uses. + """ + + names: tuple[str, ...] + + def __init__(self, *names: str) -> None: + object.__setattr__(self, "names", tuple(names)) + + +RecomputeTarget: TypeAlias = KeptOps | KeptCallables +"""What a single :class:`SaveUnit` resolves to for one architecture.""" + + +class RecomputeConfig(PydanticBaseModel): + """All recompute settings in one place: which layers, and what stays saved + inside them. + + ``ratio`` and ``vision_ratio`` decide *which layers* are recomputed at all; ``save`` decides + *what inside a recomputed layer* is kept resident anyway. They used to live in two different + configs -- the ratios on ``FSDPConfig``, the unit selection on the model config -- which made + the pair hard to reason about, since neither means much without the other. + + Args: + ratio (float): Fraction of decoder layers to recompute, counted from the end. Defaults to + 1.0, i.e. recompute every layer. + vision_ratio (float): The same, for a compose model's vision tower. Defaults to 1.0. + save (list[SaveUnit] | bool | None): What stays resident inside the recomputed layers. + ``None`` and ``False`` save nothing, ``True`` saves every unit the model declares, and a + list saves exactly those. Defaults to None. + + ``None`` deliberately does not mean "model default", unlike ``compile_cfg``: saving a + unit trades memory for speed, so an unset config must not silently change an existing + run's peak memory. + """ + + model_config = ConfigDict(extra="forbid") + + ratio: Annotated[float, Parameter(help="Fraction of decoder layers to recompute, from the end")] = 1.0 + vision_ratio: Annotated[float, Parameter(help="The same, for a compose model's vision tower")] = 1.0 + save: Annotated[ + list[SaveUnit] | bool | None, + Parameter(help="What stays resident inside recomputed layers: None/False none, True all declared, or a list"), + ] = None + + +RecomputeTargetMap: TypeAlias = dict[SaveUnit, RecomputeTarget] +"""Per-model declaration of what each supported :class:`SaveUnit` resolves to. + +Models expose this as a ``default_recompute_cfg`` property, mirroring ``default_compile_cfg``. A unit absent from the +mapping is not supported by that architecture. +""" + + +def active_recompute_unit() -> SaveUnit | None: + """Return the unit whose callable is currently executing, if any. + + Meaningful only while a checkpoint policy is running, which is the only caller. Units resolved by + :class:`KeptOps` never set this -- they are recognised from the op itself. + + Returns: + SaveUnit | None: The innermost open unit, or None outside one. + """ + return _ACTIVE_UNIT.get() + + +@contextmanager +def recompute_unit(unit: SaveUnit) -> Iterator[None]: + """Mark the enclosed call as belonging to ``unit``. + + Entered by the wrapper the engine installs on a :class:`KeptCallables` target, never by model + code. It is a plain ``ContextVar``, which is only readable because the wrapped callable has been + taken out of the compiled set; inside compiled code it would neither be set nor read. + + Args: + unit (SaveUnit): The unit the enclosed call belongs to. + """ + token = _ACTIVE_UNIT.set(unit) + try: + yield + finally: + _ACTIVE_UNIT.reset(token) + + +_ACTIVE_UNIT: ContextVar[SaveUnit | None] = ContextVar("xtuner_recompute_unit", default=None)