diff --git a/docs/design/expert_tp.md b/docs/design/expert_tp.md new file mode 100644 index 0000000000..482a4ae745 --- /dev/null +++ b/docs/design/expert_tp.md @@ -0,0 +1,39 @@ +# MoE Expert Tensor Parallelism + +Expert Tensor Parallelism（ETP）在每个 expert 内切分 grouped-linear 权重，并与 Expert Parallelism（EP）正交组合。配置入口是 `MoEConfig.expert_tp_size`；它与用于模型其余部分的 `FSDPConfig.tp_size` 含义不同。 + +```mermaid +flowchart LR + A["Router 输出物理 expert id"] --> B{"Dispatcher"} + B -->|"Naive / All2All"| C["ETP rank 间复制 token"] + B -->|"DeepEP"| D["映射为 EP × ETP 虚拟 expert"] + C --> E["Column-parallel w1/w3"] + D --> E + E --> F["Row-parallel w2"] + F --> G["合并 ETP partial output"] + G --> H["返回原 token 顺序"] +``` + +## Mesh 与参数布局 + +训练 mesh 使用 `(fsdp, ep, etp)` 维度，其中 `fsdp = world_size / (ep_size * expert_tp_size)`。Expert 权重位于 `(ep, etp)` 子 mesh： + +- `fused_w1w3` 同时沿 expert 维和每个 expert 的输出维切分，placement 为 `(Shard(0), InterleavedShard(0))`。融合 gate/up projection 时，stripe 数是 `local_experts * 2`。 +- `fused_w2` 沿 expert 维和输入维切分，placement 为 `(Shard(0), Shard(1))`。 +- 非 expert 参数在 EP、ETP 两个维度均为 `Replicate()`；保留二维子 mesh，使其与 FSDP mesh 共享同一个 parent。 + +`InterleavedShard` 让运行时切片保留完整全局语义，因此 HF save/load 和 DCP 可以从 DTensor placement 直接规划，无需为 MoE 参数维护另一套 checkpoint 特例。 + +## Dispatcher + +无 EP 时使用 `NaiveDispatcher`，All2All 场景使用统一的 `TorchAll2AllDispatcher`。两者在 dispatch 后执行 ETP row all-gather，在 combine 前执行 row reduce-scatter sum，并共享同步、异步实现。 + +DeepEP 使用扁平的 `(ep, etp)` process group。每个物理 expert 映射为 `expert_tp_size` 个虚拟 expert；`topk_ids` 和 `topk_weights` 在 dispatch 前同步扩展。这样 DeepEP 一次 collective 即可完成 EP 路由和 ETP token 复制，且同步事件覆盖扩展 kernel，支持多 micro-batch 重叠。 + +## 梯度与 FP8 + +Expert 梯度按 `ep_size * expert_tp_size` 缩放；非 expert DTensor 根据 `Replicate` placement 归约。梯度范数对所有 expert shard 求和，确保 clipping 与单模型基线一致。 + +BF16 与 tile-wise FP8 grouped GEMM 使用相同权重布局。FP8 padding、量化和输出 reshape 均基于本地 shard shape，checkpoint 仍保留原始全局 shape。 + +综上，ETP 只改变 expert 内部的计算与通信布局；路由语义、模型输出及 HF/DCP checkpoint 的全局表示保持不变。 diff --git a/docs/design/load_spec_refactor.md b/docs/design/load_spec_refactor.md new file mode 100644 index 0000000000..e734c2d080 --- /dev/null +++ b/docs/design/load_spec_refactor.md @@ -0,0 +1,355 @@ +# LoadSpec、DTensor 与 DCP 设计 + +## 结论 + +HF checkpoint、DTensor runtime layout 和 DCP checkpoint 描述的是同一个逻辑 +tensor，但面向的协议不同： + +- `LoadSpec` 记录 checkpoint 映射与 runtime 分片的稳定事实。 +- `HFLoadPlan` / `HFSavePlan` 把这些事实编译成当前 rank 可直接执行的 HF I/O + 程序。 +- `ShardDescriptor` 同时描述 continuous shard 和 Expert TP even interleave。 +- DCP 不经过 HF plan；它用 `compute_runs()` 把一个离散 local shard 展开成多个 + 标准全局 chunk，从而保持分布式保存、原地加载和跨拓扑 reshard。 + +~~~mermaid +flowchart LR + A["HF keys + canonical runtime layout"] --> B["LoadSpec：稳定事实"] + B --> C["HFLoadPlan：读取与 local copy"] + B --> D["HFSavePlan：逆分片与 padding trim"] + E["DTensor placements"] --> B + E --> F["compute_runs"] + F --> G["DCP WriteItem / ReadItem"] +~~~ + +核心边界是：HF plan 负责“checkpoint 格式与 runtime tensor 之间如何搬运”，DCP +planner 负责“local storage 对应哪些全局坐标”。两者复用相同的布局数学，不强行共用 +协议对象。 + +## 一、LoadSpec 与 Plan 的职责 + +### 1.1 为什么需要两层对象 + +可以把 `LoadSpec` 理解成地图，把 Plan 理解成一次具体行程：地图不会保存“今天从哪条 +路走”，而同一张地图可以生成完整保存、保留 EP 或只 gather FSDP 等不同路线。 + +| 对象 | 负责 | 不负责 | +| --- | --- | --- | +| `LoadSpec` | 参数名、HF keys、fused dim、runtime/global shape、checkpoint 可见 shape、forward shard history | 当前 rank 的 copy offsets、collective 顺序、save policy | +| `HFLoadPlan` | 选择本 rank 所需 keys、拼接、调用 canonical adapter、执行 copy regions、清零 runtime padding | 推断模型格式、执行 collective | +| `HFSavePlan` | 标记 preserved shards、逆序 gather、最终 FP8 trim、输出 keys | 保存稳定布局事实 | +| `SaveShardStep` | 一个 descriptor 的逆操作、操作前 runtime shape、是否 preserved | 模型格式转换、跨参数策略 | +| 模型 adapter | HF layout 与 XTuner canonical layout 的 transpose/reshape/flatten | EP、ETP、FSDP、rank 和 process group | + +Plan 生成后，executor 不再回读 `LoadSpec`。因此 Plan 是 self-contained 的 +rank-local 程序，也可以独立测试。 + +### 1.2 LoadSpec 保存哪些 shape + +- `global_shape`：所有 runtime shard 之前的完整 shape，可能包含 FP8 kernel 所需的 + 尾部 padding。 +- `origin_shape`：HF checkpoint 实际可见的 shape；没有 FP8 padding 时为空或等于 + `global_shape`。 +- `local_shape`：真实 runtime local tensor shape，用于校验 descriptor 推导结果。 + +`LoadSpec` 不保存旧设计中的 `explicit_owned_regions` 或 +`save_requires_full_reconstruct`。regions 属于某次 load plan 的编译结果；是否 gather +以及保留哪个 group 属于某次 save plan 的策略。 + +## 二、统一的分片描述 + +### 2.1 ShardDescriptor + +```python +class ShardDescriptor(BaseModel): + dim: int + group: dist.ProcessGroup + interleave_factor: int = 1 +``` + +- `interleave_factor == 1`：普通 continuous `Shard`，支持 PyTorch 原有的 uneven + 分片。 +- `interleave_factor > 1`：Expert TP even interleave；factor 是每个 rank 持有的 + 有序连续 run 数。 +- descriptor 按 full runtime tensor 到 local tensor 的语义顺序排列，典型顺序是 + `EP continuous → ETP interleave → FSDP continuous`。 +- save 严格反向执行：`FSDP gather → ETP gather/deinterleave → EP gather`。 + +even interleave 设当前维长度为 `N`，group size 为 `M`，factor 为 `F`，要求： + +```text +N % (M * F) == 0 +run_size = N / (M * F) +``` + +rank `r` 持有第 `(j * M + r)` 个 run，`j = 0 ... F - 1`。当前不支持 uneven +interleave；不整除会明确报错。 + +### 2.2 DTensor placement 如何表达 Expert TP + +MoE column-parallel 权重需要 EP 分 expert、TP 再在每个 expert/projection 内部分列： + +```python +placements = ( + Shard(0), + InterleavedShard(0, num_local_stripes=local_experts * fused_projections), +) +``` + +`InterleavedShard` 继承 PyTorch 私有 `_StridedShard`。其 `split_factor` 不是“当前 +rank 数”，而是当前切分前已经存在的逻辑 stripe 数；在这里就是本 EP rank 的 +`expert × projection` 数。 + +DTensor 构造时 placements 按 mesh 维从左到右作用；从 placements 推导 carving +order、HF save 重建时则按相反方向撤销。FSDP2 还会在 mesh 最左侧 prepend 一个 +`_StridedShard` 标签，但它实际切的是已经完成 EP/ETP 的 local parameter。因此 +`LoadSpec.from_tensor()` 在唯一的转换边界上做两件事： + +1. 真正的 `InterleavedShard` 转成 `interleave_factor > 1`。 +2. FSDP prepend placement 归一化成最后应用的 continuous descriptor。 + +后续 plan 不再依赖 `_StridedShard` 类型或 `DTensorSpec.shard_order`。 + +### 2.3 通俗例子：两位仓库员各拿每个货架的一半 + +设 fused W1/W3 的 dim-0 共 16 行，业务顺序为： + +```text +expert 0 gate: [0,1,2,3] expert 0 up: [4,5,6,7] +expert 1 gate: [8,9,10,11] expert 1 up: [12,13,14,15] +``` + +`EP=2, TP=2` 时，每个 EP rank 只有一个 expert，因 gate/up 融合所以 +`interleave_factor=2`： + +| `(ep,tp)` | local tensor 对应的 global rows | +| --- | --- | +| `(0,0)` | `[0,1,4,5]` | +| `(0,1)` | `[2,3,6,7]` | +| `(1,0)` | `[8,9,12,13]` | +| `(1,1)` | `[10,11,14,15]` | + +这像两位仓库员分别负责每个货架的左半或右半，而不是一人搬走完整货架。若误用第二个 +`Shard(0)`，TP rank 会拿到连续的完整 projection 子集，gate/up 的列并行语义就错了。 + +这种 `(Shard, InterleavedShard)` 同维布局在部分 PyTorch 版本无法归约出 +`shard_order`，因此 `full_tensor()` / `redistribute()` 不能作为稳定实现。XTuner +训练直接使用 `DTensor.from_local()` / `to_local()`，HF 和 DCP 则使用本文的自定义 +布局编译路径。 + +## 三、HF Load + +### 3.1 Plan 生成 + +`LoadSpec.plan_hf_load()` 从完整 canonical tensor 开始，按 descriptors 正向计算 +当前 rank 的 ownership： + +1. continuous shard 生成一个 slice；uneven 时复用 PyTorch shard size/offset + 语义。 +2. even interleave 生成多个有序 runs。 +3. 后续 FSDP continuous shard 切的是已拼接的 ETP-local tensor，segment compiler + 再把它映射回 global runs。 +4. 用 `origin_shape` 裁掉 checkpoint 中不存在的 FP8 padding。 +5. 选择覆盖这些 regions 的最小 HF-key envelope，并编译为 + `HFLoadPlan.copy_regions`。 + +编译使用连续 segments 表达 ownership，不会按 tensor 大小展开逐元素索引。 + +### 3.2 固定执行顺序 + +~~~mermaid +flowchart LR + A["读取 plan.hf_keys"] --> B["沿 fused_dim 拼接"] + B --> C["模型 hf_tensor_to_canonical"] + C --> D["校验 canonical shape"] + D --> E["执行 source-to-local copies"] + E --> F["未写 FP8 padding 保持为 0"] +~~~ + +必须先 canonicalize 再按 ownership copy。例如 GPT-OSS 的 HF +`gate_up_proj` 先 transpose、重排 gate/up 并 flatten，之后 dim-0 才与 +`ShardDescriptor` 描述的 expert-major runtime 维一致；先 slice HF tensor 会切错 +语义维。 + +### 3.3 模型差异被限制在 adapter + +| 模型 | HF → canonical 的主要差异 | +| --- | --- | +| GPT-OSS | `gate_up_proj` transpose、gate/up 重排并 flatten；`down_proj` transpose；bias 同步重排 | +| Qwen3.5 | 非 MTP expert 权重主要 flatten expert 维；MTP 保留自己的布局 | +| Qwen3-VL | `gate_up_proj` 的 HF 维序不同，需要 transpose 后 flatten | +| GLM | 一个 fused XTuner 参数对应多个 per-expert HF keys，先拼接再 flatten | + +新增模型只需提供 key mapping 和 canonical adapter，不复制 EP/ETP/FSDP load +流程。 + +## 四、HF Save、padding 与重建 + +### 4.1 逆分片 + +`plan_hf_save()` 为每个 descriptor 生成 `SaveShardStep`，executor 逆序执行所有未 +preserve 的 step： + +- continuous：把各 rank local shard 补到 collective 所需等长，all-gather 后按 + rank concat，再 trim 到该 step 的 `shape_before_shard`。 +- even interleave：各 rank/run 已等长；all-gather 后把 rank-major 数据重排成 + `[run][rank]`，不需要 collective padding。 + +例如 `TP=2, F=2`，gather 得到： + +```text +rank0: [A0, A2] +rank1: [A1, A3] +``` + +deinterleave 后必须为 `[A0, A1, A2, A3]`，不能像 continuous shard 一样直接按 +rank concat。 + +### 4.2 collective padding 与 FP8 padding + +两类 padding 分层处理： + +| padding | 原因 | 处理位置 | +| --- | --- | --- | +| collective padding | uneven continuous ranks 的 collective 输入必须等长 | 当前 continuous `SaveShardStep` merge 后 trim 到 runtime `shape_before_shard` | +| FP8 padding | runtime kernel 对齐，HF checkpoint 不存在这些尾部元素 | 所有目标 gather 完成后，从 `runtime_output_shape` 一次 trim 到 `output_shape` | + +通俗例子：runtime 长度 10、HF 可见长度 9，以 3 ranks continuous shard。三个 +local 长度是 `4, 4, 2`，第三个 rank 先临时补成 4；gather 得到长度 12 后先 trim +到 runtime 长度 10，这一步只去 collective padding；所有逆分片结束后再 trim 到 +HF 长度 9，这一步才去 FP8 padding。这样外层 gather 不会把已经提前裁掉的 FP8 +尾部重新当作 collective padding 补回来。 + +load 侧没有 collective trim：plan 只复制 `origin_shape` 内存在的数据，并先把未写 +runtime padding 清零。 + +### 4.3 Save policy + +| 场景 | preserved shards | 实际逆分片 | 输出语义 | +| --- | --- | --- | --- | +| full / distributed HF save | 无 | FSDP → ETP → EP | 完整 checkpoint-visible canonical tensor | +| preserve EP | EP | FSDP → ETP | 连续的 EP-local canonical tensor，供 RL expert-key sync | +| only gather FSDP | EP、ETP | 仅 FSDP | 完整 EP/ETP-local runtime tensor，供 layer-wise IPC | + +`distributed_save` 是写出分配策略，不是部分重建策略；当前仍先正确重建完整 +canonical tensor，再决定哪些 rank 写哪些 keys。 + +### 4.4 reconstruct_full_tensor + +`reconstruct_full_tensor(dt)` 仍是 public convenience API，因为 PyTorch +`DTensor.full_tensor()` 不稳定支持 `(Shard, InterleavedShard)`。它现在只是薄封装： + +1. 从 DTensor placements 构造 descriptor history。 +2. 生成不 preserve 任何 shard 的 full-runtime save steps。 +3. 调用与 `HFSavePlan` 相同的逆分片 executor。 +4. 返回 `global_shape` 的 runtime tensor。 + +它不做模型格式转换，也不 trim `origin_shape`；`distributed_save` 和 preserve group +属于 `HFSavePlan` 策略，不进入这个 API。生产 HF save 不再通过 +`save_requires_full_reconstruct` 走特殊旁路。 + +## 五、DCP 的独立协议 + +### 5.1 为什么默认 DCP 不够 + +普通 DTensor rank 可由一个 `local_shape + global_offset` 描述；interleaved rank +对应多个不连续区间。若把 `[0,1,4,5]` 错当成从 offset 0 开始的连续四行，DCP 会把 +后两行写到 global rows 2、3，checkpoint 看似成功但内容错误。 + +`compute_runs()` 只根据 global shape、mesh coordinate 和 placements 做几何计算， +不读取 tensor value，也不发 collective。它把 local tensor 拆为若干“local 连续且 +global 也连续”的 `Run`： + +```python +Run( + global_offset=(global_row, 0), + sizes=(num_rows, hidden_size), + local_start=local_row, + local_size=num_rows, +) +``` + +### 5.2 通俗例子：给一本交错装订的书编页码 + +沿用 `(ep=0,tp=0)` 的 local rows `[0,1,4,5]`。它在内存中是连续四行，但 +`compute_runs()` 返回两段： + +```text +Run(global rows [0,2), local rows [0,2)) +Run(global rows [4,6), local rows [2,4)) +``` + +就像一本小册子依次装订了原书第 0、1、4、5 页；保存时必须在目录中记录成两段， +不能声称它是原书第 0 到 3 页。 + +Save planner 为它生成： + +```text +Chunk(offset=(0,0), size=(2,H)) <- local.narrow(0, 0, 2) +Chunk(offset=(4,0), size=(2,H)) <- local.narrow(0, 2, 2) +``` + +Load planner 在当前目标拓扑重新计算 runs，让 DCP 用 saved chunk 与 destination +chunk 的全局坐标求交，并把数据直接写入对应 local `narrow()` view。 + +~~~mermaid +flowchart LR + A["Interleaved DTensor local storage"] --> B["compute_runs：多个全局连续区间"] + B --> C["SavePlanner：每个 run 一个 WriteItem"] + C --> D["DCP metadata + distributed storage"] + D --> E["LoadPlanner：目标拓扑 runs"] + E --> F["saved/destination chunk 求交"] + F --> G["原地写入 local views"] +~~~ + +### 5.3 分布式保存与跨拓扑恢复 + +- 每个 rank 只写自己已有的 local views，不 materialize full tensor，也不把权重集中 + 到 coordinator。 +- metadata 保存 FQN、global shape、offset 和 size，不保存“rank N 的业务含义”。 +- 因此可执行 `save EP=2,TP=2 → load EP=1,TP=4`：新 TP ranks 按当前 placements + 声明目标 runs，DCP 从旧 chunks 中读取相交区域。 +- model 参数和 optimizer 的 `exp_avg` / `exp_avg_sq` 都使用相同 planner，才能完整 + resume。 +- `TrainEngine.save_dcp/load_dcp` 显式传入 + `InterleavedShardSavePlanner/LoadPlanner`；正确性不依赖 Torch 2.7 的可选 monkey + patch。 + +`InterleavedShardSavePlanner` 继承 `XtunerCacheSavePlanner`，但 cache 只复用 plan、 +metadata、`WriteResult` 等 control-plane 信息，不比较权重值，也不会省略本次 tensor +写盘。显式构造 planner 的默认路径当前以正确性为主，并未开启 plan cache。 + +### 5.4 与 HF plan 的边界 + +HF 要处理 key 拼接、canonical adapter、save policy 和 FP8 trim；DCP 已有自己的 +`WriteItem` / `ReadItem`、global metadata 和 storage 协议。因此 DCP 不调用 +`reconstruct_full_tensor`，也不经过 `HFSavePlan`。二者只共享 even interleave 的 +几何语义。 + +## 六、约束与验证 + +当前有意保留以下边界： + +- HF `ShardDescriptor` 支持 uneven continuous，但 ETP 只支持 even interleave。 +- `compute_runs()` 当前只支持 tensor dim-0 sharding，strided run 要求均匀切分。 +- `InterleavedShard` 和 DCP planner 使用部分 PyTorch 私有 API，升级 PyTorch 时需要 + 回归。 +- FSDP prepend placement 必须在 EP/ETP 后应用；不能按 placements 的表面顺序直接 + 推导数据 ownership。 + +行为测试应覆盖：continuous even/uneven、ETP even interleave、EP→ETP→FSDP、 +FP8 load padding、collective+FP8 save padding、full/preserve EP/only gather FSDP、 +`reconstruct_full_tensor`、DCP 同拓扑 round-trip 与跨拓扑 reshard，以及 GPT-OSS、 +Qwen3.5、Qwen3-VL、GLM 的 canonical adapter。 + +## 总结 + +当前设计把 Expert TP 的特殊性收敛到三个清晰位置： + +1. `ShardDescriptor.interleave_factor > 1` 描述 HF I/O 所需的 even runs。 +2. HF load segment compiler 与 save deinterleave 分支负责 rank-local 数据搬运。 +3. DCP `compute_runs()` 把离散 storage 转成标准全局 chunks。 + +模型与 `BaseModel` 主流程不再判断 `InterleavedShard`，HF save 和 +`reconstruct_full_tensor` 也不再维护两套重建算法。continuous 与 interleave 共用 +布局描述、计划调度和 padding 边界，但保留各自最直观的 merge 算法。 diff --git a/tests/engine/test_moe_train_engine_deepep_expert_tp.py b/tests/engine/test_moe_train_engine_deepep_expert_tp.py new file mode 100644 index 0000000000..b0cbca9d7f --- /dev/null +++ b/tests/engine/test_moe_train_engine_deepep_expert_tp.py @@ -0,0 +1,413 @@ +from __future__ import annotations + +import os +import unittest +from typing import Literal, TypeAlias + +# 本测试关注 DeepEP + ExpertTP 的真实 grouped-GEMM 训练路径； +# 与既有 engine ExpertTP 测试一致，用 Cutlass 后端规避本地 Triton TMA 兼容性差异。 +os.environ.setdefault("XTUNER_USE_CUTLASS_GROUP_GEMM", "1") + +import torch +import torch.distributed as dist +from mmengine.utils import is_installed +from torch.testing._comparison import default_tolerances + +from xtuner._testing import DeterministicDDPTestCase +from xtuner.v1.config import AdamWConfig, FSDPConfig +from xtuner.v1.engine.train_engine import TrainEngine +from xtuner.v1.loss.ce_loss import CELossConfig +from xtuner.v1.module.dispatcher.deepep import DeepEPDispatcher +from xtuner.v1.module.dispatcher.torch_all2all import TorchAll2AllDispatcher + +from .test_moe_train_engine_tpep import ( + _build_tiny_moe_cfg, + _copy_matching_engine_weights, + _get_local_param_grad, + _get_param_grad, + _get_tpep_grouped_linear, + _make_engine_input, + _run_train_step_items_without_clip, + _run_one_step_with_norm, + _run_train_step_without_clip, + _slice_tpep_weight, + _sync_engine_weights, + _zero_non_expert_grads, +) + +BF16_RTOL, BF16_ATOL = default_tolerances(torch.bfloat16) +BF16_GRAD_ATOL = BF16_ATOL * 2 +TopKExpansion: TypeAlias = tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], bool, int] + + +def _assert_bf16_training_close(actual: torch.Tensor, expected: torch.Tensor) -> None: + # 中文注释：梯度矩阵经过 grouped-GEMM 与 TP/EP 规约，近 0 元素会出现极小累加顺序差异； + # 这里仍以 torch.testing 的 bf16 默认精度为基准，只给梯度绝对误差留 2 倍余量。 + torch.testing.assert_close( + actual.to(torch.bfloat16), + expected.to(torch.bfloat16), + atol=BF16_GRAD_ATOL, + rtol=BF16_RTOL, + ) + + +def _build_engine( + *, + dispatcher: Literal["all2all", "deepep"], + ep_size: int, + expert_tp_size: int, + intra_layer_micro_batch: int = 1, +) -> TrainEngine: + moe_cfg = _build_tiny_moe_cfg(ep_size=ep_size, expert_tp_size=expert_tp_size) + moe_cfg.dispatcher = dispatcher + optim_cfg = AdamWConfig() + fsdp_cfg = FSDPConfig( + ep_size=ep_size, + cpu_offload=False, + ) + return TrainEngine( + model_cfg=moe_cfg, + optim_cfg=optim_cfg, + fsdp_cfg=fsdp_cfg, + intra_layer_micro_batch=intra_layer_micro_batch, + ) + + +def _record_deepep_expert_tp_domino_stages( + engine: TrainEngine, +) -> tuple[dict[str, list[str]], list[TopKExpansion]]: + stages: dict[str, list[str]] = { + "async_op_true": [], + } + topk_expansions: list[TopKExpansion] = [] + + for layer in engine.model.layers.values(): + dispatcher = layer.dispatcher + assert isinstance(dispatcher, DeepEPDispatcher) + assert dispatcher._tp_size > 1 + assert dispatcher._process_group is not None + assert dispatcher._process_group.size() == dispatcher._ep_size * dispatcher._tp_size + assert dispatcher._virtual_n_experts == dispatcher._n_routed_experts * dispatcher._tp_size + + for stage_name in ( + "dispatch_preprocess", + "dispatch", + "dispatch_postprocess", + "combine_preprocess", + "combine", + "combine_postprocess", + ): + original_stage = getattr(dispatcher, stage_name) + + def stage_wrapper( + *args, + _original_stage=original_stage, + _stage_name=stage_name, + _dispatcher=dispatcher, + **kwargs, + ): + if kwargs.get("async_op", False): + stages["async_op_true"].append(_stage_name) + original_topk_ids = kwargs.get("topk_ids") + result = _original_stage(*args, **kwargs) + + if _stage_name == "dispatch_preprocess": + assert isinstance(original_topk_ids, torch.Tensor) + # 中文注释：当前 DeePEP ExpertTP 不再有独立 _expert_tp helper； + # dispatch_preprocess 会把物理 expert topK 扩展成 virtual expert topK。 + topk_expansions.append( + ( + tuple(original_topk_ids.shape), + tuple(result["topk_ids"].shape), + tuple(result["topk_weights"].shape), + result["topk_weights"].requires_grad, + _dispatcher._tp_size, + ) + ) + return result + + setattr(dispatcher, stage_name, stage_wrapper) + + return stages, topk_expansions + + +def _assert_domino_deepep_expert_tp_collective_stages( + stages: dict[str, list[str]], + topk_expansions: list[TopKExpansion], +) -> None: + assert set(stages["async_op_true"]) == { + "dispatch_preprocess", + "dispatch", + "dispatch_postprocess", + "combine_preprocess", + "combine", + "combine_postprocess", + } + assert topk_expansions + for expansion in topk_expansions: + original_topk_shape, expanded_topk_shape, expanded_weight_shape, _, tp_size = expansion + assert expanded_topk_shape[:-1] == original_topk_shape[:-1] + assert expanded_topk_shape[-1] == original_topk_shape[-1] * tp_size + assert expanded_weight_shape == expanded_topk_shape + # 中文注释：checkpoint wrapper 的第一次 forward 可能在 no_grad 下记录到不可微扩展； + # 只要重算 forward 存在可微 topK weight 扩展，就覆盖了 DeepEP virtual TP 的 backward 路径。 + assert any(expansion[3] for expansion in topk_expansions) + + +@unittest.skipIf( + not torch.cuda.is_available() or not is_installed("deep_ep"), + "CUDA/NCCL and DeepEP are required for real DeepEP ExpertTP TrainEngine validation.", +) +class TestMoETrainEngineDeepEPExpertTP(DeterministicDDPTestCase): + def test_deepep_matches_all2all_with_same_expert_tp_topology(self) -> None: + pg = self.create_pg("cuda") + + ep_size = 2 + expert_tp_size = 2 + engine_all2all = _build_engine( + dispatcher="all2all", + ep_size=ep_size, + expert_tp_size=expert_tp_size, + ) + engine_all2all.init_model_weights() + + engine_deepep = _build_engine( + dispatcher="deepep", + ep_size=ep_size, + expert_tp_size=expert_tp_size, + ) + engine_deepep.init_model_weights() + _copy_matching_engine_weights(engine_all2all, engine_deepep) + dist.barrier() + + assert isinstance(engine_all2all.model.layers["0"].dispatcher, TorchAll2AllDispatcher) + assert isinstance(engine_deepep.model.layers["0"].dispatcher, DeepEPDispatcher) + assert engine_all2all.model.ep_mesh is not None + assert engine_deepep.model.ep_mesh is not None + assert engine_all2all.model.expert_tp_mesh is not None + assert engine_deepep.model.expert_tp_mesh is not None + assert engine_all2all.model.ep_mesh.size() == engine_deepep.model.ep_mesh.size() == ep_size + assert ( + engine_all2all.model.expert_tp_mesh.size() + == engine_deepep.model.expert_tp_mesh.size() + == expert_tp_size + ) + assert type(engine_all2all.optimizer) is type(engine_deepep.optimizer) + assert len(engine_all2all.optimizer.param_groups) == len(engine_deepep.optimizer.param_groups) + assert [ + len(group["params"]) for group in engine_all2all.optimizer.param_groups + ] == [len(group["params"]) for group in engine_deepep.optimizer.param_groups] + + device = torch.device("cuda", dist.get_rank() % torch.cuda.device_count()) + input_ids, labels = _make_engine_input(device=device, seed_offset=dist.get_rank()) + loss_cfg = CELossConfig() + + loss_deepep, _, norm_deepep = _run_one_step_with_norm(engine_deepep, loss_cfg, input_ids, labels) + loss_all2all, _, norm_all2all = _run_one_step_with_norm(engine_all2all, loss_cfg, input_ids, labels) + + torch.testing.assert_close( + torch.tensor(loss_deepep), + torch.tensor(loss_all2all), + atol=BF16_ATOL, + rtol=BF16_RTOL, + ) + + gate_grad_deepep = _get_param_grad(engine_deepep, "layers.0.gate.weight") + gate_grad_all2all = _get_param_grad(engine_all2all, "layers.0.gate.weight") + _assert_bf16_training_close(gate_grad_deepep, gate_grad_all2all) + torch.testing.assert_close( + norm_deepep, + norm_all2all, + atol=BF16_ATOL, + rtol=BF16_RTOL, + ) + + dist.barrier() + torch.cuda.empty_cache() + try: + dist.destroy_process_group(pg) + except Exception: + pass + + def test_deepep_expert_tp_matches_single_model_baseline(self) -> None: + pg = self.create_pg("cuda") + + ep_size = 2 + expert_tp_size = 2 + engine_ref = _build_engine( + dispatcher="all2all", + ep_size=1, + expert_tp_size=1, + ) + engine_ref.init_model_weights() + + engine_deepep = _build_engine( + dispatcher="deepep", + ep_size=ep_size, + expert_tp_size=expert_tp_size, + ) + engine_deepep.init_model_weights() + _sync_engine_weights(engine_ref, engine_deepep) + dist.barrier() + + assert isinstance(engine_deepep.model.layers["0"].dispatcher, DeepEPDispatcher) + assert engine_deepep.model.ep_mesh is not None + assert engine_deepep.model.expert_tp_mesh is not None + assert engine_deepep.model.ep_mesh.size() == ep_size + assert engine_deepep.model.expert_tp_mesh.size() == expert_tp_size + + device = torch.device("cuda", dist.get_rank() % torch.cuda.device_count()) + input_ids, labels = _make_engine_input(device=device, seed_offset=dist.get_rank()) + loss_cfg = CELossConfig() + + loss_deepep, _, norm_deepep = _run_one_step_with_norm(engine_deepep, loss_cfg, input_ids, labels) + loss_ref, _, norm_ref = _run_one_step_with_norm(engine_ref, loss_cfg, input_ids, labels) + + torch.testing.assert_close( + torch.tensor(loss_deepep), + torch.tensor(loss_ref), + atol=BF16_ATOL, + rtol=BF16_RTOL, + ) + + gate_grad_deepep = _get_param_grad(engine_deepep, "layers.0.gate.weight") + gate_grad_ref = _get_param_grad(engine_ref, "layers.0.gate.weight") + _assert_bf16_training_close(gate_grad_deepep, gate_grad_ref) + + for module_suffix, fused_gate_up in ( + ("layers.0.experts.fused_w1w3", True), + ("layers.0.experts.fused_w2", False), + ): + ref_grad = _get_param_grad(engine_ref, f"{module_suffix}.weight") + deepep_grad = _get_local_param_grad(engine_deepep, f"{module_suffix}.weight") + deepep_module = _get_tpep_grouped_linear(engine_deepep, module_suffix) + expected_deepep_grad = _slice_tpep_weight(deepep_module, ref_grad, fused_gate_up=fused_gate_up) + _assert_bf16_training_close(deepep_grad, expected_deepep_grad) + + torch.testing.assert_close( + norm_deepep, + norm_ref, + atol=BF16_ATOL, + rtol=BF16_RTOL, + ) + + dist.barrier() + torch.cuda.empty_cache() + try: + dist.destroy_process_group(pg) + except Exception: + pass + + def test_deepep_expert_tp_expert_only_grad_norm_matches_single_model_baseline(self) -> None: + pg = self.create_pg("cuda") + + ep_size = 2 + expert_tp_size = 2 + engine_ref = _build_engine( + dispatcher="all2all", + ep_size=1, + expert_tp_size=1, + ) + engine_ref.init_model_weights() + + engine_deepep = _build_engine( + dispatcher="deepep", + ep_size=ep_size, + expert_tp_size=expert_tp_size, + ) + engine_deepep.init_model_weights() + _sync_engine_weights(engine_ref, engine_deepep) + dist.barrier() + + device = torch.device("cuda", dist.get_rank() % torch.cuda.device_count()) + input_ids, labels = _make_engine_input(device=device, seed_offset=dist.get_rank()) + loss_cfg = CELossConfig() + + _run_train_step_without_clip(engine_deepep, loss_cfg, input_ids, labels) + _run_train_step_without_clip(engine_ref, loss_cfg, input_ids, labels) + # 中文注释：expert-only norm 单独验证 EP 和 ExpertTP shard 的 norm-square 汇总语义。 + _zero_non_expert_grads(engine_deepep) + _zero_non_expert_grads(engine_ref) + expert_norm_deepep = engine_deepep.clip_grad_norm(do_clip=False).detach().float().cpu() + expert_norm_ref = engine_ref.clip_grad_norm(do_clip=False).detach().float().cpu() + torch.testing.assert_close( + expert_norm_deepep, + expert_norm_ref, + atol=BF16_ATOL, + rtol=BF16_RTOL, + ) + + dist.barrier() + torch.cuda.empty_cache() + try: + dist.destroy_process_group(pg) + except Exception: + pass + + def test_deepep_expert_tp_domino_micro_batch_matches_sync_baseline(self) -> None: + pg = self.create_pg("cuda") + + ep_size = 2 + expert_tp_size = 2 + engine_ref = _build_engine( + dispatcher="deepep", + ep_size=ep_size, + expert_tp_size=expert_tp_size, + ) + engine_ref.init_model_weights() + + engine_domino = _build_engine( + dispatcher="deepep", + ep_size=ep_size, + expert_tp_size=expert_tp_size, + intra_layer_micro_batch=2, + ) + engine_domino.init_model_weights() + _copy_matching_engine_weights(engine_ref, engine_domino) + stages, topk_expansions = _record_deepep_expert_tp_domino_stages(engine_domino) + dist.barrier() + + device = torch.device("cuda", dist.get_rank() % torch.cuda.device_count()) + batches = [ + _make_engine_input(device=device, seed_offset=dist.get_rank() * 2), + _make_engine_input(device=device, seed_offset=dist.get_rank() * 2 + 1), + ] + loss_cfg = CELossConfig() + + loss_domino = _run_train_step_items_without_clip(engine_domino, loss_cfg, batches) + norm_domino = engine_domino.clip_grad_norm(do_clip=False).detach().float().cpu() + gate_grad_domino = _get_param_grad(engine_domino, "layers.0.gate.weight") + + loss_ref = _run_train_step_items_without_clip(engine_ref, loss_cfg, batches) + norm_ref = engine_ref.clip_grad_norm(do_clip=False).detach().float().cpu() + gate_grad_ref = _get_param_grad(engine_ref, "layers.0.gate.weight") + + _assert_domino_deepep_expert_tp_collective_stages(stages, topk_expansions) + torch.testing.assert_close( + torch.tensor(loss_domino), + torch.tensor(loss_ref), + atol=BF16_ATOL, + rtol=BF16_RTOL, + ) + torch.testing.assert_close( + norm_domino, + norm_ref, + atol=BF16_ATOL, + rtol=BF16_RTOL, + ) + _assert_bf16_training_close(gate_grad_domino, gate_grad_ref) + + dist.barrier() + torch.cuda.empty_cache() + try: + dist.destroy_process_group(pg) + except Exception: + pass + + @property + def world_size(self) -> int: + return 4 + + @property + def destroy_pg_upon_exit(self) -> bool: + return False diff --git a/tests/engine/test_moe_train_engine_float8.py b/tests/engine/test_moe_train_engine_float8.py index 8a20c1515e..ec6df8ba9c 100644 --- a/tests/engine/test_moe_train_engine_float8.py +++ b/tests/engine/test_moe_train_engine_float8.py @@ -113,7 +113,7 @@ def warmup_fn(x): torch.cuda.empty_cache() try: dist.destroy_process_group(pg) - except: + except Exception: pass @parametrize.parametrize( @@ -192,7 +192,99 @@ def warmup_fn(x): torch.cuda.empty_cache() try: dist.destroy_process_group(pg) - except: + except Exception: + pass + + @parametrize.parametrize( + "device,ep_size,expert_tp_size", + [ + ("cuda", 2, 2), + ], + ) + def test_fp8_ep2_etp2_fsdp2_train(self, device, ep_size, expert_tp_size): + pg = self.create_pg(device) + assert dist.get_world_size() == 8 + + moe_cfg = Qwen3MoE30BA3Config( + ep_size=ep_size, + expert_tp_size=expert_tp_size, + dispatcher="all2all", + balancing_loss_cfg=BalancingLossConfig(), + float8_cfg=Float8Config( + scaling_granularity_gemm=ScalingGranularity.TILEWISE, + scaling_granularity_grouped_gemm=ScalingGranularity.TILEWISE, + ), + ) + optim_cfg: AdamWConfig = AdamWConfig() + lr_cfg: LRConfig = LRConfig() + fsdp_cfg: FSDPConfig = FSDPConfig( + cpu_offload=False, + ep_size=ep_size, + ) + engine = TrainEngine( + model_cfg=moe_cfg, + optim_cfg=optim_cfg, + fsdp_cfg=fsdp_cfg, + ) + assert engine.model.fsdp_mesh is not None + assert engine.model.fsdp_mesh.size() == 2 + assert engine.model.ep_mesh is not None + assert engine.model.ep_mesh.size() == ep_size + assert engine.model.expert_tp_mesh is not None + assert engine.model.expert_tp_mesh.size() == expert_tp_size + + engine.from_hf(hf_path=QWEN3_MOE_PATH) + + loss_cfg = CELossConfig() + total_steps = 1000 + warmup_steps = total_steps * lr_cfg.warmup_ratio + + def warmup_fn(x): + return x / warmup_steps if x < warmup_steps else 1 + + lr_scheduler = LambdaLR(engine.optimizer, warmup_fn) + + tok = AutoTokenizer.from_pretrained(QWEN3_MOE_PATH) + txt = "根据国际地球自转和参考系服务机构的数据，今年夏天是自2020年以来第六次地球自转加速。7月9日将成为有史以来最短的一天，比平时短1.3到1.6毫秒。 " + input_ids = tok.encode(txt, return_tensors="pt").view(1, -1) + labels = input_ids.clone() + input_ids = input_ids[:, :-1] + labels = labels[:, 1:] + pack_len = 8192 - input_ids.shape[1] + input_ids = pad_to_max_length(input_ids, 0, max_length=8192) + labels = pad_to_max_length(labels, -100, max_length=8192).to(DEVICE) + + losses = [] + for _ in range(10): + seq_ctx = SequenceContext.from_input_ids((input_ids,), device=DEVICE) + seq_ctx.num_padding = pack_len + LossContext = loss_cfg.loss_ctx_cls + loss_ctx = loss_cfg.build(data={"shifted_labels": labels}, sp_mesh=None) + loss_ctx = LossContext.build_batches([loss_ctx])[0] + engine_input = [ModelItem(seq_ctx=seq_ctx, loss_ctx={"lm": loss_ctx})] + + loss_log = engine.train_step(engine_input)["logs_info"] + grad_norm = engine.clip_grad_norm() + engine.step_optimizer(grad_norm) + lr_scheduler.step() + assert torch.isfinite(grad_norm) + assert grad_norm.item() > 0 + losses.append(loss_log["reduced_llm_loss"]) + + losses = torch.tensor(losses) + if dist.get_rank() == 0: + print(f"fp8_ep2_etp2_fsdp2 losses: {losses.tolist()}", flush=True) + losses_ref = torch.tensor([2.41, 2.41, 1.79, 1.39, 1.02, 0.68, 0.52, 0.31, 0.18, 0.12]) + # 2026-07-08 实测 EP2/ETP2/FSDP2 full tilewise FP8 loss: + # [2.4088690281, 2.4088690281, 1.7859514952, 1.3338183165, 0.9629249573, + # 0.6677960753, 0.4466774166, 0.2764163315, 0.1643507332, 0.1138044521] + # 相对原 tilewise ref: cosine similarity = 0.999698, avg relative diff = 0.05054. + self._check_loss_curve(losses, losses_ref, sim_tol=0.02, rtol=0.2) + + torch.cuda.empty_cache() + try: + dist.destroy_process_group(pg) + except Exception: pass @parametrize.parametrize( @@ -292,7 +384,7 @@ def warmup_fn(x): torch.cuda.empty_cache() try: dist.destroy_process_group(pg) - except: + except Exception: pass @parametrize.parametrize( diff --git a/tests/engine/test_moe_train_engine_tpep.py b/tests/engine/test_moe_train_engine_tpep.py new file mode 100644 index 0000000000..6a4d1b78e0 --- /dev/null +++ b/tests/engine/test_moe_train_engine_tpep.py @@ -0,0 +1,1084 @@ +"""Validate that EP+TP training produces the same forward loss and backward +gradients as a pure single-GPU (EP=1, TP=1) run. + +Test topology: world_size = EP * TP * DP = 2 * 2 * 1 = 4 GPUs. + +Strategy +-------- +1. Build a tiny Qwen3MoE model with EP=2, TP=2. +2. Build the same model with EP=1, TP=1 (4 identical DP replicas). +3. Init both engines with ``init_model_weights()``. Because weights for EP+TP + models are Shard(0) on ep_mesh for experts and Replicate for non-experts, + and ``init_params`` always initialises the *full* tensor before sharding, + the underlying full weight values are identical when the same RNG seed is + active on all ranks. +4. Sync expert weights from EP=1 engine to EP=2 engine via DCP so the two + models start from the exact same checkpoint. +5. Run one ``train_step`` + ``clip_grad_norm`` on both engines with the same + input. +6. Assert: + - losses agree within tolerance + - gate (router) gradients agree within tolerance (non-expert, replicated + on all ranks in both configs) +""" + +from __future__ import annotations + +import gc +import os + +# 本测试关注 FSDP + EP + expert TP 的 loss/梯度校准。 +# Triton TMA grouped-GEMM 在部分本地 Triton/LLVM 组合下会编译失败， +# 因此沿用 .dev_scripts 的做法，用 Cutlass 后端跑真实 grouped-GEMM。 +os.environ.setdefault("XTUNER_USE_CUTLASS_GROUP_GEMM", "1") + +import parametrize +import torch +import torch.distributed as dist +from torch.distributed.tensor import DTensor, distribute_tensor + +from xtuner._testing import DeterministicDDPTestCase +from xtuner.v1.config import AdamWConfig, FSDPConfig +from xtuner.v1.engine.train_engine import TrainEngine +from xtuner.v1.loss.ce_loss import CELossConfig +from xtuner.v1.module.attention import MHAConfig +from xtuner.v1.module.dispatcher.base import NaiveDispatcher +from xtuner.v1.module.dispatcher.torch_all2all import TorchAll2AllDispatcher +from xtuner.v1.module.grouped_linear.moe_group_linear import GroupedLinear +from xtuner.v1.module.router.greedy import GreedyRouterConfig +from xtuner.v1.model.base import ModelItem +from xtuner.v1.model.moe.moe import SequenceContext +from xtuner.v1.model.moe.qwen3 import Qwen3MoEConfig +from xtuner.v1.utils.device import get_device + +DEVICE = get_device() + +# 本测试的模型参数和主要计算是 bf16，容忍度对齐 torch.testing 的 +# bf16 默认值，避免过宽阈值掩盖 expert TP 维度缺失这类校准错误。 +BF16_ATOL = 1e-5 +BF16_RTOL = 1.6e-2 +# grouped-GEMM 和 TP 分片规约会改变 bf16 的累加顺序；逐元素梯度矩阵 +# 在接近 0 的位置会有数个 ulp 的差异，不能用它承载 loss/norm 校准红灯。 +BF16_GEMM_ATOL = 1e-4 +BF16_GEMM_RTOL = BF16_RTOL + +# Use a very small model to keep test runtime manageable. +_TINY_LAYERS = 2 +_SEQ_LEN = 32 +_VOCAB_SIZE = 128 + + +def _build_tiny_moe_cfg(ep_size: int = 1, expert_tp_size: int = 1) -> Qwen3MoEConfig: + return Qwen3MoEConfig( + vocab_size=_VOCAB_SIZE, + max_position_embeddings=128, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + num_hidden_layers=_TINY_LAYERS, + hidden_size=128, + intermediate_size=256, + rms_norm_eps=1e-6, + rope_theta=1e6, + hidden_act="silu", + attention=MHAConfig(num_attention_heads=4, num_key_value_heads=2, head_dim=32, qk_norm=True), + tie_word_embeddings=False, + n_routed_experts=4, + n_shared_experts=0, + num_experts_per_tok=2, + first_k_dense_replace=0, + hidden_factor=1.0, + moe_intermediate_size=64, + router=GreedyRouterConfig(scoring_func="softmax", norm_topk_prob=True, router_scaling_factor=1.0), + ep_size=ep_size, + expert_tp_size=expert_tp_size, + dispatcher="all2all" if ep_size > 1 else None, + compile_cfg=False, + # Disable auxiliary losses to keep the comparison clean. + balancing_loss_cfg=None, + z_loss_cfg=None, + ) + + +def _build_engine( + ep_size: int, + expert_tp_size: int, + data_tp_size: int = 1, + intra_layer_micro_batch: int = 1, +) -> TrainEngine: + moe_cfg = _build_tiny_moe_cfg(ep_size, expert_tp_size) + optim_cfg = AdamWConfig() + fsdp_cfg = FSDPConfig( + ep_size=ep_size, + tp_size=data_tp_size, + cpu_offload=False, + ) + return TrainEngine( + model_cfg=moe_cfg, + optim_cfg=optim_cfg, + fsdp_cfg=fsdp_cfg, + intra_layer_micro_batch=intra_layer_micro_batch, + ) + + +def _make_engine_input(device: torch.device, seed_offset: int = 0) -> tuple[torch.Tensor, torch.Tensor]: + """Return (input_ids [1, SEQ_LEN-1], shifted_labels [1, SEQ_LEN-1]) on *device*.""" + torch.manual_seed(12345 + seed_offset) + full_ids = torch.randint(0, _VOCAB_SIZE, (1, _SEQ_LEN), dtype=torch.long, device=device) + input_ids = full_ids[:, :-1] # [1, SEQ_LEN-1] + labels = full_ids[:, 1:] # [1, SEQ_LEN-1] already shifted + return input_ids, labels + + +def _run_one_step( + engine: TrainEngine, + loss_cfg: CELossConfig, + input_ids: torch.Tensor, + labels: torch.Tensor, +) -> tuple[float, dict[str, torch.Tensor]]: + """Run one train step; return (loss_value, {param_name: grad_tensor}).""" + loss_val, grads, _ = _run_one_step_with_norm(engine, loss_cfg, input_ids, labels) + return loss_val, grads + + +def _run_one_step_with_norm( + engine: TrainEngine, + loss_cfg: CELossConfig, + input_ids: torch.Tensor, + labels: torch.Tensor, +) -> tuple[float, dict[str, torch.Tensor], torch.Tensor]: + """Run one train step; return loss, gate grads and un-clipped grad norm.""" + loss_val = _run_train_step_without_clip(engine, loss_cfg, input_ids, labels) + grad_norm = engine.clip_grad_norm(do_clip=False) + + # Collect gradients from gate (router) parameters; these are non-expert + # parameters replicated on all ranks in both configs, so they're easy to + # compare directly. + grads: dict[str, torch.Tensor] = {} + for name, param in engine.model.named_parameters(): + if "gate.weight" in name and param.grad is not None: + grad = param.grad + if hasattr(grad, "full_tensor"): + grad = grad.full_tensor() # type: ignore[attr-defined] + grads[name] = grad.detach().float().cpu() + break # one gate layer is sufficient + + return loss_val, grads, grad_norm.detach().float().cpu() + + +def _run_train_step_without_clip( + engine: TrainEngine, + loss_cfg: CELossConfig, + input_ids: torch.Tensor, + labels: torch.Tensor, +) -> float: + engine_input = _make_engine_items(loss_cfg, [(input_ids, labels)]) + step_info = engine.train_step(engine_input) + return step_info["logs_info"]["reduced_llm_loss"] + + +def _make_engine_items( + loss_cfg: CELossConfig, + batches: list[tuple[torch.Tensor, torch.Tensor]], +) -> list[ModelItem]: + loss_ctx_list = [] + seq_ctx_list = [] + for input_ids, labels in batches: + seq_ctx_list.append(SequenceContext.from_input_ids((input_ids,), device=DEVICE)) + shifted_labels = labels.to(DEVICE) + loss_ctx_list.append(loss_cfg.build(data={"shifted_labels": shifted_labels}, sp_mesh=None)) + + LossContext = loss_cfg.loss_ctx_cls + loss_ctx_list = LossContext.build_batches(loss_ctx_list) + return [ + ModelItem(seq_ctx=seq_ctx, loss_ctx={"lm": loss_ctx}) + for seq_ctx, loss_ctx in zip(seq_ctx_list, loss_ctx_list) + ] + + +def _run_train_step_items_without_clip( + engine: TrainEngine, + loss_cfg: CELossConfig, + batches: list[tuple[torch.Tensor, torch.Tensor]], +) -> float: + engine_input = _make_engine_items(loss_cfg, batches) + step_info = engine.train_step(engine_input) + return step_info["logs_info"]["reduced_llm_loss"] + + +def _record_expert_tp_collective_stages(engine: TrainEngine) -> dict[str, list[str]]: + stages: dict[str, list[str]] = { + "async_op_true": [], + "async_all_gather_rows": [], + "async_all_gather_row_metadata": [], + "async_all_gather_per_rank_metadata": [], + "async_reduce_scatter_rows_sum": [], + } + current_stage: list[str] = [] + + for layer in engine.model.layers.values(): + dispatcher = layer.dispatcher + expert_tp = dispatcher._expert_tp + if expert_tp is None: + continue + + for stage_name in ( + "dispatch_preprocess", + "dispatch", + "dispatch_postprocess", + "combine_preprocess", + "combine", + "combine_postprocess", + ): + original_stage = getattr(dispatcher, stage_name) + + def stage_wrapper(*args, _original_stage=original_stage, _stage_name=stage_name, **kwargs): + if kwargs.get("async_op", False): + stages["async_op_true"].append(_stage_name) + current_stage.append(_stage_name) + try: + return _original_stage(*args, **kwargs) + finally: + current_stage.pop() + + setattr(dispatcher, stage_name, stage_wrapper) + + for collective_name in ( + "async_all_gather_rows", + "async_all_gather_row_metadata", + "async_all_gather_per_rank_metadata", + "async_reduce_scatter_rows_sum", + ): + original_collective = getattr(expert_tp, collective_name) + + def collective_wrapper( + *args, + _original_collective=original_collective, + _collective_name=collective_name, + **kwargs, + ): + stages[_collective_name].append(current_stage[-1] if current_stage else "") + return _original_collective(*args, **kwargs) + + setattr(expert_tp, collective_name, collective_wrapper) + + return stages + + +def _assert_domino_expert_tp_collective_stages(stages: dict[str, list[str]]) -> None: + assert set(stages["async_op_true"]) == { + "dispatch_preprocess", + "dispatch", + "dispatch_postprocess", + "combine_preprocess", + "combine", + "combine_postprocess", + } + assert stages["async_all_gather_rows"] + assert stages["async_all_gather_row_metadata"] + assert stages["async_reduce_scatter_rows_sum"] + assert set(stages["async_all_gather_rows"]) == {"dispatch"} + assert set(stages["async_all_gather_row_metadata"]) == {"dispatch"} + assert set(stages["async_reduce_scatter_rows_sum"]) == {"combine"} + + +def _assert_domino_all2all_expert_tp_collective_stages(stages: dict[str, list[str]]) -> None: + assert set(stages["async_op_true"]) == { + "dispatch_preprocess", + "dispatch", + "dispatch_postprocess", + "combine_preprocess", + "combine", + "combine_postprocess", + } + assert stages["async_all_gather_rows"] + assert stages["async_all_gather_per_rank_metadata"] + assert stages["async_reduce_scatter_rows_sum"] + assert set(stages["async_all_gather_rows"]) == {"dispatch"} + assert set(stages["async_all_gather_per_rank_metadata"]) == {"dispatch"} + assert set(stages["async_reduce_scatter_rows_sum"]) == {"combine"} + + +def _assert_rank_inputs_are_distinct(batches: list[tuple[torch.Tensor, torch.Tensor]]) -> None: + local_input_ids = tuple(tuple(input_ids.detach().cpu().reshape(-1).tolist()) for input_ids, _ in batches) + gathered_input_ids: list[tuple[tuple[int, ...], ...] | None] = [None for _ in range(dist.get_world_size())] + dist.all_gather_object(gathered_input_ids, local_input_ids) + # ExpertTP-only 下每个 TP rank 使用不同样本，避免重复输入掩盖 shard 问题。 + assert len(set(gathered_input_ids)) == len(gathered_input_ids) + + +def _get_param_grad(engine: TrainEngine, name_suffix: str) -> torch.Tensor: + for name, param in engine.model.named_parameters(): + if _canonical_name(name).endswith(name_suffix): + grad = param.grad + assert grad is not None, f"Missing gradient for {name}" + if hasattr(grad, "full_tensor"): + grad = grad.full_tensor() # type: ignore[attr-defined] + return grad.detach().float().cpu() + raise AssertionError(f"Cannot find parameter ending with {name_suffix}") + + +def _get_local_param_grad(engine: TrainEngine, name_suffix: str) -> torch.Tensor: + for name, param in engine.model.named_parameters(): + if _canonical_name(name).endswith(name_suffix): + grad = param.grad + assert grad is not None, f"Missing gradient for {name}" + if isinstance(grad, DTensor): + grad = grad.to_local() + return grad.detach().float().cpu() + raise AssertionError(f"Cannot find parameter ending with {name_suffix}") + + +def _get_tpep_grouped_linear(engine: TrainEngine, module_suffix: str) -> GroupedLinear: + for name, module in engine.model.named_modules(): + if _canonical_name(name).endswith(module_suffix): + assert isinstance(module, GroupedLinear) + return module + raise AssertionError(f"Cannot find grouped linear module ending with {module_suffix}") + + +def _canonical_name(name: str) -> str: + # 第一层会被 activation checkpoint wrapper 包一层，比较逻辑不关心该包装。 + return name.replace("._checkpoint_wrapped_module", "") + + +def _zero_non_expert_grads(engine: TrainEngine) -> None: + with torch.no_grad(): + for name, param in engine.model.named_parameters(): + if ".experts" not in _canonical_name(name) and param.grad is not None: + param.grad.zero_() + + +def _full_tensor(tensor: torch.Tensor) -> torch.Tensor: + if isinstance(tensor, DTensor): + return tensor.full_tensor() + return tensor + + +def _copy_param_from_full(param: torch.nn.Parameter, full_tensor: torch.Tensor) -> None: + if isinstance(param, DTensor): + param.copy_(distribute_tensor(full_tensor, param.device_mesh, param.placements)) + else: + param.copy_(full_tensor) + + +def _copy_param_from_local_shard(param: torch.nn.Parameter, local_shard: torch.Tensor) -> None: + if isinstance(param, DTensor): + # ExpertTP GroupedLinear 的 DTensor 参数以本 rank local shard 为真实写入单元； + # 避免对 InterleavedShard 走 full_tensor/redistribute 路径。 + param.copy_(DTensor.from_local(local_shard, param.device_mesh, param.placements, run_check=False)) + else: + param.copy_(local_shard) + + +def _sync_engine_weights(engine_ref: TrainEngine, engine_tpep: TrainEngine) -> None: + """Synchronize a non-TP reference model into the EP+TP model layout.""" + ref_params = dict(engine_ref.model.named_parameters()) + ref_modules = dict(engine_ref.model.named_modules()) + tpep_modules = dict(engine_tpep.model.named_modules()) + + with torch.no_grad(): + for name, param in engine_tpep.model.named_parameters(): + ref_param = ref_params[name] + full_param = _full_tensor(ref_param.detach()).to(device=param.device, dtype=param.dtype) + + module_name, _, param_name = name.rpartition(".") + module = tpep_modules[module_name] + ref_module = ref_modules[module_name] + if isinstance(module, GroupedLinear) and getattr(module, "tp_enabled", False): + if param_name == "weight": + shard = _slice_tpep_weight(module, full_param, fused_gate_up="fused_w1w3" in module_name) + _copy_param_from_local_shard(param, shard) + elif param_name == "bias": + shard = _slice_tpep_bias(module, full_param) + _copy_param_from_local_shard(param, shard) + else: + raise RuntimeError(f"Unexpected GroupedLinear parameter: {name}.") + else: + ref_full = _full_tensor(getattr(ref_module, param_name).detach()).to(device=param.device, dtype=param.dtype) + _copy_param_from_full(param, ref_full) + + +def _copy_matching_engine_weights(engine_src: TrainEngine, engine_dst: TrainEngine) -> None: + """Copy weights between engines that already use the same parameter layout.""" + src_params = dict(engine_src.model.named_parameters()) + + with torch.no_grad(): + for name, dst_param in engine_dst.model.named_parameters(): + src_param = src_params[name].detach() + if isinstance(dst_param, DTensor): + assert isinstance(src_param, DTensor), f"Parameter layout mismatch for {name}" + # 两个 engine 的并行布局相同，直接拷贝本 rank 的 DTensor shard。 + dst_param.copy_(src_param.to(dtype=dst_param.dtype)) + else: + src_tensor = _full_tensor(src_param).to(device=dst_param.device, dtype=dst_param.dtype) + dst_param.copy_(src_tensor) + + +def _snapshot_local_engine_weights(engine: TrainEngine) -> dict[str, torch.Tensor]: + snapshot: dict[str, torch.Tensor] = {} + with torch.no_grad(): + for name, param in engine.model.named_parameters(): + local_param = param.to_local() if isinstance(param, DTensor) else param + snapshot[name] = local_param.detach().cpu().clone() + return snapshot + + +def _copy_local_engine_weight_snapshot(snapshot: dict[str, torch.Tensor], engine: TrainEngine) -> None: + with torch.no_grad(): + for name, param in engine.model.named_parameters(): + local_param = snapshot[name].to(device=param.device, dtype=param.dtype) + if isinstance(param, DTensor): + param.copy_(DTensor.from_local(local_param, param.device_mesh, param.placements, run_check=False)) + else: + param.copy_(local_param) + + +def _slice_tpep_weight(grouped_linear: GroupedLinear, full_weight: torch.Tensor, *, fused_gate_up: bool) -> torch.Tensor: + num_experts = grouped_linear.num_routed_experts + out_features = grouped_linear.out_features + in_features = grouped_linear.in_features + expert_weight = full_weight.view(num_experts, out_features, in_features) + expert_weight = expert_weight[grouped_linear.local_expert_start : grouped_linear.local_expert_end] + + tp_rank = grouped_linear.tp_rank + tp_size = grouped_linear.tp_size + if grouped_linear.parallel_style == "column": + if fused_gate_up: + intermediate_size = out_features // 2 + local_intermediate_size = intermediate_size // tp_size + gate_start = tp_rank * local_intermediate_size + gate_end = gate_start + local_intermediate_size + up_start = intermediate_size + gate_start + up_end = intermediate_size + gate_end + expert_weight = torch.cat( + [ + expert_weight[:, gate_start:gate_end, :], + expert_weight[:, up_start:up_end, :], + ], + dim=1, + ) + else: + local_out_features = out_features // tp_size + out_start = tp_rank * local_out_features + out_end = out_start + local_out_features + expert_weight = expert_weight[:, out_start:out_end, :] + elif grouped_linear.parallel_style == "row": + local_in_features = in_features // tp_size + in_start = tp_rank * local_in_features + in_end = in_start + local_in_features + expert_weight = expert_weight[:, :, in_start:in_end] + else: + raise RuntimeError(f"Unexpected grouped linear parallel style: {grouped_linear.parallel_style}.") + + weight_shape = ( + grouped_linear.weight.to_local().shape + if isinstance(grouped_linear.weight, DTensor) + else grouped_linear.weight.shape + ) + return expert_weight.reshape(weight_shape) + + +def _slice_tpep_bias(grouped_linear: GroupedLinear, full_bias: torch.Tensor) -> torch.Tensor: + expert_bias = full_bias[grouped_linear.local_expert_start : grouped_linear.local_expert_end] + if grouped_linear.parallel_style == "column": + local_out_features = grouped_linear.out_features // grouped_linear.tp_size + out_start = grouped_linear.tp_rank * local_out_features + out_end = out_start + local_out_features + expert_bias = expert_bias[:, out_start:out_end] + bias_shape = ( + grouped_linear.bias.to_local().shape + if isinstance(grouped_linear.bias, DTensor) + else grouped_linear.bias.shape + ) + return expert_bias.reshape(bias_shape) + + +class TestMoETrainEngineExpertTPOnly(DeterministicDDPTestCase): + """Verify ExpertTP-only training matches the non-ExpertTP baseline.""" + + @parametrize.parametrize( + "device,expert_tp_size", + [ + ("cuda", 2), + ], + ) + def test_expert_tp_only_engine_constructs_and_trains(self, device: str, expert_tp_size: int) -> None: + pg = self.create_pg(device) + + engine = _build_engine(ep_size=1, expert_tp_size=expert_tp_size) + engine.init_model_weights() + + assert engine.model.ep_mesh is not None + assert engine.model.expert_tp_mesh is not None + assert engine.model.ep_mesh.size() == 1 + assert engine.model.expert_tp_mesh.size() == expert_tp_size + assert engine.model.expert_tp_mesh.mesh_dim_names == (f"{engine.model.config.mesh_prefix}.etp",) + assert isinstance(engine.model.layers["0"].dispatcher, NaiveDispatcher) + + input_ids, labels = _make_engine_input( + torch.device(device, dist.get_rank() % torch.cuda.device_count()), + seed_offset=dist.get_rank(), + ) + loss_cfg = CELossConfig() + + loss_val = _run_train_step_without_clip(engine, loss_cfg, input_ids, labels) + grad_norm = engine.clip_grad_norm() + engine.step_optimizer(grad_norm) + + assert torch.isfinite(torch.tensor(loss_val)) + assert torch.isfinite(grad_norm) + + dist.barrier() + torch.cuda.empty_cache() + try: + dist.destroy_process_group(pg) + except Exception: + pass + + @parametrize.parametrize( + "device,expert_tp_size", + [ + ("cuda", 2), + ], + ) + def test_expert_tp_only_matches_single_with_distinct_source_slices( + self, device: str, expert_tp_size: int + ) -> None: + pg = self.create_pg(device) + + engine_ref = _build_engine(ep_size=1, expert_tp_size=1) + engine_ref.init_model_weights() + + engine_etp = _build_engine(ep_size=1, expert_tp_size=expert_tp_size) + engine_etp.init_model_weights() + _sync_engine_weights(engine_ref, engine_etp) + dist.barrier() + + input_ids, labels = _make_engine_input( + torch.device(device, dist.get_rank() % torch.cuda.device_count()), + seed_offset=dist.get_rank(), + ) + loss_cfg = CELossConfig() + + loss_etp, _, norm_etp = _run_one_step_with_norm(engine_etp, loss_cfg, input_ids, labels) + loss_ref, _, norm_ref = _run_one_step_with_norm(engine_ref, loss_cfg, input_ids, labels) + + torch.testing.assert_close( + torch.tensor(loss_etp), + torch.tensor(loss_ref), + atol=BF16_ATOL, + rtol=BF16_RTOL, + ) + + gate_grad_ref = _get_param_grad(engine_ref, "layers.0.gate.weight") + gate_grad_etp = _get_param_grad(engine_etp, "layers.0.gate.weight") + torch.testing.assert_close( + gate_grad_etp, + gate_grad_ref, + atol=BF16_GEMM_ATOL, + rtol=BF16_GEMM_RTOL, + ) + + for module_suffix, fused_gate_up in ( + ("layers.0.experts.fused_w1w3", True), + ("layers.0.experts.fused_w2", False), + ): + ref_grad = _get_param_grad(engine_ref, f"{module_suffix}.weight") + etp_grad = _get_local_param_grad(engine_etp, f"{module_suffix}.weight") + etp_module = _get_tpep_grouped_linear(engine_etp, module_suffix) + expected_etp_grad = _slice_tpep_weight(etp_module, ref_grad, fused_gate_up=fused_gate_up) + torch.testing.assert_close( + etp_grad, + expected_etp_grad, + atol=BF16_GEMM_ATOL, + rtol=BF16_GEMM_RTOL, + ) + + torch.testing.assert_close( + norm_etp, + norm_ref, + atol=BF16_ATOL, + rtol=BF16_RTOL, + ) + + dist.barrier() + torch.cuda.empty_cache() + try: + dist.destroy_process_group(pg) + except Exception: + pass + + @parametrize.parametrize( + "device,expert_tp_size", + [ + ("cuda", 2), + ], + ) + def test_expert_tp_only_expert_grad_norm_matches_single_with_distinct_source_slices( + self, device: str, expert_tp_size: int + ) -> None: + pg = self.create_pg(device) + + engine_ref = _build_engine(ep_size=1, expert_tp_size=1) + engine_ref.init_model_weights() + + engine_etp = _build_engine(ep_size=1, expert_tp_size=expert_tp_size) + engine_etp.init_model_weights() + _sync_engine_weights(engine_ref, engine_etp) + dist.barrier() + + input_ids, labels = _make_engine_input( + torch.device(device, dist.get_rank() % torch.cuda.device_count()), + seed_offset=dist.get_rank(), + ) + loss_cfg = CELossConfig() + + _run_train_step_without_clip(engine_etp, loss_cfg, input_ids, labels) + _run_train_step_without_clip(engine_ref, loss_cfg, input_ids, labels) + _zero_non_expert_grads(engine_etp) + _zero_non_expert_grads(engine_ref) + + norm_etp = engine_etp.clip_grad_norm(do_clip=False).detach().float().cpu() + norm_ref = engine_ref.clip_grad_norm(do_clip=False).detach().float().cpu() + torch.testing.assert_close( + norm_etp, + norm_ref, + atol=BF16_ATOL, + rtol=BF16_RTOL, + ) + + dist.barrier() + torch.cuda.empty_cache() + try: + dist.destroy_process_group(pg) + except Exception: + pass + + @parametrize.parametrize( + "device,expert_tp_size", + [ + ("cuda", 2), + ], + ) + def test_expert_tp_only_domino_micro_batch_matches_sync_baseline( + self, device: str, expert_tp_size: int + ) -> None: + pg = self.create_pg(device) + + engine_ref = _build_engine(ep_size=1, expert_tp_size=expert_tp_size) + engine_ref.init_model_weights() + + engine_domino = _build_engine( + ep_size=1, + expert_tp_size=expert_tp_size, + intra_layer_micro_batch=2, + ) + engine_domino.init_model_weights() + _copy_matching_engine_weights(engine_ref, engine_domino) + collective_stages = _record_expert_tp_collective_stages(engine_domino) + dist.barrier() + + device_obj = torch.device(device, dist.get_rank() % torch.cuda.device_count()) + batches = [ + _make_engine_input(device_obj, seed_offset=dist.get_rank() * 2), + _make_engine_input(device_obj, seed_offset=dist.get_rank() * 2 + 1), + ] + _assert_rank_inputs_are_distinct(batches) + loss_cfg = CELossConfig() + + loss_domino = _run_train_step_items_without_clip(engine_domino, loss_cfg, batches) + norm_domino = engine_domino.clip_grad_norm(do_clip=False).detach().float().cpu() + + loss_ref = _run_train_step_items_without_clip(engine_ref, loss_cfg, batches) + norm_ref = engine_ref.clip_grad_norm(do_clip=False).detach().float().cpu() + + _assert_domino_expert_tp_collective_stages(collective_stages) + torch.testing.assert_close( + torch.tensor(loss_domino), + torch.tensor(loss_ref), + atol=BF16_ATOL, + rtol=BF16_RTOL, + ) + torch.testing.assert_close( + norm_domino, + norm_ref, + atol=BF16_ATOL, + rtol=BF16_RTOL, + ) + assert torch.isfinite(torch.tensor(loss_domino)) + assert torch.isfinite(norm_domino) + + dist.barrier() + torch.cuda.empty_cache() + try: + dist.destroy_process_group(pg) + except Exception: + pass + + @property + def world_size(self) -> int: + # ExpertTP-only topology: EP=1, TP=2, DP=1. + return 2 + + @property + def destroy_pg_upon_exit(self) -> bool: + return False + + +class TestMoETrainEngineTPEP(DeterministicDDPTestCase): + """Verify EP+TP training matches single-GPU (EP=1, TP=1) forward and backward.""" + + @parametrize.parametrize( + "device,ep_size,expert_tp_size", + [ + ("cuda", 2, 2), + ], + ) + def test_tpep_forward_backward_matches_single( + self, device: str, ep_size: int, expert_tp_size: int + ) -> None: + """Loss and gate gradients with EP+TP must match the EP=1, TP=1 baseline.""" + pg = self.create_pg(device) + + # ------------------------------------------------------------------ + # Build reference engine: EP=1, TP=1 (world acts as pure DP). + # ------------------------------------------------------------------ + engine_ref = _build_engine(ep_size=1, expert_tp_size=1) + engine_ref.init_model_weights() + + # ------------------------------------------------------------------ + # Build EP+TP engine. + # ------------------------------------------------------------------ + engine_tpep = _build_engine(ep_size=ep_size, expert_tp_size=expert_tp_size) + engine_tpep.init_model_weights() + + # ------------------------------------------------------------------ + # Sync weights by explicitly slicing full expert weights into the real + # TP column/row shards used by GroupedLinear. + # ------------------------------------------------------------------ + _sync_engine_weights(engine_ref, engine_tpep) + dist.barrier() + + # ------------------------------------------------------------------ + # Prepare shared input (identical on all ranks – no SP). + # ------------------------------------------------------------------ + input_ids, labels = _make_engine_input(torch.device(device, dist.get_rank() % torch.cuda.device_count())) + loss_cfg = CELossConfig() + + # Run EP+TP step. + loss_tpep, grads_tpep = _run_one_step(engine_tpep, loss_cfg, input_ids, labels) + + # Run reference step. + loss_ref, grads_ref = _run_one_step(engine_ref, loss_cfg, input_ids, labels) + + # ------------------------------------------------------------------ + # Assert losses match. + # ------------------------------------------------------------------ + if dist.get_rank() == 0: + torch.testing.assert_close( + torch.tensor(loss_tpep), + torch.tensor(loss_ref), + atol=BF16_ATOL, + rtol=BF16_RTOL, + ) + + # ------------------------------------------------------------------ + # Assert gate gradients match (key non-expert parameter). + # ------------------------------------------------------------------ + if grads_tpep and grads_ref: + for name in grads_ref: + if name not in grads_tpep: + continue + g_tpep = grads_tpep[name] + g_ref = grads_ref[name] + if dist.get_rank() == 0: + try: + torch.testing.assert_close( + g_tpep, + g_ref, + atol=BF16_GEMM_ATOL, + rtol=BF16_GEMM_RTOL, + ) + except AssertionError as exc: + max_diff = (g_tpep - g_ref).abs().max().item() + raise AssertionError( + f"Gate gradient mismatch for '{name}': " + f"max_abs_diff={max_diff:.4e}, EP+TP shape={g_tpep.shape}, ref shape={g_ref.shape}" + ) from exc + + dist.barrier() + torch.cuda.empty_cache() + try: + dist.destroy_process_group(pg) + except Exception: + pass + + @parametrize.parametrize( + "device,ep_size,expert_tp_size", + [ + ("cuda", 2, 2), + ], + ) + def test_tpep_expert_gradients_match_single_with_distinct_expert_tp_data( + self, device: str, ep_size: int, expert_tp_size: int + ) -> None: + """Expert TP shards should match the corresponding single-model expert gradients.""" + pg = self.create_pg(device) + + engine_ref = _build_engine(ep_size=1, expert_tp_size=1) + engine_ref.init_model_weights() + + engine_tpep = _build_engine(ep_size=ep_size, expert_tp_size=expert_tp_size) + engine_tpep.init_model_weights() + _sync_engine_weights(engine_ref, engine_tpep) + dist.barrier() + + input_ids, labels = _make_engine_input( + torch.device(device, dist.get_rank() % torch.cuda.device_count()), + seed_offset=dist.get_rank(), + ) + loss_cfg = CELossConfig() + + _run_one_step(engine_tpep, loss_cfg, input_ids, labels) + _run_one_step(engine_ref, loss_cfg, input_ids, labels) + + ref_grad = _get_param_grad(engine_ref, "layers.0.experts.fused_w1w3.weight") + tpep_grad = _get_local_param_grad(engine_tpep, "layers.0.experts.fused_w1w3.weight") + tpep_module = _get_tpep_grouped_linear(engine_tpep, "layers.0.experts.fused_w1w3") + expected_tpep_grad = _slice_tpep_weight(tpep_module, ref_grad, fused_gate_up=True) + + torch.testing.assert_close( + tpep_grad, + expected_tpep_grad, + atol=BF16_GEMM_ATOL, + rtol=BF16_GEMM_RTOL, + ) + + dist.barrier() + torch.cuda.empty_cache() + try: + dist.destroy_process_group(pg) + except Exception: + pass + + @parametrize.parametrize( + "device,ep_size,expert_tp_size", + [ + ("cuda", 2, 2), + ], + ) + def test_tpep_replicated_gradients_and_norm_match_single_with_distinct_expert_tp_data( + self, device: str, ep_size: int, expert_tp_size: int + ) -> None: + """Non-expert replicas and grad norm should match the single-model baseline.""" + pg = self.create_pg(device) + + engine_ref = _build_engine(ep_size=1, expert_tp_size=1) + engine_ref.init_model_weights() + + engine_tpep = _build_engine(ep_size=ep_size, expert_tp_size=expert_tp_size) + engine_tpep.init_model_weights() + _sync_engine_weights(engine_ref, engine_tpep) + dist.barrier() + + input_ids, labels = _make_engine_input( + torch.device(device, dist.get_rank() % torch.cuda.device_count()), + seed_offset=dist.get_rank(), + ) + loss_cfg = CELossConfig() + + _, _, norm_tpep = _run_one_step_with_norm(engine_tpep, loss_cfg, input_ids, labels) + _, _, norm_ref = _run_one_step_with_norm(engine_ref, loss_cfg, input_ids, labels) + + gate_grad_ref = _get_param_grad(engine_ref, "layers.0.gate.weight") + gate_grad_tpep = _get_param_grad(engine_tpep, "layers.0.gate.weight") + + torch.testing.assert_close( + gate_grad_tpep, + gate_grad_ref, + atol=BF16_GEMM_ATOL, + rtol=BF16_GEMM_RTOL, + ) + torch.testing.assert_close( + norm_tpep, + norm_ref, + atol=BF16_ATOL, + rtol=BF16_RTOL, + ) + + dist.barrier() + torch.cuda.empty_cache() + try: + dist.destroy_process_group(pg) + except Exception: + pass + + @parametrize.parametrize( + "device,ep_size,expert_tp_size", + [ + ("cuda", 2, 2), + ], + ) + def test_tpep_expert_only_grad_norm_matches_single_with_distinct_expert_tp_data( + self, device: str, ep_size: int, expert_tp_size: int + ) -> None: + """Expert-only grad norm must sum norm square across EP and expert TP shards.""" + pg = self.create_pg(device) + + engine_ref = _build_engine(ep_size=1, expert_tp_size=1) + engine_ref.init_model_weights() + + engine_tpep = _build_engine(ep_size=ep_size, expert_tp_size=expert_tp_size) + engine_tpep.init_model_weights() + _sync_engine_weights(engine_ref, engine_tpep) + dist.barrier() + + input_ids, labels = _make_engine_input( + torch.device(device, dist.get_rank() % torch.cuda.device_count()), + seed_offset=dist.get_rank(), + ) + loss_cfg = CELossConfig() + + _run_train_step_without_clip(engine_tpep, loss_cfg, input_ids, labels) + _run_train_step_without_clip(engine_ref, loss_cfg, input_ids, labels) + _zero_non_expert_grads(engine_tpep) + _zero_non_expert_grads(engine_ref) + + norm_tpep = engine_tpep.clip_grad_norm(do_clip=False).detach().float().cpu() + norm_ref = engine_ref.clip_grad_norm(do_clip=False).detach().float().cpu() + + torch.testing.assert_close( + norm_tpep, + norm_ref, + atol=BF16_ATOL, + rtol=BF16_RTOL, + ) + + dist.barrier() + torch.cuda.empty_cache() + try: + dist.destroy_process_group(pg) + except Exception: + pass + + @parametrize.parametrize( + "device,ep_size,expert_tp_size", + [ + ("cuda", 2, 2), + ], + ) + def test_tpep_domino_micro_batch_matches_sync_baseline( + self, device: str, ep_size: int, expert_tp_size: int + ) -> None: + pg = self.create_pg(device) + + engine_ref = _build_engine(ep_size=ep_size, expert_tp_size=expert_tp_size) + engine_ref.init_model_weights() + initial_weights = _snapshot_local_engine_weights(engine_ref) + + device_obj = torch.device(device, dist.get_rank() % torch.cuda.device_count()) + batches = [ + _make_engine_input(device_obj, seed_offset=dist.get_rank() * 2), + _make_engine_input(device_obj, seed_offset=dist.get_rank() * 2 + 1), + ] + _assert_rank_inputs_are_distinct(batches) + loss_cfg = CELossConfig() + + loss_ref = _run_train_step_items_without_clip(engine_ref, loss_cfg, batches) + norm_ref = engine_ref.clip_grad_norm(do_clip=False).detach().float().cpu() + del engine_ref + gc.collect() + torch.cuda.empty_cache() + dist.barrier() + + engine_domino = _build_engine( + ep_size=ep_size, + expert_tp_size=expert_tp_size, + intra_layer_micro_batch=2, + ) + engine_domino.init_model_weights() + _copy_local_engine_weight_snapshot(initial_weights, engine_domino) + + for layer in engine_domino.model.layers.values(): + assert isinstance(layer.dispatcher, TorchAll2AllDispatcher) + collective_stages = _record_expert_tp_collective_stages(engine_domino) + dist.barrier() + + loss_domino = _run_train_step_items_without_clip(engine_domino, loss_cfg, batches) + norm_domino = engine_domino.clip_grad_norm(do_clip=False).detach().float().cpu() + + _assert_domino_all2all_expert_tp_collective_stages(collective_stages) + torch.testing.assert_close( + torch.tensor(loss_domino), + torch.tensor(loss_ref), + atol=BF16_ATOL, + rtol=BF16_RTOL, + ) + torch.testing.assert_close( + norm_domino, + norm_ref, + atol=BF16_ATOL, + rtol=BF16_RTOL, + ) + assert torch.isfinite(torch.tensor(loss_domino)) + assert torch.isfinite(norm_domino) + + dist.barrier() + torch.cuda.empty_cache() + try: + dist.destroy_process_group(pg) + except Exception: + pass + + @parametrize.parametrize( + "device,ep_size,expert_tp_size", + [ + ("cuda", 2, 2), + ], + ) + def test_tpep_training_stability(self, device: str, ep_size: int, expert_tp_size: int) -> None: + """EP+TP training should produce finite losses and decreasing trend.""" + pg = self.create_pg(device) + + engine = _build_engine(ep_size=ep_size, expert_tp_size=expert_tp_size) + engine.init_model_weights() + + input_ids, labels = _make_engine_input(torch.device(device, dist.get_rank() % torch.cuda.device_count())) + loss_cfg = CELossConfig() + + losses: list[float] = [] + for _ in range(4): + seq_ctx = SequenceContext.from_input_ids((input_ids,), device=DEVICE) + shifted_labels = labels.to(DEVICE) + LossContext = loss_cfg.loss_ctx_cls + loss_ctx = loss_cfg.build(data={"shifted_labels": shifted_labels}, sp_mesh=None) + loss_ctx_list = LossContext.build_batches([loss_ctx]) + engine_input = [ModelItem(seq_ctx=seq_ctx, loss_ctx={"lm": loss_ctx_list[0]})] + step_info = engine.train_step(engine_input) + grad_norm = engine.clip_grad_norm() + engine.step_optimizer(grad_norm) + losses.append(step_info["logs_info"]["reduced_llm_loss"]) + + if dist.get_rank() == 0: + for i, loss_val in enumerate(losses): + self.assertTrue( + torch.isfinite(torch.tensor(loss_val)), + f"Loss at step {i} is not finite: {loss_val}", + ) + + dist.barrier() + torch.cuda.empty_cache() + try: + dist.destroy_process_group(pg) + except Exception: + pass + + @property + def world_size(self) -> int: + # EP=2, TP=2, DP=1 → 4 GPUs + return 4 + + @property + def destroy_pg_upon_exit(self) -> bool: + return False diff --git a/tests/model/test_gpt_oss_moe.py b/tests/model/test_gpt_oss_moe.py index 47c3cd5c18..24c13e3320 100644 --- a/tests/model/test_gpt_oss_moe.py +++ b/tests/model/test_gpt_oss_moe.py @@ -1,19 +1,21 @@ +import json import os +import tempfile from functools import wraps -import torch.distributed as dist -from safetensors import safe_open -import json +from pathlib import Path import parametrize import torch +import torch.distributed as dist +from safetensors import safe_open + +from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer from xtuner._testing import DeterministicDDPTestCase, patch_hf_rms_norm -from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig -import tempfile -from pathlib import Path -from xtuner.v1.model.moe.moe import SequenceContext -from xtuner.v1.model.moe.gpt_oss import GptOss21BA3P6Config from xtuner.v1.config import FSDPConfig from xtuner.v1.loss.ce_loss import CELossConfig +from xtuner.v1.model.moe.gpt_oss import GptOss21BA3P6Config +from xtuner.v1.model.moe.moe import SequenceContext + GPT_OSS_MINI_PATH = os.environ["GPT_OSS_MINI_PATH"] @@ -94,13 +96,16 @@ def test_gpt_oss_run(self, device, dispatcher, ep_size, compile, tol, loss_class self.assertTrue(torch.allclose(loss, expected_loss.to(loss.dtype), atol=tol, rtol=tol)) @parametrize.parametrize( - "device,dispatcher,ep_size", + "device,dispatcher,ep_size,expert_tp_size", [ - ("cuda", "all2all", 4), - ("cuda", None, 1), + ("cuda", "all2all", 4, 1), + ("cuda", None, 1, 1), + # Packed expert weights and biases must be canonicalized before + # applying the FSDP + EP + Expert TP ownership map. + ("cuda", "all2all", 2, 2), ], ) - def test_fsdp_accuracy(self, device, dispatcher, ep_size): + def test_fsdp_accuracy(self, device, dispatcher, ep_size, expert_tp_size): self.create_pg(device) hf_config = AutoConfig.from_pretrained(GPT_OSS_MINI_PATH) @@ -128,6 +133,7 @@ def test_fsdp_accuracy(self, device, dispatcher, ep_size): with torch.device("meta"): cfg = GptOss21BA3P6Config(compile_cfg=False) cfg.ep_size = ep_size + cfg.expert_tp_size = expert_tp_size cfg.dispatcher = dispatcher gpt_oss_model = cfg.build()._to_device_dtype(dtype=torch.bfloat16, skip_buffers_dtype=True) @@ -159,18 +165,20 @@ def test_fsdp_accuracy(self, device, dispatcher, ep_size): self.assertTrue(torch.allclose(loss, expected_loss.to(loss.dtype), atol=5e-2, rtol=5e-2)) @parametrize.parametrize( - "device,dispatcher,ep_size", + "device,dispatcher,ep_size,expert_tp_size", [ - ("cuda", None, 1), - ("cuda", "all2all", 4), + ("cuda", None, 1, 1), + ("cuda", "all2all", 4, 1), + ("cuda", "all2all", 2, 2), ], ) - def test_save_hf(self, device, dispatcher, ep_size): + def test_save_hf(self, device, dispatcher, ep_size, expert_tp_size): self.create_pg(device) with torch.device("meta"): cfg = GptOss21BA3P6Config() cfg.dispatcher = dispatcher cfg.ep_size = ep_size + cfg.expert_tp_size = expert_tp_size gpt_oss_model = cfg.build()._to_device_dtype(dtype=torch.bfloat16, skip_buffers_dtype=True) fsdp_config = FSDPConfig( diff --git a/tests/model/test_moe_expert_tp_without_ep.py b/tests/model/test_moe_expert_tp_without_ep.py new file mode 100644 index 0000000000..94bcc2cc12 --- /dev/null +++ b/tests/model/test_moe_expert_tp_without_ep.py @@ -0,0 +1,72 @@ +import os +import unittest + +import torch +import torch.distributed as dist + +from xtuner._testing import DeterministicDDPTestCase +from xtuner.v1.module.attention import MHAConfig +from xtuner.v1.module.dispatcher.base import NaiveDispatcher +from xtuner.v1.module.router.greedy import GreedyRouterConfig +from xtuner.v1.model.moe.qwen3 import Qwen3MoEConfig + + +def _tiny_moe_cfg() -> Qwen3MoEConfig: + return Qwen3MoEConfig( + vocab_size=32, + max_position_embeddings=32, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + num_hidden_layers=1, + hidden_size=16, + intermediate_size=32, + rms_norm_eps=1e-6, + rope_theta=1e6, + hidden_act="silu", + attention=MHAConfig(num_attention_heads=2, num_key_value_heads=1, head_dim=8, qk_norm=True), + tie_word_embeddings=False, + n_routed_experts=4, + n_shared_experts=0, + num_experts_per_tok=2, + first_k_dense_replace=0, + hidden_factor=1.0, + moe_intermediate_size=8, + router=GreedyRouterConfig(scoring_func="softmax", norm_topk_prob=True, router_scaling_factor=1.0), + ep_size=1, + expert_tp_size=2, + dispatcher=None, + compile_cfg=False, + balancing_loss_cfg=None, + z_loss_cfg=None, + ) + + +@unittest.skipIf(not torch.cuda.is_available(), "CUDA/NCCL is required for real ExpertTP mesh validation.") +class TestMoEExpertTPWithoutEP(DeterministicDDPTestCase): + def test_builds_real_ep_ownership_mesh_for_expert_tp_without_ep(self) -> None: + pg = self.create_pg("cuda") + rank = dist.get_rank() + torch.cuda.set_device(rank % torch.cuda.device_count()) + + model = _tiny_moe_cfg().build() + layer = model.layers["0"] + + # 中文注释：不开 EP 但开启 expert TP 时，EP ownership 维度仍然真实存在，只是 size=1。 + assert model.ep_mesh is not None + assert model.expert_tp_mesh is not None + assert model.ep_mesh.size() == 1 + assert model.expert_tp_mesh.size() == 2 + assert model.expert_tp_mesh.mesh_dim_names == (f"{model.config.mesh_prefix}.etp",) + assert layer.experts.fused_w1w3.ep_size == 1 + assert layer.experts.fused_w1w3.tp_size == 2 + assert layer.experts.fused_w1w3.expert_tp_mesh is not None + assert layer.experts.fused_w1w3.expert_tp_mesh.mesh_dim_names == (f"{model.config.mesh_prefix}.etp",) + assert isinstance(layer.dispatcher, NaiveDispatcher) + + dist.barrier() + dist.destroy_process_group(pg) + + @property + def world_size(self) -> int: + return int(os.getenv("XTUNER_TEST_WORLD_SIZE", "2")) diff --git a/tests/model/test_qwen3_moe.py b/tests/model/test_qwen3_moe.py index a9e64446c1..179ec55cff 100644 --- a/tests/model/test_qwen3_moe.py +++ b/tests/model/test_qwen3_moe.py @@ -277,19 +277,22 @@ def test_sliding_windows(self, use_sliding_window, max_window_layers, sliding_wi assert "loss" in output @parametrize.parametrize( - "device,dispatcher,ep_size", + "device,dispatcher,ep_size,expert_tp_size", [ - ("cuda", None, 1), - ("cuda", "all2all", 4), - ("cuda", "all2all", 8), + ("cuda", None, 1, 1), + ("cuda", "all2all", 4, 1), + ("cuda", "all2all", 8, 1), + # 覆盖 post-FSDP 的 EP + expert TP HF load/save 路径。 + ("cuda", "all2all", 2, 2), ], ) - def test_save_hf(self, device, dispatcher, ep_size): + def test_save_hf(self, device, dispatcher, ep_size, expert_tp_size): self.create_pg(device) with torch.device("meta"): cfg = Qwen3MoE30BA3Config() cfg.dispatcher = dispatcher cfg.ep_size = ep_size + cfg.expert_tp_size = expert_tp_size qwen_model = cfg.build()._to_device_dtype(dtype=torch.bfloat16, skip_buffers_dtype=True) fsdp_config = FSDPConfig( diff --git a/tests/module/attention/test_dsa_mla.py b/tests/module/attention/test_dsa_mla.py index 289a6ec8c4..3f3fde1b79 100644 --- a/tests/module/attention/test_dsa_mla.py +++ b/tests/module/attention/test_dsa_mla.py @@ -236,74 +236,9 @@ def test_reentrant_checkpoint_reuses_and_releases_topk(self): assert seq_ctx.dsa_topk_cache.offloaded == {} -class TestAcceleratedSparseMLA: - @pytest.mark.skipif( - not _tilelang_sparse_mla_available(), - reason="requires CUDA and importable TileLang runtime", - ) - def test_tilelang_forward_backward_matches_torch(self): - # 验证 TileLang SparseMLA 的输出、LSE、dQ 和 dKV 与 PyTorch oracle 一致。 - q, kv, indices = _tilelang_sparse_mla_inputs() - scaling = 1 / math.sqrt(q.shape[-1]) - q_ref = q.detach().clone().requires_grad_() - kv_ref = kv.detach().clone().requires_grad_() - q_tilelang = q.detach().clone().requires_grad_() - kv_tilelang = kv.detach().clone().requires_grad_() - - expected = sparse_mla(q_ref, kv_ref, indices, scaling=scaling, value_dim=512, backend="torch") - actual = sparse_mla( - q_tilelang, - kv_tilelang, - indices.to(torch.int32), - scaling=scaling, - value_dim=512, - backend="tilelang", - ) - grad_output = torch.randn_like(expected.raw_output) - expected.raw_output.backward(grad_output) - actual.raw_output.backward(grad_output) - - torch.testing.assert_close(actual.raw_output, expected.raw_output, atol=BF16_ATOL, rtol=BF16_RTOL) - torch.testing.assert_close(actual.softmax_lse, expected.softmax_lse, atol=BF16_ATOL, rtol=BF16_RTOL) - torch.testing.assert_close(q_tilelang.grad, q_ref.grad, atol=BF16_ATOL, rtol=BF16_RTOL) - torch.testing.assert_close(kv_tilelang.grad, kv_ref.grad, atol=DKV_ATOL, rtol=DKV_RTOL) - - @pytest.mark.skipif( - not (_tilelang_sparse_mla_available() and _cudnn_dsa_sparse_mla_available()), - reason="requires CUDA, TileLang, and cuDNN DSA sparse attention backward", - ) - def test_compiled_cudnn_backward_matches_tilelang(self): - # 验证 torch.compile 下 cuDNN DSA 的输出与梯度仍和 TileLang oracle 一致。 - q, kv, indices = _cudnn_dsa_sparse_mla_inputs() - scaling = 1 / math.sqrt(q.shape[-1]) - - def compiled_sparse_mla(q: torch.Tensor, kv: torch.Tensor, backend: str) -> torch.Tensor: - return sparse_mla( - q, - kv, - indices, - scaling=scaling, - value_dim=512, - backend=backend, - ).raw_output - - compiled_sparse_mla = torch.compile(compiled_sparse_mla, fullgraph=False) - q_tilelang = q.detach().clone().requires_grad_() - kv_tilelang = kv.detach().clone().requires_grad_() - q_cudnn = q.detach().clone().requires_grad_() - kv_cudnn = kv.detach().clone().requires_grad_() - - expected = compiled_sparse_mla(q_tilelang, kv_tilelang, "tilelang") - actual = compiled_sparse_mla(q_cudnn, kv_cudnn, "cudnn_dsa") - grad_output = torch.randn_like(expected) - expected.backward(grad_output) - actual.backward(grad_output) - - torch.testing.assert_close(actual, expected, atol=BF16_ATOL, rtol=BF16_RTOL) - torch.testing.assert_close(q_cudnn.grad, q_tilelang.grad, atol=CUDNN_DQ_ATOL, rtol=CUDNN_DQ_RTOL) - torch.testing.assert_close(kv_cudnn.grad, kv_tilelang.grad, atol=DKV_ATOL, rtol=DKV_RTOL) - - +# The multiprocess cases must run before TileLang JIT is initialized in the +# pytest parent. With TileLang 0.1.11, spawning them afterwards crashes rank 0 +# during child-process bootstrap, before the indexer kernel is invoked. class TestDSASequenceParallel(DeterministicDDPTestCase): def test_packed_attention_matches_full_sequence(self): # 验证 SP2 packed attention 的输出、top-k 与输入梯度拼回后等同完整序列。 @@ -446,3 +381,71 @@ def test_cudnn_local_query_global_kv_matches_full_sequence(self): @property def world_size(self) -> int: return 2 + + +class TestAcceleratedSparseMLA: + @pytest.mark.skipif( + not _tilelang_sparse_mla_available(), + reason="requires CUDA and importable TileLang runtime", + ) + def test_tilelang_forward_backward_matches_torch(self): + # 验证 TileLang SparseMLA 的输出、LSE、dQ 和 dKV 与 PyTorch oracle 一致。 + q, kv, indices = _tilelang_sparse_mla_inputs() + scaling = 1 / math.sqrt(q.shape[-1]) + q_ref = q.detach().clone().requires_grad_() + kv_ref = kv.detach().clone().requires_grad_() + q_tilelang = q.detach().clone().requires_grad_() + kv_tilelang = kv.detach().clone().requires_grad_() + + expected = sparse_mla(q_ref, kv_ref, indices, scaling=scaling, value_dim=512, backend="torch") + actual = sparse_mla( + q_tilelang, + kv_tilelang, + indices.to(torch.int32), + scaling=scaling, + value_dim=512, + backend="tilelang", + ) + grad_output = torch.randn_like(expected.raw_output) + expected.raw_output.backward(grad_output) + actual.raw_output.backward(grad_output) + + torch.testing.assert_close(actual.raw_output, expected.raw_output, atol=BF16_ATOL, rtol=BF16_RTOL) + torch.testing.assert_close(actual.softmax_lse, expected.softmax_lse, atol=BF16_ATOL, rtol=BF16_RTOL) + torch.testing.assert_close(q_tilelang.grad, q_ref.grad, atol=BF16_ATOL, rtol=BF16_RTOL) + torch.testing.assert_close(kv_tilelang.grad, kv_ref.grad, atol=DKV_ATOL, rtol=DKV_RTOL) + + @pytest.mark.skipif( + not (_tilelang_sparse_mla_available() and _cudnn_dsa_sparse_mla_available()), + reason="requires CUDA, TileLang, and cuDNN DSA sparse attention backward", + ) + def test_compiled_cudnn_backward_matches_tilelang(self): + # 验证 torch.compile 下 cuDNN DSA 的输出与梯度仍和 TileLang oracle 一致。 + q, kv, indices = _cudnn_dsa_sparse_mla_inputs() + scaling = 1 / math.sqrt(q.shape[-1]) + + def compiled_sparse_mla(q: torch.Tensor, kv: torch.Tensor, backend: str) -> torch.Tensor: + return sparse_mla( + q, + kv, + indices, + scaling=scaling, + value_dim=512, + backend=backend, + ).raw_output + + compiled_sparse_mla = torch.compile(compiled_sparse_mla, fullgraph=False) + q_tilelang = q.detach().clone().requires_grad_() + kv_tilelang = kv.detach().clone().requires_grad_() + q_cudnn = q.detach().clone().requires_grad_() + kv_cudnn = kv.detach().clone().requires_grad_() + + expected = compiled_sparse_mla(q_tilelang, kv_tilelang, "tilelang") + actual = compiled_sparse_mla(q_cudnn, kv_cudnn, "cudnn_dsa") + grad_output = torch.randn_like(expected) + expected.backward(grad_output) + actual.backward(grad_output) + + torch.testing.assert_close(actual, expected, atol=BF16_ATOL, rtol=BF16_RTOL) + torch.testing.assert_close(q_cudnn.grad, q_tilelang.grad, atol=CUDNN_DQ_ATOL, rtol=CUDNN_DQ_RTOL) + torch.testing.assert_close(kv_cudnn.grad, kv_tilelang.grad, atol=DKV_ATOL, rtol=DKV_RTOL) diff --git a/tests/module/dispatcher/test_agrs_all2all.py b/tests/module/dispatcher/test_agrs_all2all.py index 61aa295093..1d0d68f530 100644 --- a/tests/module/dispatcher/test_agrs_all2all.py +++ b/tests/module/dispatcher/test_agrs_all2all.py @@ -79,6 +79,7 @@ def _dispatcher_call( pre_dispatched = dispatcher.dispatch_preprocess( hidden_states=hidden_states, topk_ids=topk_ids, + topk_weights=topk_weights, ) dispatched = dispatcher.dispatch( pre_dispatched=pre_dispatched, diff --git a/tests/module/dispatcher/test_deepep.py b/tests/module/dispatcher/test_deepep.py index 23c49ea499..4be7cec866 100644 --- a/tests/module/dispatcher/test_deepep.py +++ b/tests/module/dispatcher/test_deepep.py @@ -75,6 +75,7 @@ def _dispatcher_call( pre_dispatched = dispatcher.dispatch_preprocess( hidden_states=hidden_states, topk_ids=topk_ids, + topk_weights=topk_weights, async_op=async_op, ) dispatched = dispatcher.dispatch( diff --git a/tests/module/dispatcher/test_deepep_expert_tp.py b/tests/module/dispatcher/test_deepep_expert_tp.py new file mode 100644 index 0000000000..0fa3728ebf --- /dev/null +++ b/tests/module/dispatcher/test_deepep_expert_tp.py @@ -0,0 +1,310 @@ +import os +import unittest + +import torch +import torch.distributed as dist +from torch.testing._comparison import default_tolerances + +from xtuner._testing import DeterministicDDPTestCase +from xtuner.v1.module.dispatcher import build_dispatcher +from xtuner.v1.module.dispatcher.deepep import DeepEPDispatcher + + +BF16_RTOL, BF16_ATOL = default_tolerances(torch.bfloat16) +FLOAT32_RTOL, FLOAT32_ATOL = default_tolerances(torch.float32) + + +def _source_payload(rank: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + rows = rank + 2 + hidden_size = 128 + token_markers = torch.arange(rows, device=device, dtype=torch.float32) + rank * 10 + hidden = token_markers.unsqueeze(1) + torch.arange(hidden_size, device=device, dtype=torch.float32) / 100 + topk_ids = torch.tensor([0, 1, 2, 3], device=device, dtype=torch.int64).expand(rows, -1).contiguous() + slot_offsets = torch.tensor([0.1, 0.2, 0.3, 0.4], device=device, dtype=torch.float32) + topk_weights = token_markers.unsqueeze(1) / 1000 + slot_offsets + return hidden.to(torch.bfloat16), topk_ids, topk_weights + + +def _build_ep_tp_groups(ep_size: int, tp_size: int, backend: str = "nccl"): + ep_groups = [ + dist.new_group([ep_rank * tp_size + tp_rank for ep_rank in range(ep_size)], backend=backend) + for tp_rank in range(tp_size) + ] + tp_groups = [ + dist.new_group([ep_rank * tp_size + tp_rank for tp_rank in range(tp_size)], backend=backend) + for ep_rank in range(ep_size) + ] + return ep_groups, tp_groups + + +@unittest.skipIf(not torch.cuda.is_available(), "CUDA/NCCL is required for real DeepEP ExpertTP validation.") +class TestDeepEPExpertTPDispatcher(DeterministicDDPTestCase): + def test_sync_virtual_expert_path_preserves_output_and_gradients(self) -> None: + pg = self.create_pg("cuda") + rank = dist.get_rank() + torch.cuda.set_device(rank % torch.cuda.device_count()) + device = torch.device("cuda", rank % torch.cuda.device_count()) + + ep_size = 2 + tp_size = 2 + ep_rank = rank // tp_size + tp_rank = rank % tp_size + ep_groups, tp_groups = _build_ep_tp_groups(ep_size, tp_size) + ep_group = ep_groups[tp_rank] + tp_group = tp_groups[ep_rank] + + dispatcher = build_dispatcher( + dispatcher="deepep", + n_routed_experts=4, + ep_group=ep_group, + tp_group=tp_group, + ep_tp_group=dist.group.WORLD, + ) + assert isinstance(dispatcher, DeepEPDispatcher) + + local_hidden, local_topk_ids, local_topk_weights = _source_payload(rank, device) + hidden_leaf = local_hidden.detach().clone().requires_grad_(True) + topk_weights_leaf = local_topk_weights.detach().clone().requires_grad_(True) + + pre_dispatched = dispatcher.dispatch_preprocess( + hidden_states=hidden_leaf, + topk_ids=local_topk_ids, + topk_weights=topk_weights_leaf, + ) + expected_virtual_ids = torch.tensor( + [0, 2, 1, 3, 4, 6, 5, 7], + device=device, + dtype=torch.int64, + ).expand(local_topk_ids.shape[0], -1) + torch.testing.assert_close( + pre_dispatched["topk_ids"], + expected_virtual_ids, + ) + torch.testing.assert_close( + pre_dispatched["topk_weights"], + topk_weights_leaf.repeat_interleave(tp_size, dim=-1), + ) + + result = self._run_public_api( + dispatcher=dispatcher, + hidden_states=hidden_leaf, + topk_ids=local_topk_ids, + topk_weights=topk_weights_leaf, + tp_size=tp_size, + async_op=False, + ) + + expected_output = ( + hidden_leaf.detach().float() * topk_weights_leaf.detach().sum(dim=1, keepdim=True) + ).to(result["hidden_states"].dtype) + torch.testing.assert_close( + result["hidden_states"], + expected_output, + atol=BF16_ATOL, + rtol=BF16_RTOL, + ) + + result["hidden_states"].float().sum().backward() + assert hidden_leaf.grad is not None + assert topk_weights_leaf.grad is not None + expected_hidden_grad = topk_weights_leaf.detach().sum(dim=1, keepdim=True).expand_as(hidden_leaf) + expected_hidden_grad = expected_hidden_grad.to(hidden_leaf.grad.dtype) + expected_topk_grad = hidden_leaf.detach().float().sum(dim=1, keepdim=True).expand_as(topk_weights_leaf) + torch.testing.assert_close( + hidden_leaf.grad, + expected_hidden_grad, + atol=BF16_ATOL, + rtol=BF16_RTOL, + ) + torch.testing.assert_close( + topk_weights_leaf.grad, + expected_topk_grad, + atol=FLOAT32_ATOL, + rtol=FLOAT32_RTOL, + ) + + dist.barrier() + for group in ep_groups + tp_groups: + dist.destroy_process_group(group) + dist.destroy_process_group(pg) + + def test_async_path_matches_sync_output_and_gradients(self) -> None: + pg = self.create_pg("cuda") + rank = dist.get_rank() + torch.cuda.set_device(rank % torch.cuda.device_count()) + device = torch.device("cuda", rank % torch.cuda.device_count()) + + ep_size = 2 + tp_size = 2 + ep_rank = rank // tp_size + tp_rank = rank % tp_size + ep_groups, tp_groups = _build_ep_tp_groups(ep_size, tp_size) + ep_group = ep_groups[tp_rank] + tp_group = tp_groups[ep_rank] + + dispatcher = build_dispatcher( + dispatcher="deepep", + n_routed_experts=4, + ep_group=ep_group, + tp_group=tp_group, + ep_tp_group=dist.group.WORLD, + ) + + local_hidden, local_topk_ids, local_topk_weights = _source_payload(rank, device) + + sync_hidden_leaf = local_hidden.detach().clone().requires_grad_(True) + sync_topk_weights_leaf = local_topk_weights.detach().clone().requires_grad_(True) + sync_result = self._run_public_api( + dispatcher=dispatcher, + hidden_states=sync_hidden_leaf * 1.25, + topk_ids=local_topk_ids, + topk_weights=sync_topk_weights_leaf * 0.5, + tp_size=tp_size, + async_op=False, + ) + sync_result["hidden_states"].float().sum().backward() + + async_hidden_leaf = local_hidden.detach().clone().requires_grad_(True) + async_topk_weights_leaf = local_topk_weights.detach().clone().requires_grad_(True) + async_result = self._run_public_api( + dispatcher=dispatcher, + hidden_states=async_hidden_leaf * 1.25, + topk_ids=local_topk_ids, + topk_weights=async_topk_weights_leaf * 0.5, + tp_size=tp_size, + async_op=True, + ) + async_result["hidden_states"].float().sum().backward() + torch.cuda.synchronize() + + torch.testing.assert_close( + async_result["hidden_states"], + sync_result["hidden_states"], + atol=BF16_ATOL, + rtol=BF16_RTOL, + ) + assert sync_hidden_leaf.grad is not None + assert async_hidden_leaf.grad is not None + assert sync_topk_weights_leaf.grad is not None + assert async_topk_weights_leaf.grad is not None + torch.testing.assert_close( + async_hidden_leaf.grad, + sync_hidden_leaf.grad, + atol=BF16_ATOL, + rtol=BF16_RTOL, + ) + torch.testing.assert_close( + async_topk_weights_leaf.grad, + sync_topk_weights_leaf.grad, + atol=FLOAT32_ATOL, + rtol=FLOAT32_RTOL, + ) + + dist.barrier() + for group in ep_groups + tp_groups: + dist.destroy_process_group(group) + dist.destroy_process_group(pg) + + def test_async_path_accepts_topk_weights_without_gradients(self) -> None: + pg = self.create_pg("cuda") + rank = dist.get_rank() + torch.cuda.set_device(rank % torch.cuda.device_count()) + device = torch.device("cuda", rank % torch.cuda.device_count()) + + ep_size = 2 + tp_size = 2 + ep_rank = rank // tp_size + tp_rank = rank % tp_size + ep_groups, tp_groups = _build_ep_tp_groups(ep_size, tp_size) + ep_group = ep_groups[tp_rank] + tp_group = tp_groups[ep_rank] + + dispatcher = build_dispatcher( + dispatcher="deepep", + n_routed_experts=4, + ep_group=ep_group, + tp_group=tp_group, + ep_tp_group=dist.group.WORLD, + ) + + local_hidden, local_topk_ids, local_topk_weights = _source_payload(rank, device) + hidden_leaf = local_hidden.detach().clone().requires_grad_(True) + topk_weights = local_topk_weights.detach().clone() + assert topk_weights.requires_grad is False + + result = self._run_public_api( + dispatcher=dispatcher, + hidden_states=hidden_leaf, + topk_ids=local_topk_ids, + topk_weights=topk_weights, + tp_size=tp_size, + async_op=True, + ) + + assert result["hidden_states"].shape == local_hidden.shape + result["hidden_states"].float().sum().backward() + torch.cuda.synchronize() + assert hidden_leaf.grad is not None + + dist.barrier() + for group in ep_groups + tp_groups: + dist.destroy_process_group(group) + dist.destroy_process_group(pg) + + def _run_public_api( + self, + *, + dispatcher, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + tp_size: int, + async_op: bool, + ) -> dict[str, torch.Tensor]: + pre_dispatched = dispatcher.dispatch_preprocess( + hidden_states=hidden_states, + topk_ids=topk_ids, + topk_weights=topk_weights, + async_op=async_op, + ) + dispatched = dispatcher.dispatch( + pre_dispatched=pre_dispatched, + topk_weights=topk_weights, + decoding=False, + async_op=async_op, + ) + post_dispatched = dispatcher.dispatch_postprocess( + pre_dispatched=pre_dispatched, + dispatched=dispatched, + async_op=async_op, + ) + # 中文注释：测试 dispatcher public API，不模拟真实 row-parallel expert； + # 每个 ExpertTP rank 产出 1/tp_size partial，combine 应归约回完整输出。 + expert_output = post_dispatched["hidden_states"] / tp_size + pre_combined = dispatcher.combine_preprocess( + hidden_states=expert_output, + pre_dispatched=pre_dispatched, + dispatched=dispatched, + post_dispatched=post_dispatched, + async_op=async_op, + ) + combined = dispatcher.combine( + pre_dispatched=pre_dispatched, + dispatched=dispatched, + post_dispatched=post_dispatched, + pre_combined=pre_combined, + decoding=False, + async_op=async_op, + ) + return dispatcher.combine_postprocess( + pre_dispatched=pre_dispatched, + dispatched=dispatched, + post_dispatched=post_dispatched, + pre_combined=pre_combined, + combined=combined, + async_op=async_op, + ) + + @property + def world_size(self) -> int: + return int(os.getenv("XTUNER_TEST_WORLD_SIZE", "4")) diff --git a/tests/module/dispatcher/test_noep.py b/tests/module/dispatcher/test_noep.py index cb789ce79a..7790c96733 100644 --- a/tests/module/dispatcher/test_noep.py +++ b/tests/module/dispatcher/test_noep.py @@ -48,6 +48,7 @@ def test_dispatch_and_combine(self, dtype, device): pre_dispatched = self.dispatcher.dispatch_preprocess( hidden_states=hidden_states, topk_ids=topk_ids, + topk_weights=topk_weights, ) dispatched = self.dispatcher.dispatch( pre_dispatched=pre_dispatched, diff --git a/tests/module/dispatcher/test_noep_expert_tp.py b/tests/module/dispatcher/test_noep_expert_tp.py new file mode 100644 index 0000000000..415ca11965 --- /dev/null +++ b/tests/module/dispatcher/test_noep_expert_tp.py @@ -0,0 +1,310 @@ +import os +import unittest + +import torch +import torch.distributed as dist + +from xtuner._testing import DeterministicDDPTestCase +from xtuner.v1.module.dispatcher import build_dispatcher +from xtuner.v1.module.dispatcher.base import NaiveDispatcher + + +def _payload_for_rank(rank: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + rows = rank + 2 + hidden_size = 8 + start = sum(i + 2 for i in range(rank)) + token_ids = torch.arange(start, start + rows, device=device) + hidden = token_ids.to(torch.float32).unsqueeze(1) * 10 + torch.arange(hidden_size, device=device) + topk_ids = torch.stack((token_ids % 4, (token_ids + 1) % 4), dim=1).to(torch.int64) + topk_weights = torch.stack( + ( + torch.full((rows,), 1.0, device=device), + torch.full((rows,), 0.25 * (rank + 1), device=device), + ), + dim=1, + ) + return hidden, topk_ids, topk_weights + + +def _run_dispatcher( + dispatcher, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + expert_scale: float = 1.0, + async_op: bool = False, +): + pre_dispatched = dispatcher.dispatch_preprocess( + hidden_states=hidden_states, + topk_ids=topk_ids, + topk_weights=topk_weights, + async_op=async_op, + ) + dispatched = dispatcher.dispatch( + pre_dispatched=pre_dispatched, + topk_weights=topk_weights, + decoding=False, + async_op=async_op, + ) + post_dispatched = dispatcher.dispatch_postprocess( + pre_dispatched=pre_dispatched, + dispatched=dispatched, + async_op=async_op, + ) + # 中文注释：dispatcher 测试不跑真实 row-parallel expert； + # 每个 TP rank 提供 1/tp_size 的 partial output，真实 ReduceScatterRowsSum 后应回到 baseline。 + experts_results = post_dispatched["hidden_states"] * expert_scale + pre_combined = dispatcher.combine_preprocess( + hidden_states=experts_results, + pre_dispatched=pre_dispatched, + dispatched=dispatched, + post_dispatched=post_dispatched, + async_op=async_op, + ) + combined = dispatcher.combine( + pre_dispatched=pre_dispatched, + dispatched=dispatched, + post_dispatched=post_dispatched, + pre_combined=pre_combined, + decoding=False, + async_op=async_op, + ) + result = dispatcher.combine_postprocess( + pre_dispatched=pre_dispatched, + dispatched=dispatched, + post_dispatched=post_dispatched, + pre_combined=pre_combined, + combined=combined, + async_op=async_op, + ) + return result, dispatched, post_dispatched, pre_combined, combined + + +def _assert_cuda_event(value: torch.cuda.Event | None) -> None: + assert isinstance(value, torch.cuda.Event) + + +@unittest.skipIf(not torch.cuda.is_available(), "CUDA/NCCL is required for real ExpertTP dispatcher validation.") +class TestNaiveExpertTPDispatcher(DeterministicDDPTestCase): + def test_sync_path_uses_real_tp_collectives(self) -> None: + pg = self.create_pg("cuda") + rank = dist.get_rank() + world_size = dist.get_world_size() + torch.cuda.set_device(rank % torch.cuda.device_count()) + device = torch.device("cuda", rank % torch.cuda.device_count()) + + ep_groups = [dist.new_group([ep_rank], backend="nccl") for ep_rank in range(world_size)] + ep_group = ep_groups[rank] + + local_hidden, local_topk_ids, local_topk_weights = _payload_for_rank(rank, device) + full_payloads = [_payload_for_rank(tp_rank, device) for tp_rank in range(world_size)] + full_hidden = torch.cat([payload[0] for payload in full_payloads], dim=0) + full_topk_ids = torch.cat([payload[1] for payload in full_payloads], dim=0) + full_topk_weights = torch.cat([payload[2] for payload in full_payloads], dim=0) + + baseline = NaiveDispatcher(n_routed_experts=4) + baseline_result, _, baseline_post, _, _ = _run_dispatcher( + baseline, + full_hidden, + full_topk_ids, + full_topk_weights, + ) + + dispatcher = build_dispatcher( + dispatcher=None, + n_routed_experts=4, + ep_group=ep_group, + tp_group=dist.group.WORLD, + ) + result, dispatched, post_dispatched, pre_combined, combined = _run_dispatcher( + dispatcher, + local_hidden, + local_topk_ids, + local_topk_weights, + expert_scale=1.0 / world_size, + ) + + tp_rank_row_counts = [tp_rank + 2 for tp_rank in range(world_size)] + slice_start = sum(tp_rank_row_counts[:rank]) + slice_end = slice_start + tp_rank_row_counts[rank] + + torch.testing.assert_close(dispatched["hidden_states"], full_hidden) + torch.testing.assert_close(dispatched["topk_ids"], full_topk_ids) + torch.testing.assert_close(dispatched["topk_weights"], full_topk_weights) + torch.testing.assert_close(post_dispatched["tokens_per_expert"], baseline_post["tokens_per_expert"]) + torch.testing.assert_close(pre_combined["hidden_states"], baseline_result["hidden_states"] / world_size) + torch.testing.assert_close(combined["hidden_states"], baseline_result["hidden_states"][slice_start:slice_end]) + torch.testing.assert_close(result["hidden_states"], baseline_result["hidden_states"][slice_start:slice_end]) + + dist.barrier() + for group in ep_groups: + dist.destroy_process_group(group) + dist.destroy_process_group(pg) + + def test_async_path_exposes_events_at_stage_boundaries(self) -> None: + pg = self.create_pg("cuda") + rank = dist.get_rank() + world_size = dist.get_world_size() + torch.cuda.set_device(rank % torch.cuda.device_count()) + device = torch.device("cuda", rank % torch.cuda.device_count()) + + ep_groups = [dist.new_group([ep_rank], backend="nccl") for ep_rank in range(world_size)] + ep_group = ep_groups[rank] + dispatcher = build_dispatcher( + dispatcher=None, + n_routed_experts=4, + ep_group=ep_group, + tp_group=dist.group.WORLD, + ) + + local_hidden, local_topk_ids, local_topk_weights = _payload_for_rank(rank, device) + hidden_leaf = local_hidden.detach().clone().requires_grad_(True) + topk_weights_leaf = local_topk_weights.detach().clone().requires_grad_(True) + hidden = hidden_leaf * 1.25 + topk_weights = topk_weights_leaf * 0.5 + + pre_dispatched = dispatcher.dispatch_preprocess( + hidden_states=hidden, + topk_ids=local_topk_ids, + topk_weights=topk_weights, + async_op=True, + ) + _assert_cuda_event(pre_dispatched["forward_finished_event"]) + _assert_cuda_event(pre_dispatched["backward_previous_event"]) + + dispatched = dispatcher.dispatch( + pre_dispatched=pre_dispatched, + topk_weights=topk_weights, + decoding=False, + async_op=True, + ) + _assert_cuda_event(dispatched["forward_finished_event"]) + _assert_cuda_event(dispatched["backward_previous_event"]) + _assert_cuda_event(dispatched["topk_weights_backward_previous_event"]) + + # 中文注释：这里不手动 wait dispatch event，由 dispatch_postprocess 自己建立等待边界。 + post_dispatched = dispatcher.dispatch_postprocess( + pre_dispatched=pre_dispatched, + dispatched=dispatched, + async_op=True, + ) + + total_rows = sum(tp_rank + 2 for tp_rank in range(world_size)) + assert dispatched["hidden_states"].shape == (total_rows, local_hidden.shape[1]) + assert dispatched["topk_ids"].shape == (total_rows, local_topk_ids.shape[1]) + assert dispatched["topk_weights"].shape == (total_rows, local_topk_weights.shape[1]) + assert post_dispatched["hidden_states"].shape == ( + total_rows * local_topk_ids.shape[1], + local_hidden.shape[1], + ) + + experts_results = post_dispatched["hidden_states"] / world_size + pre_combined = dispatcher.combine_preprocess( + hidden_states=experts_results, + pre_dispatched=pre_dispatched, + dispatched=dispatched, + post_dispatched=post_dispatched, + async_op=True, + ) + _assert_cuda_event(pre_combined["forward_finished_event"]) + _assert_cuda_event(pre_combined["backward_previous_event"]) + assert pre_combined["hidden_states"].shape == (total_rows, local_hidden.shape[1]) + + combined = dispatcher.combine( + pre_dispatched=pre_dispatched, + dispatched=dispatched, + post_dispatched=post_dispatched, + pre_combined=pre_combined, + decoding=False, + async_op=True, + ) + _assert_cuda_event(combined["forward_finished_event"]) + _assert_cuda_event(combined["backward_previous_event"]) + assert combined["hidden_states"].shape == local_hidden.shape + + # 中文注释：这里同样不手动 wait combine event，由 combine_postprocess 返回本 rank source token slice。 + result = dispatcher.combine_postprocess( + pre_dispatched=pre_dispatched, + dispatched=dispatched, + post_dispatched=post_dispatched, + pre_combined=pre_combined, + combined=combined, + async_op=True, + ) + assert result["hidden_states"].shape == local_hidden.shape + + result["hidden_states"].square().sum().backward() + torch.cuda.synchronize() + assert hidden_leaf.grad is not None + assert topk_weights_leaf.grad is not None + + dist.barrier() + for group in ep_groups: + dist.destroy_process_group(group) + dist.destroy_process_group(pg) + + def test_async_sync_path_matches_output_and_gradients(self) -> None: + pg = self.create_pg("cuda") + rank = dist.get_rank() + world_size = dist.get_world_size() + torch.cuda.set_device(rank % torch.cuda.device_count()) + device = torch.device("cuda", rank % torch.cuda.device_count()) + + ep_groups = [dist.new_group([ep_rank], backend="nccl") for ep_rank in range(world_size)] + ep_group = ep_groups[rank] + dispatcher = build_dispatcher( + dispatcher=None, + n_routed_experts=4, + ep_group=ep_group, + tp_group=dist.group.WORLD, + ) + + local_hidden, local_topk_ids, local_topk_weights = _payload_for_rank(rank, device) + sync_hidden_leaf = local_hidden.detach().clone().requires_grad_(True) + sync_topk_weights_leaf = local_topk_weights.detach().clone().requires_grad_(True) + sync_hidden = sync_hidden_leaf * 1.25 + sync_topk_weights = sync_topk_weights_leaf * 0.5 + sync_result, *_ = _run_dispatcher( + dispatcher, + sync_hidden, + local_topk_ids, + sync_topk_weights, + expert_scale=1.0 / world_size, + async_op=False, + ) + sync_loss = sync_result["hidden_states"].square().sum() + sync_loss.backward() + torch.cuda.synchronize() + + async_hidden_leaf = local_hidden.detach().clone().requires_grad_(True) + async_topk_weights_leaf = local_topk_weights.detach().clone().requires_grad_(True) + async_hidden = async_hidden_leaf * 1.25 + async_topk_weights = async_topk_weights_leaf * 0.5 + async_result, *_ = _run_dispatcher( + dispatcher, + async_hidden, + local_topk_ids, + async_topk_weights, + expert_scale=1.0 / world_size, + async_op=True, + ) + async_loss = async_result["hidden_states"].square().sum() + async_loss.backward() + torch.cuda.synchronize() + + torch.testing.assert_close(async_result["hidden_states"], sync_result["hidden_states"]) + assert sync_hidden_leaf.grad is not None + assert async_hidden_leaf.grad is not None + assert sync_topk_weights_leaf.grad is not None + assert async_topk_weights_leaf.grad is not None + torch.testing.assert_close(async_hidden_leaf.grad, sync_hidden_leaf.grad) + torch.testing.assert_close(async_topk_weights_leaf.grad, sync_topk_weights_leaf.grad) + + dist.barrier() + for group in ep_groups: + dist.destroy_process_group(group) + dist.destroy_process_group(pg) + + @property + def world_size(self) -> int: + return int(os.getenv("XTUNER_TEST_WORLD_SIZE", "2")) diff --git a/tests/module/dispatcher/test_torch_all2all.py b/tests/module/dispatcher/test_torch_all2all.py index fe5c72f7f7..802542c450 100644 --- a/tests/module/dispatcher/test_torch_all2all.py +++ b/tests/module/dispatcher/test_torch_all2all.py @@ -65,6 +65,7 @@ def _dispatcher_call( pre_dispatched = dispatcher.dispatch_preprocess( hidden_states=hidden_states, topk_ids=topk_ids, + topk_weights=topk_weights, ) dispatched = dispatcher.dispatch( pre_dispatched=pre_dispatched, diff --git a/tests/module/dispatcher/test_torch_all2all_shared_expert_tp.py b/tests/module/dispatcher/test_torch_all2all_shared_expert_tp.py new file mode 100644 index 0000000000..db5528635f --- /dev/null +++ b/tests/module/dispatcher/test_torch_all2all_shared_expert_tp.py @@ -0,0 +1,262 @@ +import unittest + +import torch +import torch.distributed as dist + +from xtuner._testing import DeterministicDDPTestCase +from xtuner.v1.module.dispatcher import build_dispatcher +from xtuner.v1.module.dispatcher.base import DispacherInterface +from xtuner.v1.module.dispatcher.torch_all2all import TorchAll2AllDispatcher + + +def _build_ep_tp_groups( + ep_size: int, + tp_size: int, +) -> tuple[dist.ProcessGroup, dist.ProcessGroup, list[dist.ProcessGroup]]: + all_groups = [] + ep_groups = [] + tp_groups = [] + for tp_rank in range(tp_size): + group = dist.new_group([ep_rank * tp_size + tp_rank for ep_rank in range(ep_size)], backend="nccl") + ep_groups.append(group) + all_groups.append(group) + for ep_rank in range(ep_size): + group = dist.new_group([ep_rank * tp_size + tp_rank for tp_rank in range(tp_size)], backend="nccl") + tp_groups.append(group) + all_groups.append(group) + + rank = dist.get_rank() + return ep_groups[rank % tp_size], tp_groups[rank // tp_size], all_groups + + +def _payload_for_rank(rank: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + rows = rank + 2 + hidden_size = 8 + token_ids = torch.arange(sum(i + 2 for i in range(rank)), sum(i + 2 for i in range(rank + 1)), device=device) + hidden = token_ids.to(torch.float32).unsqueeze(1) * 10 + torch.arange(hidden_size, device=device) + topk_ids = torch.stack((token_ids % 4, (token_ids + 1) % 4), dim=1).to(torch.int64) + topk_weights = torch.stack( + ( + torch.full((rows,), 1.0, device=device), + torch.full((rows,), 0.2 * (rank + 1), device=device), + ), + dim=1, + ) + return hidden, topk_ids, topk_weights + + +def _run_dispatcher( + dispatcher: DispacherInterface, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + async_op: bool = False, +): + pre_dispatched = dispatcher.dispatch_preprocess( + hidden_states=hidden_states, + topk_ids=topk_ids, + topk_weights=topk_weights, + async_op=async_op, + ) + dispatched = dispatcher.dispatch( + pre_dispatched=pre_dispatched, + topk_weights=topk_weights, + decoding=False, + async_op=async_op, + ) + post_dispatched = dispatcher.dispatch_postprocess( + pre_dispatched=pre_dispatched, + dispatched=dispatched, + async_op=async_op, + ) + # 中文注释：dispatcher 级别不跑真实 row-parallel expert， + # 两个 TP rank 各提供一半 partial output。 + experts_results = post_dispatched["hidden_states"] / 2 + pre_combined = dispatcher.combine_preprocess( + hidden_states=experts_results, + pre_dispatched=pre_dispatched, + dispatched=dispatched, + post_dispatched=post_dispatched, + async_op=async_op, + ) + combined = dispatcher.combine( + pre_dispatched=pre_dispatched, + dispatched=dispatched, + post_dispatched=post_dispatched, + pre_combined=pre_combined, + decoding=False, + async_op=async_op, + ) + result = dispatcher.combine_postprocess( + pre_dispatched=pre_dispatched, + dispatched=dispatched, + post_dispatched=post_dispatched, + pre_combined=pre_combined, + combined=combined, + async_op=async_op, + ) + return result, dispatched, post_dispatched, pre_combined, combined + + +def _record_shared_expert_tp_stages(dispatcher: TorchAll2AllDispatcher) -> dict[str, list[str | int]]: + stages: dict[str, list[str | int]] = { + "async_op_true": [], + "async_all_gather_rows": [], + "async_all_gather_per_rank_metadata": [], + "async_reduce_scatter_rows_sum": [], + "comm_stream": [], + } + current_stage: list[str] = [] + expert_tp = dispatcher._expert_tp + assert expert_tp is not None + + for stage_name in ( + "dispatch_preprocess", + "dispatch", + "dispatch_postprocess", + "combine_preprocess", + "combine", + "combine_postprocess", + ): + original_stage = getattr(dispatcher, stage_name) + + def stage_wrapper(*args, _original_stage=original_stage, _stage_name=stage_name, **kwargs): + if kwargs.get("async_op", False): + stages["async_op_true"].append(_stage_name) + current_stage.append(_stage_name) + try: + return _original_stage(*args, **kwargs) + finally: + current_stage.pop() + + setattr(dispatcher, stage_name, stage_wrapper) + + for collective_name in ( + "async_all_gather_rows", + "async_all_gather_per_rank_metadata", + "async_reduce_scatter_rows_sum", + ): + original_collective = getattr(expert_tp, collective_name) + + def collective_wrapper( + *args, + _original_collective=original_collective, + _collective_name=collective_name, + **kwargs, + ): + stages[_collective_name].append(current_stage[-1] if current_stage else "") + stages["comm_stream"].append(kwargs["comm_stream"].cuda_stream) + return _original_collective(*args, **kwargs) + + setattr(expert_tp, collective_name, collective_wrapper) + + return stages + + +def _assert_shared_expert_tp_async_stages( + stages: dict[str, list[str | int]], + dispatcher: TorchAll2AllDispatcher, +) -> None: + assert set(stages["async_op_true"]) == { + "dispatch_preprocess", + "dispatch", + "dispatch_postprocess", + "combine_preprocess", + "combine", + "combine_postprocess", + } + assert stages["async_all_gather_rows"] == ["dispatch"] + assert stages["async_all_gather_per_rank_metadata"] == ["dispatch"] + assert stages["async_reduce_scatter_rows_sum"] == ["combine"] + assert set(stages["comm_stream"]) == {dispatcher._comm_stream.cuda_stream} + + +@unittest.skipIf(not torch.cuda.is_available(), "CUDA/NCCL is required for real All2All ExpertTP validation.") +class TestTorchAll2AllSharedExpertTP(DeterministicDDPTestCase): + def test_build_dispatcher_uses_shared_all2all_expert_tp(self) -> None: + pg = self.create_pg("cuda") + torch.cuda.set_device(dist.get_rank() % torch.cuda.device_count()) + ep_group, tp_group, all_groups = _build_ep_tp_groups(ep_size=2, tp_size=2) + + dispatcher = build_dispatcher( + dispatcher="all2all", + n_routed_experts=4, + ep_group=ep_group, + tp_group=tp_group, + ) + + assert isinstance(dispatcher, TorchAll2AllDispatcher) + assert dispatcher._expert_tp is not None + + dist.barrier() + for group in all_groups: + dist.destroy_process_group(group) + dist.destroy_process_group(pg) + + def test_async_shared_all2all_uses_dispatcher_comm_stream(self) -> None: + pg = self.create_pg("cuda") + rank = dist.get_rank() + torch.cuda.set_device(rank % torch.cuda.device_count()) + device = torch.device("cuda", rank % torch.cuda.device_count()) + ep_group, tp_group, all_groups = _build_ep_tp_groups(ep_size=2, tp_size=2) + + sync_dispatcher = build_dispatcher( + dispatcher="all2all", + n_routed_experts=4, + ep_group=ep_group, + tp_group=tp_group, + ) + async_dispatcher = build_dispatcher( + dispatcher="all2all", + n_routed_experts=4, + ep_group=ep_group, + tp_group=tp_group, + ) + assert isinstance(sync_dispatcher, TorchAll2AllDispatcher) + assert isinstance(async_dispatcher, TorchAll2AllDispatcher) + stages = _record_shared_expert_tp_stages(async_dispatcher) + + local_hidden, local_topk_ids, local_topk_weights = _payload_for_rank(rank, device) + sync_hidden_leaf = local_hidden.detach().clone().requires_grad_(True) + sync_topk_weights_leaf = local_topk_weights.detach().clone().requires_grad_(True) + sync_result, *_ = _run_dispatcher( + sync_dispatcher, + sync_hidden_leaf * 1.25, + local_topk_ids, + sync_topk_weights_leaf * 0.5, + ) + sync_result["hidden_states"].square().sum().backward() + + async_hidden_leaf = local_hidden.detach().clone().requires_grad_(True) + async_topk_weights_leaf = local_topk_weights.detach().clone().requires_grad_(True) + async_result, *_ = _run_dispatcher( + async_dispatcher, + async_hidden_leaf * 1.25, + local_topk_ids, + async_topk_weights_leaf * 0.5, + async_op=True, + ) + async_result["hidden_states"].square().sum().backward() + torch.cuda.synchronize() + + _assert_shared_expert_tp_async_stages(stages, async_dispatcher) + torch.testing.assert_close(async_result["hidden_states"], sync_result["hidden_states"]) + assert sync_hidden_leaf.grad is not None + assert async_hidden_leaf.grad is not None + assert sync_topk_weights_leaf.grad is not None + assert async_topk_weights_leaf.grad is not None + torch.testing.assert_close(async_hidden_leaf.grad, sync_hidden_leaf.grad) + torch.testing.assert_close(async_topk_weights_leaf.grad, sync_topk_weights_leaf.grad) + + dist.barrier() + for group in all_groups: + dist.destroy_process_group(group) + dist.destroy_process_group(pg) + + @property + def world_size(self) -> int: + return 4 + + @property + def destroy_pg_upon_exit(self) -> bool: + return False diff --git a/tests/patch/test_dcp_interleaved_planner.py b/tests/patch/test_dcp_interleaved_planner.py new file mode 100644 index 0000000000..fa6e869ce1 --- /dev/null +++ b/tests/patch/test_dcp_interleaved_planner.py @@ -0,0 +1,172 @@ +"""Regression tests for the InterleavedShard (tpep) DCP planners. + +These validate that ``InterleavedShardSavePlanner`` / ``InterleavedShardLoadPlanner`` round-trip +per-expert column-parallel fused MoE weights through DCP. DCP's default planner models each +DTensor as a single contiguous chunk and silently mis-maps an ``InterleavedShard`` local tensor +(which is several interleaved runs), so these params used to be dropped from DCP checkpoints. + +The layout is built with ``distribute_tensor`` on a 2D ``(ep, tp)`` mesh — the same +``(Shard, InterleavedShard)`` placement ``GroupedLinear`` produces — without ``fully_shard`` so the +test does not depend on FSDP2's support for the strided placement. + +Topology: world_size = ep * tp = 2 * 2 = 4. +""" + +from __future__ import annotations + +import shutil +import tempfile +from pathlib import Path + +import parametrize +import torch +import torch.distributed as dist +import torch.distributed.checkpoint as dcp +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.tensor import DTensor, Shard, distribute_tensor + +from xtuner._testing import DeterministicDDPTestCase +from xtuner.v1.patch import InterleavedShardLoadPlanner, InterleavedShardSavePlanner +from xtuner.v1.utils.interleaved_shard import InterleavedShard, compute_runs, reconstruct_full_tensor + + +NUM_EXPERTS = 4 +NUM_FUSED_PROJECTIONS = 2 # fused_w1w3 packs gate_proj + up_proj per expert. +PER_PROJ_OUT = 8 +IN_FEATURES = 6 +GLOBAL_ROWS = NUM_EXPERTS * NUM_FUSED_PROJECTIONS * PER_PROJ_OUT + + +def _global_source() -> torch.Tensor: + # Deterministic, rank-independent global tensor so every rank agrees on ground truth. + return torch.arange(GLOBAL_ROWS * IN_FEATURES, device="cuda", dtype=torch.float32).reshape( + GLOBAL_ROWS, IN_FEATURES + ) + + +def _build_interleaved_dtensor(ep_size: int, tp_size: int) -> DTensor: + # Build with ``from_local`` (not ``distribute_tensor``) to mirror ``GroupedLinear``: the latter + # goes through ``redistribute``, which crashes on ``(Shard, InterleavedShard)`` on torch >= 2.9. + mesh = init_device_mesh("cuda", (ep_size, tp_size), mesh_dim_names=("ep", "tp")) + local_experts = NUM_EXPERTS // ep_size + num_local_stripes = local_experts * NUM_FUSED_PROJECTIONS + placements = (Shard(0), InterleavedShard(0, num_local_stripes=num_local_stripes)) + g = _global_source() + # Scatter the deterministic global source into this rank's interleaved runs so every rank holds + # real, distinct data (from_local does not scatter — the caller must supply the local shard). + local = torch.empty(GLOBAL_ROWS // (ep_size * tp_size), IN_FEATURES, device="cuda") + dt = DTensor.from_local(local, mesh, placements, run_check=False) + for run in compute_runs(dt): + start = run.global_offset[0] + local[run.local_start : run.local_start + run.local_size] = g[start : start + run.local_size] + return dt + + +class TestDCPInterleavedPlanner(DeterministicDDPTestCase): + @parametrize.parametrize("device", [("cuda",)]) + def test_interleaved_round_trip(self, device: str) -> None: + """Save then load an InterleavedShard DTensor under the same topology; local + global match.""" + pg = self.create_pg(device) + + src = _build_interleaved_dtensor(ep_size=2, tp_size=2) + assert len(compute_runs(src)) > 1, "expected multiple interleaved runs per rank" + local_before = src._local_tensor.clone() + full_before = reconstruct_full_tensor(src).clone() + optimizer_moment = _build_interleaved_dtensor(ep_size=2, tp_size=2) + optimizer_moment._local_tensor.add_(1000) + moment_before = optimizer_moment._local_tensor.clone() + + dp_mesh = init_device_mesh("cuda", (self.world_size,), mesh_dim_names=("dp",)) + regular_global = torch.arange(32, device="cuda", dtype=torch.float32).reshape(8, 4) + regular = distribute_tensor(regular_global, dp_mesh, (Shard(0),)) + regular_before = regular._local_tensor.clone() + plain = torch.tensor([3.0, 5.0], device="cuda") + + tmp = [tempfile.mkdtemp()] if dist.get_rank() == 0 else [None] + dist.broadcast_object_list(tmp, src=0) + ckpt = Path(tmp[0]) + + state = { + "model": {"w": src, "regular": regular, "plain": plain}, + "optimizer": { + "state": {"w": {"exp_avg": optimizer_moment}}, + "param_groups": [{"params": ["w"], "lr": 1e-3}], + }, + } + dcp.save(state, checkpoint_id=ckpt, planner=InterleavedShardSavePlanner()) + dist.barrier() + + dst = _build_interleaved_dtensor(ep_size=2, tp_size=2) + dst._local_tensor.zero_() + dst_moment = _build_interleaved_dtensor(ep_size=2, tp_size=2) + dst_moment._local_tensor.zero_() + regular_dst = distribute_tensor(torch.zeros_like(regular_global), dp_mesh, (Shard(0),)) + plain_dst = torch.zeros_like(plain) + loaded_state = { + "model": {"w": dst, "regular": regular_dst, "plain": plain_dst}, + "optimizer": { + "state": {"w": {"exp_avg": dst_moment}}, + "param_groups": [{"params": ["w"], "lr": 0.0}], + }, + } + dcp.load(loaded_state, checkpoint_id=ckpt, planner=InterleavedShardLoadPlanner()) + dist.barrier() + + self.assertTrue(torch.equal(local_before, dst._local_tensor), "local shard mismatch after DCP round-trip") + self.assertTrue(torch.equal(full_before, reconstruct_full_tensor(dst)), "global tensor mismatch") + self.assertTrue(torch.equal(moment_before, dst_moment._local_tensor), "optimizer state mismatch") + self.assertTrue(torch.equal(regular_before, regular_dst._local_tensor), "regular DTensor mismatch") + self.assertTrue(torch.equal(plain, plain_dst), "plain tensor mismatch") + self.assertEqual(loaded_state["optimizer"]["param_groups"][0]["lr"], 1e-3) + + dist.barrier() + if dist.get_rank() == 0: + shutil.rmtree(ckpt) + torch.cuda.empty_cache() + try: + dist.destroy_process_group(pg) + except Exception: + pass + + @parametrize.parametrize("device", [("cuda",)]) + def test_interleaved_reshard_across_topology(self, device: str) -> None: + """Checkpoint saved at (ep=2, tp=2) reloads correctly at (ep=1, tp=4) — global-coordinate storage.""" + pg = self.create_pg(device) + + src = _build_interleaved_dtensor(ep_size=2, tp_size=2) + full_before = reconstruct_full_tensor(src).clone() + + tmp = [tempfile.mkdtemp()] if dist.get_rank() == 0 else [None] + dist.broadcast_object_list(tmp, src=0) + ckpt = Path(tmp[0]) + + dcp.save({"w": src}, checkpoint_id=ckpt, planner=InterleavedShardSavePlanner()) + dist.barrier() + + dst = _build_interleaved_dtensor(ep_size=1, tp_size=4) + dst._local_tensor.zero_() + dcp.load({"w": dst}, checkpoint_id=ckpt, planner=InterleavedShardLoadPlanner()) + dist.barrier() + + self.assertTrue( + torch.equal(full_before, reconstruct_full_tensor(dst)), + "global tensor mismatch after resharding (ep=2,tp=2) -> (ep=1,tp=4)", + ) + + dist.barrier() + if dist.get_rank() == 0: + shutil.rmtree(ckpt) + torch.cuda.empty_cache() + try: + dist.destroy_process_group(pg) + except Exception: + pass + + @property + def world_size(self) -> int: + # (ep, tp) = (2, 2) → 4 GPUs. + return 4 + + @property + def destroy_pg_upon_exit(self) -> bool: + return False diff --git a/tests/rl/test_weight_iterator.py b/tests/rl/test_weight_iterator.py new file mode 100644 index 0000000000..57b3b6fc73 --- /dev/null +++ b/tests/rl/test_weight_iterator.py @@ -0,0 +1,49 @@ +from types import SimpleNamespace +from typing import Any, cast + +import torch +from torch import nn + +from xtuner.v1.model.base import BaseModel, HFSaveCfg, XTunerBaseModelConfig +from xtuner.v1.rl.weight_update.data import RolloutWeightUpdateInfo +from xtuner.v1.rl.weight_update.weight_iterator import WeightIterator +from xtuner.v1.utils import get_device + + +class MixedDtypeModel(BaseModel): + def __init__(self) -> None: + config = XTunerBaseModelConfig( + hf_save_cfg=HFSaveCfg(fp32_keys_pattern=[r"fp32_weight"]), + ) + super().__init__(config) + self.bf16_weight = nn.Parameter(torch.ones(2, device=get_device(), dtype=torch.bfloat16)) + self.fp32_weight = nn.Parameter(torch.ones(2, device=get_device(), dtype=torch.float32)) + self._init_load_spec() + + def to_hf_key_list(self, key: str) -> list[str]: + return [key] + + +def test_hf_weight_update_batches_have_one_dtype() -> None: + model = MixedDtypeModel() + rollout_info = RolloutWeightUpdateInfo( + rollout_config=cast(Any, SimpleNamespace()), + weight_update_targets=(), + train_rank=0, + transport_type="ipc", + backend="pytorch", + ) + iterator = WeightIterator( + config=SimpleNamespace(update_weight_bucket_size_in_gb=1, model_cfg=None), + engine=SimpleNamespace(model=model), + rollout_info=rollout_info, + global_hf_keys_mapping_cache={}, + ) + + batches = list(iterator.iter_hf_batches()) + + assert all(len({tensor.dtype for tensor in batch.state_dict.values()}) == 1 for batch in batches) + state_dict = {name: tensor for batch in batches for name, tensor in batch.state_dict.items()} + assert set(state_dict) == {"bf16_weight", "fp32_weight"} + assert state_dict["bf16_weight"].dtype == torch.bfloat16 + assert state_dict["fp32_weight"].dtype == torch.float32 diff --git a/tests/utils/test_compile.py b/tests/utils/test_compile.py index 567a5113a0..c91a3450f7 100644 --- a/tests/utils/test_compile.py +++ b/tests/utils/test_compile.py @@ -1,11 +1,32 @@ -from xtuner.v1.model import Qwen3Dense8BConfig, Qwen3MoE30BA3Config, Qwen3VLMoE30BA3Config, GptOss21BA3P6Config, DeepSeekV3Config, InternVL3P5Dense1BConfig, XTunerBaseModelConfig -import torch -from xtuner.v1.utils import get_logger -from xtuner._testing.utils import LogCapture from ast import literal_eval +import re + import pytest +import torch -import re +from xtuner._testing.utils import LogCapture +from xtuner.v1.model import ( + DeepSeekV3Config, + GptOss21BA3P6Config, + InternVL3P5Dense1BConfig, + Qwen3Dense8BConfig, + Qwen3MoE30BA3Config, + Qwen3VLMoE30BA3Config, +) +from xtuner.v1.model.moe.moe import MOE_EP_COMPILE_CFG, MOE_NON_EP_COMPILE_CFG, MoE +from xtuner.v1.model.moe.qwen3_5_text import ( + MOE_EP_COMPILE_CFG as QWEN35_MOE_EP_COMPILE_CFG, + MOE_NON_EP_COMPILE_CFG as QWEN35_MOE_NON_EP_COMPILE_CFG, + Qwen3_5_VLTextMoE, + Qwen3_5_VLTextMoE35BA3BConfig, +) +from xtuner.v1.module.dispatcher.base import ( + NaiveCombineResult, + NaiveDispatchResult, + NaivePreCombineResult, + NaivePreDispatchResult, +) +from xtuner.v1.utils import get_logger logger = get_logger() @@ -60,3 +81,41 @@ def test_compile_model_exception(): with pytest.raises(Exception): with torch.device("meta"): Qwen3Dense8BConfig(compile_cfg={"xtuner.v1.module.decoder_layer.moe_decoder_layer.MoEBlock.fuck": {}}).build() + + +@pytest.mark.parametrize( + "ep_size,expert_tp_size,expected_compile_cfg", + [ + (1, 1, MOE_NON_EP_COMPILE_CFG), + (2, 1, MOE_EP_COMPILE_CFG), + (1, 2, MOE_EP_COMPILE_CFG), + (2, 2, MOE_EP_COMPILE_CFG), + ], +) +def test_moe_compile_cfg_treats_expert_tp_like_ep(ep_size, expert_tp_size, expected_compile_cfg): + model = object.__new__(MoE) + model.config = Qwen3MoE30BA3Config(ep_size=ep_size, expert_tp_size=expert_tp_size) + assert model.default_compile_cfg == expected_compile_cfg + + +@pytest.mark.parametrize( + "ep_size,expert_tp_size,expected_compile_cfg", + [ + (1, 1, QWEN35_MOE_NON_EP_COMPILE_CFG), + (2, 1, QWEN35_MOE_EP_COMPILE_CFG), + (1, 2, QWEN35_MOE_EP_COMPILE_CFG), + (2, 2, QWEN35_MOE_EP_COMPILE_CFG), + ], +) +def test_qwen35_moe_compile_cfg_treats_expert_tp_like_ep(ep_size, expert_tp_size, expected_compile_cfg): + model = object.__new__(Qwen3_5_VLTextMoE) + model.config = Qwen3_5_VLTextMoE35BA3BConfig(ep_size=ep_size, expert_tp_size=expert_tp_size) + assert model.default_compile_cfg == expected_compile_cfg + + +def test_naive_dispatcher_compile_result_typeddicts_have_no_optional_keys(): + # 中文注释：non-EP 默认会 compile MoEDecoderLayer.forward，Dynamo 不支持 optional-key TypedDict。 + assert NaivePreDispatchResult.__optional_keys__ == frozenset() + assert NaiveDispatchResult.__optional_keys__ == frozenset() + assert NaivePreCombineResult.__optional_keys__ == frozenset() + assert NaiveCombineResult.__optional_keys__ == frozenset() diff --git a/tests/utils/test_interleaved_shard.py b/tests/utils/test_interleaved_shard.py new file mode 100644 index 0000000000..57ce140cac --- /dev/null +++ b/tests/utils/test_interleaved_shard.py @@ -0,0 +1,198 @@ +"""Distributed behavior tests for ``xtuner.v1.utils.interleaved_shard``. + +The 2D cases cover the ``(Shard, InterleavedShard)`` layout produced by +``GroupedLinear``. The 3D case prepends FSDP, matching the runtime layout seen +by HF save/load. +""" + +from __future__ import annotations + +import shutil +import tempfile + +import torch +import torch.distributed as dist +import torch.nn as nn +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard +from torch.distributed.tensor import DTensor, Shard, distribute_tensor + +from xtuner._testing import DeterministicDDPTestCase +from xtuner.v1.utils.interleaved_shard import ( + InterleavedShard, + has_interleaved_placement, + reconstruct_full_tensor, +) + + +NUM_EXPERTS = 4 +OUT_PER_EXPERT = 4 +IN_FEATURES = 8 +GLOBAL_ROWS = NUM_EXPERTS * OUT_PER_EXPERT + + +def _build_expected_local( + global_tensor: torch.Tensor, + ep_rank: int, + tp_rank: int, + ep_size: int, + tp_size: int, +) -> torch.Tensor: + """Return the hand-computed per-expert column-parallel slice.""" + experts_per_ep = NUM_EXPERTS // ep_size + rows_per_expert = global_tensor.shape[0] // NUM_EXPERTS + rows_per_tp_per_expert = rows_per_expert // tp_size + chunks = [] + for local_expert in range(experts_per_ep): + global_expert = ep_rank * experts_per_ep + local_expert + expert_start = global_expert * rows_per_expert + row_start = expert_start + tp_rank * rows_per_tp_per_expert + chunks.append(global_tensor[row_start : row_start + rows_per_tp_per_expert]) + return torch.cat(chunks, dim=0) + + +class TestInterleavedShard2D(DeterministicDDPTestCase): + def test_layout_and_reconstruct(self) -> None: + self.create_pg("cuda") + mesh = init_device_mesh("cuda", (2, 2), mesh_dim_names=("ep", "tp")) + ep_rank = mesh.get_local_rank("ep") + tp_rank = mesh.get_local_rank("tp") + + global_tensor = torch.arange( + GLOBAL_ROWS * IN_FEATURES, + device="cuda", + dtype=torch.float32, + ).reshape(GLOBAL_ROWS, IN_FEATURES) + dist.broadcast(global_tensor, src=0) + + placements = (Shard(0), InterleavedShard(0, num_local_stripes=NUM_EXPERTS // 2)) + tensor = distribute_tensor(global_tensor, mesh, placements) + + expected = _build_expected_local(global_tensor, ep_rank, tp_rank, 2, 2) + torch.testing.assert_close(tensor.to_local(), expected) + assert has_interleaved_placement(tensor) + # ``shard_order`` only exists on torch >= 2.10; no explicit order is + # needed for the standard EP/ETP layout on either version. + assert getattr(tensor._spec, "shard_order", None) is None + torch.testing.assert_close(reconstruct_full_tensor(tensor), global_tensor) + + def test_hf_round_trip(self) -> None: + """Exercise the public BaseModel HF save/load path.""" + from transformers import PretrainedConfig + + from xtuner.v1.model.base import BaseModel, XTunerBaseModelConfig + + class _ToyConfig(XTunerBaseModelConfig): + @property + def hf_config(self) -> PretrainedConfig: + return PretrainedConfig() + + class _ToyModel(BaseModel): + def __init__(self, weight: DTensor): + super().__init__(_ToyConfig()) + self.weight = nn.Parameter(weight) + self._init_load_spec() + + def to_hf_key_list(self, key: str) -> list[str]: + return [key] + + self.create_pg("cuda") + mesh = init_device_mesh("cuda", (2, 2), mesh_dim_names=("ep", "tp")) + placements = (Shard(0), InterleavedShard(0, num_local_stripes=NUM_EXPERTS // 2)) + global_weight = torch.arange( + GLOBAL_ROWS * IN_FEATURES, + device="cuda", + dtype=torch.bfloat16, + ).reshape(GLOBAL_ROWS, IN_FEATURES) + dist.broadcast(global_weight, src=0) + model = _ToyModel(distribute_tensor(global_weight, mesh, placements)) + + checkpoint_dirs = [tempfile.mkdtemp() if dist.get_rank() == 0 else None] + dist.broadcast_object_list(checkpoint_dirs, src=0) + checkpoint_dir = checkpoint_dirs[0] + assert checkpoint_dir is not None + + try: + model.save_hf(checkpoint_dir) + restored_weight = distribute_tensor(torch.zeros_like(global_weight), mesh, placements) + restored = _ToyModel(restored_weight) + restored.from_hf(checkpoint_dir) + torch.testing.assert_close(restored.weight.to_local(), model.weight.to_local(), rtol=0, atol=0) + finally: + dist.barrier() + if dist.get_rank() == 0: + shutil.rmtree(checkpoint_dir) + + @property + def world_size(self) -> int: + return 4 + + +class _ToyGroupedLinear(nn.Module): + def __init__(self, weight: DTensor): + super().__init__() + self.weight = nn.Parameter(weight) + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + weight = self.weight.to_local() if isinstance(self.weight, DTensor) else self.weight + return torch.nn.functional.linear(inputs, weight) + + +class TestInterleavedShardPostFSDP(DeterministicDDPTestCase): + def test_reconstruct_and_load(self) -> None: + self.create_pg("cuda") + mesh = init_device_mesh("cuda", (2, 2, 2), mesh_dim_names=("fsdp", "ep", "tp")) + ep_tp_mesh = mesh["ep", "tp"] + fsdp_mesh = mesh["fsdp"] + + global_tensor = torch.arange( + GLOBAL_ROWS * IN_FEATURES, + device="cuda", + dtype=torch.float32, + ).reshape(GLOBAL_ROWS, IN_FEATURES) + dist.broadcast(global_tensor, src=0) + + placements = (Shard(0), InterleavedShard(0, num_local_stripes=NUM_EXPERTS // 2)) + tensor = distribute_tensor(global_tensor, ep_tp_mesh, placements) + model = _ToyGroupedLinear(tensor).cuda() + fully_shard( + model, + mesh=fsdp_mesh, + mp_policy=MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32), + reshard_after_forward=True, + ) + + inputs = torch.randn(6, IN_FEATURES, device="cuda", dtype=torch.bfloat16) + dist.broadcast(inputs, src=0) + output = model(inputs) + expected_local = _build_expected_local( + global_tensor, + mesh.get_local_rank("ep"), + mesh.get_local_rank("tp"), + 2, + 2, + ).to(torch.bfloat16) + expected_output = torch.nn.functional.linear(inputs, expected_local) + torch.testing.assert_close(output.detach(), expected_output, atol=1e-2, rtol=1e-2) + output.sum().backward() + + assert has_interleaved_placement(model.weight) + torch.testing.assert_close(reconstruct_full_tensor(model.weight), global_tensor) + + # LoadSpec compiles the same post-FSDP layout into source-to-local copy + # runs, so loading does not need model-specific Expert TP branches. + from xtuner.v1.utils.load_spec import LoadSpec + + local = model.weight._local_tensor + loaded_local = torch.empty_like(local, dtype=global_tensor.dtype) + load_spec = LoadSpec.from_tensor(name="weight", hf_keys=["weight"], tensor=model.weight) + load_spec.plan_hf_load().load_into( + [global_tensor], + loaded_local, + lambda _, checkpoint_tensor: checkpoint_tensor, + ) + torch.testing.assert_close(loaded_local, local.to(global_tensor.dtype)) + + @property + def world_size(self) -> int: + return 8 diff --git a/tests/utils/test_load_spec.py b/tests/utils/test_load_spec.py new file mode 100644 index 0000000000..cf2c0e9cc5 --- /dev/null +++ b/tests/utils/test_load_spec.py @@ -0,0 +1,708 @@ +import os +from collections.abc import Callable + +import pytest +import torch +import torch.distributed as dist +from pydantic import ValidationError +from torch.distributed.device_mesh import DeviceMesh +from torch.distributed.tensor import Shard as DTensorShard +from torch.distributed.tensor import distribute_tensor +from torch.distributed.tensor.placement_types import _StridedShard + +from xtuner.v1.model.base import BaseModel, XTunerBaseModelConfig +from xtuner.v1.utils import load_spec as load_spec_module +from xtuner.v1.utils.load_spec import LoadSpec, ShardDescriptor, unshard_tensors_for_hf_save + + +@pytest.fixture(scope="module") +def single_rank_group() -> dist.ProcessGroup: + if not dist.is_initialized(): + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29555") + dist.init_process_group(backend="gloo", rank=0, world_size=1) + group = dist.group.WORLD + assert group is not None + return group + + +@pytest.fixture(scope="module") +def local_groups(single_rank_group: dist.ProcessGroup): + groups = tuple(dist.new_group([0]) for _ in range(3)) + yield groups + for group in groups: + dist.destroy_process_group(group) + + +@pytest.fixture +def set_shard_rank(monkeypatch: pytest.MonkeyPatch) -> Callable[[int, int], None]: + def configure(world_size: int, rank: int) -> None: + monkeypatch.setattr(dist, "get_world_size", lambda group=None: world_size) + monkeypatch.setattr(dist, "get_rank", lambda group=None: rank) + + return configure + + +def set_group_layout( + monkeypatch: pytest.MonkeyPatch, + layouts: dict[dist.ProcessGroup, tuple[int, int, list[int]]], +) -> None: + monkeypatch.setattr( + dist, + "get_world_size", + lambda group=None: layouts[group][0] if group in layouts else 1, + ) + monkeypatch.setattr( + dist, + "get_rank", + lambda group=None: layouts[group][1] if group in layouts else 0, + ) + monkeypatch.setattr( + dist, + "get_process_group_ranks", + lambda group: layouts[group][2], + ) + + +class TestLoadSpecSchema: + def test_same_unsharded_spec(self) -> None: + spec = LoadSpec( + name="layers.0.mlp.gate.weight", + global_hf_keys=["model.layers.0.mlp.gate.weight"], + global_shape=(128, 64), + ) + + assert spec.is_fused is False + assert spec.is_sharded is False + assert spec.fused_dim is None + assert spec.shards == [] + assert spec.origin_shape is None + assert spec.unpadded_global_shape == spec.global_shape + + def test_from_tensor_derives_plain_tensor_layout(self) -> None: + spec = LoadSpec.from_tensor( + name="layers.0.experts.fused_w1w3.weight", + hf_keys=["k0", "k1"], + tensor=torch.empty(128, 64), + origin_shape=(120, 64), + ) + + assert spec.global_hf_keys == ["k0", "k1"] + assert spec.global_shape == (128, 64) + assert spec.fused_dim == 0 + assert spec.shards == [] + assert spec.origin_shape == (120, 64) + + def test_from_tensor_derives_dtensor_shards(self, single_rank_group: dist.ProcessGroup) -> None: + mesh = DeviceMesh("cpu", [0]) + tensor = distribute_tensor(torch.empty(128, 64), mesh, [DTensorShard(0)]) + + spec = LoadSpec.from_tensor(name="layers.0.mlp.gate.weight", hf_keys=["gate"], tensor=tensor) + + assert spec.global_hf_keys == ["gate"] + assert spec.global_shape == (128, 64) + assert spec.fused_dim is None + assert [(shard.dim, shard.interleave_factor) for shard in spec.shards] == [(0, 1)] + + def test_unordered_strided_placement_becomes_interleave_descriptor( + self, + single_rank_group: dist.ProcessGroup, + ) -> None: + class FakeDeviceMesh: + shape = (2, 2) + + def size(self, mesh_dim: int) -> int: + return self.shape[mesh_dim] + + def get_group(self, mesh_dim: int) -> dist.ProcessGroup: + return single_rank_group + + class FakeDTensor: + placements = (DTensorShard(0), _StridedShard(0, split_factor=2)) + device_mesh = FakeDeviceMesh() + + shards = load_spec_module._dtensor_shards(FakeDTensor()) # type: ignore[arg-type] + + assert [(shard.dim, shard.interleave_factor) for shard in shards] == [(0, 1), (0, 2)] + + def test_fused_spec_requires_fused_dim(self) -> None: + with pytest.raises(ValidationError, match="fused_dim"): + LoadSpec( + name="layers.0.mlp.fused_w1w3.weight", + global_hf_keys=["gate", "up"], + global_shape=(256, 64), + ) + + def test_descriptor_order_is_preserved(self, single_rank_group: dist.ProcessGroup) -> None: + spec = LoadSpec( + name="layers.0.experts.fused_w1w3.weight", + global_hf_keys=["gate", "up"], + global_shape=(256, 64), + fused_dim=0, + shards=[ + ShardDescriptor(dim=0, group=single_rank_group), + ShardDescriptor(dim=0, group=single_rank_group, interleave_factor=2), + ], + ) + + assert [shard.interleave_factor for shard in spec.shards] == [1, 2] + assert spec.is_fused is True + assert spec.is_sharded is True + + def test_uneven_interleave_is_rejected( + self, + single_rank_group: dist.ProcessGroup, + set_shard_rank: Callable[[int, int], None], + ) -> None: + set_shard_rank(2, 0) + + with pytest.raises(NotImplementedError, match="Even interleave requires"): + LoadSpec( + name="layers.0.experts.fused_w1w3.weight", + global_hf_keys=["gate"], + global_shape=(10, 4), + shards=[ShardDescriptor(dim=0, group=single_rank_group, interleave_factor=3)], + ) + + def test_zero_size_continuous_shard_is_valid( + self, + single_rank_group: dist.ProcessGroup, + set_shard_rank: Callable[[int, int], None], + ) -> None: + set_shard_rank(2, 1) + spec = LoadSpec( + name="embeddings.cls_embedding", + global_hf_keys=["embeddings.cls_embedding"], + global_shape=(1, 1, 1024), + shards=[ShardDescriptor(dim=0, group=single_rank_group)], + ) + + plan = spec.plan_hf_load() + + assert plan.hf_keys == [] + target = torch.empty(0, 1, 1024) + plan.load_into([], target, lambda _, tensor: tensor) + assert target.numel() == 0 + + +class TestHFLoadPlan: + def test_interleave_selects_smallest_hf_key_envelope( + self, + single_rank_group: dist.ProcessGroup, + set_shard_rank: Callable[[int, int], None], + ) -> None: + set_shard_rank(2, 0) + spec = LoadSpec( + name="layers.0.experts.fused_w1w3.weight", + global_hf_keys=["k0", "k1", "k2", "k3"], + global_shape=(400, 8), + fused_dim=0, + shards=[ShardDescriptor(dim=0, group=single_rank_group, interleave_factor=2)], + ) + + plan = spec.plan_hf_load() + + assert plan.hf_keys == ["k0", "k1", "k2"] + full = torch.arange(400 * 8).reshape(400, 8) + target = torch.empty(200, 8, dtype=full.dtype) + plan.load_into([full[:100], full[100:200], full[200:300]], target, lambda _, tensor: tensor) + torch.testing.assert_close(target, torch.cat((full[:100], full[200:300]))) + + def test_non_fused_continuous_shard( + self, + single_rank_group: dist.ProcessGroup, + set_shard_rank: Callable[[int, int], None], + ) -> None: + set_shard_rank(2, 1) + spec = LoadSpec( + name="layers.0.self_attn.q_proj.weight", + global_hf_keys=["q_proj"], + global_shape=(128, 256), + shards=[ShardDescriptor(dim=1, group=single_rank_group)], + ) + + plan = spec.plan_hf_load() + full = torch.arange(128 * 256).reshape(128, 256) + target = torch.empty(128, 128, dtype=full.dtype) + plan.load_into([full], target, lambda _, tensor: tensor) + + torch.testing.assert_close(target, full[:, 128:]) + + def test_origin_shape_clips_runtime_padding( + self, + single_rank_group: dist.ProcessGroup, + set_shard_rank: Callable[[int, int], None], + ) -> None: + set_shard_rank(4, 3) + spec = LoadSpec( + name="layers.0.experts.fused_w1w3.weight", + global_hf_keys=["k0", "k1", "k2", "k3"], + global_shape=(480, 8), + fused_dim=0, + shards=[ShardDescriptor(dim=0, group=single_rank_group)], + origin_shape=(400, 8), + ) + + plan = spec.plan_hf_load() + + assert plan.hf_keys == ["k3"] + full = torch.arange(400 * 8).reshape(400, 8) + target = torch.full((120, 8), -1) + plan.load_into([full[300:400]], target, lambda _, tensor: tensor) + torch.testing.assert_close(target[:40], full[360:400]) + torch.testing.assert_close(target[40:], torch.zeros_like(target[40:])) + + def test_origin_shape_zeroes_pad_only_rank( + self, + single_rank_group: dist.ProcessGroup, + set_shard_rank: Callable[[int, int], None], + ) -> None: + set_shard_rank(8, 7) + spec = LoadSpec( + name="layers.0.experts.fused_w1w3.weight", + global_hf_keys=["k0", "k1", "k2", "k3"], + global_shape=(480, 8), + fused_dim=0, + shards=[ShardDescriptor(dim=0, group=single_rank_group)], + origin_shape=(400, 8), + ) + + plan = spec.plan_hf_load() + + assert plan.hf_keys == [] + target = torch.ones(60, 8) + plan.load_into([], target, lambda _, tensor: tensor) + torch.testing.assert_close(target, torch.zeros_like(target)) + + def test_model_canonicalization_runs_before_interleave_copy( + self, + single_rank_group: dist.ProcessGroup, + set_shard_rank: Callable[[int, int], None], + ) -> None: + set_shard_rank(2, 0) + spec = LoadSpec( + name="layers.0.experts.fused_w1w3.weight", + global_hf_keys=["packed_gate_up"], + global_shape=(8, 2), + shards=[ShardDescriptor(dim=0, group=single_rank_group, interleave_factor=4)], + ) + packed = torch.arange(16).reshape(2, 2, 4) + canonical = packed.permute(0, 2, 1).reshape(8, 2) + target = torch.empty(4, 2, dtype=packed.dtype) + + spec.plan_hf_load().load_into( + [packed], + target, + lambda name, tensor: tensor.permute(0, 2, 1).reshape(8, 2), + ) + + torch.testing.assert_close(target, canonical[(0, 2, 4, 6), :]) + + def test_qwen35_packed_adapter_runs_before_expert_tp_copy( + self, + single_rank_group: dist.ProcessGroup, + set_shard_rank: Callable[[int, int], None], + ) -> None: + from xtuner.v1.model.moe.qwen3_5_text import Qwen3_5_VLTextMoE + + set_shard_rank(2, 0) + model = object.__new__(Qwen3_5_VLTextMoE) + packed = torch.arange(2 * 4 * 3).reshape(2, 4, 3) + canonical = packed.flatten(0, 1) + spec = LoadSpec( + name="layers.0.experts.fused_w1w3.weight", + global_hf_keys=["model.layers.0.mlp.experts.gate_up_proj"], + global_shape=tuple(canonical.shape), + shards=[ShardDescriptor(dim=0, group=single_rank_group, interleave_factor=4)], + ) + target = torch.empty(4, 3, dtype=packed.dtype) + + spec.plan_hf_load().load_into([packed], target, model.hf_tensor_to_canonical) + + torch.testing.assert_close(target, canonical[(0, 2, 4, 6), :]) + + def test_qwen3vl_packed_adapter_runs_before_expert_tp_copy( + self, + single_rank_group: dist.ProcessGroup, + set_shard_rank: Callable[[int, int], None], + ) -> None: + from xtuner.v1.model.moe.qwen3vl_text import Qwen3VLTextMoE + + set_shard_rank(2, 0) + model = object.__new__(Qwen3VLTextMoE) + packed = torch.arange(2 * 3 * 4).reshape(2, 3, 4) + canonical = packed.transpose(1, 2).reshape(8, 3) + spec = LoadSpec( + name="layers.0.experts.fused_w1w3.weight", + global_hf_keys=["model.layers.0.mlp.experts.gate_up_proj"], + global_shape=tuple(canonical.shape), + shards=[ShardDescriptor(dim=0, group=single_rank_group, interleave_factor=4)], + ) + target = torch.empty(4, 3, dtype=packed.dtype) + + spec.plan_hf_load().load_into([packed], target, model.hf_tensor_to_canonical) + + torch.testing.assert_close(target, canonical[(0, 2, 4, 6), :]) + + def test_gpt_oss_packed_adapter_runs_before_expert_tp_copy( + self, + single_rank_group: dist.ProcessGroup, + set_shard_rank: Callable[[int, int], None], + ) -> None: + from xtuner.v1.model.moe.gpt_oss import GptOss + + set_shard_rank(2, 0) + model = object.__new__(GptOss) + packed = torch.arange(2 * 3 * 4).reshape(2, 3, 4) + canonical = model.hf_tensor_to_canonical("layers.0.experts.fused_w1w3.weight", packed) + spec = LoadSpec( + name="layers.0.experts.fused_w1w3.weight", + global_hf_keys=["model.layers.0.mlp.experts.gate_up_proj"], + global_shape=tuple(canonical.shape), + shards=[ShardDescriptor(dim=0, group=single_rank_group, interleave_factor=4)], + ) + target = torch.empty(4, 3, dtype=packed.dtype) + + spec.plan_hf_load().load_into([packed], target, model.hf_tensor_to_canonical) + + torch.testing.assert_close(target, canonical[(0, 2, 4, 6), :]) + + def test_glm_per_expert_adapter_runs_before_expert_tp_copy( + self, + single_rank_group: dist.ProcessGroup, + set_shard_rank: Callable[[int, int], None], + ) -> None: + from xtuner.v1.model.moe.glm52 import Glm52MoE + + set_shard_rank(2, 0) + model = object.__new__(Glm52MoE) + packed = torch.arange(2 * 4 * 3).reshape(2, 4, 3) + canonical = packed.flatten(0, 1) + spec = LoadSpec( + name="layers.0.experts.fused_w1w3.weight", + global_hf_keys=["gate", "up"], + global_shape=tuple(canonical.shape), + fused_dim=0, + shards=[ShardDescriptor(dim=0, group=single_rank_group, interleave_factor=4)], + ) + target = torch.empty(4, 3, dtype=packed.dtype) + + spec.plan_hf_load().load_into([packed[:1], packed[1:]], target, model.hf_tensor_to_canonical) + + torch.testing.assert_close(target, canonical[(0, 2, 4, 6), :]) + + def test_ep_etp_fsdp_composition_maps_back_to_global_runs( + self, + monkeypatch: pytest.MonkeyPatch, + local_groups: tuple[dist.ProcessGroup, dist.ProcessGroup, dist.ProcessGroup], + ) -> None: + ep_group, etp_group, fsdp_group = local_groups + set_group_layout( + monkeypatch, + { + ep_group: (2, 1, [0, 1]), + etp_group: (2, 1, [0, 2]), + fsdp_group: (2, 0, [0, 3]), + }, + ) + spec = LoadSpec( + name="weight", + global_hf_keys=["weight"], + global_shape=(32, 2), + shards=[ + ShardDescriptor(dim=0, group=ep_group), + ShardDescriptor(dim=0, group=etp_group, interleave_factor=2), + ShardDescriptor(dim=0, group=fsdp_group), + ], + ) + full = torch.arange(64).reshape(32, 2) + target = torch.empty(4, 2, dtype=full.dtype) + + spec.plan_hf_load().load_into([full], target, lambda _, tensor: tensor) + + torch.testing.assert_close(target, full[20:24]) + + +class TestHFSavePolicy: + def test_fused_keys_are_split_across_save_ranks(self, monkeypatch: pytest.MonkeyPatch) -> None: + model = BaseModel(XTunerBaseModelConfig()) + model.config.hf_save_cfg.max_save_rank = 4 + spec = LoadSpec( + name="layers.0.experts.fused_w1w3.weight", + global_hf_keys=[f"k{i}" for i in range(8)], + global_shape=(800, 64), + fused_dim=0, + ) + + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "get_world_size", lambda group=None: 8) + + expected_ranges = {0: (0, 2), 1: (2, 4), 2: (4, 6), 3: (6, 8), 4: (0, 0)} + for rank, expected_range in expected_ranges.items(): + monkeypatch.setattr(dist, "get_rank", lambda group=None, rank=rank: rank) + assert model._hf_save_key_range(spec.plan_hf_save(distributed_save=True)) == expected_range + + def test_preserved_fused_shard_exposes_local_hf_keys( + self, + single_rank_group: dist.ProcessGroup, + set_shard_rank: Callable[[int, int], None], + ) -> None: + set_shard_rank(4, 1) + spec = LoadSpec( + name="layers.0.experts.fused_w1w3.weight", + global_hf_keys=["k0", "k1", "k2", "k3"], + global_shape=(400, 64), + fused_dim=0, + shards=[ShardDescriptor(dim=0, group=single_rank_group)], + ) + + save_plan = spec.plan_hf_save(preserve_process_group=single_rank_group) + + assert save_plan.preserves_shards is True + assert save_plan.hf_keys == ["k1"] + assert save_plan.runtime_output_shape == (100, 64) + assert save_plan.output_shape == (100, 64) + + def test_preserved_fused_shard_must_align_with_hf_key_boundary( + self, + single_rank_group: dist.ProcessGroup, + set_shard_rank: Callable[[int, int], None], + ) -> None: + set_shard_rank(8, 1) + spec = LoadSpec( + name="layers.0.experts.fused_w1w3.weight", + global_hf_keys=["k0", "k1", "k2", "k3"], + global_shape=(400, 64), + fused_dim=0, + shards=[ShardDescriptor(dim=0, group=single_rank_group)], + ) + + with pytest.raises(AssertionError, match="must align with HF key size"): + spec.plan_hf_save(preserve_process_group=single_rank_group) + + +class TestHFSaveUnshard: + @staticmethod + def _patch_foreach_all_gather( + monkeypatch: pytest.MonkeyPatch, + responses: list[list[list[torch.Tensor]]] | None = None, + ) -> list[dict[str, object]]: + calls: list[dict[str, object]] = [] + + def fake_foreach_all_gather( + tensor_list: list[torch.Tensor], + group: dist.ProcessGroup, + ) -> list[list[torch.Tensor]]: + calls.append( + { + "group": group, + "shapes": [tuple(tensor.shape) for tensor in tensor_list], + "dtypes": [tensor.dtype for tensor in tensor_list], + } + ) + if responses is not None: + return responses.pop(0) + return [[tensor] for tensor in tensor_list] + + monkeypatch.setattr(load_spec_module, "foreach_all_gather", fake_foreach_all_gather) + return calls + + def test_scheduler_batches_same_group_and_respects_dependencies( + self, + monkeypatch: pytest.MonkeyPatch, + single_rank_group: dist.ProcessGroup, + ) -> None: + calls = self._patch_foreach_all_gather(monkeypatch) + specs = [ + LoadSpec( + name="experts", + global_hf_keys=["k0", "k1"], + global_shape=(8, 2), + fused_dim=0, + shards=[ + ShardDescriptor(dim=0, group=single_rank_group), + ShardDescriptor(dim=0, group=single_rank_group), + ], + ), + LoadSpec( + name="gate", + global_hf_keys=["gate"], + global_shape=(4, 2), + shards=[ShardDescriptor(dim=0, group=single_rank_group)], + ), + ] + + output = unshard_tensors_for_hf_save( + [torch.ones(8, 2), torch.ones(4, 2)], + [spec.plan_hf_save() for spec in specs], + ) + + assert [tuple(tensor.shape) for tensor in output] == [(8, 2), (4, 2)] + assert [call["shapes"] for call in calls] == [[(8, 2), (4, 2)], [(8, 2)]] + + def test_scheduler_splits_different_dtypes( + self, + monkeypatch: pytest.MonkeyPatch, + single_rank_group: dist.ProcessGroup, + ) -> None: + calls = self._patch_foreach_all_gather(monkeypatch) + specs = [ + LoadSpec( + name=name, + global_hf_keys=[name], + global_shape=(4, 2), + shards=[ShardDescriptor(dim=0, group=single_rank_group)], + ) + for name in ("gate", "up") + ] + + output = unshard_tensors_for_hf_save( + [torch.ones(4, 2, dtype=torch.float32), torch.ones(4, 2, dtype=torch.float64)], + [spec.plan_hf_save() for spec in specs], + ) + + assert [tuple(tensor.shape) for tensor in output] == [(4, 2), (4, 2)] + assert [call["dtypes"] for call in calls] == [[torch.float32], [torch.float64]] + + def test_continuous_collective_padding_restores_runtime_shape( + self, + monkeypatch: pytest.MonkeyPatch, + single_rank_group: dist.ProcessGroup, + set_shard_rank: Callable[[int, int], None], + ) -> None: + set_shard_rank(3, 2) + calls = self._patch_foreach_all_gather( + monkeypatch, + responses=[[[torch.tensor([0, 1]), torch.tensor([2, 3]), torch.tensor([4, 0])]]], + ) + spec = LoadSpec( + name="weight", + global_hf_keys=["weight"], + global_shape=(5,), + shards=[ShardDescriptor(dim=0, group=single_rank_group)], + ) + + [output] = unshard_tensors_for_hf_save([torch.tensor([4])], [spec.plan_hf_save()]) + + torch.testing.assert_close(output, torch.arange(5)) + assert calls[0]["shapes"] == [(2,)] + + def test_even_interleave_deinterleaves_then_trims_fp8_padding( + self, + monkeypatch: pytest.MonkeyPatch, + single_rank_group: dist.ProcessGroup, + set_shard_rank: Callable[[int, int], None], + ) -> None: + set_shard_rank(2, 0) + calls = self._patch_foreach_all_gather( + monkeypatch, + responses=[[[torch.tensor([0, 1, 4, 5]), torch.tensor([2, 3, 6, 7])]]], + ) + spec = LoadSpec( + name="weight", + global_hf_keys=["weight"], + global_shape=(8,), + origin_shape=(6,), + shards=[ShardDescriptor(dim=0, group=single_rank_group, interleave_factor=2)], + ) + + [output] = unshard_tensors_for_hf_save( + [torch.tensor([0, 1, 4, 5])], + [spec.plan_hf_save()], + ) + + torch.testing.assert_close(output, torch.arange(6)) + assert calls[0]["shapes"] == [(4,)] + + def test_preserve_ep_still_deinterleaves_etp( + self, + monkeypatch: pytest.MonkeyPatch, + local_groups: tuple[dist.ProcessGroup, dist.ProcessGroup, dist.ProcessGroup], + ) -> None: + ep_group, etp_group, _ = local_groups + set_group_layout( + monkeypatch, + { + ep_group: (2, 1, [0, 1]), + etp_group: (2, 0, [0, 2]), + }, + ) + self._patch_foreach_all_gather( + monkeypatch, + responses=[[[torch.tensor([4, 6]), torch.tensor([5, 7])]]], + ) + spec = LoadSpec( + name="experts", + global_hf_keys=["k0", "k1", "k2", "k3"], + global_shape=(8,), + fused_dim=0, + shards=[ + ShardDescriptor(dim=0, group=ep_group), + ShardDescriptor(dim=0, group=etp_group, interleave_factor=2), + ], + ) + plan = spec.plan_hf_save(preserve_process_group=ep_group) + + [output] = unshard_tensors_for_hf_save([torch.tensor([4, 6])], [plan]) + + torch.testing.assert_close(output, torch.tensor([4, 5, 6, 7])) + assert plan.hf_keys == ["k2", "k3"] + + def test_only_gather_fsdp_preserves_etp_and_trims_local_fp8_tail( + self, + monkeypatch: pytest.MonkeyPatch, + local_groups: tuple[dist.ProcessGroup, dist.ProcessGroup, dist.ProcessGroup], + ) -> None: + etp_group, fsdp_group, _ = local_groups + set_group_layout( + monkeypatch, + { + etp_group: (2, 1, [0, 1]), + fsdp_group: (2, 1, [0, 2]), + }, + ) + self._patch_foreach_all_gather( + monkeypatch, + responses=[[[torch.tensor([4, 5, 6, 7]), torch.tensor([12, 13, 14, 15])]]], + ) + spec = LoadSpec( + name="weight", + global_hf_keys=["weight"], + global_shape=(16,), + origin_shape=(14,), + shards=[ + ShardDescriptor(dim=0, group=etp_group, interleave_factor=2), + ShardDescriptor(dim=0, group=fsdp_group), + ], + ) + plan = spec.plan_hf_save(gather_process_group=fsdp_group) + + [output] = unshard_tensors_for_hf_save([torch.tensor([12, 13, 14, 15])], [plan]) + + torch.testing.assert_close(output, torch.tensor([4, 5, 6, 7, 12, 13])) + assert plan.runtime_output_shape == (8,) + assert plan.output_shape == (6,) + + +class TestBaseModelHFSave: + def test_non_dtensor_buffers_keep_runtime_dtype(self) -> None: + class BufferModel(BaseModel): + def __init__(self) -> None: + super().__init__(XTunerBaseModelConfig()) + self.register_buffer("rotary_coef", torch.tensor([1.25], dtype=torch.float32), persistent=True) + self._init_load_spec() + + def to_hf_key_list(self, key: str) -> list[str]: + return [key] + + model = BufferModel() + + [(names, tensors)] = list( + model._get_hf_param(model._load_spec_params(), dtype=torch.bfloat16, distributed_save=True) + ) + + assert names == ["rotary_coef"] + assert tensors[0].dtype == torch.float32 + assert torch.equal(tensors[0], model.rotary_coef) diff --git a/xtuner/v1/engine/train_engine.py b/xtuner/v1/engine/train_engine.py index 61a1f58fe0..0d6e1e90c6 100644 --- a/xtuner/v1/engine/train_engine.py +++ b/xtuner/v1/engine/train_engine.py @@ -38,6 +38,7 @@ ModelOutputs, XTunerBaseModelConfig, ) +from xtuner.v1.patch import InterleavedShardLoadPlanner, InterleavedShardSavePlanner from xtuner.v1.patch.xtuner_storage import XtunerCacheWriter, _get_async_dcp_save_timeout from xtuner.v1.profiler.prober import ProberList from xtuner.v1.utils import ( @@ -47,7 +48,6 @@ log_rank0, profile_time_and_memory, ) -from xtuner.v1.utils.grad_norm import cal_grad_norm class TrainStepInfo(DataBatchInfo, BatchForwardInfo): @@ -259,7 +259,7 @@ def clip_grad_norm(self, do_clip: bool = True, dtype=torch.float32): self.model.scale_and_reduce_grad() params = self.model.trainable_parameters() grads = [p.grad for _, p in params if p.grad is not None] - grad_norm, grouped_grads = cal_grad_norm(grads, dtype=dtype) + grad_norm, grouped_grads = self.model.cal_grad_norm(grads, dtype=dtype) if do_clip: clip_coef = self.optim_cfg.max_grad_norm / (grad_norm + 1e-6) clip_coef_clamped = torch.clamp(clip_coef, max=1.0) @@ -355,6 +355,7 @@ def save_dcp( dcp.save( state_dict, checkpoint_id=weights_dir, + planner=InterleavedShardSavePlanner(), ) def _get_async_checkpoint_pg(self) -> dist.ProcessGroup: @@ -407,6 +408,7 @@ def start_async_save() -> Future: return cast(Any, dcp.async_save)( state_dict, checkpoint_id=incomplete_dir, + planner=InterleavedShardSavePlanner(), storage_writer=storage_writer, process_group=async_checkpoint_pg, **async_save_kwargs, @@ -503,7 +505,11 @@ def load_dcp( set_options = StateDictOptions(cpu_offload=True, strict=True) with profile_time_and_memory(f"[Load DCP from {weights_dir}]"): - dcp.load(state_dict=state_dict, checkpoint_id=weights_dir) + dcp.load( + state_dict=state_dict, + checkpoint_id=weights_dir, + planner=InterleavedShardLoadPlanner(), + ) set_model_state_dict(self.model, state_dict["model"], options=set_options) diff --git a/xtuner/v1/float8/float8_gmm_tile_wise.py b/xtuner/v1/float8/float8_gmm_tile_wise.py index 96a4c01daf..d9704d5180 100644 --- a/xtuner/v1/float8/float8_gmm_tile_wise.py +++ b/xtuner/v1/float8/float8_gmm_tile_wise.py @@ -1,7 +1,7 @@ # Copyright (c) OpenMMLab. All rights reserved. import math -from typing import Optional, Tuple, cast +from typing import Literal, Optional, Tuple import torch import torch.nn as nn @@ -18,6 +18,7 @@ trans_per_block_quant_expand_128x, trans_per_tile_quant_expand_128x, ) +from xtuner.v1.utils.interleaved_shard import InterleavedShard # from xtuner.v1.module.grouped_linear.moe_group_linear import GroupedLinear @@ -220,6 +221,10 @@ def __init__( num_routed_experts: int, moe_bias: bool = False, ep_mesh: DeviceMesh | None = None, + expert_tp_mesh: DeviceMesh | None = None, + parallel_style: Literal["column", "row"] | None = None, + ep_tp_mesh: DeviceMesh | None = None, + num_fused_projections: int = 1, ) -> None: super().__init__() @@ -234,38 +239,104 @@ def __init__( self.in_features = in_features self.out_features = out_features self.num_routed_experts = num_routed_experts + self.ep_mesh = ep_mesh + self.expert_tp_mesh = expert_tp_mesh + self.parallel_style = parallel_style + self.ep_size = ep_mesh.size() if ep_mesh is not None else 1 + self.tp_size = expert_tp_mesh.size() if expert_tp_mesh is not None else 1 + self.tp_enabled = self.expert_tp_mesh is not None and self.tp_size > 1 and self.parallel_style is not None + self.num_fused_projections = num_fused_projections + if self.expert_tp_mesh is not None and self.expert_tp_mesh.size() > 1 and self.parallel_style is None: + raise ValueError("parallel_style must be set when expert_tp_mesh size is greater than 1.") + if self.num_routed_experts % self.ep_size != 0: + raise ValueError( + f"num_routed_experts ({self.num_routed_experts}) must be divisible by ep_size ({self.ep_size})." + ) + + self.local_num_routed_experts = self.num_routed_experts // self.ep_size + self.local_in_features = in_features + self.local_out_features = out_features + if self.tp_enabled: + if self.parallel_style == "column": + if out_features % self.tp_size != 0: + raise ValueError(f"out_features ({out_features}) must be divisible by tp_size ({self.tp_size}).") + self.local_out_features = out_features // self.tp_size + elif self.parallel_style == "row": + if in_features % self.tp_size != 0: + raise ValueError(f"in_features ({in_features}) must be divisible by tp_size ({self.tp_size}).") + self.local_in_features = in_features // self.tp_size + else: + raise ValueError(f"Unsupported parallel_style: {self.parallel_style}.") + self.ori_shape = (num_routed_experts, out_features, in_features) self.ori_local_shape = ( - (num_routed_experts // ep_mesh.size(), out_features, in_features) - if ep_mesh is not None - else self.ori_shape + self.local_num_routed_experts, + self.local_out_features, + self.local_in_features, ) # We have padded the dim0 of GroupedLinear's weight to make fsdp compatible with block-wise fp8. - weight = WeightWithDynamicTilewiseFloat8CastTensor( - torch.empty(num_routed_experts * out_features, in_features), - torch.float8_e4m3fn, - (num_routed_experts * out_features, in_features), - ) - self.ep_mesh = ep_mesh - if ep_mesh is not None and ep_mesh.size() > 1: + local_shape = (self.local_num_routed_experts * self.local_out_features, self.local_in_features) + if ep_tp_mesh is not None and self.tp_enabled: + weight = WeightWithDynamicTilewiseFloat8CastTensor( + torch.empty(local_shape), + torch.float8_e4m3fn, + (num_routed_experts * out_features, in_features), + ) + self.weight = nn.Parameter( + DTensor.from_local(weight, ep_tp_mesh, self._weight_placements(), run_check=False) + ) + elif ep_mesh is not None and ep_mesh.size() > 1: + weight = WeightWithDynamicTilewiseFloat8CastTensor( + torch.empty(num_routed_experts * out_features, in_features), + torch.float8_e4m3fn, + (num_routed_experts * out_features, in_features), + ) self.weight = nn.Parameter(distribute_tensor(weight, ep_mesh, [Shard(0)])) else: + weight = WeightWithDynamicTilewiseFloat8CastTensor( + torch.empty(local_shape), + torch.float8_e4m3fn, + (num_routed_experts * out_features, in_features), + ) self.weight = nn.Parameter(weight) self.pad_shape: Optional[Tuple[int, int]] = None self.reset_parameters() def reset_parameters(self) -> None: - init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + weight = self.weight.to_local() if isinstance(self.weight, DTensor) else self.weight + init.kaiming_uniform_(weight, a=math.sqrt(5)) + + def _weight_placements(self): + if self.tp_enabled: + if self.parallel_style == "column": + # fused_w1w3 has one stripe per (expert, fused projection); TP cuts inside every stripe. + num_local_stripes = self.local_num_routed_experts * self.num_fused_projections + return ( + Shard(0), + InterleavedShard(0, num_local_stripes=num_local_stripes), + ) + return (Shard(0), Shard(1)) + return (Shard(0),) + + def _dim_shard_size(self, dim: int) -> int: + if dim == 0: + return self.ep_size * (self.tp_size if self.tp_enabled and self.parallel_style == "column" else 1) + if dim == 1: + return self.tp_size if self.tp_enabled and self.parallel_style == "row" else 1 + return 1 def _check_shape(self, weight): - ep_size = self.ep_mesh.size() if self.ep_mesh is not None else 1 if self.is_padded: assert weight.shape == ( - self.pad_shape[0] // ep_size, - self.pad_shape[1], - ), f"Expected weight shape {(self.pad_shape[0] // ep_size, self.pad_shape[1])}, but got {weight.shape}." + self.pad_shape[0] // self._dim_shard_size(0), + self.pad_shape[1] // self._dim_shard_size(1), + ), ( + f"Expected weight shape " + f"{(self.pad_shape[0] // self._dim_shard_size(0), self.pad_shape[1] // self._dim_shard_size(1))}, " + f"but got {weight.shape}." + ) else: assert weight.shape == ( self.ori_local_shape[0] * self.ori_local_shape[1], @@ -282,7 +353,8 @@ def forward(self, input: torch.Tensor, tokens_per_expert, decoding: bool = False if tensor_already_casted_to_fp8(weight): # If we use fsdp, the weight is already casted to fp8. - # If self.is_padded is True, ep size should be 1 + # FSDP padding only extends flattened dim0; trim it before restoring + # the local (expert, out, in) grouped-GEMM layout. weight_fp8 = slice_weight.apply(weight, self.ori_local_shape) if self.is_padded else weight weight_fp8 = view_weight.apply(weight_fp8, self.ori_local_shape) else: @@ -293,7 +365,7 @@ def forward(self, input: torch.Tensor, tokens_per_expert, decoding: bool = False num_tokens = input.numel() // input.shape[-1] input = input.view(num_tokens, input.shape[-1]) out = fp8_gmm_weight_per_block_act_per_tile.apply(input, weight_fp8, tokens_per_expert) - out = out.view(*orig_shape[:-1], self.out_features) + out = out.view(*orig_shape[:-1], self.local_out_features) return out @property @@ -314,8 +386,13 @@ def pad_for_fsdp(self, padded_out_features: int) -> None: if padded_out_features == self.weight.shape[0]: return + ori_local_weight = self.weight.to_local() if isinstance(self.weight, DTensor) else self.weight + local_shape = ( + padded_out_features // self._dim_shard_size(0), + self.in_features // self._dim_shard_size(1), + ) weight = torch.empty( - (padded_out_features, self.in_features), + local_shape if isinstance(self.weight, DTensor) else (padded_out_features, self.in_features), dtype=self.weight.dtype, layout=self.weight.layout, device=self.weight.device, @@ -325,9 +402,8 @@ def pad_for_fsdp(self, padded_out_features: int) -> None: torch.float8_e4m3fn, (self.num_routed_experts * self.out_features, self.in_features), ) - if self.ep_mesh is not None and self.ep_mesh.size() > 1: - weight = distribute_tensor(weight, self.ep_mesh, [Shard(0)]) - ori_local_weight = cast(DTensor, self.weight)._local_tensor + if isinstance(self.weight, DTensor): + weight = DTensor.from_local(weight, self.weight.device_mesh, self._weight_placements(), run_check=False) local_weight = weight._local_tensor local_weight[: ori_local_weight.shape[0]].data.copy_(ori_local_weight.data) # copy the original weight local_weight[ori_local_weight.shape[0] :].data.copy_(0.0) # type: ignore # zero pad the weight diff --git a/xtuner/v1/float8/float8_handler.py b/xtuner/v1/float8/float8_handler.py index 36e9f429de..1e72bce45e 100644 --- a/xtuner/v1/float8/float8_handler.py +++ b/xtuner/v1/float8/float8_handler.py @@ -7,6 +7,8 @@ import torch.nn as nn from torch.distributed._tensor import DTensor from torch.distributed.device_mesh import DeviceMesh, init_device_mesh +from torch.distributed.tensor import Shard +from torch.distributed.tensor.placement_types import _StridedShard from xtuner.v1.float8.config import ScalingGranularity from xtuner.v1.float8.fsdp_utils import ( @@ -105,6 +107,17 @@ def get_num_features_after_pad(tensor_size, fsdp_shard_dim, num_chunks): break return chunk_size * num_chunks + @staticmethod + def get_shard_size_on_dim(tensor: torch.Tensor | DTensor, dim: int) -> int: + if not isinstance(tensor, DTensor): + return 1 + + shard_size = 1 + for mesh_dim, placement in enumerate(tensor.placements): + if isinstance(placement, (Shard, _StridedShard)) and placement.dim == dim: + shard_size *= tensor.device_mesh.size(mesh_dim) + return shard_size + @staticmethod def pad_for_fsdp(model: nn.Module, fsdp_mesh: DeviceMesh, callback_after_pad: Callable | None = None): from xtuner.v1.float8.float8_gmm_tile_wise import TileWiseFloat8GroupedLinear @@ -120,7 +133,7 @@ def pad_for_fsdp(model: nn.Module, fsdp_mesh: DeviceMesh, callback_after_pad: Ca "Currently only support even distributed TP or EP weight for float8 training." ) tensor_size = module.weight._local_tensor.size() - parallel_size = module.weight.device_mesh.size() + parallel_size = Float8Handler.get_shard_size_on_dim(module.weight, dim=0) else: tensor_size = module.weight.size() parallel_size = 1 diff --git a/xtuner/v1/model/base.py b/xtuner/v1/model/base.py index 0856a3fd7a..24a7959f6d 100644 --- a/xtuner/v1/model/base.py +++ b/xtuner/v1/model/base.py @@ -1,6 +1,5 @@ import importlib import json -import math import multiprocessing as py_mp import os import pydoc @@ -12,12 +11,11 @@ from itertools import chain from pathlib import Path from shutil import copy, copytree, rmtree -from typing import Annotated, Any, Generator, Iterable, Literal, Mapping, Sequence, cast +from typing import Annotated, Any, Generator, Iterable, Literal, Mapping, NamedTuple, Sequence, cast import torch import torch.distributed as dist 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 @@ -31,10 +29,7 @@ MixedPrecisionPolicy, fully_shard, ) -from torch.distributed.tensor import DTensor, Placement, Replicate, Shard, distribute_tensor -from torch.distributed.tensor._utils import ( - compute_local_shape_and_global_offset as _compute_local_shape_and_global_offset, -) +from torch.distributed.tensor import DTensor, Replicate, distribute_tensor from torch.utils import _pytree from typing_extensions import NotRequired, Self, TypedDict, overload @@ -50,10 +45,13 @@ from xtuner.v1.loss import BaseLossConfig, BaseLossContext, CELossConfig from xtuner.v1.module.attention import GatedDeltaNetConfig, MHAConfig, MLAConfig from xtuner.v1.module.rope import RopeParametersConfig, RopeScalingConfig -from xtuner.v1.ops.comm.foreach_allgather import foreach_all_gather from xtuner.v1.utils import get_device, get_logger, get_torch_device_module, log_rank0, profile_time_and_memory from xtuner.v1.utils.compile import MaybeCompile, is_compiled_function, maybe_compile -from xtuner.v1.utils.load_spec import LoadEnum, LoadSpec +from xtuner.v1.utils.load_spec import ( + HFSavePlan, + LoadSpec, + unshard_tensors_for_hf_save, +) from xtuner.v1.utils.loader import HFCheckpointLoader from xtuner.v1.utils.misc import FunctionEnum, FunctionType, get_function_full_qualname, get_function_type from xtuner.v1.utils.process import ( @@ -71,12 +69,6 @@ DEVICE = get_device() -def compute_local_shape_and_global_offset(*args, **kwargs): - "wrapper of _compute_local_shape_and_global_offset avoiding meta tensor error" - with torch.device(DEVICE): - return _compute_local_shape_and_global_offset(*args, **kwargs) - - class DataBatchInfo(TypedDict): step_consumed_tokens: int step_seqlen_tokens: int @@ -538,7 +530,28 @@ def _save_file( save_file(tensors, filename, metadata=metadata) +class _HFSaveBucketItem(NamedTuple): + tensor: torch.Tensor + save_plan: HFSavePlan + runtime_is_float8: bool + byte_size: int + + class BaseModel(nn.Module): + """Base class for all xtuner training models with HF checkpoint I/O + support. + + Subclass ``__init__`` **must** call ``self._init_load_spec()`` at the end, + once every parameter and submodule has been constructed (including any + ``__init__``-time sharding such as MoE EP via ``distribute_tensor``). This + populates ``self.load_spec_mapping`` so that ``from_hf`` / ``save_hf`` and + the RL weight-sync path can translate between local params and HF + checkpoint keys. ``fully_shard`` and ``Float8Handler.pad_for_fsdp`` may + re-invoke ``_init_load_spec`` afterwards to keep the mapping in sync with + the current layout. See ``docs/design/load_spec_refactor.md`` §5.2 for the + full contract. + """ + load_spec_mapping: dict[str, LoadSpec] = {} fsdp_mesh: DeviceMesh | None = None hsdp_mesh: DeviceMesh | None = None @@ -567,6 +580,10 @@ def from_hf( ) -> tuple[ Annotated[set[str], "loaded keys"], Annotated[set[str], "unloaded keys"], Annotated[set[str], "missing keys"] ]: + # Recompute from the complete HF key list and the current runtime layout. + # `__init__` still initializes the mapping for consumers that read it before checkpoint I/O. + self._init_load_spec() + self._assert_load_spec_initialized() self._hf_path = Path(hf_path) if isinstance(hf_path, Path): @@ -579,6 +596,11 @@ def from_hf( def scale_and_reduce_grad(self): return + def cal_grad_norm(self, grads: list[DTensor], dtype=torch.float32): + from xtuner.v1.utils.grad_norm import cal_grad_norm + + return cal_grad_norm(grads, dtype=dtype) + def to_hf_key_list(self, key: str) -> list[str]: raise NotImplementedError() @@ -622,6 +644,7 @@ def fully_shard( reshard_after_forward=fsdp_config.reshard_after_forward, offload_policy=CPUOffloadPolicy() if self.fsdp_config.cpu_offload else None, ) + self._init_load_spec() return self def _fully_shard( @@ -644,7 +667,7 @@ def traverse(module): load_spec = self.load_spec_mapping.get(full_name) if load_spec is None: raise ValueError(f"Internal Error. Parameter {full_name} not found in load_spec_mapping.") - hf_name_list = load_spec.hf_keys + hf_name_list = load_spec.global_hf_keys for hf_name in hf_name_list: if any(re.search(p, hf_name) for p in patterns): # type: ignore @@ -698,6 +721,9 @@ def traverse(module): ) def save_hf(self, hf_dir: Path | str, save_dtype: torch.dtype = torch.bfloat16, safetensors_prefix: str = "model"): + # Save may be called without `fully_shard`; refresh from the current runtime layout. + self._init_load_spec() + self._assert_load_spec_initialized() with profile_time_and_memory(f"[Saving HF to [{safetensors_prefix}]{hf_dir} cost]"): self._save_hf(hf_dir=hf_dir, save_dtype=save_dtype, safetensors_prefix=safetensors_prefix) @@ -845,6 +871,11 @@ def _prepare_async_hf_snapshot( safetensors_prefix: str = "model", device: torch.device | str = DEVICE, ) -> tuple[list[tuple[str, list[str]]], dict[str, str]]: + # Compose models enter async save through their child modules directly. + # Refresh here so every snapshot is planned from the post-FSDP runtime layout, + # matching the public synchronous save path. + self._init_load_spec() + self._assert_load_spec_initialized() file_to_names: list[tuple[str, list[str]]] = [] weight_map: dict[str, str] = {} for safetensor_name, name_list, hf_tensor_list in self._iter_hf_save_chunks( @@ -916,40 +947,15 @@ def _commit_async_hf_save( log_rank0.info(f"[Async saving HF to {hf_dir}] finalized") return hf_dir - def safetensors_to_params( - self, - safetensors: list[torch.Tensor], - local_tensor: torch.Tensor, - param_name: str, - start: int | None, - end: int | None, - dim: int | None, - ): - if len(safetensors) > 1: - assert dim is not None, "Internal Error dim must not be None when len(safetensors) > 1" - loaded_tensor = torch.cat(safetensors, dim=dim) - else: - loaded_tensor = safetensors[0] - - if start is not None and end is not None: - assert self.fsdp_config is not None, ( - "Internal Error. fsdp_config must not be None when start and end is not None" - ) - start = min(start, loaded_tensor.shape[self.FSDP_SHARD_DIM]) - end = min(end, loaded_tensor.shape[self.FSDP_SHARD_DIM]) - loaded_tensor_slice = loaded_tensor.index_select( - dim=self.FSDP_SHARD_DIM, index=torch.arange(start, end, dtype=torch.int64, device=loaded_tensor.device) - ) - non_pad_len = end - start - local_tensor[:non_pad_len].copy_(loaded_tensor_slice) + def hf_tensor_to_canonical(self, name: str, loaded_tensor: torch.Tensor) -> torch.Tensor: + """Convert one checkpoint tensor to XTuner's unsharded canonical + layout. - if non_pad_len < local_tensor.shape[self.FSDP_SHARD_DIM]: - assert self.config.float8_cfg is not None, ( - f"Shape mismatched! xtuner param shape: {param_name}, hf param {loaded_tensor.shape}" - ) - local_tensor[non_pad_len:].copy_(0.0) # type: ignore # padded part must be set to 0 - else: - local_tensor.copy_(loaded_tensor) + ``HFLoadPlan`` owns HF-key concatenation, rank-local slicing, interleaved + copies, and runtime padding. Model subclasses override only this format + adapter when their HF tensor shape/order differs from XTuner's layout. + """ + return loaded_tensor def param_to_safetensor( self, @@ -1028,30 +1034,6 @@ def build_rotary_embedding(self, config): return get_rope_embedding(config=config) def _init_load_spec(self) -> None: - # NOTE: (yehaochen) This is a workaround to distinguish between different parameter HF loading methods - # and model partitioning methods. Although PyTorch provides Shard, Replicate and other Placements, in - # MoE models, we need to handle both how to load HF weights and how to calculate gradients for partitioned - # parameters during the backward phase, so a more complex ParallelParamSpec is defined to describe these: - # Specifically: - # - For model loading and saving: - # From a computational efficiency perspective, we have to make the model parameter layout different from the - # HF model, resulting in a one-to-one or many-to-many mapping relationship, and we need a specification to - # describe this mapping. - # - For gradient computation: - # In MoE models, we need to divide the gradients of EP-partitioned parameters by ep_size (this is another - # complex issue not elaborated here), and although ep and ep both belong to Shard, their processing logic - # is different, so we need a specification to express the partitioning method in a more fine-grained way. - - def get_shard_placement(placements: tuple[Placement, ...]) -> Shard | None: - ret = None - for p in placements: - if isinstance(p, Shard): - if ret is None: - ret = p - else: - raise RuntimeError("Multiple Shard placements found, please report this issue") - return ret - if self.__class__.to_hf_key_list is BaseModel.to_hf_key_list: self.load_spec_mapping = {} return @@ -1085,82 +1067,15 @@ def get_shard_placement(placements: tuple[Placement, ...]) -> Shard | None: repl = self.config.hf_key_mapping[max_matched_pattern] hf_keys.append(re.sub(max_matched_pattern, repl, key)) - if isinstance(param, DTensor) and (placement := get_shard_placement(param.placements)) is not None: - dim = placement.dim - _, _offset = compute_local_shape_and_global_offset(param.shape, param.device_mesh, param.placements) - start = _offset[dim] - end = _offset[dim] + param._local_tensor.shape[dim] - local_shape = param._local_tensor.shape - global_size = param.shape[dim] - - if len(hf_keys) > 1: - start_hf_key_idx = start / global_size * len(hf_keys) - - assert start_hf_key_idx.is_integer(), "Internal xtuner error, please report this issue" - start_hf_key_idx = int(start_hf_key_idx) - - end_hf_key_idx = end / global_size * len(hf_keys) - # TODO: (yehaochen) Support TP - assert end_hf_key_idx.is_integer(), "Internal xtuner error, please report this issue" - load_type = LoadEnum.FUSED - end_hf_key_idx = int(end_hf_key_idx) - elif len(hf_keys) == 1: - start_hf_key_idx = start / global_size - end_hf_key_idx = end / global_size - if start_hf_key_idx == 0 and end_hf_key_idx == 1: - load_type = LoadEnum.SAME - else: - load_type = LoadEnum.SHARD - else: - raise RuntimeError - - # TP shard - if load_type is LoadEnum.SHARD: - load_spec = LoadSpec( - name=name, - hf_keys=hf_keys, - shape=local_shape, - dim=dim, - load_enum=LoadEnum.SHARD, - shard_start=start, - shard_end=end, - group=param.device_mesh.get_group(), - ) - # Replicate - elif load_type == LoadEnum.SAME: - load_spec = LoadSpec( - name=name, - hf_keys=hf_keys, - shape=local_shape, - dim=dim, - load_enum=LoadEnum.SAME, - group=param.device_mesh.get_group(), - ) - # EPSHard - else: - load_spec = LoadSpec( - name=name, - hf_keys=hf_keys[start_hf_key_idx:end_hf_key_idx], - shape=local_shape, - dim=dim, - load_enum=LoadEnum.FUSED, - group=param.device_mesh.get_group(), - ) - else: - if len(hf_keys) == 1: - load_spec = LoadSpec( - name=name, - hf_keys=hf_keys, - shape=param.shape, - load_enum=LoadEnum.SAME, - ) - else: - load_spec = LoadSpec( - name=name, - hf_keys=hf_keys, - shape=param.shape, - load_enum=LoadEnum.FUSED, - ) + runtime_tensor = param._local_tensor if isinstance(param, DTensor) else param + runtime_is_float8 = is_float8_weight(runtime_tensor) + origin_shape = tuple(runtime_tensor._ori_shape) if runtime_is_float8 else None # type: ignore[attr-defined] + load_spec = LoadSpec.from_tensor( + name=name, + hf_keys=hf_keys, + tensor=param, + origin_shape=origin_shape, + ) load_spec_mapping[name] = load_spec if hf_key_mapping_missing: @@ -1169,16 +1084,28 @@ def get_shard_placement(placements: tuple[Placement, ...]) -> Shard | None: self.load_spec_mapping = load_spec_mapping + def _assert_load_spec_initialized(self) -> None: + # `load_spec_mapping` defaults to the class-level empty dict; `_init_load_spec` + # always assigns an instance attribute (possibly empty), so presence on + # `self.__dict__` is the reliable signal that the subclass contract was honored. + assert "load_spec_mapping" in self.__dict__, ( + f"{type(self).__name__}.__init__ must call self._init_load_spec() at the end. " + "See docs/design/load_spec_refactor.md §5.2." + ) + def _to_float8( self, gathered_tensor_list: list[torch.Tensor], name_list: list[str], - ori_tensor_list: list[torch.Tensor], + runtime_is_float8_list: list[bool], dtype: torch.dtype, ) -> tuple[list[torch.Tensor], list[str]]: + assert len(gathered_tensor_list) == len(name_list) == len(runtime_is_float8_list), ( + "Internal error: float8 conversion metadata length does not match tensor list" + ) gathered_tensor_list_new, name_list_new = [], [] - for gathered_tensor, name, ori_tensor in zip(gathered_tensor_list, name_list, ori_tensor_list): - if not is_float8_weight(ori_tensor): + for gathered_tensor, name, runtime_is_float8 in zip(gathered_tensor_list, name_list, runtime_is_float8_list): + if not runtime_is_float8: gathered_tensor_list_new.append(gathered_tensor) name_list_new.append(name) continue @@ -1352,344 +1279,221 @@ def _get_save_dtype(self, name: str, dtype: torch.dtype) -> torch.dtype: return torch.float32 return dtype - def _get_shard_hf_param( + def _get_hf_param( self, params: list[tuple[torch.Tensor, LoadSpec]], dtype: torch.dtype, - device="cpu", - bucket_size=None, + device: torch.device | str = "cpu", + bucket_size: int | None = None, + distributed_save: bool = False, + preserved_fused_shard_group: dist.ProcessGroup | None = None, ) -> Generator[tuple[list[str], list[torch.Tensor]], None, None]: - if not params: - return + """Yield HF checkpoint tensors for the given runtime params. - ignored_params, params = self._split_ignored_params(params) - if ignored_params: - name_list: list[str] = [load_spec.hf_keys[0] for _, load_spec in ignored_params] - hf_params = [param._local_tensor if isinstance(param, DTensor) else param for param, _ in ignored_params] - yield name_list, hf_params + Args: + params (list[tuple[torch.Tensor, LoadSpec]]): Runtime tensors and their new-schema LoadSpecs. + dtype (torch.dtype): Target checkpoint dtype, currently bfloat16 or float8_e4m3fn. + device (torch.device | str): Device to move yielded tensors to. + bucket_size (int | None): Approximate bucket size in bytes. + distributed_save (bool): Whether to apply the HF save write policy. When enabled, non-fused tensors are + yielded only on rank0 and fused HF keys are divided across save ranks. + preserved_fused_shard_group (dist.ProcessGroup | None): Communication group whose fused-dim shard should + stay local instead of being all-gathered. RL weight sync uses this to stream EP-local expert slices. + Returns: + Generator[tuple[list[str], list[torch.Tensor]], None, None]: HF key names and tensors to save. + """ + assert not (distributed_save and preserved_fused_shard_group is not None), ( + "distributed_save writes checkpoint files, while preserved_fused_shard_group streams local fused shards " + "for RL." + ) if not params: return - if dtype != torch.bfloat16: - raise NotImplementedError - - load_spec0 = params[0][1] - assert load_spec0.group is not None - - def _get_hf_params(fsdp_tensor_list: list[tuple[torch.Tensor, LoadSpec]]) -> list[torch.Tensor]: - # Get fsdp unsharded params - _tensor_list, _spec_list = list(zip(*fsdp_tensor_list)) - if self.fsdp_mesh is not None: - fsdp_unsharded_tensor_list = self._fsdp_foreach_allgather(_tensor_list, _spec_list) # type: ignore - else: - fsdp_unsharded_tensor_list = _tensor_list # type: ignore - - # Get unsharded params - _unsharded_tensor_list = foreach_all_gather(fsdp_unsharded_tensor_list, load_spec0.group) - unsharded_tensor_list = [ - torch.cat(list(tensors), dim=load_spec0.dim) for tensors in _unsharded_tensor_list - ] - name_list = [spec.hf_keys[0] for _, spec in fsdp_tensor_list] - unsharded_tensor_list = [ - self.param_to_safetensor(safetensor, name) - for safetensor, name in zip(unsharded_tensor_list, name_list) - ] - unsharded_tensor_list = [t.to(device) for t in unsharded_tensor_list] - return unsharded_tensor_list - if bucket_size is None: bucket_size = self.config.hf_save_cfg.bucket_size - safetensor_size = 0 - tensor_list: list[tuple[torch.Tensor, LoadSpec]] = [] - name_list = [] + bucket_bytes = 0 + bucket: list[_HFSaveBucketItem] = [] + buffer_names = {self._clean_param_name(name) for name, _ in self.named_buffers()} for param, load_spec in params: - local_tensor = param._local_tensor if isinstance(param, DTensor) else param - local_tensor = local_tensor.to(dtype=self._get_save_dtype(load_spec.hf_keys[0], torch.bfloat16)) - tensor_size = self._get_tensor_size(param, dtype) - if safetensor_size + tensor_size > bucket_size and tensor_list: - hf_params = _get_hf_params(tensor_list) - - yield name_list, hf_params - safetensor_size = tensor_size - name_list = load_spec.hf_keys.copy() - tensor_list = [(local_tensor, load_spec)] - continue - safetensor_size += tensor_size - tensor_list.append((local_tensor, load_spec)) - name_list.append(load_spec.hf_keys[0]) - - if tensor_list: - hf_params = _get_hf_params(tensor_list) - yield name_list, hf_params - - def _get_fused_hf_param( - self, - params: list[tuple[torch.Tensor, LoadSpec]], - dtype: torch.dtype, - device="cpu", - bucket_size=None, - update_weights_for_rl: bool = False, - ) -> Generator[tuple[list[str], list[torch.Tensor]], None, None]: - if not params: - return - - ignored_params, params = self._split_ignored_params(params) - if ignored_params: - fp32_name_list: list[str] = [load_spec.hf_keys[0] for _, load_spec in ignored_params] - fp32_params = [param._local_tensor if isinstance(param, DTensor) else param for param, _ in ignored_params] - yield fp32_name_list, fp32_params - - def _get_hf_params( - fsdp_tensor_list: list[tuple[torch.Tensor, LoadSpec]], - name_list: list[str], - ) -> tuple[list[torch.Tensor], list[str]]: - # Get fsdp unsharded params - spec_list: list[LoadSpec] - tensor_list: list[torch.Tensor] - - tensor_list, spec_list = list(zip(*fsdp_tensor_list)) # type: ignore[assignment] - if self.fsdp_mesh is not None: - fsdp_unshard_tensor_list = self._fsdp_foreach_allgather(tensor_list, spec_list) # type: ignore - else: - fsdp_unshard_tensor_list = tensor_list # type: ignore - - saved_fused_tensor_list: list[torch.Tensor] = [] - hf_keys_list: list[list[str]] = [] - - for load_spec, fsdp_unshared_tensor in zip(spec_list, fsdp_unshard_tensor_list): - hf_keys = load_spec.hf_keys - - if update_weights_for_rl: - hf_keys_list.append(hf_keys) - saved_fused_tensor_list.append(fsdp_unshared_tensor) - else: - if load_spec.group is not None: - all_hf_keys_list: list[None] | list[list[str]] = [None for _ in range(load_spec.group.size())] - dist.all_gather_object(all_hf_keys_list, hf_keys, group=load_spec.group) - all_hf_keys_list = cast(list[list[str]], all_hf_keys_list) - all_hf_keys = list(chain(*all_hf_keys_list)) - else: - all_hf_keys = hf_keys - - current_rank = dist.get_rank() - - expected_fused_save_ranks = self._get_ranks_to_save_fused_tensor(len(all_hf_keys)) - hardcode_fused_save_ranks = list( - range(min((dist.get_world_size(), self.config.hf_save_cfg.max_save_rank))) - ) - - key_per_rank = len(all_hf_keys) / len(hardcode_fused_save_ranks) - # assert key_per_rank.is_integer(), ( - # f"XTuner Internal Error, size of all_hf_keys: {len(all_hf_keys)}, " - # f"size of `fused_save_ranks` {len(fused_save_ranks)}" - # ) - if not key_per_rank.is_integer(): - key_per_rank = len(all_hf_keys) / len(expected_fused_save_ranks) - - start = int(current_rank * key_per_rank) - end = int(start + key_per_rank) - - _hf_key_list = all_hf_keys[start:end] - - if not _hf_key_list: - continue - - hf_keys_list.append(_hf_key_list) - - assert load_spec.dim is not None - if load_spec.group is not None: - assert load_spec.dim is not None - _gathered_tensor_list = [ - torch.zeros_like(fsdp_unshared_tensor) for _ in range(load_spec.group.size()) - ] - dist.all_gather(_gathered_tensor_list, fsdp_unshared_tensor, group=load_spec.group) - _gathered_tensor = torch.cat(_gathered_tensor_list, dim=load_spec.dim) - else: - _gathered_tensor = fsdp_unshared_tensor - hf_tensor_size = _gathered_tensor.shape[load_spec.dim] / len(all_hf_keys) - _saved_fused_tensor = torch.index_select( - _gathered_tensor, - dim=load_spec.dim, - index=torch.arange( - int(start * hf_tensor_size), - int(end * hf_tensor_size), - dtype=torch.int64, - device=_gathered_tensor.device, - ), - ) - saved_fused_tensor_list.append(_saved_fused_tensor) - - # Split the fused tensor into hf tensors - hf_tensor_list: list[torch.Tensor] = [] - # used in self._to_float8 to determine whether to convert a unshard hf_tensor to fp8 - fsdp_shard_tensor_list: list[torch.Tensor] = [] - # `origin_tensor_list` is only used to mark, which tensors are float8 weights for the - # `_to_float8` function - origin_tensor_list: list[torch.Tensor] = [] - - for saved_tensor, load_spec, hf_keys, origin_tensor in zip( - saved_fused_tensor_list, spec_list, hf_keys_list, tensor_list - ): - dim = cast(int, load_spec.dim) - hf_tensor_size = saved_tensor.shape[dim] / len(hf_keys) - assert hf_tensor_size.is_integer(), "Internal Error, hf_tensor_size is not integer" - hf_tensor_size = int(hf_tensor_size) - hf_tensor = saved_tensor.split([hf_tensor_size] * len(hf_keys), dim=dim) - hf_tensor_list.extend(hf_tensor) - fsdp_shard_tensor_list.extend([saved_tensor] * len(hf_tensor)) - origin_tensor_list.extend([origin_tensor] * len(hf_tensor)) - - name_list = list(chain.from_iterable(hf_keys_list)) - hf_tensor_list = [ - self.param_to_safetensor(safetensor, name) for safetensor, name in zip(hf_tensor_list, name_list) - ] - - if dtype == torch.float8_e4m3fn: - hf_tensor_list_new, name_list_new = self._to_float8( - hf_tensor_list, name_list, origin_tensor_list, dtype + save_item = self._make_hf_save_item( + param, + load_spec, + checkpoint_dtype=dtype, + is_buffer=load_spec.name in buffer_names, + distributed_save=distributed_save, + preserved_fused_shard_group=preserved_fused_shard_group, + ) + if bucket_bytes + save_item.byte_size > bucket_size and bucket: + yield self._build_hf_param_bucket( + bucket, + dtype=dtype, + device=device, ) - return hf_tensor_list_new, name_list_new - - hf_tensor_list = [t.to(device=device) for t in hf_tensor_list] + bucket_bytes = 0 + bucket = [] - return hf_tensor_list, name_list + bucket_bytes += save_item.byte_size + bucket.append(save_item) - if bucket_size is None: - bucket_size = self.config.hf_save_cfg.bucket_size - safetensor_size = 0 - tensor_list: list[tuple[torch.Tensor, LoadSpec]] = [] - name_list: list[str] = [] + if bucket: + yield self._build_hf_param_bucket( + bucket, + dtype=dtype, + device=device, + ) - for param, load_spec in params: - local_tensor = param._local_tensor if isinstance(param, DTensor) else param - local_tensor = local_tensor.to(dtype=self._get_save_dtype(load_spec.hf_keys[0], torch.bfloat16)) - tensor_size = self._get_tensor_size(param, dtype) - if safetensor_size + tensor_size > bucket_size and tensor_list: - hf_params, name_list = _get_hf_params(tensor_list, name_list) - yield name_list, hf_params - safetensor_size = tensor_size - name_list = load_spec.hf_keys.copy() - tensor_list = [(local_tensor, load_spec)] - continue - safetensor_size += tensor_size - tensor_list.append((local_tensor, load_spec)) - name_list.extend(load_spec.hf_keys) + def _make_hf_save_item( + self, + param: torch.Tensor, + load_spec: LoadSpec, + *, + checkpoint_dtype: torch.dtype, + is_buffer: bool, + distributed_save: bool, + preserved_fused_shard_group: dist.ProcessGroup | None, + ) -> _HFSaveBucketItem: + """Normalize one runtime parameter and compile its checkpoint save + policy.""" + # LoadSpec records both continuous and interleaved shard history, so every + # parameter enters the same SavePlan path from its runtime-local tensor. + runtime_tensor = param._local_tensor if isinstance(param, DTensor) else param + if runtime_tensor.is_floating_point() and not is_buffer: + save_dtype = self._get_save_dtype(load_spec.global_hf_keys[0], torch.bfloat16) + checkpoint_tensor = runtime_tensor.to(dtype=save_dtype) + else: + # Persistent buffers, e.g. FoPE rotary coefficients, keep their runtime dtype. + checkpoint_tensor = runtime_tensor + + return _HFSaveBucketItem( + tensor=checkpoint_tensor, + save_plan=load_spec.plan_hf_save( + distributed_save=distributed_save, + preserve_process_group=preserved_fused_shard_group, + ), + runtime_is_float8=is_float8_weight(runtime_tensor), + byte_size=self._get_tensor_size(runtime_tensor, checkpoint_dtype), + ) - if tensor_list: - hf_params, name_list = _get_hf_params(tensor_list, name_list) - yield name_list, hf_params + def _load_spec_params(self) -> list[tuple[torch.Tensor, LoadSpec]]: + ret: list[tuple[torch.Tensor, LoadSpec]] = [] + for name, param in self.state_dict().items(): + name = self._clean_param_name(name) + load_spec = self.load_spec_mapping.get(name) + if load_spec is None: + raise ValueError(f"Internal Error. Parameter {name} not found in load_spec_mapping.") + ret.append((param, load_spec)) + return ret - def _get_same_hf_param( + def _build_hf_param_bucket( self, - params: list[tuple[torch.Tensor, LoadSpec]], + bucket: list[_HFSaveBucketItem], dtype: torch.dtype, - device: torch.device | str = "cpu", - bucket_size: int | None = None, - ) -> Generator[tuple[list[str], list[torch.Tensor]], None, None]: - if not params: - return + device: torch.device | str, + ) -> tuple[list[str], list[torch.Tensor]]: + name_list: list[str] = [] + tensor_list: list[torch.Tensor] = [] + runtime_is_float8_list: list[bool] = [] - ignored_params, params = self._split_ignored_params(params) - if ignored_params: - fp32_name_list: list[str] = [load_spec.hf_keys[0] for _, load_spec in ignored_params] - fp32_tensor_list: list[torch.Tensor] = [ - param._local_tensor if isinstance(param, DTensor) else param for param, _ in ignored_params - ] - yield fp32_name_list, fp32_tensor_list + full_tensor_list = unshard_tensors_for_hf_save( + [item.tensor for item in bucket], + [item.save_plan for item in bucket], + ) + for full_tensor, save_item in zip(full_tensor_list, bucket, strict=True): + hf_names, hf_tensors = self._split_hf_tensors_for_save(full_tensor, save_item.save_plan) + name_list.extend(hf_names) + tensor_list.extend(hf_tensors) + runtime_is_float8_list.extend([save_item.runtime_is_float8] * len(hf_tensors)) - if bucket_size is None: - bucket_size = self.config.hf_save_cfg.bucket_size - safetensor_size = 0 - tensor_list: list[torch.Tensor] = [] - load_spec_list: list[LoadSpec] = [] - name_list: list[str] = [] - buffer_tensor_list: list[torch.Tensor] = [] - buffer_name_list: list[str] = [] + if dtype == torch.float8_e4m3fn: + tensor_list, name_list = self._to_float8(tensor_list, name_list, runtime_is_float8_list, dtype) - for param, load_spec in params: - if not isinstance(param, DTensor): - # in case, param is a buffer of module, FSDP will not shard it, so it's not a DTensor - buffer_tensor_list.append(param) - buffer_name_list.append(load_spec.hf_keys[0]) - continue - local_tensor = param._local_tensor if isinstance(param, DTensor) else param + tensor_list = [tensor.to(device=device) for tensor in tensor_list] + return name_list, tensor_list + + def _split_hf_tensors_for_save( + self, + full_tensor: torch.Tensor, + save_plan: HFSavePlan, + ) -> tuple[list[str], list[torch.Tensor]]: + if not save_plan.hf_keys: + return [], [] + + if len(save_plan.hf_keys) == 1: if ( - self.fsdp_config is not None - and self.fsdp_config.fp32_lm_head - and load_spec.hf_keys[0] == "lm_head.weight" + not save_plan.preserves_shards + and save_plan.distributed_save + and dist.is_initialized() + and dist.get_rank() != 0 ): - log_rank0.info(f"handling same hf param: {load_spec.hf_keys} separately") - lm_head_tensor_list = self._fsdp_foreach_allgather([local_tensor], [load_spec]) - lm_head_tensor_list = [ - self.param_to_safetensor(safetensor, name) - for safetensor, name in zip(lm_head_tensor_list, load_spec.hf_keys.copy()) - ] - lm_head_tensor_list = [t.to(device=device) for t in lm_head_tensor_list] - yield load_spec.hf_keys.copy(), lm_head_tensor_list - del lm_head_tensor_list, local_tensor - continue - else: - local_tensor = local_tensor.to(dtype=self._get_save_dtype(load_spec.hf_keys[0], torch.bfloat16)) - tensor_size = self._get_tensor_size(param, dtype) - if safetensor_size + tensor_size > bucket_size and tensor_list: - if self.fsdp_mesh is not None: - gathered_tensor_list = self._fsdp_foreach_allgather(tensor_list, load_spec_list) - else: - gathered_tensor_list = tensor_list - gathered_tensor_list = [ - self.param_to_safetensor(safetensor, name) - for safetensor, name in zip(gathered_tensor_list, name_list) - ] - if dtype == torch.float8_e4m3fn: - gathered_tensor_list, name_list = self._to_float8( - gathered_tensor_list, name_list, tensor_list, dtype - ) - gathered_tensor_list = [t.to(device=device) for t in gathered_tensor_list] - yield name_list, gathered_tensor_list - safetensor_size = tensor_size - name_list = load_spec.hf_keys.copy() - tensor_list = [local_tensor] - load_spec_list = [load_spec] - continue - safetensor_size += tensor_size - tensor_list.append(local_tensor) - name_list.append(load_spec.hf_keys[0]) - load_spec_list.append(load_spec) + return [], [] + hf_name = save_plan.hf_keys[0] + return [hf_name], [self.param_to_safetensor(full_tensor, hf_name)] + + dim = save_plan.fused_dim + assert dim is not None, "fused_dim must be set when saving fused HF tensors" + if save_plan.preserves_shards: + hf_names = save_plan.hf_keys + tensor_to_split = full_tensor + else: + hf_names = save_plan.hf_keys.copy() + key_start, key_end = ( + self._hf_save_key_range(save_plan) + if save_plan.distributed_save + else ( + 0, + len(hf_names), + ) + ) + if key_start == key_end: + return [], [] + hf_names = hf_names[key_start:key_end] + key_size = full_tensor.shape[dim] / len(save_plan.hf_keys) + assert key_size.is_integer(), ( + f"Fused dim size {full_tensor.shape[dim]} is not divisible by " + f"{len(save_plan.hf_keys)} HF keys for {save_plan.name}" + ) + key_size = int(key_size) + # Keep the legacy save behavior here: fp8 per-block quant kernels have had correctness issues with + # non-zero-storage-offset views, so materialize the save-rank slice before splitting HF keys. + index = torch.arange( + key_start * key_size, + key_end * key_size, + dtype=torch.int64, + device=full_tensor.device, + ) + tensor_to_split = torch.index_select(full_tensor, dim=dim, index=index) - if tensor_list: - if self.fsdp_mesh is not None: - gathered_tensor_list = self._fsdp_foreach_allgather(tensor_list, load_spec_list) - else: - gathered_tensor_list = tensor_list + if not hf_names: + return [], [] - gathered_tensor_list = [ - self.param_to_safetensor(safetensor, name) for safetensor, name in zip(gathered_tensor_list, name_list) - ] - if dtype == torch.float8_e4m3fn: - gathered_tensor_list, name_list = self._to_float8(gathered_tensor_list, name_list, tensor_list, dtype) - gathered_tensor_list = [t.to(device=device) for t in gathered_tensor_list] - yield name_list, gathered_tensor_list + hf_tensor_size = tensor_to_split.shape[dim] / len(hf_names) + assert hf_tensor_size.is_integer(), ( + f"Fused dim size {tensor_to_split.shape[dim]} is not divisible by " + f"{len(hf_names)} HF keys for {save_plan.name}" + ) + split_size = int(hf_tensor_size) + hf_tensors = tensor_to_split.split([split_size] * len(hf_names), dim=dim) + return ( + hf_names, + [self.param_to_safetensor(safetensor, name) for safetensor, name in zip(hf_tensors, hf_names)], + ) - if buffer_tensor_list: - yield buffer_name_list, buffer_tensor_list + def _hf_save_key_range(self, save_plan: HFSavePlan) -> tuple[int, int]: + if not dist.is_initialized(): + return 0, len(save_plan.hf_keys) - def _is_ignored_params(self, key: str): - patterns = self.config.hf_save_cfg.fp32_keys_pattern - if patterns is None: - return False - return any(re.search(p, key) for p in patterns) - - def _split_ignored_params( - self, params: list[tuple[torch.Tensor, LoadSpec]] - ) -> tuple[list[tuple[torch.Tensor, LoadSpec]], list[tuple[torch.Tensor, LoadSpec]]]: - if not self.config.hf_save_cfg.fp32_keys_pattern: - return [], params - ignored_params = [(p, l) for p, l in params if self._is_ignored_params(l.hf_keys[0])] - remaining = [(p, l) for p, l in params if not self._is_ignored_params(l.hf_keys[0])] - return ignored_params, remaining + current_rank = dist.get_rank() + save_ranks = self._get_fused_save_ranks(len(save_plan.hf_keys)) + if current_rank not in save_ranks: + return 0, 0 + + key_per_rank = len(save_plan.hf_keys) // len(save_ranks) + rank_index = save_ranks.index(current_rank) + start = rank_index * key_per_rank + return start, start + key_per_rank # TODO: Using `xtuenr.v1.utils.misc.clean_param_name` def _clean_param_name(self, name: str) -> str: @@ -1699,46 +1503,11 @@ def _clean_param_name(self, name: str) -> str: name = name.replace("_orig_mod.", "") return name - def _group_param_by_load_spec(self, load_enum: LoadEnum): - """Group the parameters by load spec.""" - ret = [] - for name, param in self.state_dict().items(): - load_spec = self.load_spec_mapping.get(name) - if load_spec is None: - raise ValueError(f"Internal Error. Parameter {name} not found in load_spec_mapping.") - if load_spec.load_enum == load_enum: - ret.append((param, load_spec)) - else: - continue - return ret - def _get_tensor_size(self, tensor: torch.Tensor, dtype: torch.dtype) -> int: """Get the size of the tensor in bytes.""" # return tensor.element_size() * tensor.numel() return dtype.itemsize * tensor.numel() - def _get_safe_tensor_num(self, dtype: torch.dtype) -> int: - """Get the size of the model in bytes.""" - bucket_size = self.config.hf_save_cfg.bucket_size - shard_size = 0 - same_size = 0 - fused_size = 0 - for name, param in self.state_dict().items(): - load_spec = self.load_spec_mapping.get(name) - if load_spec is None: - raise ValueError(f"Internal Error. Parameter {name} not found in load_spec_mapping.") - if load_spec.load_enum == LoadEnum.SHARD: - shard_size += self._get_tensor_size(param, dtype) - elif load_spec.load_enum == LoadEnum.SAME: - same_size += self._get_tensor_size(param, dtype) - elif load_spec.load_enum == LoadEnum.FUSED: - fused_size += self._get_tensor_size(param, dtype) - return ( - math.ceil(shard_size / bucket_size) - + math.ceil(same_size / bucket_size) - + math.ceil(fused_size / bucket_size) - ) - def _iter_hf_save_chunks( self, save_dtype: torch.dtype = torch.bfloat16, @@ -1746,43 +1515,19 @@ def _iter_hf_save_chunks( device: torch.device | str = "cpu", ) -> Generator[tuple[str, list[str], list[torch.Tensor]], None, None]: assert save_dtype in [torch.float8_e4m3fn, torch.bfloat16], f"save_dtype {save_dtype} is not supported" - - shard_gen = self._get_shard_hf_param( - self._group_param_by_load_spec(LoadEnum.SHARD), - dtype=save_dtype, - device=device, - ) - same_gen = self._get_same_hf_param( - self._group_param_by_load_spec(LoadEnum.SAME), - dtype=save_dtype, - device=device, - ) - fused_gen = self._get_fused_hf_param( - self._group_param_by_load_spec(LoadEnum.FUSED), - dtype=save_dtype, - device=device, - ) - - is_others_save_rank = not dist.is_initialized() or dist.get_rank() == 0 save_rank = dist.get_rank() if dist.is_initialized() else 0 - saved_names: set[str] = set() safetensor_index = 0 - for name_list, hf_tensor_list in fused_gen: + param_gen = self._get_hf_param( + self._load_spec_params(), + dtype=save_dtype, + device=device, + distributed_save=True, + ) + for name_list, hf_tensor_list in param_gen: if not name_list: continue - safetensor_index += 1 - safetensor_name = f"{safetensors_prefix}-{safetensor_index:04d}-fused-save_rank{save_rank}.safetensors" - saved_names.update(name_list) - yield safetensor_name, name_list, hf_tensor_list - - safetensor_index = 0 - for name_list, hf_tensor_list in chain(same_gen, shard_gen): - safetensor_index += 1 - safetensor_name = f"{safetensors_prefix}-{safetensor_index:04d}-others-save_rank{save_rank}.safetensors" - if not is_others_save_rank: - continue unique_name_list: list[str] = [] unique_hf_tensor_list: list[torch.Tensor] = [] @@ -1793,6 +1538,8 @@ def _iter_hf_save_chunks( unique_name_list.append(name) unique_hf_tensor_list.append(hf_tensor) if unique_name_list: + safetensor_index += 1 + safetensor_name = f"{safetensors_prefix}-{safetensor_index:04d}-save_rank{save_rank}.safetensors" yield safetensor_name, unique_name_list, unique_hf_tensor_list def _write_hf_save_plan(self, save_plan: _HFSavePlan) -> list[str]: @@ -1920,12 +1667,7 @@ def _save_hf( DEVICE_MODULE.empty_cache() assert save_dtype in [torch.float8_e4m3fn, torch.bfloat16], f"save_dtype {save_dtype} is not supported" - # TODO: Support fp8 saving - shard_gen = self._get_shard_hf_param(self._group_param_by_load_spec(LoadEnum.SHARD), dtype=save_dtype) - same_gen = self._get_same_hf_param(self._group_param_by_load_spec(LoadEnum.SAME), dtype=save_dtype) - fused_gen = self._get_fused_hf_param(self._group_param_by_load_spec(LoadEnum.FUSED), dtype=save_dtype) - - is_others_save_rank = not dist.is_initialized() or dist.get_rank() == 0 + param_gen = self._get_hf_param(self._load_spec_params(), dtype=save_dtype, distributed_save=True) # Tell me why! why! old cao! @HIT-cwh # mp_context = multiprocessing.get_context("fork") @@ -1937,53 +1679,38 @@ def _save_hf( else: save_rank = 0 - # Sepreately save fused parameters and others to make sure each saving rank will not save - # dupilicated keys - # save_futures = [] weight_map = {} safetensor_index = 0 - for name_list, hf_tensor_list in fused_gen: + for name_list, hf_tensor_list in param_gen: if not name_list: continue + # Tied weights may map multiple runtime tensors to the same HF key; keep the first one. + unique_name_list = [] + unique_hf_tensor_list = [] + for name, hf_tensor in zip(name_list, hf_tensor_list): + if name in weight_map: + continue + unique_name_list.append(name) + unique_hf_tensor_list.append(hf_tensor) + + if not unique_name_list: + continue + safetensor_index += 1 - safetensor_name = f"{safetensors_prefix}-{safetensor_index:04d}-fused-save_rank{save_rank}.safetensors" - weight_map.update(dict.fromkeys(name_list, safetensor_name)) + safetensor_name = f"{safetensors_prefix}-{safetensor_index:04d}-save_rank{save_rank}.safetensors" + weight_map.update(dict.fromkeys(unique_name_list, safetensor_name)) assert save_executor is not None, "Internal Error, save_executor should not be None" future = save_executor.submit( _save_file, - dict(zip(name_list, hf_tensor_list)), + dict(zip(unique_name_list, unique_hf_tensor_list)), hf_dir / safetensor_name, ) save_futures.append(future) self._wait_save_task(save_futures) - safetensor_index = 0 - for name_list, hf_tensor_list in chain(same_gen, shard_gen): - safetensor_index += 1 - safetensor_name = f"{safetensors_prefix}-{safetensor_index:04d}-others-save_rank{save_rank}.safetensors" - - if is_others_save_rank: - # for tie_word_embeddings, we need to make sure each key is only saved once - unique_name_list = [] - unique_hf_tensor_list = [] - for name, hf_tensor in zip(name_list, hf_tensor_list): - if name not in weight_map: - unique_name_list.append(name) - unique_hf_tensor_list.append(hf_tensor) - weight_map[name] = safetensor_name - - assert save_executor is not None, "Internal Error, save_executor should not be None" - future = save_executor.submit( - _save_file, - dict(zip(unique_name_list, unique_hf_tensor_list)), - hf_dir / safetensor_name, - ) - save_futures.append(future) - self._wait_save_task(save_futures) - if save_futures: wait(save_futures) for future in save_futures: @@ -2061,14 +1788,7 @@ def _load_params_from_module(module: nn.Module, module_prefix: str): if load_spec is None: raise RuntimeError(f"Internal Error. Parameter {name} not found in load_spec_mapping.") - if load_spec.load_enum == LoadEnum.SAME: - _missing_keys = self._load_same_hf_param(param, load_spec, checkpoint_loader) - elif load_spec.load_enum == LoadEnum.FUSED: - _missing_keys = self._load_fused_hf_param(param, load_spec, checkpoint_loader) - elif load_spec.load_enum == LoadEnum.SHARD: - _missing_keys = self._load_shard_hf_param(param, load_spec, checkpoint_loader) - else: - raise RuntimeError(f"Unsupported load_enum: {load_spec.load_enum}") + _missing_keys = self._load_hf_param(param, load_spec, checkpoint_loader) missing_keys.update(_missing_keys) if not _missing_keys: @@ -2091,199 +1811,81 @@ def _load_params_from_module(module: nn.Module, module_prefix: str): return loaded_keys, unloaded_keys, missing_keys - def _is_loaded_param_fp8(self, hf_key: str, checkpoint_loader: HFCheckpointLoader) -> bool: - hf_key_scale_inv = hf_key + "_scale_inv" - return checkpoint_loader.is_key_exist(hf_key) and checkpoint_loader.is_key_exist(hf_key_scale_inv) - - def _load_fp8(self, hf_key: str, checkpoint_loader: HFCheckpointLoader) -> torch.Tensor | None: - hf_key_scale_inv = hf_key + "_scale_inv" - loaded_tensor_fp8 = checkpoint_loader.load(hf_key) - loaded_tensor_scales = checkpoint_loader.load(hf_key_scale_inv) - if loaded_tensor_fp8 is None or loaded_tensor_scales is None: - return None - - from xtuner.v1.float8.triton_kernels import per_block_dequant_gemm - - loaded_tensor = per_block_dequant_gemm( - loaded_tensor_fp8.to(DEVICE), - loaded_tensor_scales.to(DEVICE), - ) - return loaded_tensor - - def _load_same_hf_param( + def _load_hf_param( self, param: torch.Tensor, load_spec: LoadSpec, checkpoint_loader: HFCheckpointLoader - ) -> list[str]: # return missing key - local_tensor = param._local_tensor if isinstance(param, DTensor) else param - hf_key = load_spec.hf_keys[0] - if self._is_loaded_param_fp8(hf_key, checkpoint_loader): - if not _is_float8_available(): - raise RuntimeError( - f"Float8 is not available on {DEVICE}. Please convert the checkpoint from float8 to bfloat16 on SM89 or later (H100+ GPUs)." - ) - loaded_tensor = self._load_fp8(hf_key, checkpoint_loader) - else: - loaded_tensor = checkpoint_loader.load(hf_key) - if loaded_tensor is None: - return [hf_key] - - loaded_tensor = loaded_tensor.to(local_tensor.device) + ) -> list[str]: + """Unified HF load path for a single parameter / buffer. - if ( - self.fsdp_mesh is not None - and isinstance(param, nn.Parameter) - and isinstance(param, DTensor) - and any(isinstance(p, Shard) for p in param.placements) - ): - shape_before_fsdp = load_spec.shape - _, _offset = compute_local_shape_and_global_offset( - shape_before_fsdp, self.fsdp_mesh, [Shard(self.FSDP_SHARD_DIM)] - ) - fsdp_start = _offset[self.FSDP_SHARD_DIM] - fsdp_end = fsdp_start + local_tensor.shape[self.FSDP_SHARD_DIM] + ``LoadSpec.plan_hf_load`` computes this rank's HF keys and canonical source-to-local copy regions. This + method reads and dequantizes those keys, then lets the plan drive model canonicalization and local writes. - start = fsdp_start - end = fsdp_end - else: - start = None - end = None + Returns the list of hf_keys that were expected but missing from the + checkpoint; callers aggregate these for strict-mode reporting. + """ + local_tensor = param._local_tensor if isinstance(param, DTensor) else param + load_plan = load_spec.plan_hf_load() + loaded_tensors, missing_keys = self._read_hf_tensors( + load_plan.hf_keys, + checkpoint_loader, + device=local_tensor.device, + ) + if missing_keys: + return missing_keys - self.safetensors_to_params( - [loaded_tensor], local_tensor, param_name=load_spec.name, start=start, end=end, dim=load_spec.dim + load_plan.load_into( + loaded_tensors, + local_tensor, + canonicalize=self.hf_tensor_to_canonical, ) return [] - def _load_fused_hf_param( - self, param: torch.Tensor, load_spec: LoadSpec, checkpoint_loader: HFCheckpointLoader - ) -> list[str]: - # For expert parallel - # NOTE: - # 1. Get `hf-keys` required by sharded param (sharded by ep group) - # 2. Asumming FSDP sharding the tensor at the same dim as ep group, Get the twice sharded - # `hf-keys`. For example, if we have 128 experts with ep-size 8 and fsdp-size 16. The - # the param sharded by ep group will have 128/8 = 16 `hf-keys`, and the param further sharded - # by FSDP will only have 128/8/16 = 1 `hf-keys` - # 3. Calculating the `offset` and `size` of FSDP param base on the ep sharded params, and fill - # the FSDP param with the loaded tensor. - - hf_keys = load_spec.hf_keys - local_tensor = param._local_tensor if isinstance(param, DTensor) else param - - assert load_spec.dim == self.FSDP_SHARD_DIM, "Only support FSDP and model parallel sharding at the same dim!" - if self.fsdp_mesh is not None: - shape_before_fsdp = load_spec.shape - if is_float8_weight(local_tensor): - # fp8 weights may be padded, so we need to calculate the hf_key_size base on local_tensor._ori_shape - if load_spec.group is None: - hf_key_size = local_tensor._ori_shape[self.FSDP_SHARD_DIM] / len(hf_keys) # type: ignore - else: - hf_key_size = ( - local_tensor._ori_shape[self.FSDP_SHARD_DIM] # type: ignore - / dist.get_world_size(group=load_spec.group) - / len(hf_keys) - ) - else: - # shape_before_fsdp[self.FSDP_SHARD_DIM] == local_tensor.shape[self.FSDP_SHARD_DIM] / dist.get_world_size(group=load_spec.group) - hf_key_size = shape_before_fsdp[self.FSDP_SHARD_DIM] / len(hf_keys) - assert hf_key_size.is_integer(), ( - "Model parallel sharding size should be divisible by fused huggingface tensors!" - ) - hf_key_size = int(hf_key_size) - _, _offset = compute_local_shape_and_global_offset( - shape_before_fsdp, self.fsdp_mesh, [Shard(self.FSDP_SHARD_DIM)] - ) - fsdp_start = _offset[self.FSDP_SHARD_DIM] - fsdp_end = fsdp_start + local_tensor.shape[self.FSDP_SHARD_DIM] - - hf_keys_start = int(fsdp_start / hf_key_size) - hf_keys_end = math.ceil(fsdp_end / hf_key_size) - - # Empty pad by fsdp - if hf_keys_start == hf_keys_end: - return [] - - hf_keys = hf_keys[hf_keys_start:hf_keys_end] - - start = fsdp_start % hf_key_size - end = start + local_tensor.shape[self.FSDP_SHARD_DIM] - else: - start = None - end = None - + def _read_hf_tensors( + self, + hf_keys: list[str], + checkpoint_loader: HFCheckpointLoader, + *, + device: torch.device, + ) -> tuple[list[torch.Tensor], list[str]]: + """Read one plan's HF keys, dequantizing checkpoint FP8 when + present.""" + loaded_tensors: list[torch.Tensor] = [] missing_keys: list[str] = [] - _loaded_tensor: list[torch.Tensor] = [] for hf_key in hf_keys: - weight = self._load_fp8(hf_key, checkpoint_loader) - if weight is None: - weight = checkpoint_loader.load(hf_key) - if weight is None: + tensor = self._read_hf_tensor(hf_key, checkpoint_loader) + if tensor is None: missing_keys.append(hf_key) - continue - _loaded_tensor.append(weight.to(local_tensor.device)) - - if not _loaded_tensor: - return missing_keys - - if not hf_keys: - # fp8 pad - assert self.config.float8_cfg is not None - # assert self.fsdp_config is not None and self.fsdp_config.ep_size == 1, ( - # "Only support fp8 pad for MoE with ep_size == 1" - # ) - local_tensor.zero_() # type: ignore # padded part must be set to 0 - return missing_keys - - if missing_keys: - return missing_keys - - self.safetensors_to_params( - _loaded_tensor, local_tensor, param_name=load_spec.name, start=start, end=end, dim=load_spec.dim - ) - return missing_keys - - def _load_shard_hf_param( - self, param: torch.Tensor, load_spec: LoadSpec, checkpoint_loader: HFCheckpointLoader - ) -> list[str]: - # For tensor parallel - # NOTE: - # 1. Get `hf-keys` required by sharded param (sharded by tp group, only 1 key) - # 2. all gather the sharded param across tp group - # 3 Fill the sharded param with the sliced gathered tensor. - hf_key = load_spec.hf_keys[0] - local_tensor = param._local_tensor if isinstance(param, DTensor) else param - - loaded_tensor = checkpoint_loader.load(hf_key) - if loaded_tensor is None: - return [hf_key] + else: + loaded_tensors.append(tensor.to(device)) + return loaded_tensors, missing_keys - loaded_tensor = loaded_tensor.to(local_tensor.device) + def _read_hf_tensor( + self, + hf_key: str, + checkpoint_loader: HFCheckpointLoader, + ) -> torch.Tensor | None: + scale_key = hf_key + "_scale_inv" + is_checkpoint_fp8 = checkpoint_loader.is_key_exist(hf_key) and checkpoint_loader.is_key_exist(scale_key) + if not is_checkpoint_fp8: + return checkpoint_loader.load(hf_key) + if not _is_float8_available(): + raise RuntimeError( + f"Float8 is not available on {DEVICE}. Please convert the checkpoint from float8 " + "to bfloat16 on SM89 or later (H100+ GPUs)." + ) + return self._load_fp8(hf_key, checkpoint_loader) - assert load_spec.shard_start is not None and load_spec.shard_end is not None, ( - "load_spec.shard_start and load_spec.shard_end should not be None for sharded params" - ) + def _load_fp8(self, hf_key: str, checkpoint_loader: HFCheckpointLoader) -> torch.Tensor | None: + loaded_tensor_fp8 = checkpoint_loader.load(hf_key) + loaded_tensor_scales = checkpoint_loader.load(hf_key + "_scale_inv") + if loaded_tensor_fp8 is None or loaded_tensor_scales is None: + return None - if self.fsdp_mesh is not None: - shape_before_fsdp = load_spec.shape - _, _offset = compute_local_shape_and_global_offset( - shape_before_fsdp, self.fsdp_mesh, [Shard(self.FSDP_SHARD_DIM)] - ) - fsdp_start = _offset[self.FSDP_SHARD_DIM] - fsdp_end = fsdp_start + local_tensor.shape[self.FSDP_SHARD_DIM] + from xtuner.v1.float8.triton_kernels import per_block_dequant_gemm - start = fsdp_start + load_spec.shard_start - end = fsdp_end + load_spec.shard_start - else: - start = load_spec.shard_start - end = load_spec.shard_end - - self.safetensors_to_params( - safetensors=[loaded_tensor], - local_tensor=local_tensor, - param_name=load_spec.name, - start=start, - end=end, - dim=load_spec.dim, + return per_block_dequant_gemm( + loaded_tensor_fp8.to(DEVICE), + loaded_tensor_scales.to(DEVICE), ) - return [] def _has_meta_param(self, module: nn.Module, recurse: bool = False) -> bool: """Check if the module has meta parameters.""" @@ -2295,125 +1897,26 @@ def _has_meta_param(self, module: nn.Module, recurse: bool = False) -> bool: def _fsdp_foreach_allgather( self, tensor_list: list[torch.Tensor], load_spec_list: list[LoadSpec] ) -> list[torch.Tensor]: - assert self.fsdp_mesh is not None, "Internal Error, fsdp_mesh should not be None" - origin_fsdp_size = [] - padded_tensor_list = [] - - for param, load_spec in zip(tensor_list, load_spec_list): - shape_before_fsdp = load_spec.shape[self.FSDP_SHARD_DIM] - padded_size = math.ceil(shape_before_fsdp / self.fsdp_mesh.size()) - pad_list = [0] * (2 * param.dim()) - pad_idx = 2 * (param.dim() - 1 - self.FSDP_SHARD_DIM) - pad_list[pad_idx + 1] = padded_size - param.shape[self.FSDP_SHARD_DIM] - padded_tensor = F.pad(param, pad_list) - padded_tensor_list.append(padded_tensor) - if is_float8_weight(param): - dim_before_fsdp: int - if load_spec.group is None: - dim_before_fsdp = param._ori_shape[self.FSDP_SHARD_DIM] # type: ignore - else: - dim_before_fsdp = param._ori_shape[self.FSDP_SHARD_DIM] // dist.get_world_size( # type: ignore - group=load_spec.group - ) - origin_fsdp_size.append(dim_before_fsdp) - else: - origin_fsdp_size.append(load_spec.shape[self.FSDP_SHARD_DIM]) - - _fsdp_unsharded_tensor_list = foreach_all_gather(padded_tensor_list, self.fsdp_mesh.get_group()) - fsdp_unsharded_tensor_list = [] - - # Concatenate the tensors along the FSDP shard dim - fuse_without_alloc = self.FSDP_SHARD_DIM == 0 and len(_fsdp_unsharded_tensor_list) == 1 - for tensors, size in zip(_fsdp_unsharded_tensor_list, origin_fsdp_size): - if fuse_without_alloc: - # In the case of only one big tensor in tensor_list, the partition of tensors are contiguous. - # Therefore the cat and index_select operation can be omitted, - # and use _fuse_contiguous_chunks_without_alloc instead to reduce device peak memory. - # e.g. When a fused MoE weight exceeds bucket_size given, len(tensor_list) would be 1, - # and tensor is not None reducing peak device memory. - tensor = self._fuse_contiguous_chunks_without_alloc(tensors) - else: - tensor = torch.cat(tensors, dim=self.FSDP_SHARD_DIM) - unpaded_tensor = tensor.narrow(self.FSDP_SHARD_DIM, 0, size) - pad_tensor = tensor.narrow(self.FSDP_SHARD_DIM, size, tensor.shape[self.FSDP_SHARD_DIM] - size) - assert (pad_tensor == 0).all(), f"Internal Error, padded tensor is not zero {pad_tensor}!" - # when self.FSDP_SHARD_DIM != 0, narrow operation may lead to non-contiguous tensor - fsdp_unsharded_tensor_list.append(unpaded_tensor.contiguous()) + if self.fsdp_mesh is None: + return tensor_list - return fsdp_unsharded_tensor_list + fsdp_group = self.fsdp_mesh.get_group() + save_plan_list = [load_spec.plan_hf_save(gather_process_group=fsdp_group) for load_spec in load_spec_list] + return unshard_tensors_for_hf_save(list(tensor_list), save_plan_list) @staticmethod - def _fuse_contiguous_chunks_without_alloc(tensors: list[torch.Tensor]) -> torch.Tensor: - """Fuse contiguous chunks without extra memory allocation. + def _is_same_process_group(left: dist.ProcessGroup, right: dist.ProcessGroup) -> bool: + if left is right: + return True + return dist.get_process_group_ranks(left) == dist.get_process_group_ranks(right) - Return None if not possible. - """ - if not tensors: - raise ValueError("tensors should not be empty") - base = tensors[0] - storage = base.untyped_storage() - dtype = base.dtype - device = base.device - stride = base.stride() - - inner_stride = stride[1:] - inner_elems = math.prod(base.shape[1:]) if base.dim() > 1 else 1 - - chunks = [] - for t in tensors: - # we should check both storage and stride to ensure contiguity - # regardless of the implementation of foreach_all_gather - if t.untyped_storage().data_ptr() != storage.data_ptr(): - raise RuntimeError("Tensors are not sharing the same storage.") - if t.stride()[1:] != inner_stride: - raise RuntimeError("Tensors have mismatched strides.") - chunks.append((t.storage_offset(), t.shape[0], t)) - chunks.sort(key=lambda x: x[0]) - - expected_offset = chunks[0][0] - total_rows = 0 - for offset, rows, _ in chunks: - if offset != expected_offset: - raise RuntimeError("Tensors are not contiguous in the storage") - expected_offset += rows * inner_elems - total_rows += rows - - size = (total_rows, *base.shape[1:]) - flat = torch.empty(0, dtype=dtype, device=device) - flat.set_(storage, chunks[0][0], size, stride) - return flat - - def _get_ranks_to_save_fused_tensor(self, fused_size: int) -> list[int]: - # Goal: decide how many ranks are used to store model/expert parameters. - # Policy: choose d such that: - # 1) d is a positive divisor of world_size, - # 2) d <= num_experts, - # 3) d is as close to num_experts as possible under (1)(2). - # This is equivalent to: pick the largest divisor of world_size that does not exceed num_experts. - # Rationale: ensures feasibility under expert count, maximizes utilization, and yields balanced groups. - # Implementation hint: enumerate divisor pairs (i, world_size // i) for i up to sqrt(world_size) and keep the max d <= num_experts. - # Complexity: O(sqrt(world_size)). + def _get_fused_save_ranks(self, hf_key_count: int) -> list[int]: world_size = dist.get_world_size() - - if world_size >= fused_size: - return list(range(fused_size)) - - num_ranks_to_save = None - best_diff = None - - i = 1 - while i * i <= fused_size: - if fused_size % i == 0: - for d in (i, fused_size // i): - diff = abs(d - world_size) - if ( - num_ranks_to_save is None - or (diff < best_diff) # type: ignore - or (diff == best_diff and d < num_ranks_to_save) - ): - num_ranks_to_save, best_diff = d, diff - i += 1 - return list(range(cast(int, num_ranks_to_save))) + max_save_ranks = min(world_size, self.config.hf_save_cfg.max_save_rank, hf_key_count) + for save_rank_count in range(max_save_ranks, 0, -1): + if hf_key_count % save_rank_count == 0: + return list(range(save_rank_count)) + raise RuntimeError(f"Unable to choose save ranks for {hf_key_count} fused HF keys") def _to_device_dtype( self, @@ -2609,7 +2112,17 @@ def _collect_full_state_dict(self, module: nn.Module): ret = {} for name, param in module.state_dict().items(): # type: ignore[attr-defined] if isinstance(param, DTensor): - param = param.full_tensor() + from xtuner.v1.utils.interleaved_shard import ( + has_interleaved_placement, + reconstruct_full_tensor, + ) + + if has_interleaved_placement(param): + # `(Shard, InterleavedShard)`-style placements can't be redistributed; use the + # explicit reconstruct path instead of `.full_tensor()`. + param = reconstruct_full_tensor(param) + else: + param = param.full_tensor() ret[name] = param return ret diff --git a/xtuner/v1/model/dense/dense.py b/xtuner/v1/model/dense/dense.py index 71fd1c461d..47952f0c6d 100644 --- a/xtuner/v1/model/dense/dense.py +++ b/xtuner/v1/model/dense/dense.py @@ -295,6 +295,7 @@ def fully_shard( ) self.set_modules_to_forward_prefetch([self.embed_tokens, self.layers["0"]]) # type: ignore + self._init_load_spec() self._to_empty_meta() # Make sure it works properly when using fsdp diff --git a/xtuner/v1/model/moe/glm52.py b/xtuner/v1/model/moe/glm52.py index 838546677e..1c71eb4c54 100644 --- a/xtuner/v1/model/moe/glm52.py +++ b/xtuner/v1/model/moe/glm52.py @@ -157,40 +157,10 @@ def to_hf_key_list(self, key: str) -> list[str]: else: return [key] - def safetensors_to_params( - self, - safetensors: list[torch.Tensor], - local_tensor: torch.Tensor, - param_name: str, - start: int | None, - end: int | None, - dim: int | None, - ): - if len(safetensors) > 1: - assert dim is not None, "Internal Error dim must not be None when len(safetensors) > 1" - loaded_tensor = torch.cat(safetensors, dim=dim) - else: - loaded_tensor = safetensors[0] - - if ( - "fused_w1w3.weight" in param_name or "fused_w2.weight" in param_name - ) and loaded_tensor.ndim == local_tensor.ndim + 1: + def hf_tensor_to_canonical(self, name: str, loaded_tensor: torch.Tensor) -> torch.Tensor: + if ("fused_w1w3.weight" in name or "fused_w2.weight" in name) and loaded_tensor.ndim == 3: loaded_tensor = loaded_tensor.flatten(0, 1) - - if start is not None and end is not None: - start = min(start, loaded_tensor.shape[self.FSDP_SHARD_DIM]) - end = min(end, loaded_tensor.shape[self.FSDP_SHARD_DIM]) - loaded_tensor_slice = loaded_tensor.index_select( - dim=self.FSDP_SHARD_DIM, index=torch.arange(start, end, dtype=torch.int64, device=loaded_tensor.device) - ) - non_pad_len = end - start - local_tensor[:non_pad_len].copy_(loaded_tensor_slice) - - if non_pad_len < local_tensor.shape[self.FSDP_SHARD_DIM]: - assert self.config.float8_cfg is not None - local_tensor[non_pad_len:].zero_() - else: - local_tensor.copy_(loaded_tensor) + return loaded_tensor def param_to_safetensor( self, diff --git a/xtuner/v1/model/moe/gpt_oss.py b/xtuner/v1/model/moe/gpt_oss.py index 5f9e7bc083..1e484942c1 100644 --- a/xtuner/v1/model/moe/gpt_oss.py +++ b/xtuner/v1/model/moe/gpt_oss.py @@ -40,22 +40,8 @@ def to_hf_key_list(self, key: str) -> list[str]: else: return [key] - def safetensors_to_params( - self, - safetensors: list[torch.Tensor], - local_tensor: torch.Tensor, - param_name: str, - start: int | None, - end: int | None, - dim: int | None, - ): - if len(safetensors) > 1: - assert dim is not None, "Internal Error dim must not be None when len(safetensors) > 1" - loaded_tensor = torch.cat(safetensors, dim=dim) - else: - loaded_tensor = safetensors[0] - - if "fused_w1w3.weight" in param_name: + def hf_tensor_to_canonical(self, name: str, loaded_tensor: torch.Tensor) -> torch.Tensor: + if "fused_w1w3.weight" in name: # hf: num_experts, hidden_size, expert_dim * 2 # xtuner: num_experts * 2 * expert_dim, hidden_size num_experts, hidden_size = loaded_tensor.shape[:2] @@ -64,32 +50,19 @@ def safetensors_to_params( # # num_experts *2 * expert_dim, hidden_size loaded_tensor = loaded_tensor.transpose(1, 2).reshape(-1, hidden_size) - elif "fused_w2.weight" in param_name: + elif "fused_w2.weight" in name: # hf: num_experts, expert_dim, hidden_size # xtuner: num_experts * hidden_size, expert_dim loaded_tensor = loaded_tensor.transpose(1, 2).flatten(0, 1) - if "fused_w1w3.bias" in param_name: + if "fused_w1w3.bias" in name: # hf: num_experts, expert_dim * 2 - # xtuner: num_experts, 2 * expert_dim + # xtuner: num_experts * 2 * expert_dim (flattened so Expert TP can shard each projection stripe) num_experts = loaded_tensor.size(0) loaded_tensor = loaded_tensor.reshape(num_experts, -1, 2) - loaded_tensor = loaded_tensor.transpose(1, 2).reshape(num_experts, -1) - - if start is not None and end is not None: - start = min(start, loaded_tensor.shape[self.FSDP_SHARD_DIM]) - end = min(end, loaded_tensor.shape[self.FSDP_SHARD_DIM]) - loaded_tensor_slice = loaded_tensor.index_select( - dim=self.FSDP_SHARD_DIM, index=torch.arange(start, end, dtype=torch.int64, device=loaded_tensor.device) - ) - non_pad_len = end - start - local_tensor[:non_pad_len].copy_(loaded_tensor_slice) + loaded_tensor = loaded_tensor.transpose(1, 2).reshape(-1) - if non_pad_len < local_tensor.shape[self.FSDP_SHARD_DIM]: - assert self.config.float8_cfg is not None - local_tensor[non_pad_len:].copy_(0.0) # type: ignore # padded part must be set to 0 - else: - local_tensor.copy_(loaded_tensor) + return loaded_tensor def param_to_safetensor( self, diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py index f27e0a2dbc..50fbfa3869 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/xtuner/v1/model/moe/moe.py @@ -18,7 +18,7 @@ CPUOffloadPolicy, MixedPrecisionPolicy, ) -from torch.distributed.tensor import DTensor, Replicate, distribute_tensor +from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor from tqdm import tqdm from typing_extensions import overload, override @@ -147,6 +147,7 @@ class MoEConfig(TransformerConfig): hidden_factor: Annotated[float, Parameter(group="moe")] = 1.0 moe_intermediate_size: Annotated[int, Parameter(group="moe")] ep_size: Annotated[int, Parameter(group="moe")] = 1 + expert_tp_size: Annotated[int, Parameter(group="moe")] = 1 dispatcher: Annotated[Literal["deepep", "all2all", "agrs"] | None, Parameter(group="moe")] = None router: GreedyRouterConfig | NoAuxRouterConfig balancing_loss_cfg: BalancingLossConfig | None = BalancingLossConfig() @@ -178,6 +179,11 @@ def build(self) -> "MoE": return MoE(self) +def use_moe_ep_compile_cfg(config: MoEConfig) -> bool: + # 中文注释：ExpertTP 也会跨 rank 进入 dispatcher 通信段，compile 边界应和 EP 路径一致。 + return config.ep_size > 1 or config.expert_tp_size > 1 + + class MoE(BaseModel): """Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`InternLM3DecoderLayer`] @@ -188,18 +194,46 @@ class MoE(BaseModel): config: MoEConfig ep_mesh: DeviceMesh | None = None + expert_tp_mesh: DeviceMesh | None = None + ep_tp_mesh: DeviceMesh | None = None def __init__(self, config: MoEConfig): super().__init__(config) - if config.ep_size is not None and config.ep_size > 1: + ep_size = config.ep_size if config.ep_size is not None else 1 + expert_tp_size = config.expert_tp_size if config.expert_tp_size > 1 else 1 + if ep_size > 1 or expert_tp_size > 1: world_size = dist.get_world_size() - self.ep_mesh = init_device_mesh( - DEVICE, - (world_size // config.ep_size, config.ep_size), - mesh_dim_names=(f"{self.config.mesh_prefix}.dp", f"{self.config.mesh_prefix}.ep"), - )[f"{self.config.mesh_prefix}.ep"] + fsdp_size = world_size // (ep_size * expert_tp_size) + if expert_tp_size > 1: + # 中文注释：即使不开 EP，也保留 size=1 的 expert ownership 维度， + # 这样 routed experts 和 expert TP 仍然使用同一套 mesh 语义。 + _init_mesh = init_device_mesh( + DEVICE, + (fsdp_size, ep_size, expert_tp_size), + mesh_dim_names=( + f"{self.config.mesh_prefix}.dp", + f"{self.config.mesh_prefix}.ep", + f"{self.config.mesh_prefix}.etp", + ), + ) + self.ep_mesh = _init_mesh[f"{self.config.mesh_prefix}.ep"] + self.expert_tp_mesh = _init_mesh[f"{self.config.mesh_prefix}.etp"] + # 2D (ep, etp) sub-mesh used by GroupedLinear for per-expert column-parallel weights. + # LoadSpec records both placements so HF plans can load and reconstruct the layout. + self.ep_tp_mesh = _init_mesh[f"{self.config.mesh_prefix}.ep", f"{self.config.mesh_prefix}.etp"] + else: + _init_mesh = init_device_mesh( + DEVICE, + (fsdp_size, ep_size), + mesh_dim_names=(f"{self.config.mesh_prefix}.dp", f"{self.config.mesh_prefix}.ep"), + ) + self.ep_mesh = _init_mesh[f"{self.config.mesh_prefix}.ep"] + self.expert_tp_mesh = None + self.ep_tp_mesh = None else: self.ep_mesh = None + self.expert_tp_mesh = None + self.ep_tp_mesh = None self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, type=config.rms_norm_type) self.lm_head = LMHead(config.hidden_size, config.vocab_size, bias=False) @@ -1000,6 +1034,8 @@ def build_layers(self, config: MoEConfig) -> nn.ModuleDict: layer_idx=layer_idx, dispatcher=config.dispatcher, ep_mesh=self.ep_mesh, + expert_tp_mesh=self.expert_tp_mesh, + ep_tp_mesh=self.ep_tp_mesh, ) if self.config.freeze_routers: layers[str(layer_idx)].gate.requires_grad_(False) @@ -1065,6 +1101,8 @@ def build_mtp_block(self, config: MoEConfig) -> MTPBlock: layer_idx=config.num_hidden_layers + i, dispatcher=config.dispatcher, ep_mesh=self.ep_mesh, + expert_tp_mesh=self.expert_tp_mesh, + ep_tp_mesh=self.ep_tp_mesh, ) # Wrap decoder layer in MTPLayer @@ -1134,7 +1172,10 @@ def fully_shard( for param in self.parameters(): param.requires_grad = False - if self.ep_mesh.size() > 1: + tp_enabled = self.expert_tp_mesh is not None and self.expert_tp_mesh.size() > 1 + if self.ep_mesh.size() > 1 or tp_enabled: + # 中文注释：不开 EP 但开启 expert TP 时，非 expert 参数仍是 TP rank 间的逻辑副本， + # 需要显式放到 Replicate DTensor 上，后续梯度才会跨 expert TP 平均。 self._replicate_other_params(self) # Although rotary_emb was already constructed in __init__, it was built on the meta device. @@ -1269,13 +1310,14 @@ def fully_shard( if isinstance(module, nn.Embedding): module.forward = types.MethodType(self.patched_emb_forward, module) # type: ignore + self._init_load_spec() self._to_empty_meta() return self @property @override def default_compile_cfg(self) -> dict[str, TorchCompileOption]: - if self.config.ep_size > 1: + if use_moe_ep_compile_cfg(self.config): return MOE_EP_COMPILE_CFG else: return MOE_NON_EP_COMPILE_CFG @@ -1287,8 +1329,6 @@ def need_update_bias(self) -> bool: @torch.no_grad # type: ignore def scale_and_reduce_grad(self): - ep_enabled = self.ep_mesh is not None and self.ep_mesh.size() > 1 - # Bucket gradients that need a cross-rank reduction by their target process # group. Each bucket is reduced with a single coalesced NCCL all_reduce # instead of one launch per parameter, which used to dominate latency for @@ -1299,10 +1339,13 @@ def scale_and_reduce_grad(self): if param.grad is None: continue - # Expert parameters live on a unique EP rank, so no cross-rank reduction - # is needed — just rescale by `ep_size` to keep the effective average. - if ep_enabled and ".experts" in name: - param.grad.div_(self.ep_mesh.size()) # type: ignore + expert_parallel_size = ( + self.ep_mesh.size() if self.ep_mesh is not None else 1 + ) * self.config.expert_tp_size + # 中文注释：expert 参数会在 EP 和 expert TP 维度上看到全量 token 梯度和， + # 需要按参与该 expert 计算的 rank 数平均，才能对齐普通 DP/FSDP baseline。 + if expert_parallel_size > 1 and ".experts" in name: + param.grad.div_(expert_parallel_size) # type: ignore continue if not isinstance(param, DTensor): @@ -1339,19 +1382,67 @@ def scale_and_reduce_grad(self): for grad in grads: dist.all_reduce(grad, ReduceOp.SUM, group=group) + def cal_grad_norm(self, grads: list[DTensor], dtype=torch.float32): + # Keep the established reduction order (and therefore FP8 training + # baseline) when Expert TP is disabled. The custom path below is only + # needed for Expert TP's additional interleaved shard placement. + if self.config.expert_tp_size <= 1: + return super().cal_grad_norm(grads, dtype=dtype) + + from xtuner.v1.utils.grad_norm import group_tensors_by_device_mesh_and_placements + + grouped_grads = group_tensors_by_device_mesh_and_placements(grads) + if len(grads) == 0: + return torch.tensor(0.0, dtype=dtype), grouped_grads + + total_norm_squared = torch.zeros((), dtype=dtype, device=grads[0].device) + for name, param in self.trainable_parameters(): + grad = param.grad + if grad is None: + continue + + local_grad = grad.to_local() if isinstance(grad, DTensor) else grad + local_norm_squared = torch.linalg.vector_norm(local_grad, ord=2.0, dtype=dtype) ** 2 + if isinstance(grad, DTensor): + for i, placement in enumerate(grad.placements): + if isinstance(placement, Shard): + dist.all_reduce(local_norm_squared, group=grad.device_mesh.get_group(i)) + elif isinstance(placement, Replicate): + pass + else: + raise ValueError(f"Unsupported placement type {placement} in clip_grad_norm") + + total_norm_squared += local_norm_squared + + grad_norm = total_norm_squared**0.5 + grad_norm = grad_norm.to(grads[0].dtype) + return grad_norm, grouped_grads + def _init_device_mesh(self, fsdp_config: FSDPConfig): self.fsdp_config = fsdp_config device = DEVICE world_size = dist.get_world_size() - experts_fsdp_size = world_size // self.fsdp_config.ep_size + expert_tp_size = self.config.expert_tp_size if self.config.expert_tp_size > 1 else 1 + experts_fsdp_size = world_size // (self.fsdp_config.ep_size * expert_tp_size) if self.fsdp_config.hsdp_sharding_size is None: - model_mesh = init_device_mesh( - device, - (experts_fsdp_size, self.fsdp_config.ep_size), - mesh_dim_names=(f"{self.config.mesh_prefix}.fsdp", f"{self.config.mesh_prefix}.ep"), - ) + if expert_tp_size > 1: + model_mesh = init_device_mesh( + device, + (experts_fsdp_size, self.fsdp_config.ep_size, expert_tp_size), + mesh_dim_names=( + f"{self.config.mesh_prefix}.fsdp", + f"{self.config.mesh_prefix}.ep", + f"{self.config.mesh_prefix}.etp", + ), + ) + else: + model_mesh = init_device_mesh( + device, + (experts_fsdp_size, self.fsdp_config.ep_size), + mesh_dim_names=(f"{self.config.mesh_prefix}.fsdp", f"{self.config.mesh_prefix}.ep"), + ) self._world_mesh = model_mesh if self.ep_mesh is not None: # WARN: This assertion is **VERY** important. @@ -1390,6 +1481,14 @@ def _init_device_mesh(self, fsdp_config: FSDPConfig): else: self.ep_mesh = model_mesh[f"{self.config.mesh_prefix}.ep"] + if expert_tp_size > 1: + new_expert_tp_mesh = model_mesh[f"{self.config.mesh_prefix}.etp"] + if self.expert_tp_mesh is not None: + assert new_expert_tp_mesh.mesh_dim_names == self.expert_tp_mesh.mesh_dim_names + assert torch.equal(self.expert_tp_mesh.mesh, new_expert_tp_mesh.mesh) + else: + self.expert_tp_mesh = new_expert_tp_mesh + self.fsdp_mesh = model_mesh[f"{self.config.mesh_prefix}.fsdp"] else: assert self.fsdp_config.ep_size == 1, "Currently, HSDP requires expert parallel size to be 1" @@ -1413,12 +1512,26 @@ def _init_device_mesh(self, fsdp_config: FSDPConfig): self.fsdp_mesh = self.hsdp_mesh[f"{self.config.mesh_prefix}.hsdp_shard"] def _replicate_other_params(self, model: nn.Module): - def traverse(module): + def traverse(module: nn.Module) -> None: if isinstance(module, MoEBlock): + # Expert params are already partitioned by build_grouped_linear. return for name, param in module.named_parameters(recurse=False): + assert self.ep_mesh is not None + replicate_mesh = self.ep_mesh + placements = [Replicate()] + if self.expert_tp_mesh is not None and self.expert_tp_mesh.size() > 1: + assert self._world_mesh is not None + # Keep both model-parallel dimensions attached to the FSDP + # root mesh. A separately flattened mesh is cached globally + # by PyTorch and can lose that parent on a later model build. + replicate_mesh = self._world_mesh[ + (f"{self.config.mesh_prefix}.ep", f"{self.config.mesh_prefix}.etp") + ] + placements = [Replicate(), Replicate()] dist_param = nn.Parameter( - distribute_tensor(param, self.ep_mesh, [Replicate()]), requires_grad=param.requires_grad + distribute_tensor(param, replicate_mesh, placements), + requires_grad=param.requires_grad, ) module.register_parameter(name, dist_param) for child in module.children(): diff --git a/xtuner/v1/model/moe/qwen3_5_text.py b/xtuner/v1/model/moe/qwen3_5_text.py index e4cf2e2fc3..33414ef06b 100644 --- a/xtuner/v1/model/moe/qwen3_5_text.py +++ b/xtuner/v1/model/moe/qwen3_5_text.py @@ -10,7 +10,7 @@ HFSaveCfg, TorchCompileOption, ) -from xtuner.v1.model.moe.moe import BalancingLossConfig, MoEConfig, ZLossConfig +from xtuner.v1.model.moe.moe import BalancingLossConfig, MoEConfig, ZLossConfig, use_moe_ep_compile_cfg from xtuner.v1.module.attention import GatedDeltaNetConfig, MHAConfig from xtuner.v1.module.rope import RopeParametersConfig from xtuner.v1.module.router.greedy import GreedyRouterConfig @@ -122,46 +122,19 @@ def to_hf_key_list(self, key: str) -> list[str]: else: return [key] - def safetensors_to_params( - self, - safetensors: list[torch.Tensor], - local_tensor: torch.Tensor, - param_name: str, - start: int | None, - end: int | None, - dim: int | None, - ): - if len(safetensors) > 1: - assert dim is not None, "Internal Error dim must not be None when len(safetensors) > 1" - loaded_tensor = torch.cat(safetensors, dim=dim) - else: - loaded_tensor = safetensors[0] - - if "fused_w1w3.weight" in param_name and "mtp" not in param_name: + def hf_tensor_to_canonical(self, name: str, loaded_tensor: torch.Tensor) -> torch.Tensor: + if "fused_w1w3.weight" in name and "mtp" not in name: # hf: num_experts, 2 * expert_dim, hidden_size # xtuner: num_experts * 2 * expert_dim, hidden_size # num_experts * 2 * expert_dim, hidden_size loaded_tensor = loaded_tensor.flatten(0, 1) - elif "fused_w2.weight" in param_name and "mtp" not in param_name: + elif "fused_w2.weight" in name and "mtp" not in name: # hf: num_experts, hidden_size, expert_dim # xtuner: num_experts * hidden_size, expert_dim loaded_tensor = loaded_tensor.flatten(0, 1) - if start is not None and end is not None: - start = min(start, loaded_tensor.shape[self.FSDP_SHARD_DIM]) - end = min(end, loaded_tensor.shape[self.FSDP_SHARD_DIM]) - loaded_tensor_slice = loaded_tensor.index_select( - dim=self.FSDP_SHARD_DIM, index=torch.arange(start, end, dtype=torch.int64, device=loaded_tensor.device) - ) - non_pad_len = end - start - local_tensor[:non_pad_len].copy_(loaded_tensor_slice) - - if non_pad_len < local_tensor.shape[self.FSDP_SHARD_DIM]: - assert self.config.float8_cfg is not None - local_tensor[non_pad_len:].copy_(0.0) # type: ignore # padded part must be set to 0 - else: - local_tensor.copy_(loaded_tensor) + return loaded_tensor def param_to_safetensor( self, @@ -191,7 +164,7 @@ def param_to_safetensor( @property @override def default_compile_cfg(self) -> dict[str, TorchCompileOption]: - if self.config.ep_size > 1: + if use_moe_ep_compile_cfg(self.config): return MOE_EP_COMPILE_CFG else: return MOE_NON_EP_COMPILE_CFG diff --git a/xtuner/v1/model/moe/qwen3vl_text.py b/xtuner/v1/model/moe/qwen3vl_text.py index 5451c31346..45742d8527 100644 --- a/xtuner/v1/model/moe/qwen3vl_text.py +++ b/xtuner/v1/model/moe/qwen3vl_text.py @@ -35,22 +35,8 @@ def to_hf_key_list(self, key: str) -> list[str]: else: return [key] - def safetensors_to_params( - self, - safetensors: list[torch.Tensor], - local_tensor: torch.Tensor, - param_name: str, - start: int | None, - end: int | None, - dim: int | None, - ): - if len(safetensors) > 1: - assert dim is not None, "Internal Error dim must not be None when len(safetensors) > 1" - loaded_tensor = torch.cat(safetensors, dim=dim) - else: - loaded_tensor = safetensors[0] - - if "fused_w1w3.weight" in param_name: + def hf_tensor_to_canonical(self, name: str, loaded_tensor: torch.Tensor) -> torch.Tensor: + if "fused_w1w3.weight" in name: # hf: num_experts, hidden_size, 2 * expert_dim # xtuner: num_experts * 2 * expert_dim, hidden_size num_experts, hidden_size = loaded_tensor.shape[:2] @@ -58,25 +44,12 @@ def safetensors_to_params( # num_experts * 2 * expert_dim, hidden_size loaded_tensor = loaded_tensor.reshape(-1, hidden_size) - elif "fused_w2.weight" in param_name: + elif "fused_w2.weight" in name: # hf: num_experts, expert_dim, hidden_size # xtuner: num_experts * hidden_size, expert_dim loaded_tensor = loaded_tensor.transpose(1, 2).flatten(0, 1) - if start is not None and end is not None: - start = min(start, loaded_tensor.shape[self.FSDP_SHARD_DIM]) - end = min(end, loaded_tensor.shape[self.FSDP_SHARD_DIM]) - loaded_tensor_slice = loaded_tensor.index_select( - dim=self.FSDP_SHARD_DIM, index=torch.arange(start, end, dtype=torch.int64, device=loaded_tensor.device) - ) - non_pad_len = end - start - local_tensor[:non_pad_len].copy_(loaded_tensor_slice) - - if non_pad_len < local_tensor.shape[self.FSDP_SHARD_DIM]: - assert self.config.float8_cfg is not None - local_tensor[non_pad_len:].copy_(0.0) # type: ignore # padded part must be set to 0 - else: - local_tensor.copy_(loaded_tensor) + return loaded_tensor def param_to_safetensor( self, diff --git a/xtuner/v1/module/decoder_layer/moe_decoder_layer.py b/xtuner/v1/module/decoder_layer/moe_decoder_layer.py index 00e5d6c27e..862d09c030 100644 --- a/xtuner/v1/module/decoder_layer/moe_decoder_layer.py +++ b/xtuner/v1/module/decoder_layer/moe_decoder_layer.py @@ -156,8 +156,10 @@ def __init__( n_routed_experts: int, moe_bias: bool = False, ep_mesh: DeviceMesh | None = None, + expert_tp_mesh: DeviceMesh | None = None, float8_cfg: Float8Config | None = None, moe_act_fn_cfg: MoEActFnConfig, + ep_tp_mesh: DeviceMesh | None = None, ): super().__init__() self.hidden_size = hidden_size @@ -172,7 +174,11 @@ def __init__( self.num_routed_experts, moe_bias=moe_bias, ep_mesh=self.ep_mesh, + expert_tp_mesh=expert_tp_mesh, + parallel_style="column", float8_cfg=float8_cfg, + ep_tp_mesh=ep_tp_mesh, + num_fused_projections=2, ) self.fused_w2 = build_grouped_linear( self.intermediate_size, @@ -180,7 +186,10 @@ def __init__( self.num_routed_experts, moe_bias=moe_bias, ep_mesh=self.ep_mesh, + expert_tp_mesh=expert_tp_mesh, + parallel_style="row", float8_cfg=float8_cfg, + ep_tp_mesh=ep_tp_mesh, ) self.moe_act = moe_act_fn_cfg.build() @@ -222,9 +231,12 @@ def __init__( layer_idx: int = 0, dispatcher: Literal["deepep", "all2all", "agrs"] | None, ep_mesh: DeviceMesh | None = None, + expert_tp_mesh: DeviceMesh | None = None, + ep_tp_mesh: DeviceMesh | None = None, ): super().__init__() self.ep_mesh = ep_mesh + self.ep_tp_mesh = ep_tp_mesh self.hidden_size = hidden_size self.n_routed_experts = n_routed_experts self.n_shared_experts = n_shared_experts @@ -276,15 +288,21 @@ def __init__( n_routed_experts=n_routed_experts, moe_bias=moe_bias, ep_mesh=ep_mesh, + expert_tp_mesh=expert_tp_mesh, float8_cfg=float8_cfg, moe_act_fn_cfg=moe_act_fn_cfg, + ep_tp_mesh=ep_tp_mesh, ) # TODO: (yehaochen) Maybe should be replaced by build_dispatcher process_group = ep_mesh.get_group() if ep_mesh is not None else None + tp_group = expert_tp_mesh.get_group() if expert_tp_mesh is not None else None + ep_tp_group = ep_tp_mesh._flatten().get_group() if ep_tp_mesh is not None else None self.dispatcher = build_dispatcher( dispatcher=dispatcher, n_routed_experts=n_routed_experts, ep_group=process_group, + tp_group=tp_group, + ep_tp_group=ep_tp_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", ) @@ -393,6 +411,7 @@ def _forward( pre_dispatched = self.dispatcher.dispatch_preprocess( hidden_states=hidden_states.view(-1, hidden_states.shape[-1]), topk_ids=router_results["topk_ids"], + topk_weights=router_results["topk_weights"], ) dispatched = self.dispatcher.dispatch( pre_dispatched=pre_dispatched, @@ -507,6 +526,7 @@ def _micro_batch_forward( pre_dispatched = self.dispatcher.dispatch_preprocess( hidden_states=hidden_states, topk_ids=router_results["topk_ids"], + topk_weights=router_results["topk_weights"], async_op=True, ) pre_dispatched_list.append(pre_dispatched) diff --git a/xtuner/v1/module/dispatcher/__init__.py b/xtuner/v1/module/dispatcher/__init__.py index e4981392d0..d5d3c96860 100644 --- a/xtuner/v1/module/dispatcher/__init__.py +++ b/xtuner/v1/module/dispatcher/__init__.py @@ -31,6 +31,8 @@ def build_dispatcher( dispatcher: Literal["deepep", "all2all", "agrs"] | None, n_routed_experts: int, ep_group: dist.ProcessGroup | None = None, + tp_group: dist.ProcessGroup | None = None, + ep_tp_group: dist.ProcessGroup | None = None, training_dtype: Literal["bf16", "fp8"] = "bf16", generate_dtype: Literal["bf16", "fp8"] = "bf16", ) -> DispacherInterface: @@ -40,6 +42,7 @@ def build_dispatcher( return NaiveDispatcher( n_routed_experts=n_routed_experts, process_group=ep_group, + tp_group=tp_group, training_dtype=training_dtype, generate_dtype=generate_dtype, ) # type: ignore[return-value] @@ -52,18 +55,32 @@ def build_dispatcher( from .deepep import DeepEPDispatcher # type: ignore[attr-defined] assert ep_group is not None, "DeepEPDispatcher requires a non-null process group." + # When expert TP is enabled, fuse EP dispatch + TP replication into a single DeepEP + # collective: the dispatcher operates on the combined (ep × tp) group with each + # physical expert virtualized into ``tp_size`` copies (see ``DeepEPDispatcher``). + tp_size = tp_group.size() if tp_group is not None else 1 + if tp_size > 1: + assert ep_tp_group is not None, ( + "DeepEPDispatcher with expert TP requires the combined (ep × tp) process group; " + "pass ``ep_tp_group`` from ``ep_tp_mesh._flatten().get_group()``." + ) + process_group = ep_tp_group + else: + process_group = ep_group # TODO: remove type ignore here return DeepEPDispatcher( n_routed_experts=n_routed_experts, - process_group=ep_group, + process_group=process_group, + tp_size=tp_size, training_dtype=training_dtype, generate_dtype=generate_dtype, ) # type: ignore elif dispatcher == "all2all": - assert ep_group is not None, "DeepEPDispatcher requires a non-null process group." + assert ep_group is not None, "TorchAll2AllDispatcher requires a non-null ep_group." return TorchAll2AllDispatcher( n_routed_experts=n_routed_experts, process_group=ep_group, + tp_group=tp_group, training_dtype=training_dtype, generate_dtype=generate_dtype, ) # type: ignore[return-value] diff --git a/xtuner/v1/module/dispatcher/agrs.py b/xtuner/v1/module/dispatcher/agrs.py index 664cc59d56..255070a6b9 100644 --- a/xtuner/v1/module/dispatcher/agrs.py +++ b/xtuner/v1/module/dispatcher/agrs.py @@ -258,6 +258,7 @@ def dispatch_preprocess( *, hidden_states: torch.Tensor, topk_ids: torch.Tensor, + topk_weights: torch.Tensor, # noqa: ARG002 — kept for interface compatibility; not used here async_op: bool = False, ) -> MoEAGRSPreDispatchResult: if async_op: diff --git a/xtuner/v1/module/dispatcher/base.py b/xtuner/v1/module/dispatcher/base.py index b268d75f63..81bc94d919 100644 --- a/xtuner/v1/module/dispatcher/base.py +++ b/xtuner/v1/module/dispatcher/base.py @@ -11,10 +11,26 @@ from xtuner.v1.ops import permute, unpermute +from .expert_tp import ExpertTP + HiddenStates: TypeAlias = torch.Tensor +def _get_backward_pre_hook(backward_previous_event: torch.cuda.Event): + def _backward_pre_hook(*_): + torch.cuda.current_stream().wait_event(backward_previous_event) + + return _backward_pre_hook + + +def _get_backward_hook(backward_finished_event: torch.cuda.Event): + def _backward_hook(*_): + backward_finished_event.record() + + return _backward_hook + + class PreDispatchResult(TypedDict): hidden_states: torch.Tensor topk_ids: torch.Tensor @@ -119,6 +135,7 @@ def dispatch_preprocess( *, hidden_states: torch.Tensor, topk_ids: torch.Tensor, + topk_weights: torch.Tensor, async_op: bool = False, ) -> PreDispatch: ... @@ -171,20 +188,32 @@ class DispacherInterface( ): ... -class NaivePreDispatchResult(PreDispatchResult): ... +class NaivePreDispatchResult(PreDispatchResult): + # 中文注释：这些 key 必须始终存在；torch.compile 不支持 optional-key TypedDict。 + forward_finished_event: torch.cuda.Event | None + backward_previous_event: torch.cuda.Event | None -class NaiveDispatchResult(DispatchResult): ... +class NaiveDispatchResult(DispatchResult): + topk_ids: torch.Tensor + tp_rank_row_counts: list[int] + forward_finished_event: torch.cuda.Event | None + backward_previous_event: torch.cuda.Event | None + topk_weights_backward_previous_event: torch.cuda.Event | None class NaivePostDispatchResult(PostDispatchResult): row_ids_map: torch.Tensor -class NaivePreCombineResult(PreCombineResult): ... +class NaivePreCombineResult(PreCombineResult): + forward_finished_event: torch.cuda.Event | None + backward_previous_event: torch.cuda.Event | None -class NaiveCombineResult(CombineResult): ... +class NaiveCombineResult(CombineResult): + forward_finished_event: torch.cuda.Event | None + backward_previous_event: torch.cuda.Event | None class NaivePostCombineResult(PostCombineResult): ... @@ -200,11 +229,14 @@ class NaiveDispatcher( NaivePostCombineResult, ] ): + _comm_stream: torch.cuda.Stream | None = None + def __init__( self, *, n_routed_experts: int, process_group: torch.distributed.ProcessGroup | None = None, + tp_group: torch.distributed.ProcessGroup | None = None, training_dtype: Literal["fp8", "bf16"] = "bf16", generate_dtype: Literal["fp8", "bf16"] = "bf16", ): @@ -216,6 +248,9 @@ def __init__( ) if self._process_group is not None: assert self._process_group.size() == 1, "Naive dispatcher is only for ep=1." + self._expert_tp = ExpertTP(tp_group) if tp_group is not None and tp_group.size() > 1 else None + if self._expert_tp is not None and NaiveDispatcher._comm_stream is None: + NaiveDispatcher._comm_stream = torch.cuda.Stream() @override def dispatch_preprocess( @@ -223,31 +258,120 @@ def dispatch_preprocess( *, hidden_states: torch.Tensor, topk_ids: torch.Tensor, + topk_weights: torch.Tensor, async_op: bool = False, - ) -> PreDispatchResult: + ) -> NaivePreDispatchResult: if async_op: - raise NotImplementedError("Naive dispatcher is only for ep=1.") + if self._expert_tp is None: + raise NotImplementedError("Naive dispatcher async_op=True requires ExpertTP.") + + forward_finished_event = torch.cuda.Event() + forward_finished_event.record() + backward_previous_event = torch.cuda.Event() + if hidden_states.grad_fn is not None: + hidden_states.grad_fn.register_prehook(_get_backward_pre_hook(backward_previous_event)) + + return NaivePreDispatchResult( + hidden_states=hidden_states, + topk_ids=topk_ids, + forward_finished_event=forward_finished_event, + backward_previous_event=backward_previous_event, + ) return NaivePreDispatchResult( hidden_states=hidden_states, topk_ids=topk_ids, + forward_finished_event=None, + backward_previous_event=None, ) @override def dispatch( self, *, - pre_dispatched: PreDispatchResult, + pre_dispatched: NaivePreDispatchResult, topk_weights: torch.Tensor, async_op: bool = False, decoding: bool = False, ) -> NaiveDispatchResult: if async_op: - raise NotImplementedError("Naive dispatcher is only for ep=1.") + if self._expert_tp is None: + raise NotImplementedError("Naive dispatcher async_op=True requires ExpertTP.") + + forward_previous_event = pre_dispatched["forward_finished_event"] + backward_finished_event = pre_dispatched["backward_previous_event"] + assert forward_previous_event is not None, "Use async_op=True for dispatch_preprocess!" + assert backward_finished_event is not None, "Use async_op=True for dispatch_preprocess!" + assert self._comm_stream is not None + + tp_rank_row_counts = self._expert_tp.gather_tp_rank_row_counts(pre_dispatched["hidden_states"]) + # 中文注释：dispatch 内部的 TP AllGather 都排在同一个 comm stream， + # 互相不需要 event 串行化；只在 dispatch 阶段边界记录最终完成事件。 + forward_finished_event = torch.cuda.Event() + hidden_backward_previous_event = torch.cuda.Event() + topk_weights_backward_previous_event = torch.cuda.Event() + topk_weights_backward_finished_event = torch.cuda.Event() + if topk_weights.grad_fn is not None: + topk_weights.grad_fn.register_prehook(_get_backward_pre_hook(topk_weights_backward_finished_event)) + + hidden_states = self._expert_tp.async_all_gather_rows( + pre_dispatched["hidden_states"], + tp_rank_row_counts=tp_rank_row_counts, + forward_previous_event=forward_previous_event, + forward_finished_event=None, + backward_previous_event=hidden_backward_previous_event, + backward_finished_event=backward_finished_event, + comm_stream=self._comm_stream, + ) + topk_ids = self._expert_tp.async_all_gather_row_metadata( + pre_dispatched["topk_ids"], + tp_rank_row_counts=tp_rank_row_counts, + forward_previous_event=None, + forward_finished_event=None, + comm_stream=self._comm_stream, + ) + topk_weights = self._expert_tp.async_all_gather_rows( + topk_weights, + tp_rank_row_counts=tp_rank_row_counts, + forward_previous_event=None, + forward_finished_event=forward_finished_event, + backward_previous_event=topk_weights_backward_previous_event, + backward_finished_event=topk_weights_backward_finished_event, + comm_stream=self._comm_stream, + ) + + return NaiveDispatchResult( + hidden_states=hidden_states, + topk_ids=topk_ids, + topk_weights=topk_weights, + tp_rank_row_counts=tp_rank_row_counts, + forward_finished_event=forward_finished_event, + backward_previous_event=hidden_backward_previous_event, + topk_weights_backward_previous_event=topk_weights_backward_previous_event, + ) + + if self._expert_tp is not None: + hidden_states, tp_rank_row_counts = self._expert_tp.all_gather_rows(pre_dispatched["hidden_states"]) + topk_ids = self._expert_tp.all_gather_row_metadata(pre_dispatched["topk_ids"], tp_rank_row_counts) + topk_weights = self._expert_tp.all_gather_row_metadata(topk_weights, tp_rank_row_counts) + return NaiveDispatchResult( + hidden_states=hidden_states, + topk_ids=topk_ids, + topk_weights=topk_weights, + tp_rank_row_counts=tp_rank_row_counts, + forward_finished_event=None, + backward_previous_event=None, + topk_weights_backward_previous_event=None, + ) return NaiveDispatchResult( hidden_states=pre_dispatched["hidden_states"], + topk_ids=pre_dispatched["topk_ids"], topk_weights=topk_weights, + tp_rank_row_counts=[], + forward_finished_event=None, + backward_previous_event=None, + topk_weights_backward_previous_event=None, ) @override @@ -260,14 +384,24 @@ def dispatch_postprocess( decoding: bool = False, ) -> NaivePostDispatchResult: if async_op: - raise NotImplementedError("Naive dispatcher is only for ep=1.") + if self._expert_tp is None: + raise NotImplementedError("Naive dispatcher async_op=True requires ExpertTP.") + forward_finished_event = dispatched["forward_finished_event"] + assert forward_finished_event is not None, "Use async_op=True for dispatch!" + torch.cuda.current_stream().wait_event(forward_finished_event) + topk_ids = dispatched["topk_ids"] if self._expert_tp is not None else pre_dispatched["topk_ids"] hidden_states, row_id_maps = permute( dispatched["hidden_states"], - pre_dispatched["topk_ids"].to(torch.int32), + topk_ids.to(torch.int32), ) - topk_ids = pre_dispatched["topk_ids"] tokens_per_expert = torch.histc(topk_ids, bins=self._n_routed_experts, min=0, max=self._n_routed_experts) + if async_op: + backward_previous_event = dispatched["backward_previous_event"] + assert backward_previous_event is not None, "Use async_op=True for dispatch!" + if hidden_states.grad_fn is not None: + hidden_states.grad_fn.register_hook(_get_backward_hook(backward_previous_event)) + if decoding: raise NotImplementedError else: @@ -287,19 +421,37 @@ def combine_preprocess( post_dispatched: NaivePostDispatchResult, async_op: bool = False, decoding: bool = False, - ) -> PreCombineResult: + ) -> NaivePreCombineResult: if async_op: - raise NotImplementedError("Naive dispatcher is only for ep=1.") + if self._expert_tp is None: + raise NotImplementedError("Naive dispatcher async_op=True requires ExpertTP.") hidden_states = unpermute( input_act=hidden_states, row_id_map=post_dispatched["row_ids_map"], probs=dispatched["topk_weights"], ) + if async_op: + backward_previous_event = torch.cuda.Event() + forward_finished_event = torch.cuda.Event() + forward_finished_event.record() + if hidden_states.grad_fn is not None: + hidden_states.grad_fn.register_prehook(_get_backward_pre_hook(backward_previous_event)) + topk_weights_backward_previous_event = dispatched["topk_weights_backward_previous_event"] + assert topk_weights_backward_previous_event is not None, "Use async_op=True for dispatch!" + hidden_states.grad_fn.register_hook(_get_backward_hook(topk_weights_backward_previous_event)) + else: + backward_previous_event = None + forward_finished_event = None + if decoding: raise NotImplementedError("NaiveDispatcher does not support decoding.") else: - return PreCombineResult(hidden_states=hidden_states) + return NaivePreCombineResult( + hidden_states=hidden_states, + backward_previous_event=backward_previous_event, + forward_finished_event=forward_finished_event, + ) @override def combine( @@ -313,12 +465,52 @@ def combine( decoding: bool = False, ) -> NaiveCombineResult: if async_op: - raise NotImplementedError("Naive dispatcher is only for ep=1.") + if self._expert_tp is None: + raise NotImplementedError("Naive dispatcher async_op=True requires ExpertTP.") if decoding: raise NotImplementedError else: - return NaiveCombineResult(hidden_states=pre_combined["hidden_states"]) + if self._expert_tp is not None: + if async_op: + forward_previous_event = pre_combined["forward_finished_event"] + backward_finished_event = pre_combined["backward_previous_event"] + assert forward_previous_event is not None, "Use async_op=True for combine_preprocess!" + assert backward_finished_event is not None, "Use async_op=True for combine_preprocess!" + assert self._comm_stream is not None + + forward_finished_event = torch.cuda.Event() + backward_previous_event = torch.cuda.Event() + hidden_states = self._expert_tp.async_reduce_scatter_rows_sum( + pre_combined["hidden_states"], + tp_rank_row_counts=dispatched["tp_rank_row_counts"], + forward_previous_event=forward_previous_event, + forward_finished_event=forward_finished_event, + backward_previous_event=backward_previous_event, + backward_finished_event=backward_finished_event, + comm_stream=self._comm_stream, + ) + return NaiveCombineResult( + hidden_states=hidden_states, + forward_finished_event=forward_finished_event, + backward_previous_event=backward_previous_event, + ) + + hidden_states = self._expert_tp.reduce_scatter_rows_sum( + pre_combined["hidden_states"], + dispatched["tp_rank_row_counts"], + ) + return NaiveCombineResult( + hidden_states=hidden_states, + forward_finished_event=None, + backward_previous_event=None, + ) + + return NaiveCombineResult( + hidden_states=pre_combined["hidden_states"], + forward_finished_event=None, + backward_previous_event=None, + ) @override def combine_postprocess( @@ -332,6 +524,16 @@ def combine_postprocess( async_op: bool = False, ) -> PostCombineResult: if async_op: - raise NotImplementedError("Naive dispatcher is only for ep=1.") + if self._expert_tp is None: + raise NotImplementedError("Naive dispatcher async_op=True requires ExpertTP.") + forward_finished_event = combined["forward_finished_event"] + backward_previous_event = combined["backward_previous_event"] + assert forward_finished_event is not None, "Use async_op=True for combine!" + assert backward_previous_event is not None, "Use async_op=True for combine!" + torch.cuda.current_stream().wait_event(forward_finished_event) + hidden_states = combined["hidden_states"].view_as(combined["hidden_states"]) + if hidden_states.grad_fn is not None: + hidden_states.grad_fn.register_hook(_get_backward_hook(backward_previous_event)) + return PostCombineResult(hidden_states=hidden_states) return PostCombineResult(hidden_states=combined["hidden_states"]) diff --git a/xtuner/v1/module/dispatcher/deepep.py b/xtuner/v1/module/dispatcher/deepep.py index 679253e2ba..de264bf4ab 100644 --- a/xtuner/v1/module/dispatcher/deepep.py +++ b/xtuner/v1/module/dispatcher/deepep.py @@ -42,6 +42,11 @@ # DeepEP handle include 6 tensor: # (rank_prefix_matrix, channel_prefix_matrix, recv_channel_prefix_matrix, recv_src_idx, is_token_in_rank, send_head) class DeepEPPreDispatchResult(PreDispatchResult): + # Final ``topk_weights`` fed to DeepEP. Equal to the caller's ``topk_weights`` for ep-only + # routing; for virtual expert TP (``tp_size > 1``) it is ``repeat_interleave``'d here in + # ``dispatch_preprocess`` so the expand kernel runs on the compute stream during Loop A + # (overlapping the next microbatch's attention/gate) instead of inside ``dispatch``. + topk_weights: torch.Tensor backward_previous_event: EventOverlap | None forward_finished_event: EventOverlap | None @@ -258,9 +263,32 @@ def __init__( *, n_routed_experts: int, process_group: torch.distributed.ProcessGroup, + tp_size: int = 1, training_dtype: Literal["fp8", "bf16"] = "bf16", generate_dtype: Literal["fp8", "bf16"] = "bf16", ): + """DeepEP-backed MoE dispatcher. + + When ``tp_size > 1`` the dispatcher fuses expert-parallel dispatch and tensor-parallel + token replication into a single DeepEP collective. The caller must: + + * Build the combined ``(ep × tp)`` process group via ``ep_tp_mesh._flatten().get_group()`` + (mesh dims ordered with ``tp`` as the inner/fastest dim) and pass it as + ``process_group``. ``process_group.size() == ep_size * tp_size``. + * Pass ``tp_size`` so this class can: + - Treat the expert space as ``n_routed_experts * tp_size`` *virtual* experts. + Each physical expert ``e`` gets ``tp_size`` virtual copies, one owned by each + TP rank in the EP group ``e`` belongs to. + - Expand caller-supplied ``topk_ids`` so a token routed to physical expert ``e`` + lands on **both** TP ranks within EP rank ``ep(e)`` — exactly what + column-parallel ``fused_w1w3`` needs. + + DeepEP's NVL+RDMA path encodes destination as + ``(rdma_rank, is_token_in_nvl_rank_bits)`` (see ``DeepEP/csrc/kernels/internode.cu``), + so duplicated routings landing on the same node are sent as a single RDMA transfer + with the appropriate NVL bitmask. Cross-node bandwidth therefore matches the + ep-only case; only the local intra-node fan-out is doubled. + """ if not is_installed("deep_ep"): raise RuntimeError("`DeepEP` is not installed!") super().__init__( @@ -273,6 +301,21 @@ 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." ) + self._tp_size = tp_size + assert process_group.size() % tp_size == 0, ( + f"process_group size {process_group.size()} must be a multiple of tp_size {tp_size}; " + f"the caller is expected to pass the combined (ep × tp) group." + ) + self._ep_size = process_group.size() // tp_size + assert n_routed_experts % self._ep_size == 0, ( + f"n_routed_experts {n_routed_experts} must be divisible by ep_size {self._ep_size}" + ) + self._local_experts = n_routed_experts // self._ep_size + # Virtual expert count seen by DeepEP. Per-rank count + # (= virtual_n_experts / process_group.size()) stays equal to ``_local_experts`` — + # downstream ``permute`` / ``group_gemm`` consume ``num_recv_tokens_per_expert_list`` of + # that fixed length and no aggregation is needed. + self._virtual_n_experts = n_routed_experts * tp_size @override def dispatch_preprocess( @@ -280,39 +323,96 @@ def dispatch_preprocess( *, hidden_states: torch.Tensor, topk_ids: torch.Tensor, + topk_weights: torch.Tensor, async_op: bool = False, ) -> DeepEPPreDispatchResult: if async_op: backward_previous_event = EventOverlap(None) - forward_finished_event = buffer_capture() if hidden_states.grad_fn is not None: hidden_states.grad_fn.register_prehook( get_backward_pre_hook( backward_previous_event=backward_previous_event, - name="TorchAll2AllDispatcher.dispatch_preprocess", + name="DeepEPDispatcher.dispatch_preprocess.hidden_states", debug=XTUNER_DISPATCHER_DEBUG, ) ) else: - forward_finished_event = None backward_previous_event = None + topk_ids = topk_ids.to(torch.int64) + if self._tp_size > 1: + topk_ids = self._expand_topk_ids_for_tp(topk_ids) + # ``topk_ids`` was duplicated tp_size× above; ``topk_weights`` must follow with + # the SAME value per duplicate. No 1/tp scaling — the two TP partial outputs sum + # to the full expert output, so weighting both by ``w_k`` already gives + # ``w_k * full`` after combine. + topk_weights = topk_weights.repeat_interleave(self._tp_size, dim=-1).contiguous() + if async_op and topk_weights.grad_fn is not None: + # Symmetric to the ``hidden_states`` prehook: the grad for ``topk_weights`` + # flows back through ``repeat_interleave_backward`` on the compute stream, + # while DeepEP's dispatch backward writes that grad on the comm stream and + # stamps the event into ``backward_previous_event``. Without this prehook + # the compute-stream backward starts before that event fires and reads + # stale grad memory — observed as ``grad_norm=NaN``. + topk_weights.grad_fn.register_prehook( + get_backward_pre_hook( + backward_previous_event=backward_previous_event, + name="DeepEPDispatcher.dispatch_preprocess.topk_weights", + debug=XTUNER_DISPATCHER_DEBUG, + ) + ) + + # Capture AFTER all compute-stream work above (topk_ids expand, topk_weights + # repeat_interleave) so DeepEP's ``stream_wait(previous_event)`` covers those + # kernels. Capturing before them leaves their writes outside the event, and + # DeepEP's comm-stream dispatch may read stale memory — observed as NaN / + # divergent loss under ``intra_layer_micro_batch>1`` with virtual expert TP. + forward_finished_event = buffer_capture() if async_op else None + return DeepEPPreDispatchResult( hidden_states=hidden_states, - topk_ids=topk_ids.to(torch.int64), + topk_ids=topk_ids, + topk_weights=topk_weights, backward_previous_event=backward_previous_event, forward_finished_event=forward_finished_event, ) + def _expand_topk_ids_for_tp(self, topk_ids: torch.Tensor) -> torch.Tensor: + """Map physical-expert ids to virtual-expert ids so DeepEP routes each + token to every TP rank within its owning EP group. + + Virtual id layout (rank ``r = ep * tp + t`` owns ids ``[r * local, (r + 1) * local)``):: + + virtual_id(e, t) = (ep(e) * tp + t) * local_experts + (e mod local_experts) + + The two virtuals for the same physical expert sit on adjacent ranks (same EP, t=0/t=1), + so DeepEP's NUMA layout collapses the cross-node copy to a single RDMA transfer with a + 2-bit NVL bitmask. ``-1`` (padding) is preserved. + """ + local_experts = self._local_experts + tp = self._tp_size + + ep_e = topk_ids // local_experts + local_idx = topk_ids % local_experts + tp_offsets = torch.arange(tp, device=topk_ids.device, dtype=topk_ids.dtype) + virtual = (ep_e.unsqueeze(-1) * tp + tp_offsets) * local_experts + local_idx.unsqueeze(-1) + # Preserve sentinel (-1) for padded slots after expansion. + virtual = torch.where(topk_ids.unsqueeze(-1) < 0, topk_ids.unsqueeze(-1), virtual) + out = virtual.reshape(*topk_ids.shape[:-1], topk_ids.shape[-1] * tp) + return out.contiguous() + @override def dispatch( self, *, pre_dispatched: DeepEPPreDispatchResult, - topk_weights: torch.Tensor, + topk_weights: torch.Tensor, # noqa: ARG002 — already expanded and stashed in pre_dispatched async_op: bool = False, decoding: bool = False, ) -> DeepEPDispatchResult: + # ``topk_ids`` / ``topk_weights`` expansion and the cross-stream sync setup live in + # ``dispatch_preprocess`` so they run on Loop A's compute stream and overlap with the + # next microbatch's attention/gate. ``dispatch`` itself only kicks off DeepEP. ( dispatched_hidden_states, dispatched_topk_idx, @@ -323,8 +423,8 @@ def dispatch( ) = _async_dispatch( pre_dispatched["hidden_states"], pre_dispatched["topk_ids"], - topk_weights, - self._n_routed_experts, + pre_dispatched["topk_weights"], + self._virtual_n_experts, self._process_group, pre_dispatched["forward_finished_event"], pre_dispatched["backward_previous_event"], @@ -464,7 +564,7 @@ def combine( combined_hidden_states, event = _async_combine( pre_combined["hidden_states"], - self._n_routed_experts, + self._virtual_n_experts, dispatched["handle"], self._process_group, pre_combined["forward_finished_event"], @@ -494,6 +594,7 @@ def combine_postprocess( combined: DeepEPCombineResult, async_op: bool = False, ) -> PostCombineResult: + # Restored original wait order (after view_as) to test torch_compile interaction hidden_states = combined["hidden_states"] forward_previous_event = combined["forward_finished_event"] diff --git a/xtuner/v1/module/dispatcher/expert_tp.py b/xtuner/v1/module/dispatcher/expert_tp.py new file mode 100644 index 0000000000..c0652b455a --- /dev/null +++ b/xtuner/v1/module/dispatcher/expert_tp.py @@ -0,0 +1,405 @@ +from __future__ import annotations + +from typing import Any + +import torch +import torch.distributed as dist + + +def _record_stream(value: Any, stream: torch.cuda.Stream) -> None: + if isinstance(value, torch.Tensor): + value.record_stream(stream) + elif isinstance(value, (list, tuple)): + for item in value: + _record_stream(item, stream) + + +def _tp_all_gather_rows_forward_impl( + tensor: torch.Tensor, + tp_rank_row_counts: list[int], + tp_group: dist.ProcessGroup, +) -> tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]: + tensor = tensor.contiguous() + chunks = [ + torch.empty((size, *tensor.shape[1:]), dtype=tensor.dtype, device=tensor.device) for size in tp_rank_row_counts + ] + dist.all_gather(chunks, tensor, group=tp_group) + return torch.cat(chunks, dim=0), tensor, chunks + + +def _tp_reduce_scatter_rows_sum_impl( + tensor: torch.Tensor, + tp_rank_row_counts: list[int], + tp_rank: int, + tp_group: dist.ProcessGroup, +) -> tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]: + tensor = tensor.contiguous() + assert tensor.shape[0] == sum(tp_rank_row_counts), ( + "TP ReduceScatterRowsSum input rows must match tp_rank_row_counts." + ) + + out = tensor.new_empty((tp_rank_row_counts[tp_rank], *tensor.shape[1:])) + if tensor.shape[0] == 0: + # 中文注释：所有 TP rank 都没有 token 时没有通信量，直接返回当前 rank 的 0 行 slice。 + return out, tensor, [] + + if all(size == tp_rank_row_counts[0] for size in tp_rank_row_counts): + dist.reduce_scatter_tensor(out, tensor, op=dist.ReduceOp.SUM, group=tp_group) + return out, tensor, [] + + input_chunks = list(torch.split(tensor, tp_rank_row_counts, dim=0)) + dist.reduce_scatter(out, input_chunks, op=dist.ReduceOp.SUM, group=tp_group) + return out, tensor, input_chunks + + +def _tp_all_gather_rows_backward_impl( + grad: torch.Tensor, + tp_rank_row_counts: list[int], + tp_rank: int, + tp_group: dist.ProcessGroup, +) -> tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]: + return _tp_reduce_scatter_rows_sum_impl(grad, tp_rank_row_counts, tp_rank, tp_group) + + +def _tp_reduce_scatter_rows_sum_backward_impl( + grad_slice: torch.Tensor, + tp_rank_row_counts: list[int], + tp_group: dist.ProcessGroup, +) -> tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]: + grad_slice = grad_slice.contiguous() + chunks = [ + torch.empty((size, *grad_slice.shape[1:]), dtype=grad_slice.dtype, device=grad_slice.device) + for size in tp_rank_row_counts + ] + dist.all_gather(chunks, grad_slice, group=tp_group) + return torch.cat(chunks, dim=0), grad_slice, chunks + + +class _TPAllGatherRows(torch.autograd.Function): + @staticmethod + def forward( + ctx: Any, + tensor: torch.Tensor, + tp_rank_row_counts: list[int], + tp_group: dist.ProcessGroup, + tp_size: int, + tp_rank: int, + ) -> torch.Tensor: + gathered, _, _ = _tp_all_gather_rows_forward_impl(tensor, tp_rank_row_counts, tp_group) + ctx.tp_rank_row_counts = tp_rank_row_counts + ctx.tp_group = tp_group + ctx.tp_rank = tp_rank + return gathered + + @staticmethod + def backward(ctx: Any, grad: torch.Tensor) -> tuple[torch.Tensor, None, None, None, None]: + grad_input, _, _ = _tp_all_gather_rows_backward_impl(grad, ctx.tp_rank_row_counts, ctx.tp_rank, ctx.tp_group) + return grad_input, None, None, None, None + + +class _AsyncTPAllGatherRows(torch.autograd.Function): + @staticmethod + def forward( + ctx: Any, + tensor: torch.Tensor, + tp_rank_row_counts: list[int], + tp_group: dist.ProcessGroup, + tp_size: int, + tp_rank: int, + forward_previous_event: torch.cuda.Event | None, + forward_finished_event: torch.cuda.Event | None, + backward_previous_event: torch.cuda.Event, + backward_finished_event: torch.cuda.Event, + comm_stream: torch.cuda.Stream, + ) -> torch.Tensor: + with torch.cuda.stream(comm_stream): + if forward_previous_event is not None: + comm_stream.wait_event(forward_previous_event) + gathered, tensor_for_comm, chunks = _tp_all_gather_rows_forward_impl(tensor, tp_rank_row_counts, tp_group) + # 中文注释：异步路径只增加 stream/event 管理； + # collective 核心逻辑和同步路径一致。 + _record_stream((tensor_for_comm, chunks, gathered), comm_stream) + if forward_finished_event is not None: + forward_finished_event.record(comm_stream) + + ctx.tp_rank_row_counts = tp_rank_row_counts + ctx.tp_group = tp_group + ctx.tp_rank = tp_rank + ctx.backward_previous_event = backward_previous_event + ctx.backward_finished_event = backward_finished_event + ctx.comm_stream = comm_stream + return gathered + + @staticmethod + def backward( + ctx: Any, + grad: torch.Tensor, + ) -> tuple[torch.Tensor, None, None, None, None, None, None, None, None, None]: + grad_ready_event = torch.cuda.Event() + grad_ready_event.record() + with torch.cuda.stream(ctx.comm_stream): + ctx.comm_stream.wait_event(ctx.backward_previous_event) + ctx.comm_stream.wait_event(grad_ready_event) + grad_input, grad_for_comm, chunks = _tp_all_gather_rows_backward_impl( + grad, + ctx.tp_rank_row_counts, + ctx.tp_rank, + ctx.tp_group, + ) + _record_stream((grad_for_comm, chunks, grad_input), ctx.comm_stream) + ctx.backward_finished_event.record(ctx.comm_stream) + + return grad_input, None, None, None, None, None, None, None, None, None + + +class _TPReduceScatterRowsSum(torch.autograd.Function): + @staticmethod + def forward( + ctx: Any, + tensor: torch.Tensor, + tp_rank_row_counts: list[int], + tp_group: dist.ProcessGroup, + tp_size: int, + tp_rank: int, + ) -> torch.Tensor: + out, _, _ = _tp_reduce_scatter_rows_sum_impl(tensor, tp_rank_row_counts, tp_rank, tp_group) + ctx.tp_rank_row_counts = tp_rank_row_counts + ctx.tp_group = tp_group + return out + + @staticmethod + def backward(ctx: Any, grad_slice: torch.Tensor) -> tuple[torch.Tensor, None, None, None, None]: + full_grad, _, _ = _tp_reduce_scatter_rows_sum_backward_impl(grad_slice, ctx.tp_rank_row_counts, ctx.tp_group) + return full_grad, None, None, None, None + + +class _AsyncTPReduceScatterRowsSum(torch.autograd.Function): + @staticmethod + def forward( + ctx: Any, + tensor: torch.Tensor, + tp_rank_row_counts: list[int], + tp_group: dist.ProcessGroup, + tp_size: int, + tp_rank: int, + forward_previous_event: torch.cuda.Event, + forward_finished_event: torch.cuda.Event, + backward_previous_event: torch.cuda.Event, + backward_finished_event: torch.cuda.Event, + comm_stream: torch.cuda.Stream, + ) -> torch.Tensor: + with torch.cuda.stream(comm_stream): + comm_stream.wait_event(forward_previous_event) + out, tensor_for_comm, chunks = _tp_reduce_scatter_rows_sum_impl( + tensor, + tp_rank_row_counts, + tp_rank, + tp_group, + ) + # 中文注释：TP ReduceScatterRowsSum 属于 combine 通信段； + # 输出事件交给 combine_postprocess 等待。 + _record_stream((tensor_for_comm, chunks, out), comm_stream) + forward_finished_event.record(comm_stream) + + ctx.tp_rank_row_counts = tp_rank_row_counts + ctx.tp_group = tp_group + ctx.backward_previous_event = backward_previous_event + ctx.backward_finished_event = backward_finished_event + ctx.comm_stream = comm_stream + return out + + @staticmethod + def backward( + ctx: Any, + grad_slice: torch.Tensor, + ) -> tuple[torch.Tensor, None, None, None, None, None, None, None, None, None]: + grad_ready_event = torch.cuda.Event() + grad_ready_event.record() + with torch.cuda.stream(ctx.comm_stream): + ctx.comm_stream.wait_event(ctx.backward_previous_event) + ctx.comm_stream.wait_event(grad_ready_event) + full_grad, grad_slice_for_comm, chunks = _tp_reduce_scatter_rows_sum_backward_impl( + grad_slice, + ctx.tp_rank_row_counts, + ctx.tp_group, + ) + _record_stream((grad_slice_for_comm, chunks, full_grad), ctx.comm_stream) + ctx.backward_finished_event.record(ctx.comm_stream) + + return full_grad, None, None, None, None, None, None, None, None, None + + +class ExpertTP: + """Token-sliced Expert TP collectives shared by dispatcher routing + paths.""" + + def __init__(self, tp_group: dist.ProcessGroup) -> None: + self._tp_group = tp_group + self._tp_size = tp_group.size() + + @property + def size(self) -> int: + return self._tp_size + + def gather_tp_rank_row_counts(self, tensor: torch.Tensor, stream: torch.cuda.Stream | None = None) -> list[int]: + if self._tp_size == 1: + return [tensor.shape[0]] + + if stream is None: + local_size = tensor.new_tensor([tensor.shape[0]], dtype=torch.long) + tp_rank_row_counts_t = tensor.new_empty([self._tp_size], dtype=torch.long) + dist.all_gather_into_tensor(tp_rank_row_counts_t, local_size, group=self._tp_group) + else: + # 中文注释：行数要转成 Python list；单独 stream 避免同步 + # dispatcher comm stream 上的大 tensor 通信。 + with torch.cuda.stream(stream): + local_size = tensor.new_tensor([tensor.shape[0]], dtype=torch.long) + tp_rank_row_counts_t = tensor.new_empty([self._tp_size], dtype=torch.long) + dist.all_gather_into_tensor(tp_rank_row_counts_t, local_size, group=self._tp_group) + _record_stream((local_size, tp_rank_row_counts_t), stream) + stream.synchronize() + return [int(size) for size in tp_rank_row_counts_t.tolist()] + + def all_gather_rows( + self, + tensor: torch.Tensor, + tp_rank_row_counts: list[int] | None = None, + ) -> tuple[torch.Tensor, list[int]]: + if self._tp_size == 1: + return tensor, [tensor.shape[0]] + + if tp_rank_row_counts is None: + tp_rank_row_counts = self.gather_tp_rank_row_counts(tensor) + + tp_rank = dist.get_rank(group=self._tp_group) + gathered = _TPAllGatherRows.apply(tensor, tp_rank_row_counts, self._tp_group, self._tp_size, tp_rank) + return gathered, tp_rank_row_counts + + def all_gather_row_metadata(self, tensor: torch.Tensor, tp_rank_row_counts: list[int]) -> torch.Tensor: + # 中文注释：topk_ids/topk_weights 和 hidden 使用同一份 + # tp_rank_row_counts，保证 source token 对齐。 + gathered, _ = self.all_gather_rows(tensor, tp_rank_row_counts) + return gathered + + def all_gather_per_rank_metadata(self, tensor: torch.Tensor) -> torch.Tensor: + # 中文注释：tokens_per_expert_group 这类固定形状 meta + # 不沿 token 维变长，使用独立 gather。 + if self._tp_size == 1: + return tensor.unsqueeze(0) + + gathered = tensor.new_empty((self._tp_size, *tensor.shape)) + dist.all_gather_into_tensor(gathered, tensor.contiguous(), group=self._tp_group) + return gathered + + def async_all_gather_rows( + self, + tensor: torch.Tensor, + tp_rank_row_counts: list[int], + forward_previous_event: torch.cuda.Event | None, + forward_finished_event: torch.cuda.Event | None, + backward_previous_event: torch.cuda.Event, + backward_finished_event: torch.cuda.Event, + comm_stream: torch.cuda.Stream, + ) -> torch.Tensor: + if self._tp_size == 1: + if forward_finished_event is not None: + forward_finished_event.record() + return tensor + + tp_rank = dist.get_rank(group=self._tp_group) + return _AsyncTPAllGatherRows.apply( + tensor, + tp_rank_row_counts, + self._tp_group, + self._tp_size, + tp_rank, + forward_previous_event, + forward_finished_event, + backward_previous_event, + backward_finished_event, + comm_stream, + ) + + def async_all_gather_row_metadata( + self, + tensor: torch.Tensor, + tp_rank_row_counts: list[int], + forward_previous_event: torch.cuda.Event | None, + forward_finished_event: torch.cuda.Event | None, + comm_stream: torch.cuda.Stream, + ) -> torch.Tensor: + if self._tp_size == 1: + if forward_finished_event is not None: + forward_finished_event.record() + return tensor + + with torch.cuda.stream(comm_stream): + if forward_previous_event is not None: + comm_stream.wait_event(forward_previous_event) + gathered, tensor_for_comm, chunks = _tp_all_gather_rows_forward_impl( + tensor, + tp_rank_row_counts, + self._tp_group, + ) + _record_stream((tensor_for_comm, chunks, gathered), comm_stream) + if forward_finished_event is not None: + forward_finished_event.record(comm_stream) + return gathered + + def async_all_gather_per_rank_metadata( + self, + tensor: torch.Tensor, + forward_previous_event: torch.cuda.Event | None, + forward_finished_event: torch.cuda.Event | None, + comm_stream: torch.cuda.Stream, + ) -> torch.Tensor: + if self._tp_size == 1: + if forward_finished_event is not None: + forward_finished_event.record() + return tensor.unsqueeze(0) + + gathered = tensor.new_empty((self._tp_size, *tensor.shape)) + with torch.cuda.stream(comm_stream): + if forward_previous_event is not None: + comm_stream.wait_event(forward_previous_event) + tensor_for_comm = tensor.contiguous() + dist.all_gather_into_tensor(gathered, tensor_for_comm, group=self._tp_group) + _record_stream((tensor_for_comm, gathered), comm_stream) + if forward_finished_event is not None: + forward_finished_event.record(comm_stream) + return gathered + + def reduce_scatter_rows_sum(self, tensor: torch.Tensor, tp_rank_row_counts: list[int]) -> torch.Tensor: + if self._tp_size == 1: + return tensor + + tp_rank = dist.get_rank(group=self._tp_group) + return _TPReduceScatterRowsSum.apply(tensor, tp_rank_row_counts, self._tp_group, self._tp_size, tp_rank) + + def async_reduce_scatter_rows_sum( + self, + tensor: torch.Tensor, + tp_rank_row_counts: list[int], + forward_previous_event: torch.cuda.Event, + forward_finished_event: torch.cuda.Event, + backward_previous_event: torch.cuda.Event, + backward_finished_event: torch.cuda.Event, + comm_stream: torch.cuda.Stream, + ) -> torch.Tensor: + if self._tp_size == 1: + forward_finished_event.record() + return tensor + + tp_rank = dist.get_rank(group=self._tp_group) + return _AsyncTPReduceScatterRowsSum.apply( + tensor, + tp_rank_row_counts, + self._tp_group, + self._tp_size, + tp_rank, + forward_previous_event, + forward_finished_event, + backward_previous_event, + backward_finished_event, + comm_stream, + ) diff --git a/xtuner/v1/module/dispatcher/torch_all2all.py b/xtuner/v1/module/dispatcher/torch_all2all.py index ba1d021e6a..6edc6002be 100644 --- a/xtuner/v1/module/dispatcher/torch_all2all.py +++ b/xtuner/v1/module/dispatcher/torch_all2all.py @@ -19,6 +19,7 @@ PreCombineResult, PreDispatchResult, ) +from .expert_tp import ExpertTP if get_device() == "npu": @@ -51,6 +52,7 @@ class TorchAll2AllDispatchResult(DispatchResult): tokens_per_expert_group: torch.Tensor input_splits: list[int] output_splits: list[int] + tp_rank_row_counts: list[int] forward_finished_event: torch.cuda.Event | None backward_previous_event: torch.cuda.Event | None @@ -285,6 +287,7 @@ class TorchAll2AllDispatcher( ] ): _comm_stream = None + _tp_row_count_stream: torch.cuda.Stream | None = None _process_group: dist.ProcessGroup def __init__( @@ -292,6 +295,7 @@ def __init__( *, n_routed_experts: int, process_group: torch.distributed.ProcessGroup, + tp_group: torch.distributed.ProcessGroup | None = None, training_dtype: Literal["fp8", "bf16"] = "bf16", generate_dtype: Literal["fp8", "bf16"] = "bf16", ): @@ -314,6 +318,10 @@ def __init__( ) if TorchAll2AllDispatcher._comm_stream is None: TorchAll2AllDispatcher._comm_stream = cast(torch.cuda.Stream, torch.cuda.Stream(device=DEVICE)) + self._expert_tp = ExpertTP(tp_group) if tp_group is not None and tp_group.size() > 1 else None + if self._expert_tp is not None and TorchAll2AllDispatcher._tp_row_count_stream is None: + TorchAll2AllDispatcher._tp_row_count_stream = torch.cuda.Stream(device=DEVICE) + self._tp_row_count_stream = TorchAll2AllDispatcher._tp_row_count_stream # if training_dtype == "fp8": # raise NotImplementedError @@ -323,6 +331,7 @@ def dispatch_preprocess( *, hidden_states: torch.Tensor, topk_ids: torch.Tensor, + topk_weights: torch.Tensor, # noqa: ARG002 — kept for interface compatibility; not used here async_op: bool = False, ) -> TorchAll2AllPreDispatchResult: permuted_hidden_states, row_ids_map = permute(hidden_states, topk_ids.to(torch.int32)) @@ -368,6 +377,10 @@ def dispatch( self._n_routed_experts, self._process_group, ) + tp_rank_row_counts = [hidden_states.shape[0]] + if self._expert_tp is not None: + hidden_states, tp_rank_row_counts = self._expert_tp.all_gather_rows(hidden_states) + tokens_per_expert_group = self._expert_tp.all_gather_per_rank_metadata(tokens_per_expert_group) if decoding: raise NotImplementedError else: @@ -377,6 +390,7 @@ def dispatch( tokens_per_expert_group=cast(torch.Tensor, tokens_per_expert_group), input_splits=cast(list[int], input_splits), output_splits=cast(list[int], output_splits), + tp_rank_row_counts=tp_rank_row_counts, forward_finished_event=None, backward_previous_event=None, ) @@ -400,6 +414,36 @@ def dispatch( self._comm_stream, self._process_group, ) + tp_rank_row_counts = [hidden_states.shape[0]] + if self._expert_tp is not None: + comm_stream = cast(torch.cuda.Stream, self._comm_stream) + assert self._tp_row_count_stream is not None + # 中文注释：只同步 TP 变长 tp_rank_row_counts； + # hidden/counts TP 通信继续排在 dispatcher comm stream。 + tp_rank_row_counts = self._expert_tp.gather_tp_rank_row_counts( + hidden_states, + stream=self._tp_row_count_stream, + ) + tp_hidden_finished_event = cast(torch.cuda.Event, torch.cuda.Event()) + tp_counts_finished_event = cast(torch.cuda.Event, torch.cuda.Event()) + tp_backward_previous_event = cast(torch.cuda.Event, torch.cuda.Event()) + hidden_states = self._expert_tp.async_all_gather_rows( + hidden_states, + tp_rank_row_counts=tp_rank_row_counts, + forward_previous_event=forward_finished_event, + forward_finished_event=tp_hidden_finished_event, + backward_previous_event=tp_backward_previous_event, + backward_finished_event=backward_finished_event, + comm_stream=comm_stream, + ) + tokens_per_expert_group = self._expert_tp.async_all_gather_per_rank_metadata( + tokens_per_expert_group, + forward_previous_event=tp_hidden_finished_event, + forward_finished_event=tp_counts_finished_event, + comm_stream=comm_stream, + ) + forward_finished_event = tp_counts_finished_event + backward_previous_event = tp_backward_previous_event if decoding: raise NotImplementedError else: @@ -409,6 +453,7 @@ def dispatch( tokens_per_expert_group=tokens_per_expert_group, input_splits=cast(list[int], input_splits), output_splits=cast(list[int], output_splits), + tp_rank_row_counts=tp_rank_row_counts, backward_previous_event=backward_previous_event, forward_finished_event=forward_finished_event, ) @@ -427,9 +472,20 @@ def dispatch_postprocess( self.wait_comm_stream(dispatched["forward_finished_event"]) tokens_per_expert_group = dispatched["tokens_per_expert_group"] - token_counts = tokens_per_expert_group.ravel() + token_counts = tokens_per_expert_group.ravel().to(torch.long) + if self._expert_tp is not None: + local_expert_ids = self._expert_ids_per_ep_rank.repeat(self._expert_tp.size) + output_size = dispatched["hidden_states"].shape[0] + tokens_per_expert = tokens_per_expert_group.sum(dim=(0, 1)) + else: + local_expert_ids = self._expert_ids_per_ep_rank + output_size = sum(dispatched["output_splits"]) + tokens_per_expert = tokens_per_expert_group.sum(dim=0) + global_input_tokens_local_experts_indices = torch.repeat_interleave( - self._expert_ids_per_ep_rank, token_counts, output_size=sum(dispatched["output_splits"]) + local_expert_ids, + token_counts, + output_size=output_size, ) # The dispatch result is already permuted, so we can return it directly. @@ -437,7 +493,6 @@ def dispatch_postprocess( dispatched["hidden_states"], global_input_tokens_local_experts_indices.to(torch.int32), ) - tokens_per_expert = tokens_per_expert_group.sum(dim=0) if async_op: assert dispatched["backward_previous_event"] is not None, "Please use `async_op=True` for dispatch!" @@ -513,8 +568,14 @@ def combine( decoding: bool = False, ) -> CombineResult: if not async_op: + hidden_states_for_combine = pre_combined["hidden_states"] + if self._expert_tp is not None: + hidden_states_for_combine = self._expert_tp.reduce_scatter_rows_sum( + hidden_states_for_combine, + dispatched["tp_rank_row_counts"], + ) hidden_states = all_to_all_single_autograd( - pre_combined["hidden_states"], + hidden_states_for_combine, input_split_sizes=dispatched["output_splits"], output_split_sizes=dispatched["input_splits"], group=self._process_group, @@ -530,8 +591,26 @@ def combine( assert forward_previous_event is not None, "Please use `async_op=True` for combine_preprocess!" assert backward_finished_event is not None, "Please use `async_op=True` for combine_preprocess!" + hidden_states_for_combine = pre_combined["hidden_states"] + if self._expert_tp is not None: + tp_forward_finished_event = cast(torch.cuda.Event, torch.cuda.Event()) + tp_backward_previous_event = cast(torch.cuda.Event, torch.cuda.Event()) + # 中文注释：TP ReduceScatterRowsSum 属于 combine 通信段， + # EP combine 等 TP 输出事件后再发起。 + hidden_states_for_combine = self._expert_tp.async_reduce_scatter_rows_sum( + hidden_states_for_combine, + tp_rank_row_counts=dispatched["tp_rank_row_counts"], + forward_previous_event=forward_previous_event, + forward_finished_event=tp_forward_finished_event, + backward_previous_event=tp_backward_previous_event, + backward_finished_event=backward_finished_event, + comm_stream=cast(torch.cuda.Stream, self._comm_stream), + ) + forward_previous_event = tp_forward_finished_event + backward_finished_event = tp_backward_previous_event + hidden_states = _async_combine( - pre_combined["hidden_states"], + hidden_states_for_combine, dispatched["output_splits"], dispatched["input_splits"], forward_previous_event, diff --git a/xtuner/v1/module/grouped_linear/moe_group_linear.py b/xtuner/v1/module/grouped_linear/moe_group_linear.py index 1e2653f762..e00a129e34 100644 --- a/xtuner/v1/module/grouped_linear/moe_group_linear.py +++ b/xtuner/v1/module/grouped_linear/moe_group_linear.py @@ -1,11 +1,17 @@ +from typing import Literal + import torch import torch.nn as nn from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.tensor import DTensor, Shard, distribute_tensor +from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor from xtuner.v1.float8.config import Float8Config, ScalingGranularity from xtuner.v1.float8.float8_gmm_tile_wise import TileWiseFloat8GroupedLinear from xtuner.v1.ops import group_gemm +from xtuner.v1.utils.interleaved_shard import InterleavedShard + + +GroupedLinearParallelStyle = Literal["column", "row"] class GroupedLinear(nn.Module): @@ -17,34 +23,152 @@ def __init__( num_routed_experts: int, moe_bias: bool = False, ep_mesh: DeviceMesh | None = None, + expert_tp_mesh: DeviceMesh | None = None, + parallel_style: GroupedLinearParallelStyle | None = None, + ep_tp_mesh: DeviceMesh | None = None, + num_fused_projections: int = 1, ): super().__init__() self.in_features = in_features self.out_features = out_features self.num_routed_experts = num_routed_experts - weight = torch.empty(num_routed_experts * out_features, in_features) self.ep_mesh = ep_mesh - if self.ep_mesh is not None and self.ep_mesh.size() > 1: - self.weight = nn.Parameter(distribute_tensor(weight, ep_mesh, [Shard(0)])) + self.expert_tp_mesh = expert_tp_mesh + self.parallel_style: GroupedLinearParallelStyle | None = parallel_style + self.ep_size = ep_mesh.size() if ep_mesh is not None else 1 + self.tp_size = expert_tp_mesh.size() if expert_tp_mesh is not None else 1 + self.ep_rank = ep_mesh.get_local_rank() if ep_mesh is not None else 0 + self.tp_rank = expert_tp_mesh.get_local_rank() if expert_tp_mesh is not None else 0 + self.tp_enabled = self.expert_tp_mesh is not None and self.tp_size > 1 and self.parallel_style is not None + if self.expert_tp_mesh is not None and self.expert_tp_mesh.size() > 1 and self.parallel_style is None: + raise ValueError("parallel_style must be set when expert_tp_mesh size is greater than 1.") + if self.num_routed_experts % self.ep_size != 0: + raise ValueError( + f"num_routed_experts ({self.num_routed_experts}) must be divisible by ep_size ({self.ep_size})." + ) + + self.local_num_routed_experts = self.num_routed_experts // self.ep_size + self.local_expert_start = self.ep_rank * self.local_num_routed_experts + self.local_expert_end = self.local_expert_start + self.local_num_routed_experts + self.local_in_features = in_features + self.local_out_features = out_features + use_dtensor = False + if self.tp_enabled: + if self.parallel_style == "column": + if out_features % self.tp_size != 0: + raise ValueError(f"out_features ({out_features}) must be divisible by tp_size ({self.tp_size}).") + self.local_out_features = out_features // self.tp_size + elif self.parallel_style == "row": + if in_features % self.tp_size != 0: + raise ValueError(f"in_features ({in_features}) must be divisible by tp_size ({self.tp_size}).") + self.local_in_features = in_features // self.tp_size + else: + raise ValueError(f"Unsupported parallel_style: {self.parallel_style}.") + + # When the caller provides the (ep, tp) 2D sub-mesh, wrap the weight in a DTensor so HF save / load + # know how this rank's slice maps back to the global tensor. Choice of placement depends on + # parallel_style: + # * column-parallel: TP cuts `out_features` inside every local expert → use InterleavedShard + # (per-expert column parallel). EP and TP both slice tensor dim 0. + # * row-parallel: TP cuts `in_features` → just Shard(1). EP still slices dim 0. Two different + # tensor dims, no shard_order conflict. + # Without ep_tp_mesh we fall back to a plain tensor (legacy behavior); the param stays sharded but + # cannot be unsharded for HF save. + use_dtensor = ep_tp_mesh is not None and self.tp_size > 1 + if use_dtensor: + assert ep_tp_mesh is not None # for type narrowing + assert ep_tp_mesh.ndim == 2, f"ep_tp_mesh must be a 2D (ep, tp) sub-mesh, got ndim={ep_tp_mesh.ndim}" + local = torch.empty( + self.local_num_routed_experts * self.local_out_features, + self.local_in_features, + ) + if self.parallel_style == "column": + # `from_local` (not `distribute_tensor`) — the latter goes through redistribute, which + # crashes on the `(Shard, InterleavedShard)` combo (shard_order is None). + # For a fused weight (e.g. fused_w1w3 packing gate_proj + up_proj per expert), the + # per-rank dim has `local_experts * num_fused_projections` stripes — one per (expert, + # fused projection). InterleavedShard must cut INSIDE each stripe so each TP rank ends + # up with the same half of every projection. Passing `num_experts_per_ep` here instead + # of `local_experts * num_fused_projections` swaps the fused projections between TP + # ranks and silently corrupts ``silu(gate) * up``. + num_local_stripes = self.local_num_routed_experts * num_fused_projections + placements: tuple = ( + Shard(0), + InterleavedShard(0, num_local_stripes=num_local_stripes), + ) + else: # row + placements = (Shard(0), Shard(1)) + self.weight = nn.Parameter(DTensor.from_local(local, ep_tp_mesh, placements, run_check=False)) + else: + weight = torch.empty( + self.local_num_routed_experts * self.local_out_features, + self.local_in_features, + ) + self.weight = nn.Parameter(weight) else: - self.weight = nn.Parameter(weight) + weight = torch.empty(num_routed_experts * out_features, in_features) + if self.ep_mesh is not None and self.ep_mesh.size() > 1: + self.weight = nn.Parameter(distribute_tensor(weight, ep_mesh, [Shard(0)])) + else: + self.weight = nn.Parameter(weight) self.moe_bias = moe_bias if self.moe_bias: - bias = torch.zeros(num_routed_experts, out_features) - if self.ep_mesh is not None and self.ep_mesh.size() > 1: - self.bias = nn.Parameter(distribute_tensor(bias, ep_mesh, [Shard(0)])) + if self.parallel_style == "column": + # Keep column-parallel bias flattened like the weight's output dimension. This lets EP and Expert TP + # reuse the same (Shard, InterleavedShard) ownership map, while forward restores the per-expert view. + if use_dtensor: + assert ep_tp_mesh is not None + local_bias = torch.zeros(self.local_num_routed_experts * self.local_out_features) + bias_placements = ( + Shard(0), + InterleavedShard( + 0, + num_local_stripes=self.local_num_routed_experts * num_fused_projections, + ), + ) + self.bias = nn.Parameter( + DTensor.from_local(local_bias, ep_tp_mesh, bias_placements, run_check=False) + ) + elif self.tp_enabled: + self.bias = nn.Parameter(torch.zeros(self.local_num_routed_experts * self.local_out_features)) + else: + bias = torch.zeros(num_routed_experts * out_features) + if self.ep_mesh is not None and self.ep_mesh.size() > 1: + self.bias = nn.Parameter(distribute_tensor(bias, self.ep_mesh, [Shard(0)])) + else: + self.bias = nn.Parameter(bias) else: - self.bias = nn.Parameter(torch.zeros(num_routed_experts, out_features)) + bias = torch.zeros(num_routed_experts, out_features) + if use_dtensor: + assert ep_tp_mesh is not None + local_bias = torch.zeros(self.local_num_routed_experts, out_features) + self.bias = nn.Parameter( + DTensor.from_local( + local_bias, + ep_tp_mesh, + (Shard(0), Replicate()), + run_check=False, + ) + ) + elif self.tp_enabled: + self.bias = nn.Parameter(torch.zeros(self.local_num_routed_experts, out_features)) + elif self.ep_mesh is not None and self.ep_mesh.size() > 1: + self.bias = nn.Parameter(distribute_tensor(bias, ep_mesh, [Shard(0)])) + else: + self.bias = nn.Parameter(bias) def forward(self, x: torch.Tensor, tokens_per_expert: torch.Tensor, decoding: bool = False): weight = self.weight.to_local() if isinstance(self.weight, DTensor) else self.weight - weight = weight.view(-1, self.out_features, self.in_features) + weight = weight.view(-1, self.local_out_features, self.local_in_features) out = group_gemm(x, weight, tokens_per_expert) if self.moe_bias: bias = self.bias.to_local() if isinstance(self.bias, DTensor) else self.bias + bias = bias.view(self.local_num_routed_experts, -1) + if self.tp_enabled and self.parallel_style == "row" and self.tp_rank != 0: + return out out = out + bias.repeat_interleave(tokens_per_expert, dim=0) # TODO: 无法 compile return out @@ -55,16 +179,38 @@ def build_grouped_linear( num_routed_experts: int, moe_bias: bool = False, ep_mesh: DeviceMesh | None = None, + expert_tp_mesh: DeviceMesh | None = None, + parallel_style: GroupedLinearParallelStyle | None = None, float8_cfg: Float8Config | None = None, + ep_tp_mesh: DeviceMesh | None = None, + num_fused_projections: int = 1, ): """Build a grouped linear layer with optional float8 support.""" if float8_cfg is None or float8_cfg.scaling_granularity_grouped_gemm is None: - return GroupedLinear(in_features, out_features, num_routed_experts, moe_bias=moe_bias, ep_mesh=ep_mesh) + return GroupedLinear( + in_features, + out_features, + num_routed_experts, + moe_bias=moe_bias, + ep_mesh=ep_mesh, + expert_tp_mesh=expert_tp_mesh, + parallel_style=parallel_style, + ep_tp_mesh=ep_tp_mesh, + num_fused_projections=num_fused_projections, + ) elif float8_cfg.scaling_granularity_grouped_gemm == ScalingGranularity.TILEWISE: return TileWiseFloat8GroupedLinear( - in_features, out_features, num_routed_experts, moe_bias=moe_bias, ep_mesh=ep_mesh + in_features, + out_features, + num_routed_experts, + moe_bias=moe_bias, + ep_mesh=ep_mesh, + expert_tp_mesh=expert_tp_mesh, + parallel_style=parallel_style, + ep_tp_mesh=ep_tp_mesh, + num_fused_projections=num_fused_projections, ) else: raise NotImplementedError( - f"Unsupported grouped GEMM float8 scaling granularity: {float8_cfg.scaling_granularity_grouped_gemm}" + f"Unsupported float8 grouped GEMM scaling granularity: {float8_cfg.scaling_granularity_grouped_gemm}" ) diff --git a/xtuner/v1/ops/comm/deepep_op.py b/xtuner/v1/ops/comm/deepep_op.py index 6fe92ecf6d..ac5939f93a 100644 --- a/xtuner/v1/ops/comm/deepep_op.py +++ b/xtuner/v1/ops/comm/deepep_op.py @@ -85,22 +85,35 @@ def get_low_latency_buffer( if _buffer is None: # NOTES: for best performance, the QP number **must** be equal to the number of the local experts assert num_experts % group.size() == 0 - # _buffer = Buffer(group, num_nvl_bytes, num_rdma_bytes) + num_qps_per_rank = max(num_experts // group.size(), Buffer.num_sms // 2) _buffer = Buffer( group, num_nvl_bytes, num_rdma_bytes, low_latency_mode=True, - num_qps_per_rank=max(num_experts // group.size(), Buffer.num_sms // 2), + num_qps_per_rank=num_qps_per_rank, ) logger.info( - f"{num_nvl_bytes}, {_buffer.num_nvl_bytes}, {num_max_dispatch_tokens_per_rank}, {hidden}, {num_experts}, {group.size()}" + "[DeepEP low-latency] allocated buffer: " + f"num_nvl_bytes={num_nvl_bytes} (allocated={_buffer.num_nvl_bytes}), " + f"num_rdma_bytes={num_rdma_bytes} (allocated={_buffer.num_rdma_bytes}), " + f"num_max_dispatch_tokens_per_rank={num_max_dispatch_tokens_per_rank}, " + f"hidden={hidden}, num_experts={num_experts}, ep_group_size={group.size()}, " + f"num_qps_per_rank={num_qps_per_rank}" ) else: assert num_nvl_bytes <= _buffer.num_nvl_bytes, ( - f"{num_nvl_bytes}, {_buffer.num_nvl_bytes}, {num_max_dispatch_tokens_per_rank}, {hidden}, {num_experts}, {group.size()}" + "[DeepEP low-latency] NVL buffer too small: " + f"required={num_nvl_bytes}, allocated={_buffer.num_nvl_bytes}, " + f"num_max_dispatch_tokens_per_rank={num_max_dispatch_tokens_per_rank}, " + f"hidden={hidden}, num_experts={num_experts}, ep_group_size={group.size()}" + ) + assert num_rdma_bytes <= _buffer.num_rdma_bytes, ( + "[DeepEP low-latency] RDMA buffer too small: " + f"required={num_rdma_bytes}, allocated={_buffer.num_rdma_bytes}, " + f"num_max_dispatch_tokens_per_rank={num_max_dispatch_tokens_per_rank}, " + f"hidden={hidden}, num_experts={num_experts}, ep_group_size={group.size()}" ) - assert num_rdma_bytes <= _buffer.num_rdma_bytes return _buffer diff --git a/xtuner/v1/patch/__init__.py b/xtuner/v1/patch/__init__.py index 6106a59f99..f2206080c1 100644 --- a/xtuner/v1/patch/__init__.py +++ b/xtuner/v1/patch/__init__.py @@ -1,5 +1,6 @@ from . import torch_shape_env_simplify_pt28 from .dcp_async_port import patch_dcp_async_daemon_port +from .dcp_interleaved_planner import InterleavedShardLoadPlanner, InterleavedShardSavePlanner from .torch_dcp_planner import patch_dcp_save_state_dict, patch_dcp_save_with_cache_storage, patch_default_save_plan @@ -9,4 +10,6 @@ "patch_dcp_save_state_dict", "patch_dcp_save_with_cache_storage", "patch_dcp_async_daemon_port", + "InterleavedShardSavePlanner", + "InterleavedShardLoadPlanner", ] diff --git a/xtuner/v1/patch/dcp_interleaved_planner.py b/xtuner/v1/patch/dcp_interleaved_planner.py new file mode 100644 index 0000000000..5242bc38be --- /dev/null +++ b/xtuner/v1/patch/dcp_interleaved_planner.py @@ -0,0 +1,187 @@ +"""DCP planners that make ``InterleavedShard`` (tpep) DTensors participate in +DCP. + +An ``InterleavedShard`` DTensor (per-expert column-parallel fused MoE weights, produced when +``tp_size > 1``) stores a local tensor that is **several interleaved runs** of the global +tensor rather than one contiguous slice. DCP's default planners model each DTensor as a single +contiguous chunk via ``compute_local_shape_and_global_offset`` — which cannot describe this +placement: on torch >= 2.9 it raises (``_StridedShard`` split_factor != aggregate mesh size), +and on 2.8 it silently returns a wrong ``(size, offset)``. Either way the default path is +unusable for these params, which is why they used to be dropped from DCP. + +The fix routes those DTensors through :func:`xtuner.v1.utils.interleaved_shard.compute_runs`, +which decomposes the local tensor into contiguous ``Run`` records mapped to their true global +offsets. Each run becomes one ``WriteItem`` (save) / ``ReadItem`` (load), so the checkpoint is +stored in global coordinates and reshards correctly across different tp/ep topologies. + +Every other object (plain tensors, non-interleaved DTensors, optimizer state) falls through to +the default planner unchanged. +""" + +from __future__ import annotations + +import dataclasses + +import torch +from torch.distributed._shard._utils import narrow_tensor_by_index +from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner +from torch.distributed.checkpoint.metadata import ( + ChunkStorageMetadata, + MetadataIndex, + TensorProperties, +) +from torch.distributed.checkpoint.planner import ( + LoadPlan, + ReadItem, + SavePlan, + SavePlanner, + TensorWriteData, + WriteItem, + WriteItemType, +) +from torch.distributed.checkpoint.planner_helpers import ( # type: ignore[attr-defined] + _compare_save_plans, + _create_read_items, + _create_write_items, + create_read_items_for_chunk_list, +) +from torch.distributed.tensor import DTensor + +from xtuner.v1.utils.interleaved_shard import Run, compute_runs, has_interleaved_placement + +from .xtuner_cache_planner import XtunerCacheSavePlanner + + +__all__ = ["InterleavedShardSavePlanner", "InterleavedShardLoadPlanner"] + + +class InterleavedShardSavePlanner(XtunerCacheSavePlanner): + """DCP ``SavePlanner`` that emits one ``WriteItem`` per contiguous run for + InterleavedShard DTensors, while preserving + :class:`XtunerCacheSavePlanner`'s incremental-save plan caching for every + other object.""" + + _interleaved_runs: dict[str, dict[tuple, Run]] + + def create_local_plan(self) -> SavePlan: + # Build write items directly. Interleaved DTensors must NOT reach torch's default write-item + # builder: ``compute_local_shape_and_global_offset`` raises for their ``_StridedShard`` + # placement (split_factor != aggregate mesh size) on torch >= 2.9, and silently returns a + # wrong single chunk on 2.8. Route them through per-run items; defer every other object to + # the default builder (matching DefaultSavePlanner, incl. the DTensor submesh-coordinate + # guard). Done before the plan-caching comparison so the cached plan matches what + # ``resolve_data`` streams. + requests: list[WriteItem] = [] + self._interleaved_runs = {} + for fqn, obj in self.state_dict.items(): + if isinstance(obj, DTensor) and has_interleaved_placement(obj): + if obj.device_mesh.get_coordinate() is None: + continue + items, run_map = _interleaved_write_items(fqn, obj) + requests.extend(items) + self._interleaved_runs[fqn] = run_map + elif isinstance(obj, DTensor): + if obj.device_mesh.get_coordinate() is not None: + requests.extend(_create_write_items(fqn, obj)) + else: + requests.extend(_create_write_items(fqn, obj)) + plan = SavePlan(requests) + if self.flatten_state_dict: + plan = dataclasses.replace(plan, planner_data=self.mappings) + self.plan = plan + + # Mirror DefaultSavePlanner.create_local_plan's caching short-circuit (torch 2.7.x + # incremental save): skip re-sending an unchanged local plan to the coordinator. + if self._enable_plan_caching: # type: ignore[attr-defined] + cached = SavePlanner._cached_save_plan # type: ignore[attr-defined] + if self._cached_plans_key in cached and _compare_save_plans(plan, cached[self._cached_plans_key]): + return SavePlan([], usable=False) # type: ignore[call-arg] + cached[self._cached_plans_key] = plan + return self.plan + + def resolve_data(self, write_item: WriteItem): + offset = write_item.index.offset + if offset is not None: + run = self._interleaved_runs.get(write_item.index.fqn, {}).get(tuple(offset)) + if run is not None: + local = self.state_dict[write_item.index.fqn]._local_tensor + return local.narrow(0, run.local_start, run.local_size).contiguous() + return super().resolve_data(write_item) + + +class InterleavedShardLoadPlanner(DefaultLoadPlanner): + """DCP ``LoadPlanner`` that reads InterleavedShard DTensors as one + ``ReadItem`` per contiguous run, resharding from the global-coordinate + checkpoint into this rank's interleaved runs.""" + + _interleaved_runs: dict[str, dict[tuple, Run]] + + def create_local_plan(self) -> LoadPlan: + assert self.metadata is not None + # Reimplement the default per-fqn read-item loop so interleaved DTensors never reach torch's + # default chunk builder (``compute_local_shape_and_global_offset`` raises for their + # ``_StridedShard`` placement on torch >= 2.9). Interleaved DTensors get per-run read items + # against the checkpoint metadata (so DCP reshards into this rank's runs); every other object + # uses the default builder unchanged. The pre-2.4 checkpoint version fallback in + # ``DefaultLoadPlanner.create_local_plan`` is dropped — xtuner never writes those. + self._interleaved_runs = {} + requests: list[ReadItem] = [] + strict = not self.allow_partial_load + for fqn, obj in self.state_dict.items(): + if fqn not in self.metadata.state_dict_metadata: + if strict: + raise RuntimeError(f"Missing key in checkpoint state_dict: {fqn}.") + continue + md = self.metadata.state_dict_metadata[fqn] + if isinstance(obj, DTensor) and has_interleaved_placement(obj): + if obj.device_mesh.get_coordinate() is None: + continue + runs = compute_runs(obj) + local_chunks = [ + ChunkStorageMetadata(offsets=torch.Size(run.global_offset), sizes=torch.Size(run.sizes)) + for run in runs + ] + self._interleaved_runs[fqn] = {tuple(run.global_offset): run for run in runs} + requests.extend(create_read_items_for_chunk_list(fqn, md, local_chunks)) # type: ignore[arg-type] + elif isinstance(obj, DTensor): + if obj.device_mesh.get_coordinate() is not None: + requests.extend(_create_read_items(fqn, md, obj)) + else: + requests.extend(_create_read_items(fqn, md, obj)) + return LoadPlan(requests) + + def resolve_tensor(self, read_item: ReadItem): + offset = read_item.dest_index.offset + run = ( + self._interleaved_runs.get(read_item.dest_index.fqn, {}).get(tuple(offset)) if offset is not None else None + ) + if run is not None: + local = self.state_dict[read_item.dest_index.fqn]._local_tensor + run_slice = local.narrow(0, run.local_start, run.local_size) + # ``dest_offsets`` / ``lengths`` are relative to the run (the "current shard"), so narrow + # the run slice — not the whole local tensor — to land the checkpoint bytes correctly. + return narrow_tensor_by_index(run_slice, read_item.dest_offsets, read_item.lengths) + return super().resolve_tensor(read_item) + + +def _interleaved_write_items(fqn: str, dt: DTensor) -> tuple[list[WriteItem], dict[tuple, Run]]: + properties = TensorProperties.create_from_tensor(dt._local_tensor) + global_size = torch.Size(dt.shape) + items: list[WriteItem] = [] + run_map: dict[tuple, Run] = {} + for run in compute_runs(dt): + offsets = torch.Size(run.global_offset) + sizes = torch.Size(run.sizes) + items.append( + WriteItem( + index=MetadataIndex(fqn, offsets), + type=WriteItemType.SHARD, + tensor_data=TensorWriteData( + chunk=ChunkStorageMetadata(offsets=offsets, sizes=sizes), + properties=properties, + size=global_size, + ), + ) + ) + run_map[tuple(run.global_offset)] = run + return items, run_map diff --git a/xtuner/v1/patch/torch_dcp_planner.py b/xtuner/v1/patch/torch_dcp_planner.py index 4400df9c9e..9c1b117bd1 100644 --- a/xtuner/v1/patch/torch_dcp_planner.py +++ b/xtuner/v1/patch/torch_dcp_planner.py @@ -116,16 +116,23 @@ def dcp_save_with_cache_storage(state_dict, **kwargs): planner = kwargs.get("planner", None) storage_writer = kwargs.get("storage_writer", None) - if storage_writer is None and planner is None: - from xtuner.v1.patch.xtuner_cache_planner import XtunerCacheSavePlanner + if storage_writer is None: + from xtuner.v1.patch.dcp_interleaved_planner import InterleavedShardSavePlanner from xtuner.v1.patch.xtuner_storage import XtunerCacheWriter - planner = XtunerCacheSavePlanner(enable_plan_caching=True, cache_key_prefix=checkpoint_id.stem) - storage_writer = XtunerCacheWriter( - checkpoint_id, enable_write_result_caching=True, cache_key_prefix=checkpoint_id.stem - ) - kwargs["planner"] = planner - kwargs["storage_writer"] = storage_writer + if planner is None: + # Use the interleaved-aware planner on the torch 2.7 incremental-save path too. + planner = InterleavedShardSavePlanner( + enable_plan_caching=True, + cache_key_prefix=checkpoint_id.stem, + ) + kwargs["planner"] = planner + if isinstance(planner, InterleavedShardSavePlanner): + kwargs["storage_writer"] = XtunerCacheWriter( + checkpoint_id, + enable_write_result_caching=True, + cache_key_prefix=checkpoint_id.stem, + ) return original_dcp_save(state_dict, **kwargs) diff --git a/xtuner/v1/rl/weight_update/weight_iterator.py b/xtuner/v1/rl/weight_update/weight_iterator.py index 55f236ba11..a9e92a22b8 100644 --- a/xtuner/v1/rl/weight_update/weight_iterator.py +++ b/xtuner/v1/rl/weight_update/weight_iterator.py @@ -1,18 +1,15 @@ from __future__ import annotations from itertools import chain -from typing import Any, cast +from typing import Any import torch -import torch.distributed as dist import tqdm from torch.distributed.tensor import DTensor from xtuner.v1.model.compose.base import BaseComposeConfig from xtuner.v1.model.compose.qwen3_vl import Qwen3VLForConditionalGeneration -from xtuner.v1.model.moe.moe import MoE from xtuner.v1.utils import get_device, get_torch_device_module -from xtuner.v1.utils.load_spec import LoadEnum, LoadSpec from .data import RolloutWeightUpdateInfo, WeightUpdateBatch @@ -55,123 +52,6 @@ def iter_batch_groups(self): yield self.iter_hf_batches(final_update=True) - def _get_hf_params( - self, - model, - model_ep_size: int, - target_ep_size: int, - target_ep_rank: int, - fsdp_tensor_list: list[tuple[torch.Tensor, LoadSpec]], - should_gather_train_ep_shards: bool, - ) -> tuple[list[torch.Tensor], list[str]]: - hf_keys_list: list[str] = [] - hf_tensor_list: list[torch.Tensor] = [] - - for fsdp_tensor, load_spec in fsdp_tensor_list: - hf_keys = load_spec.hf_keys - if model_ep_size > 1 and model.ep_mesh is not None: - # Each train EP rank owns only part of the HF key list; gather the global - # mapping once so rollout EP ranks can receive the right slice. - if load_spec.name not in self._global_hf_keys_mapping_cache: - global_hf_keys: list[list[str] | None] = [None] * model_ep_size - dist.all_gather_object(global_hf_keys, hf_keys, group=model.ep_mesh.get_group()) - global_hf_keys_gathered = cast(list[list[str]], global_hf_keys) - self._global_hf_keys_mapping_cache[load_spec.name] = list( - chain.from_iterable(global_hf_keys_gathered) - ) - hf_keys = self._global_hf_keys_mapping_cache[load_spec.name] - - fused_full_tensor = fsdp_tensor.bfloat16() - if isinstance(fused_full_tensor, DTensor): - fused_full_tensor = fused_full_tensor.full_tensor() - # FUSED load specs pack multiple HF tensors along load_spec.dim; split them - # back into HF tensors before selecting the target rollout EP shard. - dim = cast(int, load_spec.dim) - - if should_gather_train_ep_shards and model_ep_size > 1: - assert model.ep_mesh is not None - ep_group = model.ep_mesh.get_group() - - output = torch.empty( - *fused_full_tensor.shape[:dim], - fused_full_tensor.shape[dim] * model_ep_size, - *fused_full_tensor.shape[dim + 1 :], - dtype=fused_full_tensor.dtype, - device=fused_full_tensor.device, - ) - dist.all_gather_into_tensor(output, fused_full_tensor.contiguous(), group=ep_group) - fused_full_tensor = output - - num_split = len(hf_keys) - hf_tensor_size = fused_full_tensor.shape[dim] / num_split - assert hf_tensor_size.is_integer(), "Internal Error, hf_tensor_size is not integer" - hf_tensor_size = int(hf_tensor_size) - - hf_tensor = fused_full_tensor.split([hf_tensor_size] * num_split, dim=dim) - assert num_split % target_ep_size == 0, ( - f"len(hf_keys) of '{hf_keys}' is {num_split}, it must be divisible by target_ep_size {target_ep_size}" - ) - start_idx = (num_split // target_ep_size) * target_ep_rank - end_idx = (num_split // target_ep_size) * (target_ep_rank + 1) - - hf_keys_list.extend(hf_keys[start_idx:end_idx]) - hf_tensor_list.extend(hf_tensor[start_idx:end_idx]) - - hf_tensor_list = [ - model.param_to_safetensor(safetensor, name) for safetensor, name in zip(hf_tensor_list, hf_keys_list) - ] - - return hf_tensor_list, hf_keys_list - - def _rl_get_fused_ep_hf_param( - self, - model: MoE, - target_ep_rank: int, - target_ep_size: int, - bucket_size: int, - should_gather_train_ep_shards: bool, - ): - fused_param_groups: list[tuple[torch.Tensor, LoadSpec]] = model._group_param_by_load_spec(LoadEnum.FUSED) - model_ep_size = 1 if model.fsdp_config is None else model.fsdp_config.ep_size - if not fused_param_groups: - return - - safetensor_size = 0 - dtype = torch.bfloat16 - tensor_list: list[tuple[torch.Tensor, LoadSpec]] = [] - - for param, load_spec in fused_param_groups: - tensor_size = dtype.itemsize * param.numel() // target_ep_size - if safetensor_size + tensor_size > bucket_size and tensor_list: - hf_params, name_list = self._get_hf_params( - model, - model_ep_size=model_ep_size, - target_ep_size=target_ep_size, - target_ep_rank=target_ep_rank, - fsdp_tensor_list=tensor_list, - should_gather_train_ep_shards=should_gather_train_ep_shards, - ) - yield name_list, hf_params - safetensor_size = tensor_size - # Kept to mirror the legacy generator layout; the next iteration rebuilds - # name_list from tensor_list before yielding. - name_list = load_spec.hf_keys.copy() - tensor_list = [(param, load_spec)] - continue - safetensor_size += tensor_size - tensor_list.append((param, load_spec)) - - if tensor_list: - hf_params, name_list = self._get_hf_params( - model=model, - model_ep_size=model_ep_size, - target_ep_size=target_ep_size, - target_ep_rank=target_ep_rank, - fsdp_tensor_list=tensor_list, - should_gather_train_ep_shards=should_gather_train_ep_shards, - ) - yield name_list, hf_params - @torch.no_grad() def iter_hf_batches(self, submodule=None, final_update=False): """Update the model weights.""" @@ -182,66 +62,50 @@ def iter_hf_batches(self, submodule=None, final_update=False): dtype = torch.bfloat16 bucket_size = int(self.config.update_weight_bucket_size_in_gb * 1024**3) - same_gen = model._get_same_hf_param( - model._group_param_by_load_spec(LoadEnum.SAME), + train_enable_ep = model.fsdp_config is not None and model.fsdp_config.ep_size > 1 + params = model._load_spec_params() + ep_mesh = getattr(model, "ep_mesh", None) + ep_group = ep_mesh.get_group() if ep_mesh is not None and ep_mesh.size() > 1 else None + + # IPC maps train ranks directly to rollout ranks, so keep the fused expert EP shard local and only gather + # later shards such as FSDP. NCCL is driven by train rank 0 and therefore needs globally gathered weights. + preserve_ep_shards = self.rollout_info.transport_type == "ipc" and ep_group is not None + ep_fused_params = [] + other_params = [] + for param, load_spec in params: + is_ep_fused = ( + preserve_ep_shards + and load_spec.is_fused + and load_spec.fused_dim is not None + and any( + shard.dim == load_spec.fused_dim and model._is_same_process_group(shard.group, ep_group) + for shard in load_spec.shards + ) + ) + (ep_fused_params if is_ep_fused else other_params).append((param, load_spec)) + + ep_fused_gen = model._get_hf_param( + ep_fused_params, dtype=dtype, device=DEVICE, bucket_size=bucket_size, + preserved_fused_shard_group=ep_group, ) - - train_enable_ep = model.fsdp_config is not None and model.fsdp_config.ep_size > 1 - should_gather_train_ep_shards = self.rollout_info.transport_type == "nccl" and train_enable_ep - - if train_enable_ep: - if self.rollout_info.transport_type == "ipc" and self.rollout_info.ep > 1: - target_ep_rank = self.rollout_info.ipc_engine_parallel_rank - target_ep_size = self.rollout_info.ipc_engine_parallel_size - assert target_ep_rank is not None, "IPC rollout target for current train rank is not resolved." - assert target_ep_size is not None, "IPC rollout target size for current train rank is not resolved." - # Colocated IPC can send only the expert slice needed by the local rollout - # EP rank - fused_gen = self._rl_get_fused_ep_hf_param( - model, - target_ep_rank=target_ep_rank, - target_ep_size=target_ep_size, - bucket_size=bucket_size, - should_gather_train_ep_shards=should_gather_train_ep_shards, - ) - else: - # Disaggregated NCCL uses one trainer-side broadcast for all rollout ranks. - # Gather train EP shards first, then send the full expert tensor instead of - # slicing by rollout EP rank. - fused_gen = self._rl_get_fused_ep_hf_param( - model, - target_ep_rank=0, - target_ep_size=1, - bucket_size=bucket_size, - should_gather_train_ep_shards=should_gather_train_ep_shards, - ) - else: - fused_gen = model._get_fused_hf_param( - model._group_param_by_load_spec(LoadEnum.FUSED), - dtype=dtype, - device=DEVICE, - bucket_size=bucket_size, - update_weights_for_rl=True, - ) - shard_gen = model._get_shard_hf_param( - model._group_param_by_load_spec(LoadEnum.SHARD), + other_gen = model._get_hf_param( + other_params, dtype=dtype, device=DEVICE, bucket_size=bucket_size, ) - - for name_list, fused_param_list in fused_gen: - state_dict = {name: param.detach() for name, param in zip(name_list, fused_param_list)} - yield WeightUpdateBatch(state_dict, train_enable_ep=train_enable_ep, finished=False) - del state_dict, name_list, fused_param_list - - for name_list, param_list in chain(same_gen, shard_gen): - state_dict = {name: param.detach() for name, param in zip(name_list, param_list)} - yield WeightUpdateBatch(state_dict, train_enable_ep=train_enable_ep, finished=False) - del state_dict, name_list, param_list + for name_list, param_list in chain(ep_fused_gen, other_gen): + # FlattenedTensorBucket stores one dtype per payload. Qwen3.5 keeps + # selected norm and A_log weights in fp32, so split those from the + # otherwise bf16 HF-save bucket before handing it to the transport. + state_dicts: dict[torch.dtype, dict[str, torch.Tensor]] = {} + for name, param in zip(name_list, param_list, strict=True): + state_dicts.setdefault(param.dtype, {})[name] = param.detach() + for state_dict in state_dicts.values(): + yield WeightUpdateBatch(state_dict, train_enable_ep=train_enable_ep, finished=False) # pytorch and vLLM use an empty final update as an end marker; SGLang and # turbomind do not consume this marker. diff --git a/xtuner/v1/utils/init_weight.py b/xtuner/v1/utils/init_weight.py index ed0fce487f..c865448775 100644 --- a/xtuner/v1/utils/init_weight.py +++ b/xtuner/v1/utils/init_weight.py @@ -24,9 +24,16 @@ def init_params(param: torch.Tensor, init_fn: Callable[[torch.Tensor], torch.Ten device = param.device if isinstance(param, DTensor): - full_param = torch.empty_like(param.full_tensor(), device=device) - init_fn(full_param) - param.copy_(distribute_tensor(full_param, param.device_mesh, param.placements)) + # InterleavedShard cannot go through full_tensor/distribute_tensor because PyTorch has no + # redistribute path for that placement chain. Initialize the values owned by this rank directly. + from .interleaved_shard import has_interleaved_placement + + if has_interleaved_placement(param): + init_fn(param._local_tensor) + else: + full_param = torch.empty_like(param.full_tensor(), device=device) + init_fn(full_param) + param.copy_(distribute_tensor(full_param, param.device_mesh, param.placements)) else: init_fn(param) diff --git a/xtuner/v1/utils/interleaved_shard.py b/xtuner/v1/utils/interleaved_shard.py new file mode 100644 index 0000000000..3f2d317edb --- /dev/null +++ b/xtuner/v1/utils/interleaved_shard.py @@ -0,0 +1,324 @@ +"""Per-expert column parallel placement and helpers. + +This module introduces ``InterleavedShard``, a custom :class:`Placement` for fused MoE weights +where TP needs to cut ``out_features`` *inside* every local expert. The layout cannot be +expressed by torch's built-in ``Shard`` (which would either give each TP rank one whole expert +or break expert boundaries). ``InterleavedShard`` does exactly per-expert column parallel. + +It is intentionally a subclass of ``_StridedShard`` so: + + * FSDP2 (``fully_shard``) recognizes it via ``isinstance(..., _StridedShard)`` and prepends + its own placement on the same tensor dim correctly. + * All ``_local_shard_size_and_offset``/``_split_tensor``/``_to_replicate_tensor`` semantics + come from ``_StridedShard`` for free. + +The cost is that PyTorch cannot reduce ``(Shard, InterleavedShard)`` (i.e. the strided shard +sitting at the *rightmost* mesh dim) to a ``ShardOrder``. Any code path that relies on +``DTensorSpec.shard_order`` — most notably ``DTensor.redistribute`` / ``DTensor.full_tensor`` — +crashes on such DTensors. xtuner deliberately bypasses those paths: + + * Forward / backward read ``weight.to_local()`` so the op dispatcher is never invoked on + InterleavedShard parameters. + * HF save / load are routed through LoadSpec plans, which do not depend on ``shard_order``. + :func:`reconstruct_full_tensor` is a convenience wrapper over the same save-plan executor. + +The reconstruction algorithm and its rationale are documented inline on +:func:`reconstruct_full_tensor`. +""" + +from __future__ import annotations + +from typing import NamedTuple + +import torch +from torch.distributed.tensor import DTensor, Shard +from torch.distributed.tensor.placement_types import _StridedShard + + +__all__ = [ + "InterleavedShard", + "Run", + "compute_runs", + "has_interleaved_placement", + "reconstruct_full_tensor", +] + + +class Run(NamedTuple): + """One contiguous run of global indices that the current rank owns on the + sharded dim. + + Used by the DCP planners to build per-run WriteItems / ReadItems. HF load/save derives + equivalent run ownership from ``ShardDescriptor`` inside its plans. + + Args: + global_offset (tuple[int, ...]): Offset into the global tensor where this run begins. + All non-sharded dims are 0. + sizes (tuple[int, ...]): Chunk size on each tensor dim for this run. + local_start (int): Row in the local tensor where this run begins. + local_size (int): Number of rows in this run (== sizes on the sharded dim). + """ + + global_offset: tuple[int, ...] + sizes: tuple[int, ...] + local_start: int + local_size: int + + +class InterleavedShard(_StridedShard): + """Per-stripe column-parallel placement for fused MoE weights. + + For a fused weight whose sharded dim contains ``num_local_stripes`` equal-size logical + stripes per rank, this placement cuts the **inside** of every stripe by ``tp_size`` and + interleaves the cuts. Each ``(ep, tp)`` rank ends up holding ``num_local_stripes`` runs + of contiguous rows; consecutive runs are spaced by one full stripe. + + Two common stripe interpretations: + + * **Non-fused MoE weight** (e.g. one projection per expert): one stripe per local expert. + ``num_local_stripes == num_experts_per_ep``. + * **Fused MoE weight** (e.g. ``fused_w1w3`` packs ``gate_proj`` and ``up_proj`` per + expert): one stripe per (expert, fused projection). For ``fused_w1w3`` with 2 projections + per expert: ``num_local_stripes == num_experts_per_ep * 2``. + + Getting ``num_local_stripes`` wrong silently produces a layout that swaps data between + fused projections (e.g. ``silu(gate) * up`` becomes ``silu(gate_half) * gate_other_half``), + so callers must pass the value that matches the HF key concatenation order. + + Internally this is a ``_StridedShard(dim, split_factor=num_local_stripes)``. + + Args: + dim (int): Tensor dim to shard. For fused MoE weights this is 0. + num_local_stripes (int): Number of equal-size stripes the per-rank dim contains. + See class docstring for how to compute this. + """ + + def __init__(self, dim: int, *, num_local_stripes: int): + super().__init__(dim, split_factor=num_local_stripes) + + @property + def num_local_stripes(self) -> int: + return self.split_factor + + def __repr__(self) -> str: + return f"InterleavedShard(dim={self.dim}, num_local_stripes={self.split_factor})" + + +def has_interleaved_placement(dt: torch.Tensor) -> bool: + """True if ``dt`` is a DTensor whose placements include a strided shard + that cannot be reduced to a valid ShardOrder — i.e. our per-expert column + parallel layout. + + Detection strategy: + + * torch >= 2.10: check ``DTensorSpec.shard_order is None``. The auto-derivation returns + ``None`` whenever an internal ``_StridedShard`` placement has no consistent + ``split_factor`` insertion position (exactly our case). + * torch < 2.10: that attribute does not exist, so fall back to a structural scan — + look for any ``_StridedShard`` whose position+sf cannot match the cumulative mesh + sizes to its right. + """ + if not isinstance(dt, DTensor): + return False + shard_order = getattr(dt._spec, "shard_order", _SENTINEL) + if shard_order is not _SENTINEL: + return shard_order is None + # Fallback for torch < 2.10: replicate the carving-order insertion check. + return _placement_chain_unsupported(dt.placements, dt.device_mesh) + + +# Marker used to distinguish "attribute missing" (older torch) vs "attribute is None" +# (the case we care about on 2.10+). +_SENTINEL = object() + + +def _placement_chain_unsupported(placements, mesh) -> bool: + """Right-to-left insertion check, identical to torch 2.10's + ``_maybe_convert_StridedShard_to_shard_order``. + + Returns ``True`` iff any + ``_StridedShard`` cannot be slotted into a consistent carving order. + """ + tensor_dim_to_order: dict[int, list[int]] = {} + for mesh_dim in reversed(range(len(placements))): + p = placements[mesh_dim] + if not isinstance(p, (Shard, _StridedShard)): + continue + order = tensor_dim_to_order.setdefault(p.dim, []) + sf = p.split_factor if isinstance(p, _StridedShard) else 1 + accumulated = 1 + inserted = False + for position in range(len(order) + 1): + if accumulated == sf: + order.insert(position, mesh_dim) + inserted = True + break + if position < len(order): + accumulated *= mesh.size(order[position]) + if not inserted: + return True + return False + + +def _strided_indices(placement, curr_size: int, num_chunks: int, rank: int) -> list[int]: + """Return the list of indices the given rank owns under a ``_StridedShard`` + placement. + + Compatible with both torch 2.9 (no ``return_first_offset`` kwarg, only contiguous offset + returned) and torch 2.10+ (full index list available). For 2.9 we replicate the formula + derived from ``_StridedShard._split_tensor``: rank ``r`` owns chunks ``r, r+M, r+2M, …`` of + the ``M*sf``-way split, each chunk being ``N / (M*sf)`` elements wide. + """ + sf = placement.split_factor + total_split = num_chunks * sf + chunk_size = curr_size // total_split + if chunk_size * total_split != curr_size: + raise NotImplementedError( + f"_strided_indices: uneven sharding (curr_size={curr_size}, " + f"num_chunks={num_chunks}, split_factor={sf}) is not yet supported." + ) + indices: list[int] = [] + for j in range(sf): + chunk_start = (j * num_chunks + rank) * chunk_size + indices.extend(range(chunk_start, chunk_start + chunk_size)) + return indices + + +def _is_fsdp_prepended_strided(placement, mesh_dim: int) -> bool: + """Heuristic: a ``_StridedShard`` at mesh dim 0 is FSDP-prepended. + + ``fully_shard`` always prepends its placement at the leftmost mesh dim, and FSDP's actual + chunking is plain contiguous (``_chunk_with_empty``) despite the strided label. Position + ``0`` is the most reliable signal because the ``_StridedShard`` subclass identity does not + survive ``distribute_tensor`` / FSDP2's internal spec construction (C++ layer reconstructs + a bare ``_StridedShard``). + + This heuristic breaks if a user places an InterleavedShard at mesh dim 0 directly without + FSDP wrapping. xtuner does not do that — InterleavedShard is always at the TP position. + """ + return mesh_dim == 0 and isinstance(placement, _StridedShard) and placement.split_factor > 1 + + +def _is_real_strided(placement, mesh_dim: int) -> bool: + """True iff ``placement`` is a real strided shard whose data layout + actually requires the interleaved gather+scatter algorithm. + + Excludes FSDP-prepended labels. + """ + return ( + isinstance(placement, _StridedShard) + and placement.split_factor > 1 + and not _is_fsdp_prepended_strided(placement, mesh_dim) + ) + + +def reconstruct_full_tensor(dt: DTensor) -> torch.Tensor: + """Reconstruct a full runtime tensor through the shared SavePlan executor. + + PyTorch ``DTensor.full_tensor()`` cannot redistribute Expert TP layouts with + ``shard_order=None``. LoadSpec normalizes those placements into continuous + and even-interleave descriptors, and its save executor performs the inverse + collectives for both HF save and this convenience API. + + Returns: + torch.Tensor: the global tensor materialized on every rank. Dtype and device match + ``dt._local_tensor``. + """ + if not isinstance(dt, DTensor): + raise TypeError(f"reconstruct_full_tensor expects a DTensor, got {type(dt).__name__}") + + from xtuner.v1.utils.load_spec import LoadSpec, unshard_tensors_for_hf_save + + load_spec = LoadSpec.from_tensor( + name="__full_runtime_tensor__", + hf_keys=["__full_runtime_tensor__"], + tensor=dt, + ) + return unshard_tensors_for_hf_save( + [dt._local_tensor.contiguous()], + [load_spec.plan_hf_save()], + )[0] + + +def compute_runs(dt: DTensor) -> list[Run]: + """Compute the contiguous-run decomposition of this rank's share of the + global tensor. + + Accumulates the global indices the current rank owns on the sharded dim. Adjacent indices + are grouped into ``Run`` records so the caller can do per-run narrow + copy without ever + materializing the full index tensor. + + FSDP prepends its placement at mesh dim 0, but semantically it shards the already EP/TP-local + parameter. So for index computation we apply non-FSDP placements first and the FSDP-prepended + shard last, matching the descriptor order used by HF plans. + + Restricted to single-dim sharding (the only layout xtuner currently uses for fused MoE + weights). For multi-dim sharding a Cartesian-product extension is straightforward. + """ + if not isinstance(dt, DTensor): + raise TypeError(f"compute_runs expects a DTensor, got {type(dt).__name__}") + + mesh = dt.device_mesh + global_shape = tuple(dt.shape) + ndim = len(global_shape) + + fsdp_prepended = [] + placement_order = [] + for mesh_dim, p in enumerate(dt.placements): + item = (mesh_dim, p) + if _is_fsdp_prepended_strided(p, mesh_dim): + fsdp_prepended.append(item) + else: + placement_order.append(item) + + dim_indices: dict[int, list[int]] = {} + for mesh_dim, p in placement_order + fsdp_prepended: + if not isinstance(p, (Shard, _StridedShard)): + continue + d = p.dim + prev = dim_indices.get(d) + prev_size = len(prev) if prev is not None else global_shape[d] + if _is_real_strided(p, mesh_dim): + new_idx = _strided_indices(p, prev_size, mesh.size(mesh_dim), mesh.get_local_rank(mesh_dim)) + else: + size, offset = Shard(d)._local_shard_size_and_offset( # type: ignore[attr-defined] + prev_size, mesh.size(mesh_dim), mesh.get_local_rank(mesh_dim) + ) + new_idx = list(range(offset, offset + size)) + dim_indices[d] = new_idx if prev is None else [prev[i] for i in new_idx] + + sharded_dims = sorted(dim_indices.keys()) + assert sharded_dims == [0], f"compute_runs currently handles dim-0 sharding only, got {sharded_dims}" + + indices = dim_indices[0] + if not indices: + return [] + + runs: list[Run] = [] + run_start = indices[0] + run_len = 1 + local_start = 0 + for i in range(1, len(indices)): + if indices[i] == indices[i - 1] + 1: + run_len += 1 + continue + runs.append( + Run( + global_offset=(run_start,) + (0,) * (ndim - 1), + sizes=(run_len,) + global_shape[1:], + local_start=local_start, + local_size=run_len, + ) + ) + local_start += run_len + run_start = indices[i] + run_len = 1 + runs.append( + Run( + global_offset=(run_start,) + (0,) * (ndim - 1), + sizes=(run_len,) + global_shape[1:], + local_start=local_start, + local_size=run_len, + ) + ) + return runs diff --git a/xtuner/v1/utils/load_spec.py b/xtuner/v1/utils/load_spec.py index ef95585fe3..f849ed8dca 100644 --- a/xtuner/v1/utils/load_spec.py +++ b/xtuner/v1/utils/load_spec.py @@ -1,36 +1,986 @@ +import math +from itertools import product +from typing import Callable, NamedTuple + +import torch import torch.distributed as dist -from pydantic import BaseModel, ConfigDict +import torch.nn.functional as F +from pydantic import BaseModel, ConfigDict, Field, computed_field +from torch.distributed.tensor import DTensor, Shard +from torch.distributed.tensor.placement_types import _StridedShard -from .enum_helper import StrEnum +from xtuner.v1.ops.comm.foreach_allgather import foreach_all_gather +from xtuner.v1.utils.device import get_device -class LoadEnum(StrEnum): - FUSED = "fused" - SAME = "same" - SHARD = "shard" +def _is_same_process_group(left: dist.ProcessGroup, right: dist.ProcessGroup) -> bool: + if left is right: + return True + return dist.get_process_group_ranks(left) == dist.get_process_group_ranks(right) -class LoadSpec(BaseModel): - # TODO: (yehaochen) Add more description +class ShardDescriptor(BaseModel): + """One runtime partition applied to the canonical full tensor. + + Descriptors are applied in forward layout order. ``interleave_factor == 1`` + follows normal ``Shard`` semantics, including uneven continuous shards. + Larger factors describe even interleave: every rank owns that many equal + runs from the current tensor dimension. + + Args: + dim (int): Tensor dim on which this partition cuts. + group (dist.ProcessGroup): Communication group that produced this partition. + interleave_factor (int): Number of ordered runs owned by each rank. One means a continuous shard. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + dim: int + group: dist.ProcessGroup + interleave_factor: int = Field(default=1, ge=1) + + def local_intervals(self, dim_size: int) -> list[tuple[int, int]]: + """Return this rank's intervals in the current tensor coordinate.""" + world_size = dist.get_world_size(group=self.group) + rank = dist.get_rank(group=self.group) + assert rank >= 0, "ShardDescriptor process group must contain the current rank" + + if self.interleave_factor == 1: + # XTuner may initialize modules while the default device is meta. PyTorch's + # placement helper inherits that default for temporary shape arithmetic. + with torch.device(get_device()): + local_size, offset = Shard(self.dim)._local_shard_size_and_offset( # type: ignore[attr-defined] + dim_size, + world_size, + rank, + ) + return [(offset, offset + local_size)] if local_size else [] + + split_count = world_size * self.interleave_factor + if dim_size % split_count != 0: + raise NotImplementedError( + "Even interleave requires size_before_shard to be divisible by " + f"group_size * interleave_factor, got {dim_size} % " + f"({world_size} * {self.interleave_factor}) != 0" + ) + run_size = dim_size // split_count + return [ + ((run_index * world_size + rank) * run_size, (run_index * world_size + rank + 1) * run_size) + for run_index in range(self.interleave_factor) + ] + + +def _dtensor_shards(tensor: DTensor) -> list[ShardDescriptor]: + try: + # Normal DTensors already have a valid carving order. The helper also + # normalizes FSDP's bookkeeping _StridedShard to a continuous Shard. + ordered_placements = _ordered_dtensor_placements(tensor) + except RuntimeError: + # Expert TP intentionally creates a placement chain with shard_order=None. + # Its semantic order is the model-parallel placements followed by the + # FSDP-prepended placement, even though FSDP occupies mesh dim 0. + fsdp_placements: list[tuple[int, object]] = [] + ordered_placements = [] + for mesh_dim, raw_placement in enumerate(tensor.placements): + if not isinstance(raw_placement, (Shard, _StridedShard)): + continue + is_fsdp_prepended = ( + mesh_dim == 0 and isinstance(raw_placement, _StridedShard) and raw_placement.split_factor > 1 + ) + (fsdp_placements if is_fsdp_prepended else ordered_placements).append((mesh_dim, raw_placement)) + ordered_placements.extend(fsdp_placements) + + shards: list[ShardDescriptor] = [] + for mesh_dim, ordered_placement in ordered_placements: + assert isinstance(ordered_placement, (Shard, _StridedShard)) + is_fsdp_prepended = ( + mesh_dim == 0 and isinstance(ordered_placement, _StridedShard) and ordered_placement.split_factor > 1 + ) + shards.append( + ShardDescriptor( + dim=ordered_placement.dim, + group=tensor.device_mesh.get_group(mesh_dim), + interleave_factor=( + ordered_placement.split_factor + if isinstance(ordered_placement, _StridedShard) and not is_fsdp_prepended + else 1 + ), + ) + ) + return shards + + +def _ordered_dtensor_placements(tensor: DTensor) -> list[tuple[int, object]]: + # Return placements expanded into carving order: for each tensor dim that is sharded, emit one + # (mesh_dim, Shard(tensor_dim)) entry per mesh dim, listed in the order each mesh dim cuts the + # tensor. `_StridedShard` placements are normalized to plain `Shard` so downstream code can + # treat every entry as a contiguous slice on its current sub-tensor. + # + # Algorithm mirrors torch 2.10's `_maybe_convert_StridedShard_to_shard_order`: process + # placements right-to-left and, for each `_StridedShard(d, split_factor=sf)`, insert it into + # the carving order for its tensor dim at the position where the product of mesh sizes of + # already-inserted entries on its right equals `sf`. Plain `Shard` is treated as `sf == 1`, + # so it always slots into the outermost free position. + # + # We re-implement the algorithm here instead of importing torch's helper because the relevant + # PyTorch symbol changed across versions (`_explicit_order_placements` in <2.10, + # `DTensorSpec._normalize_placements_into_shard_order` in >=2.10) and both are private. Keeping + # the math local insulates LoadSpec from future PyTorch refactors. + mesh = tensor.device_mesh + placements = tensor.placements + + tensor_dim_to_carving_order: dict[int, list[int]] = {} + for mesh_dim in reversed(range(len(placements))): + placement = placements[mesh_dim] + if not isinstance(placement, (Shard, _StridedShard)): + continue + tensor_dim = placement.dim + split_factor = placement.split_factor if isinstance(placement, _StridedShard) else 1 + carving_order = tensor_dim_to_carving_order.setdefault(tensor_dim, []) + + # Walk the existing carving order from outermost (index 0) inward, accumulating mesh sizes. + # The current placement slots in at the position where accumulated size equals split_factor. + accumulated = 1 + inserted = False + for position in range(len(carving_order) + 1): + if accumulated == split_factor: + carving_order.insert(position, mesh_dim) + inserted = True + break + if position < len(carving_order): + accumulated *= mesh.size(carving_order[position]) + if not inserted: + # No insertion point matched: split_factor is inconsistent with the cumulative mesh + # sizes of the other placements on this tensor dim. The placement is malformed for this + # mesh (PyTorch's algorithm would also reject it). + raise RuntimeError( + f"Cannot place {placement} at mesh dim {mesh_dim} into carving order for tensor " + f"dim {tensor_dim}: split_factor {split_factor} does not match any cumulative " + f"mesh size produced by the other placements on this dim." + ) + + ordered: list[tuple[int, object]] = [] + for tensor_dim in sorted(tensor_dim_to_carving_order): + for mesh_dim in tensor_dim_to_carving_order[tensor_dim]: + ordered.append((mesh_dim, Shard(tensor_dim))) + return ordered + + +class _OwnedRegion(BaseModel): + """One contiguous region of the global tensor owned by this rank. + + ``global_offsets`` locate the region in XTuner's canonical global tensor, + while ``local_offsets`` locate the same data in the runtime local tensor. + A regular FSDP/EP shard has one region; an interleaved Expert TP layout has + one region per contiguous run. + """ + + model_config = ConfigDict(extra="forbid") + global_offsets: tuple[int, ...] + local_offsets: tuple[int, ...] + sizes: tuple[int, ...] + + +class LoadCopyRegion(BaseModel): + """One source-to-target copy executed by :class:`HFLoadPlan`.""" + + model_config = ConfigDict(extra="forbid") + source_offsets: tuple[int, ...] + target_offsets: tuple[int, ...] + sizes: tuple[int, ...] + + +class HFLoadPlan(BaseModel): + """Rank-local program for loading HF tensors into one runtime tensor. + + Args: + name (str): Fully-qualified parameter or buffer name on the xtuner side. + hf_keys (list[str]): HF keys that must be read for this rank. + fused_dim (int | None): Concatenation dimension when multiple HF keys are loaded. + canonical_source_shape (tuple[int, ...] | None): Expected shape after the model adapter converts the loaded + HF tensor to XTuner's canonical layout. ``None`` means the plan has no checkpoint-backed copy work. + target_shape (tuple[int, ...]): Expected runtime local-tensor shape. + copy_regions (list[LoadCopyRegion]): Source-to-target copies in canonical coordinates. + zero_unwritten_target (bool): Whether to zero the target before executing copies, used for runtime padding. + """ + + model_config = ConfigDict(extra="forbid") name: str + hf_keys: list[str] + fused_dim: int | None = None + canonical_source_shape: tuple[int, ...] | None = None + target_shape: tuple[int, ...] + copy_regions: list[LoadCopyRegion] = Field(default_factory=list) + zero_unwritten_target: bool = False + + @torch.no_grad() + def load_into( + self, + checkpoint_tensors: list[torch.Tensor], + local_tensor: torch.Tensor, + canonicalize: Callable[[str, torch.Tensor], torch.Tensor], + ) -> None: + """Convert loaded HF tensors and execute this rank's copy program. + + Model adapters own only the HF-layout-to-canonical transformation. This method owns the ordering around that + adapter and all runtime-layout details, including regular slices, interleaved runs, and padding. + """ + assert tuple(local_tensor.shape) == self.target_shape, ( + f"Load target shape {tuple(local_tensor.shape)} does not match planned shape " + f"{self.target_shape} for {self.name}" + ) + assert len(checkpoint_tensors) == len(self.hf_keys), ( + f"Loaded {len(checkpoint_tensors)} tensors for {len(self.hf_keys)} planned HF keys of {self.name}" + ) + + if self.zero_unwritten_target: + local_tensor.zero_() + if not self.copy_regions: + return + + canonical_tensor = self._canonicalize_source(checkpoint_tensors, canonicalize) + for region in self.copy_regions: + source = self._narrow_region(canonical_tensor, region.source_offsets, region.sizes) + target = self._narrow_region(local_tensor, region.target_offsets, region.sizes) + target.copy_(source) + + def _canonicalize_source( + self, + checkpoint_tensors: list[torch.Tensor], + canonicalize: Callable[[str, torch.Tensor], torch.Tensor], + ) -> torch.Tensor: + """Build and validate the canonical source before any rank-local + copy.""" + assert checkpoint_tensors, f"Internal Error. No safetensors were loaded for {self.name}" + if len(checkpoint_tensors) == 1: + loaded_tensor = checkpoint_tensors[0] + else: + assert self.fused_dim is not None, ( + f"Internal Error. fused_dim must be set when loading multiple HF keys for {self.name}" + ) + loaded_tensor = torch.cat(checkpoint_tensors, dim=self.fused_dim) + + canonical_tensor = canonicalize(self.name, loaded_tensor) + assert self.canonical_source_shape is not None + assert tuple(canonical_tensor.shape) == self.canonical_source_shape, ( + f"Canonical HF tensor shape {tuple(canonical_tensor.shape)} does not match planned shape " + f"{self.canonical_source_shape} for {self.name}" + ) + return canonical_tensor + + @staticmethod + def _narrow_region( + tensor: torch.Tensor, + offsets: tuple[int, ...], + sizes: tuple[int, ...], + ) -> torch.Tensor: + assert tensor.dim() == len(offsets) == len(sizes) + for dim, (offset, size) in enumerate(zip(offsets, sizes)): + tensor = tensor.narrow(dim, offset, size) + return tensor + + +def _layout_segments( + global_shape: tuple[int, ...], + shards: list[ShardDescriptor], +) -> list[list[tuple[int, int]]]: + """Map each local tensor dimension to ordered global segments. + + A dimension starts as one full global segment. Each descriptor slices the *current local order*, so a later FSDP + shard can cut across multiple ETP runs without materializing an index for every tensor row. + """ + segments_by_dim = [[(0, dim_size)] if dim_size else [] for dim_size in global_shape] + for shard in shards: + assert 0 <= shard.dim < len(global_shape), f"Invalid shard dim {shard.dim} for shape {global_shape}" + source_segments = segments_by_dim[shard.dim] + current_size = sum(size for _, size in source_segments) + selected_intervals = shard.local_intervals(current_size) + selected_segments: list[tuple[int, int]] = [] + + for selected_start, selected_end in selected_intervals: + local_start = 0 + for global_start, segment_size in source_segments: + local_end = local_start + segment_size + overlap_start = max(selected_start, local_start) + overlap_end = min(selected_end, local_end) + if overlap_start < overlap_end: + mapped_start = global_start + overlap_start - local_start + mapped_size = overlap_end - overlap_start + previous_start, previous_size = selected_segments[-1] if selected_segments else (0, 0) + if selected_segments and previous_start + previous_size == mapped_start: + selected_segments[-1] = (previous_start, previous_size + mapped_size) + else: + selected_segments.append((mapped_start, mapped_size)) + local_start = local_end + + segments_by_dim[shard.dim] = selected_segments + return segments_by_dim + + +def _shape_from_segments( + segments_by_dim: list[list[tuple[int, int]]], + *, + visible_shape: tuple[int, ...] | None = None, +) -> tuple[int, ...]: + if visible_shape is None: + return tuple(sum(size for _, size in segments) for segments in segments_by_dim) + + assert len(visible_shape) == len(segments_by_dim) + return tuple( + sum( + max(0, min(global_start + size, visible_size) - min(global_start, visible_size)) + for global_start, size in segments + ) + for segments, visible_size in zip(segments_by_dim, visible_shape, strict=True) + ) + + +def _visible_regions_from_segments( + segments_by_dim: list[list[tuple[int, int]]], + *, + visible_shape: tuple[int, ...], +) -> list[_OwnedRegion]: + """Compile rank ownership into rectangular global-to-local copies.""" + if any(not segments for segments in segments_by_dim): + return [] + + located_segments: list[list[tuple[int, int, int]]] = [] + for segments in segments_by_dim: + local_offset = 0 + current: list[tuple[int, int, int]] = [] + for global_offset, size in segments: + current.append((global_offset, local_offset, size)) + local_offset += size + located_segments.append(current) + + regions: list[_OwnedRegion] = [] + for segment_tuple in product(*located_segments): + global_offsets: list[int] = [] + local_offsets: list[int] = [] + sizes: list[int] = [] + for dim, (global_offset, local_offset, size) in enumerate(segment_tuple): + clipped_start = min(global_offset, visible_shape[dim]) + clipped_end = min(global_offset + size, visible_shape[dim]) + clipped_size = max(0, clipped_end - clipped_start) + if clipped_size == 0: + break + global_offsets.append(clipped_start) + local_offsets.append(local_offset + clipped_start - global_offset) + sizes.append(clipped_size) + else: + regions.append( + _OwnedRegion( + global_offsets=tuple(global_offsets), + local_offsets=tuple(local_offsets), + sizes=tuple(sizes), + ) + ) + return regions + + +class SaveShardStep(BaseModel): + """Save-time work item derived from one ``LoadSpec.shards`` entry. + + ``LoadSpec.shards`` is a layout description: each descriptor says how the previous tensor was partitioned. + Saving needs the inverse operation. ``LoadSpec._save_shard_steps`` converts every shard descriptor into a work + item that contains the shard itself plus the tensor shapes that existed immediately before that shard was applied. + The save path then executes these work items in reverse order and batches compatible all-gathers by process group. + + Example: + ``LoadSpec.shards == [ep_shard, fsdp_shard]`` means the full HF tensor was first cut by EP, then the + EP-local tensor was cut by FSDP. Normal HF save executes ``[fsdp_step, ep_step]`` to rebuild the full tensor. + RL weight sync can mark ``ep_step`` as preserved, so only the FSDP work item is executed and the result stays + EP-local. + + Args: + shard (ShardDescriptor): Shard descriptor this save step reverses. + shape_before_shard (tuple[int, ...]): Runtime tensor shape immediately before ``shard`` was applied. + preserved (bool): Whether this shard should remain applied instead of being all-gathered. + """ + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + shard: ShardDescriptor + shape_before_shard: tuple[int, ...] + preserved: bool = False + + +class HFSavePlan(BaseModel): + """Execution plan for preparing one runtime tensor for HF safetensors save. + + Args: + name (str): Fully-qualified parameter or buffer name on the xtuner side. + hf_keys (list[str]): HF keys represented by the tensor after this plan's pending unshard steps finish. + runtime_output_shape (tuple[int, ...]): Shape after pending gathers, before removing FP8 runtime padding. + output_shape (tuple[int, ...]): Checkpoint-visible shape after pending gathers and final padding trim. + fused_dim (int | None): HF key concatenation dim when the underlying ``LoadSpec`` is fused; ``None`` + otherwise. + distributed_save (bool): Whether non-fused tensors are written only on rank0 and fused keys are split across + save ranks. + unshard_steps (list[SaveShardStep]): Forward-order shard history with save-time preserved flags. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + name: str hf_keys: list[str] - shape: tuple[int, ...] - dim: int | None = None - load_enum: LoadEnum - shard_start: int | None = None - shard_end: int | None = None - group: dist.ProcessGroup | None = None + runtime_output_shape: tuple[int, ...] + output_shape: tuple[int, ...] + fused_dim: int | None = None + distributed_save: bool = False + unshard_steps: list[SaveShardStep] = Field(default_factory=list) + + @computed_field # type: ignore[prop-decorator] + @property + def preserves_shards(self) -> bool: + return any(step.preserved for step in self.unshard_steps) + + +class _SaveUnshardGroup(NamedTuple): + """One compatible foreach all-gather batch in the save unshard loop. + + ``tensors`` and ``shard_steps`` are the grouped work payload. ``tensor_indices`` is kept only because the gathered + tensors must be written back to their original positions in the bucket after the collective finishes. + """ + + tensor_indices: list[int] + tensors: list[torch.Tensor] + shard_steps: list[SaveShardStep] + + +def unshard_tensors_for_hf_save( + tensors: list[torch.Tensor], + save_plans: list[HFSavePlan], +) -> list[torch.Tensor]: + """Run the all-gathers needed to turn local runtime tensors into + checkpoint-visible save tensors. + + Args: + tensors (list[torch.Tensor]): Local runtime tensors to unshard. + save_plans (list[HFSavePlan]): HF save plans corresponding to ``tensors``. + + Returns: + list[torch.Tensor]: Tensors after all pending save unshard steps have been executed. + """ + assert len(tensors) == len(save_plans), "Internal error: save tensor and plan count mismatch" + if not tensors: + return [] + + tensor_list = list(tensors) + # Each tensor reverses its shard history one step at a time. Ready heads from + # independent tensors are batched when process group and dtype match. + pending_steps = [ + [step for step in reversed(save_plan.unshard_steps) if not step.preserved] for save_plan in save_plans + ] + while unshard_groups := _take_ready_save_unshard_groups(tensor_list, pending_steps): + for unshard_group in unshard_groups: + gathered_tensors = _foreach_all_gather_save_shards( + unshard_group.tensors, + unshard_group.shard_steps, + ) + for index, gathered_tensor in zip(unshard_group.tensor_indices, gathered_tensors, strict=True): + tensor_list[index] = gathered_tensor + + # Collectives reconstruct runtime shapes; checkpoint-invisible FP8 tail + # padding is removed only after the requested shard history is complete. + return [ + _finalize_hf_save_tensor(tensor, save_plan) for tensor, save_plan in zip(tensor_list, save_plans, strict=True) + ] + + +def _take_ready_save_unshard_groups( + tensor_list: list[torch.Tensor], + pending_shard_steps_list: list[list[SaveShardStep]], +) -> list[_SaveUnshardGroup]: + """Build foreach all-gather groups for the save unshard steps that are + ready to run now.""" + unshard_groups: list[_SaveUnshardGroup] = [] + group_list: list[dist.ProcessGroup] = [] + dtype_list: list[torch.dtype] = [] + + for index, pending_shard_steps in enumerate(pending_shard_steps_list): + if not pending_shard_steps: + # This tensor has no gather work in the current save context. Common cases are unsharded tensors or + # tensors whose remaining shards are intentionally preserved, e.g. an EP-only tensor when this pass is + # only gathering FSDP shards. + continue + + # Consume one dependency-ready head step from this tensor and place it into a compatible foreach group. + shard_step = pending_shard_steps.pop(0) + shard_group = shard_step.shard.group + tensor_dtype = tensor_list[index].dtype + for group_index, (existing_group, existing_dtype) in enumerate(zip(group_list, dtype_list, strict=True)): + if tensor_dtype == existing_dtype and _is_same_process_group(existing_group, shard_group): + unshard_groups[group_index].tensor_indices.append(index) + unshard_groups[group_index].tensors.append(tensor_list[index]) + unshard_groups[group_index].shard_steps.append(shard_step) + break + else: + group_list.append(shard_group) + dtype_list.append(tensor_dtype) + unshard_groups.append( + _SaveUnshardGroup( + tensor_indices=[index], + tensors=[tensor_list[index]], + shard_steps=[shard_step], + ) + ) + + return unshard_groups + + +def _finalize_hf_save_tensor(tensor: torch.Tensor, save_plan: HFSavePlan) -> torch.Tensor: + """Validate the reconstructed runtime shape and trim FP8 tail padding.""" + assert tuple(tensor.shape) == save_plan.runtime_output_shape, ( + f"Save reconstruction produced shape {tuple(tensor.shape)}, expected runtime shape " + f"{save_plan.runtime_output_shape} for {save_plan.name}" + ) + assert all(output_size <= tensor.shape[dim] for dim, output_size in enumerate(save_plan.output_shape)) + + output = tensor[tuple(slice(0, size) for size in save_plan.output_shape)].contiguous() + assert tuple(output.shape) == save_plan.output_shape, ( + f"Saved tensor shape {tuple(output.shape)} is incompatible with HFSavePlan output_shape=" + f"{save_plan.output_shape} for {save_plan.name}" + ) + return output + + +def _foreach_all_gather_save_shards( + tensor_list: list[torch.Tensor], + shard_steps: list[SaveShardStep], +) -> list[torch.Tensor]: + assert len(tensor_list) == len(shard_steps), "Internal error: tensor and shard-step count mismatch" + assert tensor_list, "Internal error: empty save all-gather group" + group = shard_steps[0].shard.group + assert all(_is_same_process_group(group, shard_step.shard.group) for shard_step in shard_steps), ( + "Internal error: save all-gather group contains different process groups" + ) + padded_tensor_list = [ + _pad_tensor_for_save_shard(tensor, shard_step) + for tensor, shard_step in zip(tensor_list, shard_steps, strict=True) + ] + gathered_chunks_list = foreach_all_gather(padded_tensor_list, group) + return [ + _merge_gathered_save_shard(gathered_chunks, shard_step) + for gathered_chunks, shard_step in zip(gathered_chunks_list, shard_steps, strict=True) + ] + + +def _pad_tensor_for_save_shard(tensor: torch.Tensor, shard_step: SaveShardStep) -> torch.Tensor: + world_size = dist.get_world_size(group=shard_step.shard.group) + dim = shard_step.shard.dim + shard_dim_size = shard_step.shape_before_shard[dim] + + expected_local_size = sum(end - start for start, end in shard_step.shard.local_intervals(shard_dim_size)) + assert tensor.shape[dim] == expected_local_size, ( + f"Local tensor shape {tuple(tensor.shape)} does not match descriptor-local size " + f"{expected_local_size} for {shard_step.shard}" + ) + if shard_step.shard.interleave_factor > 1: + # Even interleave guarantees equal local tensors, so collective padding + # would only hide an invalid layout. + return tensor + + padded_local_size = math.ceil(shard_dim_size / world_size) + pad_len = padded_local_size - tensor.shape[dim] + assert pad_len >= 0, ( + f"Local tensor shape {tuple(tensor.shape)} exceeds padded shard size {padded_local_size} " + f"for {shard_step.shard} in save path" + ) + if not pad_len: + return tensor + + pad_list = [0] * (2 * tensor.dim()) + pad_idx = 2 * (tensor.dim() - 1 - dim) + pad_list[pad_idx + 1] = pad_len + return F.pad(tensor, pad_list) + + +def _merge_gathered_save_shard( + gathered_chunks: list[torch.Tensor], + shard_step: SaveShardStep, +) -> torch.Tensor: + dim = shard_step.shard.dim + runtime_dim_size = shard_step.shape_before_shard[dim] + if shard_step.shard.interleave_factor == 1: + gathered_tensor = torch.cat(gathered_chunks, dim=dim) + return gathered_tensor.narrow(dim, 0, runtime_dim_size).contiguous() + + world_size = dist.get_world_size(group=shard_step.shard.group) + interleave_factor = shard_step.shard.interleave_factor + assert len(gathered_chunks) == world_size + assert runtime_dim_size % (world_size * interleave_factor) == 0 + run_size = runtime_dim_size // (world_size * interleave_factor) + ordered_runs = [ + gathered_chunks[rank].narrow(dim, run_index * run_size, run_size) + for run_index in range(interleave_factor) + for rank in range(world_size) + ] + return torch.cat(ordered_runs, dim=dim).contiguous() + + +class LoadSpec(BaseModel): + """Mapping between a local param / buffer and its HF checkpoint keys. + + Args: + name (str): Fully-qualified parameter or buffer name on the xtuner side. + global_hf_keys (list[str]): Full HF key list. Concatenating these keys along ``fused_dim`` produces the + full tensor before local sharding. + global_shape (tuple[int, ...]): Shape of the fused full tensor before any ``shards`` partition is applied. + This is the runtime shape and may include padding introduced by XTuner float8 weights. + fused_dim (int | None): HF key concatenation dim when ``len(global_hf_keys) > 1``; ``None`` otherwise. + shards (list[ShardDescriptor]): Partitions applied to the full tensor in outer-to-inner order. + origin_shape (tuple[int, ...] | None): Checkpoint-visible global shape after trimming runtime-only padding. + The current caller sets it from fp8 tensor metadata; ``None`` means the runtime shape is already the + checkpoint shape. + local_shape (tuple[int, ...] | None): Runtime local-tensor shape. It is recorded explicitly for layouts such + as InterleavedShard and checked against the shape derived from ``global_shape`` and ``shards``. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + name: str + global_hf_keys: list[str] + global_shape: tuple[int, ...] + fused_dim: int | None = None + shards: list[ShardDescriptor] = Field(default_factory=list) + origin_shape: tuple[int, ...] | None = None + local_shape: tuple[int, ...] | None = None + + @computed_field # type: ignore[prop-decorator] + @property + def is_fused(self) -> bool: + return len(self.global_hf_keys) > 1 + + @computed_field # type: ignore[prop-decorator] + @property + def is_sharded(self) -> bool: + return bool(self.shards) + + @computed_field # type: ignore[prop-decorator] + @property + def unpadded_global_shape(self) -> tuple[int, ...]: + return tuple(self.origin_shape or self.global_shape) + + @classmethod + def from_tensor( + cls, + *, + name: str, + hf_keys: list[str], + tensor: torch.Tensor | DTensor, + origin_shape: tuple[int, ...] | None = None, + ) -> "LoadSpec": + """Build a load spec from a runtime tensor and its HF key mapping. + + This is the conversion boundary from PyTorch runtime layout to ``LoadSpec``. It derives the fused HF + dimension from ``hf_keys`` and converts DTensor ``Shard`` placements into ``ShardDescriptor`` entries. It does + not inspect XTuner fp8 wrapper types; callers should pass ``origin_shape`` when runtime-only padding makes the + checkpoint-visible shape smaller than the runtime shape. + + Args: + name (str): Fully-qualified parameter or buffer name on the xtuner side. + hf_keys (list[str]): HF key list corresponding to ``tensor``. + tensor (torch.Tensor | DTensor): Runtime tensor whose DTensor placements should be captured. + origin_shape (tuple[int, ...] | None): Optional checkpoint-visible shape after trimming runtime-only + padding. + + Returns: + LoadSpec: Spec derived from the runtime tensor layout. + """ + global_hf_keys = list(hf_keys) + shards = _dtensor_shards(tensor) if isinstance(tensor, DTensor) else [] + local_tensor = tensor._local_tensor if isinstance(tensor, DTensor) else tensor + return cls( + name=name, + global_hf_keys=global_hf_keys, + global_shape=tuple(tensor.shape), + fused_dim=0 if len(global_hf_keys) > 1 else None, + shards=shards, + origin_shape=origin_shape, + local_shape=tuple(local_tensor.shape), + ) + + def plan_hf_load(self) -> HFLoadPlan: + """Build a safetensors read plan from this layout spec. + + Runtime-only padding currently comes from XTuner float8 weights. In that case, ``origin_shape`` is used as + the checkpoint-visible full tensor shape, while ``global_shape`` and ``shards`` still describe the padded + runtime layout that this rank owns. + + Returns: + HFLoadPlan: The selected HF keys and canonical source-to-local copy program for this rank. + """ + segments_by_dim = _layout_segments(self.global_shape, self.shards) + target_shape = _shape_from_segments(segments_by_dim) + owned_regions = _visible_regions_from_segments( + segments_by_dim, + visible_shape=self.unpadded_global_shape, + ) + if not owned_regions: + return HFLoadPlan( + name=self.name, + hf_keys=[], + fused_dim=self.fused_dim, + target_shape=target_shape, + zero_unwritten_target=math.prod(target_shape) > 0, + ) + + hf_keys, source_offsets, source_shape = self._hf_load_source(owned_regions) + copy_regions = [ + LoadCopyRegion( + source_offsets=tuple( + region.global_offsets[dim] - source_offsets[dim] for dim in range(len(self.global_shape)) + ), + target_offsets=region.local_offsets, + sizes=region.sizes, + ) + for region in owned_regions + ] + copied_numel = sum(math.prod(region.sizes) for region in copy_regions) + target_numel = math.prod(target_shape) + assert copied_numel <= target_numel, ( + f"Owned regions for {self.name} copy {copied_numel} values into a target with {target_numel} values" + ) + + return HFLoadPlan( + name=self.name, + hf_keys=hf_keys, + fused_dim=self.fused_dim, + canonical_source_shape=source_shape, + target_shape=target_shape, + copy_regions=copy_regions, + zero_unwritten_target=copied_numel < target_numel, + ) + + def _hf_load_source( + self, + owned_regions: list[_OwnedRegion], + ) -> tuple[list[str], tuple[int, ...], tuple[int, ...]]: + """Select HF keys and express their canonical tensor in global + coordinates.""" + key_start, key_end = self._hf_key_range_for_regions(owned_regions) + source_offsets = [0 for _ in self.global_shape] + source_shape = list(self.unpadded_global_shape) + if self.is_fused: + assert self.fused_dim is not None + key_size = self._fused_key_size() + source_offsets[self.fused_dim] = key_start * key_size + source_shape[self.fused_dim] = (key_end - key_start) * key_size + + return ( + self.global_hf_keys[key_start:key_end], + tuple(source_offsets), + tuple(source_shape), + ) + + def plan_hf_save( + self, + *, + distributed_save: bool = False, + preserve_process_group: dist.ProcessGroup | None = None, + gather_process_group: dist.ProcessGroup | None = None, + ) -> HFSavePlan: + """Build a safetensors save plan from this layout spec. + + Args: + distributed_save (bool): Whether non-fused tensors are written only on rank0 and fused HF keys are split + across save ranks. + preserve_process_group (dist.ProcessGroup | None): Fused-dim shard group that should remain sharded, + used by RL weight sync to stream EP-local expert slices. + gather_process_group (dist.ProcessGroup | None): If set, only shards from this group are gathered and + all other shards are preserved. This is used by callers that need an FSDP-only all-gather. + + Returns: + HFSavePlan: Save-time unshard and HF key planning information. + """ + assert not (preserve_process_group is not None and gather_process_group is not None), ( + "preserve_process_group and gather_process_group describe different save policies and cannot be combined" + ) + preserved_shard_indices = self._preserved_shard_indices( + preserve_process_group=preserve_process_group, + gather_process_group=gather_process_group, + ) + unshard_steps = self._save_shard_steps(preserved_shard_indices) + preserved_shards = [step.shard for step in unshard_steps if step.preserved] + if preserve_process_group is not None and preserved_shards: + hf_keys = self._local_hf_keys_for_shards(preserved_shards, require_fused_key_aligned=True) + else: + # FSDP-only gather keeps ETP's runtime layout and does not produce HF + # tensors, so its key list is informational and needs no key alignment. + hf_keys = list(self.global_hf_keys) + + output_segments = _layout_segments(self.global_shape, preserved_shards) + runtime_output_shape = _shape_from_segments(output_segments) + output_shape = _shape_from_segments( + output_segments, + visible_shape=self.unpadded_global_shape, + ) + + return HFSavePlan( + name=self.name, + hf_keys=hf_keys, + runtime_output_shape=runtime_output_shape, + output_shape=output_shape, + fused_dim=self.fused_dim, + distributed_save=distributed_save, + unshard_steps=unshard_steps, + ) def model_post_init(self, _) -> None: - if self.load_enum == LoadEnum.SAME: - assert len(self.hf_keys) == 1, "hf_keys should have exactly one key when load_enum is SAME" - elif self.load_enum == LoadEnum.FUSED: - if self.dim is None: - self.dim = 0 - assert self.dim == 0, "dim should be 0 when load_enum is FUSED" - elif self.load_enum == LoadEnum.SHARD: - assert self.dim is not None, "dim should not be None when load_enum is SHARD" - assert len(self.hf_keys) == 1, "hf_keys should have more than one key when load_enum is SHARD" - assert self.shard_start is not None, "shard_start should not be None when load_enum is SHARD" - assert self.shard_end is not None, "shard_end should not be None when load_enum is SHARD" + if self.is_fused: + assert self.fused_dim is not None, "fused_dim must be set when global_hf_keys has multiple entries" + else: + assert self.fused_dim is None, "fused_dim must be None when global_hf_keys has one entry" + self._validate_origin_shape() + self._validate_shards() + + def _fused_key_size(self) -> int: + assert self.fused_dim is not None, "fused_dim must be set when global_hf_keys has multiple entries" + key_size = self.unpadded_global_shape[self.fused_dim] / len(self.global_hf_keys) + assert key_size.is_integer(), ( + f"Fused dim size {self.unpadded_global_shape[self.fused_dim]} is not divisible by " + f"{len(self.global_hf_keys)} HF keys for {self.name}" + ) + return int(key_size) + + def _hf_key_range_for_regions( + self, + regions: list[_OwnedRegion], + *, + require_fused_key_aligned: bool = False, + ) -> tuple[int, int]: + if not self.is_fused: + return 0, len(self.global_hf_keys) + + assert self.fused_dim is not None + key_size = self._fused_key_size() + fused_start = min(region.global_offsets[self.fused_dim] for region in regions) + fused_end = max(region.global_offsets[self.fused_dim] + region.sizes[self.fused_dim] for region in regions) + if require_fused_key_aligned: + assert fused_start % key_size == 0 and fused_end % key_size == 0, ( + f"Preserved fused shard range [{fused_start}, {fused_end}) for {self.name} must align with " + f"HF key size {key_size}" + ) + + # Shards may start or end inside a fused HF key, e.g. FSDP slicing an EP-local expert tensor. + # floor/ceil keeps every overlapping key; the load plan's copy regions trim to the exact local range. + key_start = fused_start // key_size + key_end = math.ceil(fused_end / key_size) + assert 0 <= key_start < key_end <= len(self.global_hf_keys), ( + f"Invalid fused key range [{key_start}, {key_end}) for {self.name}" + ) + return key_start, key_end + + def _local_hf_keys_for_shards( + self, + shards: list[ShardDescriptor], + *, + require_fused_key_aligned: bool = False, + ) -> list[str]: + segments_by_dim = _layout_segments(self.global_shape, shards) + regions = _visible_regions_from_segments( + segments_by_dim, + visible_shape=self.unpadded_global_shape, + ) + if not regions: + return [] + key_start, key_end = self._hf_key_range_for_regions( + regions, + require_fused_key_aligned=require_fused_key_aligned, + ) + return self.global_hf_keys[key_start:key_end] + + def _validate_origin_shape(self) -> None: + if self.origin_shape is None: + return + + assert len(self.origin_shape) == len(self.global_shape), ( + f"origin_shape={self.origin_shape} must have the same rank as global_shape={self.global_shape}" + ) + assert all(origin <= global_ for origin, global_ in zip(self.origin_shape, self.global_shape)), ( + f"origin_shape={self.origin_shape} must not exceed global_shape={self.global_shape}" + ) + + def _validate_shards(self) -> None: + segments_by_dim = _layout_segments(self.global_shape, self.shards) + derived_shape = _shape_from_segments(segments_by_dim) + + assert self.local_shape is None or derived_shape == self.local_shape, ( + f"Recorded local_shape={self.local_shape} does not match descriptor-derived shape " + f"{derived_shape} for {self.name}" + ) + + def _preserved_shard_indices( + self, + *, + preserve_process_group: dist.ProcessGroup | None, + gather_process_group: dist.ProcessGroup | None, + ) -> set[int]: + """Return ``self.shards`` indices that should remain sharded in this + save plan. + + ``preserve_process_group`` is only used when a fused HF tensor has an additional runtime partition on + ``fused_dim``. For example, MoE expert parallel may shard the concatenated expert keys on the same dim that + HF uses for fused keys, and FSDP may further shard that EP-local tensor on the same dim. RL weight sync wants + to preserve the EP shard so it can derive the local HF key range from that shard, while all remaining shards + such as FSDP must still be all-gathered to recover a complete weight for that preserved EP slice. + + ``gather_process_group`` is the inverse policy used by FSDP-only all-gather callers: gather shards from this + group and preserve every other shard. + + Example: + Suppose ``global_hf_keys`` represents experts ``[0..7]`` concatenated on dim 0, and the runtime layout is + ``shards=[ep_shard(dim=0, group=ep_group), fsdp_shard(dim=0, group=fsdp_group)]``. Passing ``ep_group`` as + ``preserve_process_group`` returns ``{0}``: the EP shard is preserved for local HF key planning, while the + FSDP shard at index 1 is still all-gathered so the local EP expert slice becomes complete. Passing + ``fsdp_group`` as ``gather_process_group`` produces the same preserved index set for an FSDP-only gather. + + Returns: + set[int]: Indices into ``self.shards``, not tensor dimensions. + """ + if gather_process_group is not None: + return { + shard_index + for shard_index, shard in enumerate(self.shards) + if not _is_same_process_group(shard.group, gather_process_group) + } + + if preserve_process_group is None or not self.is_fused: + return set() + + assert self.fused_dim is not None, ( + f"Internal error: fused LoadSpec {self.name} has no fused_dim. " + "LoadSpec.model_post_init should reject this layout before save planning." + ) + return { + shard_index + for shard_index, shard in enumerate(self.shards) + if shard.dim == self.fused_dim and _is_same_process_group(shard.group, preserve_process_group) + } + + def _save_shard_steps(self, preserved_shard_indices: set[int]) -> list[SaveShardStep]: + """Convert ``LoadSpec.shards`` into save-time reverse-unshard work + items. + + ``LoadSpec.shards`` is ordered in the forward partitioning direction. Each step snapshots only the runtime + shape before its shard. Save executes the steps in reverse and restores those runtime shapes; checkpoint-only + FP8 trimming happens once at the final ``HFSavePlan`` output boundary. + + Args: + preserved_shard_indices (set[int]): Original ``LoadSpec.shards`` indices that should remain sharded. + + Returns: + list[SaveShardStep]: Work items in the same largest-to-smallest order as ``LoadSpec.shards``. + """ + current_shape = list(self.global_shape) + steps: list[SaveShardStep] = [] + + for shard_index, shard in enumerate(self.shards): + steps.append( + SaveShardStep( + shard=shard, + shape_before_shard=tuple(current_shape), + preserved=shard_index in preserved_shard_indices, + ) + ) + current_shape[shard.dim] = sum( + end - start for start, end in shard.local_intervals(current_shape[shard.dim]) + ) + return steps