From 3540ed16ebc580b0ffc3f553f1c7d07bced7d8e7 Mon Sep 17 00:00:00 2001 From: zhaopenghao Date: Mon, 10 Aug 2026 07:16:52 +0000 Subject: [PATCH 1/7] [Refactor] Make HF load/save planning DTensor-placement aware --- docs/design/load_spec_refactor.md | 403 ++++++ tests/rl/test_weight_iterator.py | 49 + tests/utils/test_load_spec.py | 423 ++++++ xtuner/v1/model/base.py | 1182 +++++------------ xtuner/v1/model/dense/dense.py | 1 + xtuner/v1/model/moe/glm52.py | 32 +- xtuner/v1/model/moe/gpt_oss.py | 36 +- xtuner/v1/model/moe/moe.py | 1 + xtuner/v1/model/moe/qwen3_5_text.py | 34 +- xtuner/v1/model/moe/qwen3vl_text.py | 34 +- xtuner/v1/rl/weight_update/weight_iterator.py | 210 +-- xtuner/v1/utils/load_spec.py | 782 ++++++++++- 12 files changed, 2054 insertions(+), 1133 deletions(-) create mode 100644 docs/design/load_spec_refactor.md create mode 100644 tests/rl/test_weight_iterator.py create mode 100644 tests/utils/test_load_spec.py diff --git a/docs/design/load_spec_refactor.md b/docs/design/load_spec_refactor.md new file mode 100644 index 0000000000..785a4bec17 --- /dev/null +++ b/docs/design/load_spec_refactor.md @@ -0,0 +1,403 @@ +# LoadSpec 设计 + +> 面向 `xtuner/v1/utils/load_spec.py` 与 `xtuner/v1/model/base.py` 的加载/保存路径。 +> TP 设计（`dense_tp.md`）依赖本文档描述的抽象。 + +## TL;DR + +LoadSpec 描述 xtuner 运行时 tensor 与 HF safetensors 之间的**纯布局映射**。对一个 +param，它回答两件事： + +1. 这个 param 由哪些 HF key 组成？怎么拼？ — `global_hf_keys` + `fused_dim` +2. 本 rank 持有全量 tensor 的哪一块？ — `shards`（按外到内顺序施加） + +加载/保存执行路径不直接读 LoadSpec，而是调用 `plan_hf_load()` / +`plan_hf_save(...)` 拿到一份**不可变的 plan**，按 plan 驱动 IO 与通信。 + +**核心约束**：LoadSpec 只承担"同 dtype 下的形状/索引映射"。fp8 的量化反量化、 +padding 的 zero-fill 等 dtype 语义都住在 `base.py` 的 load/save 路径里， +LoadSpec 不感知。 + +--- + +## 1. 设计理念 + +### 1.1 单一抽象，两条正交轴 + +原先三类映射（SAME / FUSED / SHARD）统一成一个 schema 上的两个正交维度： + +| 问题 | 表达 | +| --- | --- | +| 这个 param 对应几个 HF key？怎么拼？ | `len(global_hf_keys)`；多 key 时 `fused_dim` 指定拼接维 | +| 本 rank 持有哪一块？ | `shards`（可为空；按施加顺序排列） | + +消费方用派生属性 `is_fused` / `is_sharded` 查询，**不需要**任何枚举分支。 + +### 1.2 多维切分按顺序叠加 + +`shards` 是列表，原生支持 TP × FSDP、EP × FSDP 等多轴组合。每条 +`ShardDescriptor.start/end` 的含义是"在**前面所有** descriptor 切完之后的 +子 tensor 上的偏移"。这条规则完全对齐 DTensor `placements` 从 `mesh_dim=0` 到 +`mesh_dim=N-1` 逐步施加的语义 —— 你可以把 `shards[i]` 理解成 `placements[i]` +在"此刻本 rank 实际持有"这个问题上的等价形式。 + +### 1.3 Plan 是冻结快照 + +`plan_hf_load()` / `plan_hf_save(...)` 返回的是 Pydantic dataclass： + +- 一次性从当前 LoadSpec 状态计算出执行所需的全部信息； +- 不持有对 LoadSpec 的引用； +- 执行器（`_load_hf_param` / `unshard_tensors_for_hf_save` / `_split_hf_tensors_for_save`） + 只读 plan，**不读** LoadSpec。 + +这条边界保证"布局规划"和"IO/通信执行"解耦。未来要接入新的持久化格式（例如 +DCP），只需要替换 plan 的消费者，不牵涉 LoadSpec 内部结构。 + +### 1.4 fp8 与 LoadSpec 解耦 + +LoadSpec 是"同 dtype 下的布局描述"。fp8 涉及的两件事 —— 量化/反量化、运行时 +padding —— 归属如下： + +- **运行时 padding**：用 `LoadSpec.origin_shape` 表达 checkpoint-visible shape + （剥掉运行时 padding 之后）。今天这个字段的唯一来源是 fp8 tensor metadata； + 它只记录 shape，不记录 dtype / wrapper 类型。 +- **量化/反量化**：只在 `base.py._to_float8` / 反量化分支里现场判断（通过 + `is_float8_weight(tensor)`）。LoadSpec 不包含 `runtime_is_float8` 这类 + dtype-specific 字段。 + +### 1.5 Spec → Plan → Executor 的分层 + +``` +┌────────────────────┐ plan_hf_load() ┌──────────────┐ +│ │ ──────────────────▶│ HFLoadPlan │──▶ _load_hf_param +│ LoadSpec │ └──────────────┘ +│ (pure layout) │ plan_hf_save(...) ┌──────────────┐ +│ │ ──────────────────▶│ HFSavePlan │──▶ unshard_tensors_for_hf_save +└────────────────────┘ └──────────────┘ │ + ▼ + _split_hf_tensors_for_save +``` + +"Spec 是源、Plan 是派生、Executor 只依赖 Plan"。这条线保持单向。 + +--- + +## 2. 数据模型 + +### 2.1 `ShardDescriptor` + +```python +class ShardDescriptor(BaseModel): + dim: int # 被切的维 + start: int # 在"前面切完的 sub-tensor"上的起点 + end: int # 在"前面切完的 sub-tensor"上的终点 + group: dist.ProcessGroup # 产生这次切分的通信组 +``` + +`group` 是 load/save 双向通信域。load 时只需要知道本 rank 的范围；save 时需要沿 +`group` 做 all-gather 复原全量 tensor。 + +### 2.2 `LoadSpec` + +```python +class LoadSpec(BaseModel): + name: str # xtuner 侧 fully-qualified param name + global_hf_keys: list[str] # 对应的 HF key 列表（按 fused_dim 拼接顺序） + global_shape: tuple[int, ...] # 全量 tensor（fused 之后）的 runtime shape + # 可能包含运行时 padding（例如 fp8 的 FSDP 对齐 pad） + fused_dim: int | None = None # 多 HF key 时的拼接维；单 key 时必须为 None + shards: list[ShardDescriptor] = [] # 从外到内的切分列表 + origin_shape: tuple[int, ...] | None = None # checkpoint-visible shape after runtime padding is trimmed + # None 表示"runtime shape 就是 checkpoint shape" +``` + +派生属性： + +```python +is_fused # len(global_hf_keys) > 1 +is_sharded # bool(shards) +unpadded_global_shape # origin_shape or global_shape +``` + +**不变量**（`model_post_init` 强制）： + +- `is_fused` ⇔ `fused_dim is not None`； +- 每条 shard 的 `start/end` 必须落在"前面切完之后的 sub-tensor"范围内； +- 若 `origin_shape` 给定，它的秩与 `global_shape` 相同，且每维 `≤ global_shape`。 + +### 2.3 `HFLoadPlan` + +`plan_hf_load()` 的产出： + +```python +class HFLoadPlan(BaseModel): + name: str + hf_keys: list[str] # 本 rank 实际需要读的 HF key + fused_dim: int | None = None # 多 key 时的拼接维 + slices: list[LoadSlice] = [] # 读完拼接后，再做的 narrow 列表 + zero_fill: bool = False # 本 rank 完全落在运行时 padding 区，跳过 IO +``` + +`slices` 的 start/end 是**相对已加载 tensor 的坐标**，不是相对 `global_shape`。 +zero_fill=True 时 `hf_keys` 和 `slices` 都为空。 + +### 2.4 `HFSavePlan` + +`plan_hf_save(...)` 的产出，承载两类信息： + +```python +class HFSavePlan(BaseModel): + name: str + hf_keys: list[str] # 当前 save tensor 最终要写/同步的 HF keys + global_shape: tuple[int, ...] + unpadded_global_shape: tuple[int, ...] + fused_dim: int | None = None + distributed_save: bool = False + preserves_shards: bool = False # True 表示 hf_keys 来自保留 shard 后的局部 tensor + unshard_steps: list[SaveShardStep] = [] # 所有 shard 的逆操作 + preserved 标记 +``` + +`SaveShardStep` 记录一次 shard 在"施加前的 runtime shape / checkpoint-visible +shape"两个快照 —— save 执行时倒序跑每一步、all-gather 还原、narrow 回 +checkpoint-visible shape。`preserved` 标记把某些 shard 排除在 all-gather 之外 +（见 §3.3）。`HFSavePlan.hf_keys` 始终是执行器要处理的 key 集合：普通 save 下 +它是完整 HF key list，preserved shard save 下它是当前局部 shard 覆盖的 key list。 + +--- + +## 3. 计划生成 + +### 3.1 `plan_hf_load()` + +不接受参数 —— 本 rank 的所有信息已经在 LoadSpec 里。步骤： + +1. 计算本 rank 最终持有的区间 `final_intervals`（顺序应用 `shards`）； +2. 用 `unpadded_global_shape` 裁剪掉运行时 padding 部分；若裁完为空，返回 + `zero_fill=True`； +3. 若 `is_fused`，按 `fused_dim` 上的区间算出需要的 HF key 下标范围（floor/ceil + 支持 mid-key shard，例如 FSDP 在 EP-local 专家 tensor 内部再切）； +4. 对每个 dim，如果"最终区间"比"加载后的 tensor 区间"窄，生成一条 `LoadSlice`。 + +### 3.2 `plan_hf_save(distributed_save=, preserve_process_group=, gather_process_group=)` + +三个参数对应三种 save 策略，互斥使用： + +| 参数 | 用途 | +| --- | --- | +| `distributed_save=True` | HF save：非 fused tensor 只在 rank0 写；fused tensor 的 HF key 在 save rank 间分配 | +| `preserve_process_group=ep_group` | RL 权重同步：保留 EP 在 `fused_dim` 上的 shard，每个 EP rank 只流自己的 expert key；其他 shard 照常 all-gather | +| `gather_process_group=fsdp_group` | FSDP-only all-gather：只 gather 这个 group 的 shard，其他 shard 保留 | + +策略统一落到 `_preserved_shard_indices` 这一步上 —— 决定哪些 `LoadSpec.shards` +需要保留。之后 `_save_shard_steps` 给每个 shard 生成带 `preserved` 标记的 +`SaveShardStep`。若有 preserved shard，`LoadSpec` 直接从这些 shard 推导 +`HFSavePlan.hf_keys`；save plan 只暴露最终要写/同步的 HF keys，以及 +`preserves_shards` 说明这些 keys 来自局部 tensor 还是完整 tensor。 + +### 3.3 preserve vs gather 的正交性 + +`preserve_process_group` 是"显式保留某个 group"的策略，`gather_process_group` +是"显式 gather 某个 group（其余保留）"的策略。两者不能同时使用（assert 拦截）。 +在今天的代码里： + +- 普通 HF save：两者都不传，全部 all-gather； +- RL 权重同步：传 `preserve_process_group=ep_group`； +- `_fsdp_foreach_allgather`：传 `gather_process_group=fsdp_group`，只做 FSDP + 层的 all-gather，不动 EP / TP。 + +--- + +## 4. 执行 + +### 4.1 加载路径 + +```python +def _load_hf_param(self, param, load_spec, loader): + plan = load_spec.plan_hf_load() + if plan.zero_fill: + # 本 rank 只持有运行时 padding，写 0 返回 + local_tensor.zero_() + return [] + # 按 plan.hf_keys 逐个读（fp8 走 dequant 分支，这里 base.py 现场处理） + loaded_tensors = self._load_hf_keys(plan, loader, ...) + # 拼接 + narrow 全部交给 safetensors_to_params + self.safetensors_to_params(loaded_tensors, local_tensor, plan) +``` + +`safetensors_to_params` 的签名是 `(safetensors, local_tensor, plan)`。三个 MoE +子类（`gpt_oss`、`qwen3_5_text`、`qwen3vl_text`）按 `plan.name` 做 reshape / +transpose 等模型特有变换后，调通用的 `_apply_load_slices` + `_copy_loaded_tensor_to_local`。 + +### 4.2 保存路径 + +所有 save 场景（HF save、RL 权重同步、FSDP-only gather）共用一条管道： + +```python +save_items = [HFSaveItem(tensor, load_spec.plan_hf_save(...)) for ...] +full_tensors = unshard_tensors_for_hf_save(save_items) +for full_tensor, item in zip(full_tensors, save_items): + names, tensors = self._split_hf_tensors_for_save(full_tensor, item.save_plan) +``` + +`unshard_tensors_for_hf_save` 自带**依赖感知的批量 foreach all-gather**： + +- 同一个 tensor 的多个 step 必须串行（例如 "先还原 FSDP，再还原 EP"）； +- 不同 tensor 的 step 如果 `(group, dtype)` 兼容，可以 foreach 批到同一次 NCCL 调用。 + +每一轮由 `_build_ready_save_unshard_groups` 从每个 pending 队列取头部 step，按 +group + dtype 分桶；`_foreach_all_gather_save_shards` 跑一次批量 gather；下一轮 +再消费队列的下一层。MoE EP+FSDP 的 save 就是这样两轮跑完的。 + +### 4.3 `HFSaveItem` + +```python +class HFSaveItem(NamedTuple): + tensor: torch.Tensor + save_plan: HFSavePlan +``` + +这是**跨 LoadSpec 和 BaseModel 边界**的 bundle：一边是 runtime tensor（模型侧 +概念，带 fp8 wrapper / DTensor wrapper），一边是纯布局的 `HFSavePlan`。它的 +归属地是 `base.py` ——`load_spec.py` 保持"不认识模型侧概念"。 +`unshard_tensors_for_hf_save` 的签名使用两个平行列表（`list[torch.Tensor]` + +`list[HFSavePlan]`）而不是 `list[HFSaveItem]`，避免 `load_spec.py` 反向依赖 +`base.py`。 + +--- + +## 5. 调用时机 + +`_init_load_spec` 被定位为"从当前 DTensor 布局反推 HF 映射的纯函数"。 +调用约定：**谁改 param 布局谁负责重算，后者覆盖前者**。 + +| 时机 | 调用方 | spec 代表 | +| --- | --- | --- | +| 子类 `__init__` 末尾 | 子类自己 | 构建完成时的布局（EP-only / Replicate / 其它 init-time 切分） | +| `parallelize(tp_mesh)` 结束 | `BaseModel.parallelize` | TP + 已有切分 | +| `fully_shard` 结束 | `BaseModel.fully_shard` | 叠加 FSDP（训练态） | +| `Float8Handler.pad_for_fsdp` 回调 | 回调内 | fp8 pad 后的真实 shape | + +`from_hf` / `save_hf` 入口有 assert 兜底： + +```python +assert "load_spec_mapping" in self.__dict__, ( + f"{type(self).__name__}.__init__ must call self._init_load_spec() at the end." +) +``` + +这条约定是硬契约；子类若跳过会在第一次 load/save 时被抓。 + +--- + +## 6. 示例 + +### 6.1 Dense, tp=2, fsdp=4, `q_proj.weight` + +```python +LoadSpec( + name="layers.0.self_attn.q_proj.weight", + global_hf_keys=["model.layers.0.self_attn.q_proj.weight"], + global_shape=(n*d, h), + fused_dim=None, + shards=[ + ShardDescriptor(dim=0, start=tp_start, end=tp_end, group=tp_group), + ShardDescriptor(dim=0, start=fsdp_start, end=fsdp_end, group=fsdp_group), + ], +) +``` + +`fsdp_start/end` 相对于"已经被 TP 切过的 sub-tensor"而言，不是相对 +`global_shape`。 + +### 6.2 MoE, ep=8, fsdp=4, fused expert weight + +```python +LoadSpec( + name="layers.0.experts.fused_w1w3.weight", + global_hf_keys=[f"model.layers.0.mlp.experts.{i}.gate_proj.weight" for i in range(64)] + + [f"model.layers.0.mlp.experts.{i}.up_proj.weight" for i in range(64)], + global_shape=(128 * I_padded, H), # I_padded 含 fp8 FSDP 对齐 pad + fused_dim=0, + shards=[ + ShardDescriptor(dim=0, start=ep_start, end=ep_end, group=ep_group), + ShardDescriptor(dim=0, start=fsdp_start, end=fsdp_end, group=fsdp_group), + ], + origin_shape=(128 * I, H), # 剥掉 pad 后的 checkpoint shape +) +``` + +RL 权重同步调用 `plan_hf_save(preserve_process_group=ep_group)` —— EP shard 被 +标记 preserved，保存管道只做 FSDP 还原，结果留在 EP-local 坐标系；再由 +`_request_ep_sequential_update` 按 EP rank 顺序广播。 + +### 6.3 embed_tokens, 纯 FSDP + +```python +LoadSpec( + name="embed_tokens.weight", + global_hf_keys=["model.embed_tokens.weight"], + global_shape=(V, H), + fused_dim=None, + shards=[ShardDescriptor(dim=0, start=fsdp_start, end=fsdp_end, group=fsdp_group)], +) +``` + +--- + +## 7. 为什么这样设计 + +几个关键取舍的归档。 + +### 7.1 为什么 `shards` 是列表而不是单轴四元组 + +旧的 `(dim, shard_start, shard_end, group)` 只表达一刀。TP × FSDP 或 EP × FSDP +是常见组合，旧 schema 只能靠"加载时临时推导第二刀"这种硬编码绕过（ +`FSDP_SHARD_DIM == 0` 就是这条路径的残留）。列表 + DTensor 施加顺序是最小的 +统一表达。 + +### 7.2 为什么删 `LoadEnum` + +`SAME/FUSED/SHARD` 给定 `global_hf_keys` 和 `shards` 后是可派生的。保留它相当于 +同一份状态的两种表达，下游分支要同步维护。直接用 `is_fused` / `is_sharded` 两个 +独立 bool 可以正交表达所有组合（包括原本需要新造 `FUSED_SHARD` 的情况）。 + +### 7.3 为什么 fp8 不进 LoadSpec + +LoadSpec 的定位是"同 dtype 下的映射"。fp8 涉及的反量化需要的是 tensor 的真实 +dtype / wrapper 类型，这些只有在 IO 路径里拿到 runtime tensor 才能判断。若把 +`runtime_is_float8` 放进 spec，一方面是状态重复（`is_float8_weight(tensor)` 已经 +是事实来源），另一方面污染 LoadSpec 的语义 —— 它不再是纯布局描述。 + +`origin_shape` 是 checkpoint-visible shape。它今天只服务 fp8 runtime padding， +但仍然只携带 shape 信息；fp8 的 dtype / wrapper 判断不进入 LoadSpec。 + +### 7.4 为什么 `unshard_tensors_for_hf_save` 住在 `load_spec.py` + +尽管它做的是分布式 all-gather，但它**只依赖 HFSavePlan + 一个通信原语**。把它 +放在 `load_spec.py` 让"spec → plan → 执行"三层都在一个文件里闭环，调用方 +（base.py）只需要准备 `(tensor, plan)` 对，不需要理解 shard 调度。 + +若将来 `unshard_tensors_for_hf_save` 进一步膨胀，可以拆到独立模块（例如 +`save_runner.py`），但当前规模尚不需要。 + +### 7.5 为什么保存不用 `_fuse_contiguous_chunks_without_alloc` + +旧代码对 `dim == 0` 的单 tensor all-gather 用过一个零拷贝 view 合并优化。这条 +优化只在"一次 gather 一个 tensor"时成立 —— 当前批量 foreach 把多个 tensor 交错 +塞进同一个扁平缓冲区，per-tensor chunks 不再连续，这条路径失效。换掉 NCCL 调用 +次数（O(num_tensors) → O(rounds)）比 dim=0 多一次 cat alloc 更划算。如果某个 +特定场景发现这次 trade-off 不值，可以单独给那条路走非批量路径，但默认策略保持 +批量。 + +--- + +## 8. 测试 + +核心测试都在 `tests/utils/test_load_spec.py`： + +- `TestLoadSpecSchema`：字段契约 + `shards` 顺序验证； +- `TestHFLoadPlan`：`plan_hf_load` 在 fused / non-fused / fp8 padding 下的产出； +- `TestHFSavePolicy`：`distributed_save` 的 HF key 分配规则。 + +行为等价性由 `tests/model/test_qwen3_dense.py::test_save_hf` 和 +`tests/model/test_qwen3_moe.py::test_save_hf` 的 safetensors bit-equal 保证。 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_load_spec.py b/tests/utils/test_load_spec.py new file mode 100644 index 0000000000..f26ca85498 --- /dev/null +++ b/tests/utils/test_load_spec.py @@ -0,0 +1,423 @@ +import os + +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: + # ShardDescriptor.group is typed as `dist.ProcessGroup`; Pydantic enforces + # the isinstance check even with `arbitrary_types_allowed=True`, so schema + # tests need a real (but minimal) process group. A single-rank gloo group + # is sufficient and avoids any CUDA / multi-process plumbing. + 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 + + +class TestLoadSpecSchema: + """New-schema fields should describe layout without legacy dispatch state.""" + + 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: + assert single_rank_group is not 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.start, shard.end) for shard in spec.shards] == [(0, 0, 128)] + + def test_dtensor_shards_follow_explicit_placement_order(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_local_rank(self, mesh_dim: int) -> int: + return (1, 0)[mesh_dim] + + def get_group(self, mesh_dim: int) -> dist.ProcessGroup: + return single_rank_group + + class FakeDTensor: + shape = (8,) + placements = (_StridedShard(0, split_factor=2), DTensorShard(0)) + device_mesh = FakeDeviceMesh() + + shards = load_spec_module._dtensor_shards(FakeDTensor()) # type: ignore[arg-type] + + assert [(shard.dim, shard.start, shard.end) for shard in shards] == [(0, 0, 4), (0, 2, 4)] + + 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=[ + "model.layers.0.mlp.experts.0.gate_proj.weight", + "model.layers.0.mlp.experts.0.up_proj.weight", + ], + global_shape=(256, 64), + ) + + def test_multi_axis_shards_preserve_order(self, single_rank_group: dist.ProcessGroup) -> None: + ep = ShardDescriptor(dim=0, start=64, end=128, group=single_rank_group) + fsdp = ShardDescriptor(dim=0, start=16, end=32, group=single_rank_group) + spec = LoadSpec( + name="layers.0.experts.fused_w1w3.weight", + global_hf_keys=[ + "model.layers.0.mlp.experts.0.gate_proj.weight", + "model.layers.0.mlp.experts.0.up_proj.weight", + ], + global_shape=(256, 64), + fused_dim=0, + shards=[ep, fsdp], + ) + + assert [(shard.start, shard.end) for shard in spec.shards] == [(64, 128), (16, 32)] + assert spec.is_fused is True + assert spec.is_sharded is True + + def test_ordered_shard_bounds_are_validated(self, single_rank_group: dist.ProcessGroup) -> None: + with pytest.raises(ValidationError, match="Invalid shard descriptor"): + LoadSpec( + name="layers.0.experts.fused_w1w3.weight", + global_hf_keys=["model.layers.0.mlp.experts.0.gate_proj.weight"], + global_shape=(128, 64), + shards=[ + ShardDescriptor(dim=0, start=64, end=128, group=single_rank_group), + ShardDescriptor(dim=0, start=65, end=80, group=single_rank_group), + ], + ) + + def test_zero_size_dtensor_shards_are_valid(self, single_rank_group: dist.ProcessGroup) -> None: + spec = LoadSpec( + name="embeddings.cls_embedding", + global_hf_keys=["embeddings.cls_embedding"], + global_shape=(1, 1, 1024), + shards=[ShardDescriptor(dim=0, start=1, end=1, group=single_rank_group)], + ) + + plan = spec.plan_hf_load() + + assert plan.zero_fill is True + assert plan.hf_keys == [] + + +class TestHFLoadPlan: + """LoadSpec should derive HF read plans from shards only.""" + + def test_fused_slice_selects_overlapping_hf_keys(self, single_rank_group: dist.ProcessGroup) -> None: + 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, start=150, end=260, group=single_rank_group)], + ) + + plan = spec.plan_hf_load() + + assert plan.hf_keys == ["k1", "k2"] + assert plan.fused_dim == 0 + assert [(load_slice.dim, load_slice.start, load_slice.end) for load_slice in plan.slices] == [(0, 50, 160)] + assert not hasattr(plan, "loaded_shape") + + def test_non_fused_slice_keeps_single_hf_key(self, single_rank_group: dist.ProcessGroup) -> None: + spec = LoadSpec( + name="layers.0.self_attn.q_proj.weight", + global_hf_keys=["q_proj"], + global_shape=(128, 256), + shards=[ShardDescriptor(dim=1, start=64, end=192, group=single_rank_group)], + ) + + plan = spec.plan_hf_load() + + assert plan.hf_keys == ["q_proj"] + assert plan.fused_dim is None + assert [(load_slice.dim, load_slice.start, load_slice.end) for load_slice in plan.slices] == [(1, 64, 192)] + + def test_origin_shape_clips_runtime_padding(self, single_rank_group: dist.ProcessGroup) -> None: + spec = LoadSpec( + name="layers.0.experts.fused_w1w3.weight", + global_hf_keys=["k0", "k1", "k2", "k3"], + global_shape=(480, 64), + fused_dim=0, + shards=[ShardDescriptor(dim=0, start=350, end=450, group=single_rank_group)], + origin_shape=(400, 64), + ) + + plan = spec.plan_hf_load() + + assert plan.hf_keys == ["k3"] + assert [(load_slice.dim, load_slice.start, load_slice.end) for load_slice in plan.slices] == [(0, 50, 100)] + assert plan.zero_fill is False + + def test_origin_shape_returns_zero_fill_for_pad_only_rank(self, single_rank_group: dist.ProcessGroup) -> None: + spec = LoadSpec( + name="layers.0.experts.fused_w1w3.weight", + global_hf_keys=["k0", "k1", "k2", "k3"], + global_shape=(480, 64), + fused_dim=0, + shards=[ShardDescriptor(dim=0, start=420, end=480, group=single_rank_group)], + origin_shape=(400, 64), + ) + + plan = spec.plan_hf_load() + + assert plan.zero_fill is True + assert plan.hf_keys == [] + assert plan.slices == [] + + +class TestHFSavePolicy: + """HF save should preserve the old distributed write policy from the new schema.""" + + 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) -> None: + 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, start=100, end=200, 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"] + + def test_preserved_fused_shard_must_align_with_hf_key_boundary(self, single_rank_group: dist.ProcessGroup) -> None: + 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, start=50, end=150, 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 TestHFSaveUnshardScheduler: + """Save unshard should batch independent work without violating per-tensor dependencies.""" + + @staticmethod + def _patch_foreach_all_gather(monkeypatch: pytest.MonkeyPatch) -> 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], + } + ) + return [[tensor] for tensor in tensor_list] + + monkeypatch.setattr(load_spec_module, "foreach_all_gather", fake_foreach_all_gather) + return calls + + def test_single_tensor_single_step( + self, monkeypatch: pytest.MonkeyPatch, single_rank_group: dist.ProcessGroup + ) -> None: + calls = self._patch_foreach_all_gather(monkeypatch) + spec = LoadSpec( + name="layers.0.mlp.gate.weight", + global_hf_keys=["gate"], + global_shape=(4, 2), + shards=[ShardDescriptor(dim=0, start=1, end=3, group=single_rank_group)], + ) + + output = unshard_tensors_for_hf_save( + [torch.ones(2, 2)], + [spec.plan_hf_save()], + ) + + assert [tuple(tensor.shape) for tensor in output] == [(4, 2)] + assert [call["shapes"] for call in calls] == [[(4, 2)]] + + def test_same_group_same_dtype_tensors_are_batched( + self, monkeypatch: pytest.MonkeyPatch, single_rank_group: dist.ProcessGroup + ) -> None: + calls = self._patch_foreach_all_gather(monkeypatch) + specs = [ + LoadSpec( + name="layers.0.mlp.gate.weight", + global_hf_keys=["gate"], + global_shape=(4, 2), + shards=[ShardDescriptor(dim=0, start=1, end=3, group=single_rank_group)], + ), + LoadSpec( + name="layers.0.mlp.up.weight", + global_hf_keys=["up"], + global_shape=(6, 2), + shards=[ShardDescriptor(dim=0, start=2, end=5, group=single_rank_group)], + ), + ] + + output = unshard_tensors_for_hf_save( + [torch.ones(2, 2), torch.ones(3, 2)], + [spec.plan_hf_save() for spec in specs], + ) + + assert [tuple(tensor.shape) for tensor in output] == [(4, 2), (6, 2)] + assert [call["shapes"] for call in calls] == [[(4, 2), (6, 2)]] + + def test_same_group_different_dtype_tensors_are_split( + self, monkeypatch: pytest.MonkeyPatch, single_rank_group: dist.ProcessGroup + ) -> None: + calls = self._patch_foreach_all_gather(monkeypatch) + specs = [ + LoadSpec( + name="layers.0.mlp.gate.weight", + global_hf_keys=["gate"], + global_shape=(4, 2), + shards=[ShardDescriptor(dim=0, start=1, end=3, group=single_rank_group)], + ), + LoadSpec( + name="layers.0.mlp.up.weight", + global_hf_keys=["up"], + global_shape=(4, 2), + shards=[ShardDescriptor(dim=0, start=1, end=3, group=single_rank_group)], + ), + ] + + output = unshard_tensors_for_hf_save( + [torch.ones(2, 2, dtype=torch.float32), torch.ones(2, 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_multi_step_tensor_waits_for_previous_step( + self, monkeypatch: pytest.MonkeyPatch, single_rank_group: dist.ProcessGroup + ) -> None: + calls = self._patch_foreach_all_gather(monkeypatch) + specs = [ + LoadSpec( + name="layers.0.experts.fused_w1w3.weight", + global_hf_keys=["k0", "k1"], + global_shape=(8, 2), + fused_dim=0, + shards=[ + ShardDescriptor(dim=0, start=0, end=4, group=single_rank_group), + ShardDescriptor(dim=0, start=1, end=3, group=single_rank_group), + ], + ), + LoadSpec( + name="layers.0.mlp.gate.weight", + global_hf_keys=["gate"], + global_shape=(4, 2), + shards=[ShardDescriptor(dim=0, start=1, end=3, group=single_rank_group)], + ), + ] + + output = unshard_tensors_for_hf_save( + [torch.ones(2, 2), torch.ones(2, 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] == [[(4, 2), (4, 2)], [(8, 2)]] + + +class TestBaseModelHFSave: + """BaseModel save should preserve state semantics outside LoadSpec.""" + + 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/model/base.py b/xtuner/v1/model/base.py index 0856a3fd7a..1407530d77 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,14 @@ 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 ( + HFLoadPlan, + 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 +70,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 +531,27 @@ def _save_file( save_file(tensors, filename, metadata=metadata) +class _HFSaveBucketItem(NamedTuple): + tensor: torch.Tensor + save_plan: HFSavePlan + runtime_is_float8: bool + + 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): @@ -622,6 +639,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( @@ -698,6 +716,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) @@ -920,36 +941,67 @@ 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, - ): + load_plan: HFLoadPlan, + ) -> None: + """Copy loaded HF tensors into a local parameter tensor. + + Args: + safetensors (list[torch.Tensor]): HF tensors loaded for ``load_plan.hf_keys``, in key order. + local_tensor (torch.Tensor): Destination local parameter or buffer tensor. + load_plan (HFLoadPlan): Plan whose ``slices`` are relative to ``safetensors`` after concatenation. + """ + loaded_tensor = self._cat_safetensors(safetensors, load_plan) + loaded_tensor = self._apply_load_slices(loaded_tensor, load_plan) + self._copy_loaded_tensor_to_local(loaded_tensor, local_tensor) + + def _cat_safetensors(self, safetensors: list[torch.Tensor], load_plan: HFLoadPlan) -> torch.Tensor: + assert safetensors, f"Internal Error. No safetensors were loaded for {load_plan.name}" if len(safetensors) > 1: + dim = load_plan.fused_dim 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) + return torch.cat(safetensors, dim=dim) + return safetensors[0] + + def _apply_load_slices(self, loaded_tensor: torch.Tensor, load_plan: HFLoadPlan) -> torch.Tensor: + for load_slice in load_plan.slices: + start = min(load_slice.start, loaded_tensor.shape[load_slice.dim]) + end = min(load_slice.end, loaded_tensor.shape[load_slice.dim]) + assert start <= end, f"Invalid load slice [{start}, {end}) for {load_plan.name}" + loaded_tensor = loaded_tensor.narrow(load_slice.dim, start, end - start) + return loaded_tensor - 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: + def _copy_loaded_tensor_to_local(self, loaded_tensor: torch.Tensor, local_tensor: torch.Tensor) -> None: + if loaded_tensor.shape == local_tensor.shape: local_tensor.copy_(loaded_tensor) + return + + assert loaded_tensor.dim() == local_tensor.dim(), ( + f"Loaded tensor shape {tuple(loaded_tensor.shape)} is incompatible with local tensor shape " + f"{tuple(local_tensor.shape)}" + ) + # HF checkpoints never store FSDP padding. After applying the LoadPlan slices, only the FSDP shard dim may be + # shorter than the runtime local tensor; all other dims must match exactly. + non_pad_dim_matches = all( + loaded_tensor.shape[dim] == local_tensor.shape[dim] + for dim in range(local_tensor.dim()) + if dim != self.FSDP_SHARD_DIM + ) + assert non_pad_dim_matches, ( + f"Loaded tensor shape {tuple(loaded_tensor.shape)} is incompatible with local tensor shape " + f"{tuple(local_tensor.shape)}; padding is only expected on dim {self.FSDP_SHARD_DIM}" + ) + non_pad_len = loaded_tensor.shape[self.FSDP_SHARD_DIM] + assert non_pad_len <= local_tensor.shape[self.FSDP_SHARD_DIM], ( + f"Loaded tensor shape {tuple(loaded_tensor.shape)} is larger than local tensor shape " + f"{tuple(local_tensor.shape)}" + ) + local_tensor.narrow(self.FSDP_SHARD_DIM, 0, non_pad_len).copy_(loaded_tensor) + + if non_pad_len < local_tensor.shape[self.FSDP_SHARD_DIM]: + assert self.config.float8_cfg is not None + pad_len = local_tensor.shape[self.FSDP_SHARD_DIM] - non_pad_len + # Torch casts the scalar to the destination dtype; for fp8 this writes the canonical zero value. + local_tensor.narrow(self.FSDP_SHARD_DIM, non_pad_len, pad_len).copy_(0.0) # type: ignore def param_to_safetensor( self, @@ -1028,30 +1080,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 +1113,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 +1130,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 +1325,199 @@ 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: 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 + runtime_tensor = param._local_tensor if isinstance(param, DTensor) else param + runtime_is_float8 = is_float8_weight(runtime_tensor) + is_buffer = load_spec.name in buffer_names + if runtime_tensor.is_floating_point() and not is_buffer: + save_dtype = self._get_save_dtype(load_spec.global_hf_keys[0], torch.bfloat16) + local_tensor = runtime_tensor.to(dtype=save_dtype) 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 + # Persistent buffers, e.g. FoPE rotary coefficients, are part of HF state but are not trainable + # weights. Keep the legacy behavior and write them in their runtime dtype instead of save_dtype. + local_tensor = runtime_tensor + tensor_size = self._get_tensor_size(runtime_tensor, dtype) + + if safetensor_size + tensor_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] - - return hf_tensor_list, name_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: list[str] = [] + safetensor_size = 0 + bucket = [] - 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) + save_plan = load_spec.plan_hf_save( + distributed_save=distributed_save, + preserve_process_group=preserved_fused_shard_group, + ) + bucket.append( + _HFSaveBucketItem(tensor=local_tensor, save_plan=save_plan, runtime_is_float8=runtime_is_float8) + ) - if tensor_list: - hf_params, name_list = _get_hf_params(tensor_list, name_list) - yield name_list, hf_params + if bucket: + yield self._build_hf_param_bucket( + bucket, + dtype=dtype, + device=device, + ) - def _get_same_hf_param( + 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 _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 +1527,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 +1539,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 +1562,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 +1691,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 +1703,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 +1812,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: @@ -2110,178 +1854,53 @@ def _load_fp8(self, hf_key: str, checkpoint_loader: HFCheckpointLoader) -> torch ) 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) - - 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] - - start = fsdp_start - end = fsdp_end - else: - start = None - end = None + ) -> list[str]: + """Unified HF load path for a single parameter / buffer. - self.safetensors_to_params( - [loaded_tensor], local_tensor, param_name=load_spec.name, start=start, end=end, dim=load_spec.dim - ) - return [] + ``LoadSpec.plan_hf_load`` computes this rank's HF keys and loaded-tensor-relative slices from the new + schema. This method only executes that plan: load keys, dequantize fp8 when needed, then hand off to + ``safetensors_to_params`` for cat + narrow + copy. - 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 + 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 - - 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!" + load_plan = load_spec.plan_hf_load() + if load_plan.zero_fill: + # No checkpoint key overlaps this rank. This can be fp8 runtime padding, or a legal zero-sized DTensor + # shard when a tiny tensor dimension is split across more ranks than it has elements. + assert load_spec.origin_shape is not None or local_tensor.numel() == 0, ( + "Empty load plan is only legal for runtime pad-only or zero-sized local 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 + local_tensor.zero_() # type: ignore + return [] 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: + loaded_tensors: list[torch.Tensor] = [] + for hf_key in load_plan.hf_keys: + 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)." + ) + weight = self._load_fp8(hf_key, checkpoint_loader) + else: weight = checkpoint_loader.load(hf_key) if weight 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 + loaded_tensors.append(weight.to(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 - ) - 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] - - loaded_tensor = loaded_tensor.to(local_tensor.device) - - 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" - ) - - 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] - - 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, + loaded_tensors, + local_tensor, + load_plan, ) return [] @@ -2295,125 +1914,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, 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..954ced6bd1 100644 --- a/xtuner/v1/model/moe/glm52.py +++ b/xtuner/v1/model/moe/glm52.py @@ -23,6 +23,7 @@ from xtuner.v1.module.mtp import MTPConfig, MTPLayer from xtuner.v1.module.rope import RopeParametersConfig from xtuner.v1.module.router.noaux_router import NoAuxRouterConfig +from xtuner.v1.utils.load_spec import HFLoadPlan from .moe import MoE @@ -161,36 +162,17 @@ 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] + load_plan: HFLoadPlan, + ) -> None: + loaded_tensor = self._cat_safetensors(safetensors, load_plan) if ( - "fused_w1w3.weight" in param_name or "fused_w2.weight" in param_name + "fused_w1w3.weight" in load_plan.name or "fused_w2.weight" in load_plan.name ) and loaded_tensor.ndim == local_tensor.ndim + 1: 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) + loaded_tensor = self._apply_load_slices(loaded_tensor, load_plan) + self._copy_loaded_tensor_to_local(loaded_tensor, local_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..e5c48db928 100644 --- a/xtuner/v1/model/moe/gpt_oss.py +++ b/xtuner/v1/model/moe/gpt_oss.py @@ -12,6 +12,7 @@ from xtuner.v1.module.decoder_layer.moe_decoder_layer import MoEActFnConfig from xtuner.v1.module.rope import RopeParametersConfig from xtuner.v1.module.router.greedy import GreedyRouterConfig +from xtuner.v1.utils.load_spec import HFLoadPlan from .moe import MoE @@ -44,18 +45,11 @@ 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] + load_plan: HFLoadPlan, + ) -> None: + loaded_tensor = self._cat_safetensors(safetensors, load_plan) - if "fused_w1w3.weight" in param_name: + if "fused_w1w3.weight" in load_plan.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 +58,20 @@ 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 load_plan.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 load_plan.name: # hf: num_experts, expert_dim * 2 # xtuner: num_experts, 2 * expert_dim 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) - - 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) + loaded_tensor = self._apply_load_slices(loaded_tensor, load_plan) + self._copy_loaded_tensor_to_local(loaded_tensor, local_tensor) def param_to_safetensor( self, diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py index f27e0a2dbc..225fac993c 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/xtuner/v1/model/moe/moe.py @@ -1269,6 +1269,7 @@ 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 diff --git a/xtuner/v1/model/moe/qwen3_5_text.py b/xtuner/v1/model/moe/qwen3_5_text.py index e4cf2e2fc3..6e742bbe21 100644 --- a/xtuner/v1/model/moe/qwen3_5_text.py +++ b/xtuner/v1/model/moe/qwen3_5_text.py @@ -14,6 +14,7 @@ from xtuner.v1.module.attention import GatedDeltaNetConfig, MHAConfig from xtuner.v1.module.rope import RopeParametersConfig from xtuner.v1.module.router.greedy import GreedyRouterConfig +from xtuner.v1.utils.load_spec import HFLoadPlan from .qwen3vl_text import Qwen3VLTextMoE @@ -126,42 +127,23 @@ 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] + load_plan: HFLoadPlan, + ) -> None: + loaded_tensor = self._cat_safetensors(safetensors, load_plan) - if "fused_w1w3.weight" in param_name and "mtp" not in param_name: + if "fused_w1w3.weight" in load_plan.name and "mtp" not in load_plan.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 load_plan.name and "mtp" not in load_plan.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) + loaded_tensor = self._apply_load_slices(loaded_tensor, load_plan) + self._copy_loaded_tensor_to_local(loaded_tensor, local_tensor) def param_to_safetensor( self, diff --git a/xtuner/v1/model/moe/qwen3vl_text.py b/xtuner/v1/model/moe/qwen3vl_text.py index 5451c31346..853c158a6b 100644 --- a/xtuner/v1/model/moe/qwen3vl_text.py +++ b/xtuner/v1/model/moe/qwen3vl_text.py @@ -5,6 +5,7 @@ from xtuner.v1.data_proto import SequenceContext from xtuner.v1.utils.activation_offload import async_save_on_cpu +from xtuner.v1.utils.load_spec import HFLoadPlan from .moe import MoELossContextDict, MoEModelOutputs from .qwen3 import Qwen3MoE, Qwen3MoE30BA3Config, Qwen3MoE235BA22Config @@ -39,18 +40,11 @@ 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] + load_plan: HFLoadPlan, + ) -> None: + loaded_tensor = self._cat_safetensors(safetensors, load_plan) - if "fused_w1w3.weight" in param_name: + if "fused_w1w3.weight" in load_plan.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 +52,13 @@ 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 load_plan.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) + loaded_tensor = self._apply_load_slices(loaded_tensor, load_plan) + self._copy_loaded_tensor_to_local(loaded_tensor, local_tensor) def param_to_safetensor( self, 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/load_spec.py b/xtuner/v1/utils/load_spec.py index ef95585fe3..399a5fa3e1 100644 --- a/xtuner/v1/utils/load_spec.py +++ b/xtuner/v1/utils/load_spec.py @@ -1,36 +1,768 @@ +import math +from collections.abc import Callable +from typing import Any, NamedTuple, cast + +import torch import torch.distributed as dist -from pydantic import BaseModel, ConfigDict +import torch.distributed.tensor._utils as dtensor_utils +import torch.nn.functional as F +from pydantic import BaseModel, ConfigDict, Field, computed_field +from torch.distributed.tensor import DTensor, Shard -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): + """A single partition applied to the fused full tensor. + + The full tensor is obtained by concatenating every ``LoadSpec.global_hf_keys`` along + ``LoadSpec.fused_dim`` (or taking the sole HF tensor when ``len(global_hf_keys) == 1``). + Descriptors are applied in order; later descriptors use offsets relative to the sub-tensor produced by all + earlier descriptors, matching DTensor placement semantics. + + Args: + dim (int): Tensor dim on which this partition cuts. + start (int): Inclusive start offset relative to the current sub-tensor. + end (int): Exclusive end offset relative to the current sub-tensor. + group (dist.ProcessGroup): Communication group that produced this partition. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + dim: int + start: int + end: int + group: dist.ProcessGroup + + +def _dtensor_shards(tensor: DTensor) -> list[ShardDescriptor]: + current_shape = list(tensor.shape) + shards: list[ShardDescriptor] = [] + for mesh_dim, placement in _ordered_dtensor_placements(tensor): + if not isinstance(placement, Shard): + continue + + # DTensor placement order is not always the raw mesh-dim order. FSDP2 can represent right-to-left sharding + # with _StridedShard, and PyTorch's checkpoint offset helper first expands that into the effective shard + # order. LoadSpec must preserve the same order so its descriptor intervals match DTensor local tensors. + # + # XTuner may initialize modules while the default device is "meta". PyTorch's Shard placement helpers can + # inherit that default device for temporary shape arithmetic, so force XTuner's real runtime device before + # calling the helper. + with torch.device(get_device()): + local_size, offset = placement._local_shard_size_and_offset( # type: ignore[attr-defined] + current_shape[placement.dim], + tensor.device_mesh.size(mesh_dim), + tensor.device_mesh.get_local_rank(mesh_dim), + ) + shards.append( + ShardDescriptor( + dim=placement.dim, + start=offset, + end=offset + local_size, + group=tensor.device_mesh.get_group(mesh_dim), + ) + ) + current_shape[placement.dim] = local_size + return shards + + +def _ordered_dtensor_placements(tensor: DTensor) -> list[tuple[int, object]]: + # PyTorch keeps this helper private and does not expose it in type stubs, but it is the same ordering logic used + # by `compute_local_shape_and_global_offset`. Access it dynamically so mypy does not reject the private symbol. + explicit_order_placements = cast( + Callable[[Any, Any], list[tuple[int, object]]], + getattr(dtensor_utils, "_explicit_order_placements"), + ) + return explicit_order_placements(tensor.device_mesh.shape, tensor.placements) + + +class LoadSlice(BaseModel): + """A narrow operation in the loaded HF tensor coordinate system. + + Args: + dim (int): Tensor dimension to narrow. + start (int): Inclusive start offset in the loaded tensor. + end (int): Exclusive end offset in the loaded tensor. + """ + + model_config = ConfigDict(extra="forbid") + dim: int + start: int + end: int + + +class HFLoadPlan(BaseModel): + """Execution plan for reading HF safetensors into one local 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. + slices (list[LoadSlice]): Narrow operations to apply after loading. Offsets are relative to the loaded + tensor, not the original ``LoadSpec.global_shape``. + zero_fill (bool): Whether this rank falls entirely in a padded region and should skip checkpoint reads. + """ + + model_config = ConfigDict(extra="forbid") name: str + hf_keys: list[str] + fused_dim: int | None = None + slices: list[LoadSlice] = Field(default_factory=list) + zero_fill: bool = False + + +def _final_intervals( + global_shape: tuple[int, ...], + shards: list[ShardDescriptor], +) -> list[tuple[int, int]]: + intervals = [(0, dim_size) for dim_size in global_shape] + for shard in shards: + current_start, _ = intervals[shard.dim] + intervals[shard.dim] = (current_start + shard.start, current_start + shard.end) + return intervals + + +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. + + ``load_spec_shard_index`` is only needed when some original shards should stay sharded. RL weight sync preserves + the EP shard on the fused HF dimension so each EP rank streams only its local expert keys, while later shards such + as FSDP still need to be all-gathered. Because execution reverses and groups the work items, their list positions + no longer match ``LoadSpec.shards``. The original index is the stable handle used by the save plan to decide + which work items to skip and which preserved shards should define the final expected shape. + + 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: + load_spec_shard_index (int): Index of ``shard`` in the original ``LoadSpec.shards`` list. + shard (ShardDescriptor): Shard descriptor this save step reverses. + shape_before_shard (tuple[int, ...]): Runtime tensor shape immediately before ``shard`` was applied. + unpadded_shape_before_shard (tuple[int, ...]): Checkpoint-visible shape 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") + load_spec_shard_index: int + shard: ShardDescriptor + shape_before_shard: tuple[int, ...] + unpadded_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. + global_shape (tuple[int, ...]): Runtime full tensor shape before any shard is applied. + unpadded_global_shape (tuple[int, ...]): Checkpoint-visible full tensor shape after removing runtime padding. + 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. + preserves_shards (bool): Whether the save tensor intentionally remains sharded by some original + ``LoadSpec.shards`` entries. + 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 + global_shape: tuple[int, ...] + unpadded_global_shape: tuple[int, ...] + fused_dim: int | None = None + distributed_save: bool = False + preserves_shards: bool = False + unshard_steps: list[SaveShardStep] = Field(default_factory=list) + + def _pending_unshard_steps(self) -> list[SaveShardStep]: + return [step for step in reversed(self.unshard_steps) if not step.preserved] + + def _preserved_shards(self) -> list[ShardDescriptor]: + return [step.shard for step in self.unshard_steps if step.preserved] + + def _expected_unsharded_shape(self) -> tuple[int, ...]: + """Return the save tensor shape after intentionally preserved shards + remain applied. + + The save path starts from the local tensor and all-gathers every pending shard step. If no shard is preserved, + the final shape should be ``unpadded_global_shape``. If some shards are preserved, for example an EP shard + during RL weight sync, the final tensor should still be cut by those preserved shards. This helper applies + only the preserved shard descriptors to ``unpadded_global_shape`` to compute that expected partially-unsharded + shape for the final assert. + + Example: + Suppose the runtime full tensor shape is ``(16, 8)`` because fp8 padding added rows, while + ``unpadded_global_shape == (14, 8)`` is the shape that should exist in HF. If the preserved EP shard is + ``ShardDescriptor(dim=0, start=8, end=16)``, that shard owns runtime rows ``[8, 16)``. The last two rows + are padding-only in HF coordinates, so the checkpoint-visible interval is clipped to ``[8, 14)`` and the + expected preserved tensor shape is ``(6, 8)``. If a shard were ``[14, 16)``, both boundaries would clip to + ``14`` and the expected shape on that rank would be ``(0, 8)``. + + Returns: + tuple[int, ...]: Expected shape after the preserved shards are still applied. + """ + effective_shape = list(self.unpadded_global_shape) + for shard in self._preserved_shards(): + # ShardDescriptor offsets are defined against the runtime shape, which may include XTuner-only padding. + # Clip preserved shard boundaries to the currently visible unpadded shape before computing its length. + clipped_start = min(shard.start, effective_shape[shard.dim]) + clipped_end = min(shard.end, effective_shape[shard.dim]) + effective_shape[shard.dim] = max(0, clipped_end - clipped_start) + return tuple(effective_shape) + + +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 [] + + # Shallow-copy the list, not the tensors. Entries with no gather work can be returned as-is, while entries + # that do need all-gather are overwritten in this working list with their gathered tensor. + tensor_list = list(tensors) + + # Convert each tensor's forward shard history into the save-time work queue. Save must undo shards from + # inner to outer, so the steps are reversed; preserved shards, such as an EP shard kept local for RL weight + # sync, are removed from the queue and only used later to compute the expected partially-unsharded shape. + + # Example: + # tensor A: [ep_a(index=0), fsdp_a(index=1)], preserved {0} -> pending [fsdp_a] + # tensor B: [ep_b(index=0), fsdp_b(index=1)], preserved {} -> pending [fsdp_b, ep_b] + # tensor C: [fsdp_c(index=0)], preserved {} -> pending [fsdp_c] + # tensor D: [tp_d(index=0)], preserved {} -> pending [tp_d] + # tensor E: [ep_e(index=0)], preserved {0} -> pending [] + # This produces one pending queue per tensor; the loop below consumes compatible queue heads by group. + pending_shard_steps_list = [save_plan._pending_unshard_steps() for save_plan in save_plans] + + while True: + # Build one all-gather round. For one tensor, reverse-unshard steps must run one by one: if a local + # tensor needs to undo FSDP and then EP, the EP gather must use the tensor produced by the FSDP gather. + # `_build_ready_save_unshard_groups` consumes `pending_shard_steps_list` gradually. For example, a queue + # `[fsdp_step, ep_step]` contributes `fsdp_step` in the first round; after its gathered tensor is written + # back, the next loop consumes `ep_step`. Independent tensors with compatible group/dtype can still be + # batched together in each round. + # + # With the A-E example above, round 1 consumes fsdp_a/fsdp_b/fsdp_c together if they share group/dtype, + # and consumes tp_d in another group. tensor E contributes no work. Round 2 can then consume ep_b, because + # ep_b must use tensor B after fsdp_b has been gathered and written back. + unshard_groups = _build_ready_save_unshard_groups(tensor_list, pending_shard_steps_list) + if not unshard_groups: + break + + 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 + + for tensor, save_plan in zip(tensor_list, save_plans, strict=True): + expected_shape = save_plan._expected_unsharded_shape() + assert tuple(tensor.shape) == expected_shape, ( + f"Saved tensor shape {tuple(tensor.shape)} is incompatible with HFSavePlan global_shape=" + f"{save_plan.global_shape} and unpadded_global_shape={save_plan.unpadded_global_shape} " + f"for {save_plan.name}" + ) + return tensor_list + + +def _build_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 _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] + 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 + gathered_tensor = torch.cat(gathered_chunks, dim=dim) + return gathered_tensor.narrow(dim, 0, shard_step.unpadded_shape_before_shard[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. + """ + + 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 + + @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) + 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=_dtensor_shards(tensor) if isinstance(tensor, DTensor) else [], + origin_shape=origin_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 loaded-tensor-relative slices for this rank. + """ + effective_intervals = self._effective_intervals_for_shards(self.shards) + if effective_intervals is None: + return HFLoadPlan(name=self.name, hf_keys=[], fused_dim=self.fused_dim, zero_fill=True) + + loaded_starts = [0 for _ in self.global_shape] + loaded_ends = list(self.unpadded_global_shape) + key_start, key_end = self._local_hf_key_indices(effective_intervals) + hf_keys = self.global_hf_keys[key_start:key_end] + + if self.is_fused: + key_size = self._fused_key_size() + assert self.fused_dim is not None + loaded_starts[self.fused_dim] = key_start * key_size + loaded_ends[self.fused_dim] = key_end * key_size + + slices: list[LoadSlice] = [] + for dim, (effective_start, effective_end) in enumerate(effective_intervals): + loaded_start = loaded_starts[dim] + loaded_end = loaded_ends[dim] + if effective_start == loaded_start and effective_end == loaded_end: + continue + slices.append( + LoadSlice( + dim=dim, + start=effective_start - loaded_start, + end=effective_end - loaded_start, + ) + ) + + return HFLoadPlan(name=self.name, hf_keys=hf_keys, fused_dim=self.fused_dim, slices=slices) + + 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] + hf_keys = ( + self._local_hf_keys_for_shards(preserved_shards, require_fused_key_aligned=True) + if preserved_shards + else list(self.global_hf_keys) + ) + + return HFSavePlan( + name=self.name, + hf_keys=hf_keys, + global_shape=self.global_shape, + unpadded_global_shape=self.unpadded_global_shape, + fused_dim=self.fused_dim, + distributed_save=distributed_save, + preserves_shards=bool(preserved_shards), + 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 _effective_intervals_for_shards( + self, + shards: list[ShardDescriptor], + ) -> list[tuple[int, int]] | None: + effective_shape = self.unpadded_global_shape + assert len(effective_shape) == len(self.global_shape), ( + f"origin_shape={effective_shape} must have the same rank as global_shape={self.global_shape}" + ) + assert all(effective <= global_ for effective, global_ in zip(effective_shape, self.global_shape)), ( + f"origin_shape={effective_shape} must not exceed global_shape={self.global_shape}" + ) + + final_intervals = _final_intervals(self.global_shape, shards) + effective_intervals: list[tuple[int, int]] = [] + for dim, (start, end) in enumerate(final_intervals): + effective_start = min(start, effective_shape[dim]) + effective_end = min(end, effective_shape[dim]) + if effective_start >= effective_end: + return None + effective_intervals.append((effective_start, effective_end)) + return effective_intervals + + 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 _local_hf_key_indices( + self, + effective_intervals: list[tuple[int, int]], + *, + 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, fused_end = effective_intervals[self.fused_dim] + 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; LoadSlice later trims load tensors 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]: + effective_intervals = self._effective_intervals_for_shards(shards) + if effective_intervals is None: + return [] + key_start, key_end = self._local_hf_key_indices( + effective_intervals, + 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: + current_shape = list(self.global_shape) + for shard in self.shards: + assert 0 <= shard.dim < len(current_shape), ( + f"Invalid shard dim {shard.dim} for global_shape={self.global_shape}" + ) + current_size = current_shape[shard.dim] + assert 0 <= shard.start <= shard.end <= current_size, ( + f"Invalid shard descriptor {shard} against current_shape={tuple(current_shape)}" + ) + current_shape[shard.dim] = shard.end - shard.start + + 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: start from the full runtime tensor, + apply one shard after another, and end at this rank's local tensor. The returned steps keep that same + largest-to-smallest order. Each step snapshots the runtime shape and the unpadded checkpoint-visible shape + that existed immediately before its shard was applied. + + Save executes these steps in reverse. Starting from the smallest local tensor, each reverse step all-gathers + one shard and narrows the gathered tensor back to ``unpadded_shape_before_shard``. This is how the save path + reconstructs the original shape information one partition layer at a time, while still avoiding fp8 runtime + padding in the checkpoint-visible tensor. + + Example: + Suppose ``global_shape=(16, 8)``, ``unpadded_global_shape=(14, 8)``, and + ``LoadSpec.shards == [ep(dim=0, start=8, end=16), fsdp(dim=0, start=3, end=5)]``. The returned steps are + in forward order: + + * ``ep_step`` records ``shape_before_shard=(16, 8)`` and + ``unpadded_shape_before_shard=(14, 8)``. + * ``fsdp_step`` records ``shape_before_shard=(8, 8)`` and + ``unpadded_shape_before_shard=(6, 8)``. + + A local save tensor has shape ``(2, 8)``. Save runs ``[fsdp_step, ep_step]``: gather FSDP back toward + ``(6, 8)``, then gather EP back toward ``(14, 8)``. If EP is preserved, only ``fsdp_step`` remains + pending and the result stays EP-local. + + 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) + effective_shape = list(self.unpadded_global_shape) + steps: list[SaveShardStep] = [] + + for shard_index, shard in enumerate(self.shards): + steps.append( + SaveShardStep( + load_spec_shard_index=shard_index, + shard=shard, + shape_before_shard=tuple(current_shape), + unpadded_shape_before_shard=tuple(effective_shape), + preserved=shard_index in preserved_shard_indices, + ) + ) + effective_start = min(shard.start, effective_shape[shard.dim]) + effective_end = min(shard.end, effective_shape[shard.dim]) + effective_shape[shard.dim] = max(0, effective_end - effective_start) + current_shape[shard.dim] = shard.end - shard.start + return steps From d0a176ddfb1369cb0d0ce9cc53aeef8188c331f8 Mon Sep 17 00:00:00 2001 From: zhaopenghao Date: Mon, 10 Aug 2026 07:19:46 +0000 Subject: [PATCH 2/7] [Feature] Support InterleavedShard in HF checkpoint I/O --- tests/utils/test_interleaved_shard.py | 252 ++++++++++++++++ xtuner/v1/model/base.py | 43 ++- xtuner/v1/utils/init_weight.py | 13 +- xtuner/v1/utils/interleaved_shard.py | 403 ++++++++++++++++++++++++++ xtuner/v1/utils/load_spec.py | 83 +++++- 5 files changed, 778 insertions(+), 16 deletions(-) create mode 100644 tests/utils/test_interleaved_shard.py create mode 100644 xtuner/v1/utils/interleaved_shard.py diff --git a/tests/utils/test_interleaved_shard.py b/tests/utils/test_interleaved_shard.py new file mode 100644 index 0000000000..ba923f8edb --- /dev/null +++ b/tests/utils/test_interleaved_shard.py @@ -0,0 +1,252 @@ +"""Unit tests for ``xtuner.v1.utils.interleaved_shard``. + +These tests cover the InterleavedShard placement and the ``reconstruct_full_tensor`` helper +across the layouts that XTuner actually uses: + + * Plain ``(Shard, InterleavedShard)`` on a 2D (ep, tp) mesh — the layout produced by + ``GroupedLinear`` when TP is enabled. + * The post-``fully_shard`` 3D layout with FSDP prepended on top — what HF save sees in + practice. + +Run with:: + + torchrun --nproc-per-node=8 tests/utils/test_interleaved_shard.py +""" + +from __future__ import annotations + +import os +import shutil +import sys +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 + +_HERE = os.path.dirname(os.path.abspath(__file__)) +_REPO_ROOT = os.path.dirname(os.path.dirname(_HERE)) +sys.path.insert(0, _REPO_ROOT) +# Import the module directly to avoid pulling in xtuner package's heavy deps (loguru etc.) that +# aren't required for this unit test. +import importlib.util as _ilu + +_spec = _ilu.spec_from_file_location( + "interleaved_shard", + os.path.join(_REPO_ROOT, "xtuner", "v1", "utils", "interleaved_shard.py"), +) +assert _spec is not None and _spec.loader is not None +_mod = _ilu.module_from_spec(_spec) +_spec.loader.exec_module(_mod) +InterleavedShard = _mod.InterleavedShard +compute_runs = _mod.compute_runs +has_interleaved_placement = _mod.has_interleaved_placement +reconstruct_full_tensor = _mod.reconstruct_full_tensor + + +NUM_EXPERTS = 4 +OUT_PER_EXPERT = 4 +IN_FEATURES = 8 +GLOBAL_ROWS = NUM_EXPERTS * OUT_PER_EXPERT # 16 + + +def _build_expected_local( + g: torch.Tensor, + ep_rank: int, + tp_rank: int, + ep_size: int, + tp_size: int, +) -> torch.Tensor: + """Hand-computed per-expert column parallel slice.""" + experts_per_ep = NUM_EXPERTS // ep_size + rows_per_expert = g.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(g[row_start : row_start + rows_per_tp_per_expert]) + return torch.cat(chunks, dim=0) + + +def test_2d_layout_and_reconstruct(): + """Build a DTensor on (ep, tp) with (Shard, InterleavedShard) and reconstruct.""" + 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") + + g = torch.arange(GLOBAL_ROWS * IN_FEATURES, device="cuda", dtype=torch.float32).reshape( + GLOBAL_ROWS, IN_FEATURES + ) + dist.broadcast(g, src=0) + + placements = (Shard(0), InterleavedShard(0, num_local_stripes=NUM_EXPERTS // 2)) + dt = distribute_tensor(g, mesh, placements) + + # Layout correctness: per-rank local matches hand-computed per-expert column parallel. + expected = _build_expected_local(g, ep_rank, tp_rank, 2, 2) + assert torch.allclose(dt.to_local(), expected), ( + f"rank {dist.get_rank()} local mismatch" + ) + + # Detection helper works on this placement. ``shard_order`` only exists on torch>=2.10; + # the implementation guards with ``getattr`` so the test must too. + assert has_interleaved_placement(dt), "shard_order should be None for this placement" + assert getattr(dt._spec, "shard_order", None) is None + + # Reconstruct gives back the global tensor. + full = reconstruct_full_tensor(dt) + assert torch.allclose(full, g), ( + f"reconstruct mismatch on 2D layout: max_diff={(full - g).abs().max().item()}" + ) + + +def test_hf_round_trip(): + """Exercise InterleavedShard through BaseModel's public HF save/load API.""" + 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] + + 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_dir = tempfile.mkdtemp() if dist.get_rank() == 0 else None + checkpoint_dirs = [checkpoint_dir] + 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) + assert torch.equal(restored.weight.to_local(), model.weight.to_local()) + finally: + dist.barrier() + if dist.get_rank() == 0: + shutil.rmtree(checkpoint_dir) + + +class _ToyGroupedLinear(nn.Module): + def __init__(self, weight): + super().__init__() + self.weight = nn.Parameter(weight) + + def forward(self, x): + w = self.weight.to_local() if isinstance(self.weight, DTensor) else self.weight + return torch.nn.functional.linear(x, w) + + +def test_post_fully_shard_reconstruct(): + """Layout after FSDP wraps the (ep, tp) DTensor — the case HF save actually sees.""" + mesh = init_device_mesh("cuda", (2, 2, 2), mesh_dim_names=("fsdp", "ep", "tp")) + ep_tp = mesh["ep", "tp"] + fsdp_mesh = mesh["fsdp"] + + g = torch.arange(GLOBAL_ROWS * IN_FEATURES, device="cuda", dtype=torch.float32).reshape( + GLOBAL_ROWS, IN_FEATURES + ) + dist.broadcast(g, src=0) + + placements = (Shard(0), InterleavedShard(0, num_local_stripes=NUM_EXPERTS // 2)) + dt = distribute_tensor(g, ep_tp, placements) + + model = _ToyGroupedLinear(dt).cuda() + fully_shard( + model, + mesh=fsdp_mesh, + mp_policy=MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32), + reshard_after_forward=True, + ) + + # Sanity: a forward pass through the wrapped model still produces the right output. + x = torch.randn(6, IN_FEATURES, device="cuda", dtype=torch.bfloat16) + dist.broadcast(x, src=0) + y = model(x) + ep_rank = mesh.get_local_rank("ep") + tp_rank = mesh.get_local_rank("tp") + expected_local = _build_expected_local(g, ep_rank, tp_rank, 2, 2).to(torch.bfloat16) + expected_y = torch.nn.functional.linear(x, expected_local) + assert torch.allclose(y.detach(), expected_y, atol=1e-2, rtol=1e-2) + y.sum().backward() + + # Detection helper still recognizes the wrapped DTensor. + assert has_interleaved_placement(model.weight) + + # Reconstruct from the post-FSDP local matches the original global. + full = reconstruct_full_tensor(model.weight) + assert torch.allclose(full, g), ( + f"reconstruct mismatch on post-FSDP layout: max_diff={(full - g).abs().max().item()}" + ) + + # HF load uses compute_runs to copy from the concatenated global tensor into the post-FSDP + # local tensor. This must describe FSDP's prepended shard as a contiguous cut; otherwise a + # valid HF checkpoint is loaded into the wrong local rows before training starts. + local = model.weight._local_tensor + loaded_local = torch.empty_like(local, dtype=g.dtype) + for run in compute_runs(model.weight): + loaded_slice = g.narrow(0, run.global_offset[0], run.local_size) + loaded_local.narrow(0, run.local_start, run.local_size).copy_(loaded_slice) + expected_local = local.to(g.dtype) + assert torch.allclose(loaded_local, expected_local), ( + f"compute_runs load mismatch on post-FSDP layout: " + f"max_diff={(loaded_local - expected_local).abs().max().item()}" + ) + + +def main(): + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group(backend="nccl") + world = dist.get_world_size() + + rank = dist.get_rank() + if world == 4: + test_2d_layout_and_reconstruct() + test_hf_round_trip() + if rank == 0: + print("[2d_layout_and_hf_round_trip] PASSED", flush=True) + elif world == 8: + test_post_fully_shard_reconstruct() + if rank == 0: + print("[post_fully_shard_reconstruct] PASSED", flush=True) + else: + if rank == 0: + print( + f"World size {world} not handled (expected 4 or 8). Skipping.", flush=True + ) + dist.destroy_process_group() + sys.exit(0) + + dist.barrier() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/xtuner/v1/model/base.py b/xtuner/v1/model/base.py index 1407530d77..73d055edcc 100644 --- a/xtuner/v1/model/base.py +++ b/xtuner/v1/model/base.py @@ -1364,7 +1364,20 @@ def _get_hf_param( buffer_names = {self._clean_param_name(name) for name, _ in self.named_buffers()} for param, load_spec in params: - runtime_tensor = param._local_tensor if isinstance(param, DTensor) else param + # InterleavedShard-bearing DTensors (e.g. fused MoE column-parallel weights) have + # `shard_order=None`; their layout cannot be described by per-step ShardDescriptors. + # Materialize the global tensor up-front via `reconstruct_full_tensor` and treat the + # result as already-unsharded by the rest of the save pipeline (load_spec.shards is + # empty, so `unshard_tensors_for_hf_save` becomes a no-op for these items). + if load_spec.needs_full_reconstruct: + assert isinstance(param, DTensor), ( + f"needs_full_reconstruct=True implies a DTensor param, got {type(param).__name__}" + ) + from xtuner.v1.utils.interleaved_shard import reconstruct_full_tensor + + runtime_tensor = reconstruct_full_tensor(param) + else: + runtime_tensor = param._local_tensor if isinstance(param, DTensor) else param runtime_is_float8 = is_float8_weight(runtime_tensor) is_buffer = load_spec.name in buffer_names if runtime_tensor.is_floating_point() and not is_buffer: @@ -1897,6 +1910,22 @@ def _load_hf_param( if missing_keys: return missing_keys + if load_spec.needs_full_reconstruct: + # InterleavedShard-style placements: this rank owns N contiguous "runs" of rows in + # the global tensor (one per local expert). Copy each run from the concatenated + # HF tensor to the matching slice of the local tensor. + assert isinstance(param, DTensor), ( + f"needs_full_reconstruct=True implies a DTensor param, got {type(param).__name__}" + ) + from xtuner.v1.utils.interleaved_shard import compute_runs + + loaded_tensor = self._cat_safetensors(loaded_tensors, load_plan) + local = param._local_tensor + for run in compute_runs(param): + loaded_slice = loaded_tensor.narrow(0, run.global_offset[0], run.local_size) + local.narrow(0, run.local_start, run.local_size).copy_(loaded_slice) + return [] + self.safetensors_to_params( loaded_tensors, local_tensor, @@ -2129,7 +2158,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/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..7f3e509078 --- /dev/null +++ b/xtuner/v1/utils/interleaved_shard.py @@ -0,0 +1,403 @@ +"""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. + * Save / load are routed through :func:`reconstruct_full_tensor` (this module) and the LoadSpec + machinery, neither of which depends on ``shard_order``. + +The reconstruction algorithm and its rationale are documented inline on +:func:`reconstruct_full_tensor`. +""" + +from __future__ import annotations + +from typing import NamedTuple + +import torch +import torch.distributed._functional_collectives as funcol +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 both the HF save path (build per-run WriteItems / per-run slices) and the HF load + path (per-run narrow + copy from the loaded global tensor). + + 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 the global tensor from a DTensor's local data, even when the + spec contains placements that PyTorch's ``redistribute`` cannot handle + (``shard_order=None``). + + Why a custom routine: ``DTensor.full_tensor()`` goes through ``redistribute`` which asserts + ``shard_order is not None`` in torch 2.10. For our ``(Shard, InterleavedShard)`` placement + that assert fires. We bypass redistribute by emitting collectives directly. + + Algorithm: + + 1. **Phase 1 — undo FSDP-prepended _StridedShard (mesh_dim 0) as plain Shard.** FSDP2 + actually chunks the parameter contiguously (``_chunk_with_empty``) regardless of the + strided label. So the right undo is a plain ``all_gather`` along the FSDP mesh dim. + After this phase every rank holds the pre-FSDP local. + + 2. **Phase 2 — undo remaining placements in REVERSE mesh-dim order:** + + * ``InterleavedShard`` (= real strided): ``all_gather`` along the placement's mesh dim, + then scatter the gathered chunks back to their correct global positions using + ``_local_shard_size_and_offset(return_first_offset=False)``. + * Plain ``Shard``: ``all_gather`` and concatenate. + + The reverse direction is essential because ``InterleavedShard.split_factor`` is defined + relative to the size of the tensor *after* the placements to its right have already + been undone. Doing TP undo before EP undo keeps the sf math consistent. + + 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__}") + + mesh = dt.device_mesh + placements = list(dt.placements) + # Make sure the working buffer is contiguous so all_gather copies see a well-defined layout. + result = dt._local_tensor.contiguous() + + # Phase 1: FSDP-prepended _StridedShard at mesh_dim 0 → plain gather. + for mesh_dim, placement in enumerate(placements): + if not _is_fsdp_prepended_strided(placement, mesh_dim): + continue + result = _all_gather_plain(result, placement.dim, mesh.get_group(mesh_dim)) + + # Phase 2: remaining placements in reverse mesh-dim order. + for mesh_dim in reversed(range(len(placements))): + placement = placements[mesh_dim] + if not isinstance(placement, (Shard, _StridedShard)): + continue + if _is_fsdp_prepended_strided(placement, mesh_dim): + continue # already handled in Phase 1 + if _is_real_strided(placement, mesh_dim): + result = _undo_strided(result, placement, mesh, mesh_dim) + else: + # Plain Shard or _StridedShard with sf == 1 (degenerate). + result = _all_gather_plain(result, placement.dim, mesh.get_group(mesh_dim)) + + return result + + +# --------------------------------------------------------------------------- +# Internal collective helpers +# --------------------------------------------------------------------------- + + +def _all_gather_plain(local: torch.Tensor, tensor_dim: int, group) -> torch.Tensor: + """``all_gather_tensor`` along ``tensor_dim`` then materialize the async + wrapper.""" + gathered = funcol.all_gather_tensor(local, gather_dim=tensor_dim, group=group) + if isinstance(gathered, funcol.AsyncCollectiveTensor): + gathered = gathered.wait() + return gathered + + +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, mirroring ``reconstruct_full_tensor`` which undoes FSDP first. + + 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 + + +def _undo_strided( + local: torch.Tensor, + placement, + mesh, + mesh_dim: int, +) -> torch.Tensor: + """``all_gather`` + scatter for a strided placement. + + Each rank in the mesh dim group holds a strided chunk per ``placement``'s spec. After + ``all_gather`` the result is the concatenation of those chunks in rank order. To recover + the original layout we re-index each rank's chunk back to its true positions using + ``_local_shard_size_and_offset(return_first_offset=False)`` which returns the global + indices the rank owned within the post-undo tensor. + """ + tensor_dim = placement.dim + mesh_size = mesh.size(mesh_dim) + group = mesh.get_group(mesh_dim) + + gathered = _all_gather_plain(local, tensor_dim, group) + current_size = gathered.shape[tensor_dim] + + all_indices: list[int] = [] + for r in range(mesh_size): + all_indices.extend(_strided_indices(placement, current_size, mesh_size, r)) + + indices_tensor = torch.tensor(all_indices, device=gathered.device, dtype=torch.long) + new_result = torch.empty_like(gathered) + new_result.index_copy_(tensor_dim, indices_tensor, gathered) + return new_result diff --git a/xtuner/v1/utils/load_spec.py b/xtuner/v1/utils/load_spec.py index 399a5fa3e1..991cac7198 100644 --- a/xtuner/v1/utils/load_spec.py +++ b/xtuner/v1/utils/load_spec.py @@ -1,13 +1,12 @@ import math -from collections.abc import Callable -from typing import Any, NamedTuple, cast +from typing import NamedTuple import torch import torch.distributed as dist -import torch.distributed.tensor._utils as dtensor_utils 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 xtuner.v1.ops.comm.foreach_allgather import foreach_all_gather from xtuner.v1.utils.device import get_device @@ -74,13 +73,59 @@ def _dtensor_shards(tensor: DTensor) -> list[ShardDescriptor]: def _ordered_dtensor_placements(tensor: DTensor) -> list[tuple[int, object]]: - # PyTorch keeps this helper private and does not expose it in type stubs, but it is the same ordering logic used - # by `compute_local_shape_and_global_offset`. Access it dynamically so mypy does not reject the private symbol. - explicit_order_placements = cast( - Callable[[Any, Any], list[tuple[int, object]]], - getattr(dtensor_utils, "_explicit_order_placements"), - ) - return explicit_order_placements(tensor.device_mesh.shape, tensor.placements) + # 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 LoadSlice(BaseModel): @@ -413,6 +458,7 @@ class LoadSpec(BaseModel): 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. + needs_full_reconstruct (bool): Whether HF I/O must use the explicit InterleavedShard reconstruction/run path. """ model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") @@ -422,6 +468,12 @@ class LoadSpec(BaseModel): fused_dim: int | None = None shards: list[ShardDescriptor] = Field(default_factory=list) origin_shape: tuple[int, ...] | None = None + # When True, this tensor's layout cannot be described by the ``shards`` list — typically an + # ``InterleavedShard``-bearing DTensor whose spec has ``shard_order=None``. The HF save path + # must call :func:`xtuner.v1.utils.interleaved_shard.reconstruct_full_tensor` on the param at + # save time to materialize the global tensor, and treat the result as already-unsharded + # (i.e. ``shards`` is empty, no per-step all-gather work needed). + needs_full_reconstruct: bool = False @computed_field # type: ignore[prop-decorator] @property @@ -465,13 +517,22 @@ def from_tensor( LoadSpec: Spec derived from the runtime tensor layout. """ global_hf_keys = list(hf_keys) + shards: list[ShardDescriptor] = [] + needs_full_reconstruct = False + if isinstance(tensor, DTensor): + from xtuner.v1.utils.interleaved_shard import has_interleaved_placement + + needs_full_reconstruct = has_interleaved_placement(tensor) + if not needs_full_reconstruct: + shards = _dtensor_shards(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=_dtensor_shards(tensor) if isinstance(tensor, DTensor) else [], + shards=shards, origin_shape=origin_shape, + needs_full_reconstruct=needs_full_reconstruct, ) def plan_hf_load(self) -> HFLoadPlan: From e9b6f021af50431b78826eb77da32410d962572e Mon Sep 17 00:00:00 2001 From: zhaopenghao Date: Mon, 10 Aug 2026 07:27:23 +0000 Subject: [PATCH 3/7] [Feature] Support InterleavedShard in DCP checkpoints --- tests/patch/test_dcp_interleaved_planner.py | 172 ++++++++++++++++++ xtuner/v1/engine/train_engine.py | 9 +- xtuner/v1/patch/__init__.py | 3 + xtuner/v1/patch/dcp_interleaved_planner.py | 187 ++++++++++++++++++++ xtuner/v1/patch/torch_dcp_planner.py | 23 ++- 5 files changed, 385 insertions(+), 9 deletions(-) create mode 100644 tests/patch/test_dcp_interleaved_planner.py create mode 100644 xtuner/v1/patch/dcp_interleaved_planner.py 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/xtuner/v1/engine/train_engine.py b/xtuner/v1/engine/train_engine.py index 61a1f58fe0..0c21110337 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 ( @@ -355,6 +356,7 @@ def save_dcp( dcp.save( state_dict, checkpoint_id=weights_dir, + planner=InterleavedShardSavePlanner(), ) def _get_async_checkpoint_pg(self) -> dist.ProcessGroup: @@ -407,6 +409,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 +506,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/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) From 7e1f4b84e02f106000315d599a4f9b00ece51c67 Mon Sep 17 00:00:00 2001 From: zhaopenghao Date: Mon, 10 Aug 2026 09:54:37 +0000 Subject: [PATCH 4/7] [Feature] Add MoE expert tensor parallelism --- docs/design/expert_tp.md | 39 + .../test_moe_train_engine_deepep_expert_tp.py | 413 +++++++ tests/engine/test_moe_train_engine_float8.py | 98 +- tests/engine/test_moe_train_engine_tpep.py | 1084 +++++++++++++++++ tests/model/test_gpt_oss_moe.py | 40 +- tests/model/test_moe_expert_tp_without_ep.py | 72 ++ tests/model/test_qwen3_moe.py | 13 +- tests/module/dispatcher/test_agrs_all2all.py | 1 + tests/module/dispatcher/test_deepep.py | 1 + .../dispatcher/test_deepep_expert_tp.py | 310 +++++ tests/module/dispatcher/test_noep.py | 1 + .../module/dispatcher/test_noep_expert_tp.py | 310 +++++ tests/module/dispatcher/test_torch_all2all.py | 1 + .../test_torch_all2all_shared_expert_tp.py | 262 ++++ tests/utils/test_compile.py | 69 +- xtuner/v1/engine/train_engine.py | 3 +- xtuner/v1/float8/float8_gmm_tile_wise.py | 120 +- xtuner/v1/float8/float8_handler.py | 15 +- xtuner/v1/model/base.py | 13 + xtuner/v1/model/moe/glm52.py | 18 +- xtuner/v1/model/moe/gpt_oss.py | 23 +- xtuner/v1/model/moe/moe.py | 158 ++- xtuner/v1/model/moe/qwen3_5_text.py | 21 +- xtuner/v1/model/moe/qwen3vl_text.py | 17 +- .../module/decoder_layer/moe_decoder_layer.py | 20 + xtuner/v1/module/dispatcher/__init__.py | 21 +- xtuner/v1/module/dispatcher/agrs.py | 1 + xtuner/v1/module/dispatcher/base.py | 236 +++- xtuner/v1/module/dispatcher/deepep.py | 117 +- xtuner/v1/module/dispatcher/expert_tp.py | 405 ++++++ xtuner/v1/module/dispatcher/torch_all2all.py | 89 +- .../module/grouped_linear/moe_group_linear.py | 172 ++- xtuner/v1/ops/comm/deepep_op.py | 23 +- 33 files changed, 4000 insertions(+), 186 deletions(-) create mode 100644 docs/design/expert_tp.md create mode 100644 tests/engine/test_moe_train_engine_deepep_expert_tp.py create mode 100644 tests/engine/test_moe_train_engine_tpep.py create mode 100644 tests/model/test_moe_expert_tp_without_ep.py create mode 100644 tests/module/dispatcher/test_deepep_expert_tp.py create mode 100644 tests/module/dispatcher/test_noep_expert_tp.py create mode 100644 tests/module/dispatcher/test_torch_all2all_shared_expert_tp.py create mode 100644 xtuner/v1/module/dispatcher/expert_tp.py 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/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/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/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/xtuner/v1/engine/train_engine.py b/xtuner/v1/engine/train_engine.py index 0c21110337..0d6e1e90c6 100644 --- a/xtuner/v1/engine/train_engine.py +++ b/xtuner/v1/engine/train_engine.py @@ -48,7 +48,6 @@ log_rank0, profile_time_and_memory, ) -from xtuner.v1.utils.grad_norm import cal_grad_norm class TrainStepInfo(DataBatchInfo, BatchForwardInfo): @@ -260,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) 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 73d055edcc..663b7b908f 100644 --- a/xtuner/v1/model/base.py +++ b/xtuner/v1/model/base.py @@ -596,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() @@ -951,9 +956,14 @@ def safetensors_to_params( load_plan (HFLoadPlan): Plan whose ``slices`` are relative to ``safetensors`` after concatenation. """ loaded_tensor = self._cat_safetensors(safetensors, load_plan) + loaded_tensor = self.hf_tensor_to_canonical(load_plan.name, loaded_tensor) loaded_tensor = self._apply_load_slices(loaded_tensor, load_plan) self._copy_loaded_tensor_to_local(loaded_tensor, local_tensor) + def hf_tensor_to_canonical(self, name: str, loaded_tensor: torch.Tensor) -> torch.Tensor: + """Convert one loaded HF tensor to XTuner's canonical layout.""" + return loaded_tensor + def _cat_safetensors(self, safetensors: list[torch.Tensor], load_plan: HFLoadPlan) -> torch.Tensor: assert safetensors, f"Internal Error. No safetensors were loaded for {load_plan.name}" if len(safetensors) > 1: @@ -1920,6 +1930,9 @@ def _load_hf_param( from xtuner.v1.utils.interleaved_shard import compute_runs loaded_tensor = self._cat_safetensors(loaded_tensors, load_plan) + # Interleaved runs use canonical global coordinates. Packed model + # formats such as GPT-OSS must be converted before applying them. + loaded_tensor = self.hf_tensor_to_canonical(load_plan.name, loaded_tensor) local = param._local_tensor for run in compute_runs(param): loaded_slice = loaded_tensor.narrow(0, run.global_offset[0], run.local_size) diff --git a/xtuner/v1/model/moe/glm52.py b/xtuner/v1/model/moe/glm52.py index 954ced6bd1..1c71eb4c54 100644 --- a/xtuner/v1/model/moe/glm52.py +++ b/xtuner/v1/model/moe/glm52.py @@ -23,7 +23,6 @@ from xtuner.v1.module.mtp import MTPConfig, MTPLayer from xtuner.v1.module.rope import RopeParametersConfig from xtuner.v1.module.router.noaux_router import NoAuxRouterConfig -from xtuner.v1.utils.load_spec import HFLoadPlan from .moe import MoE @@ -158,21 +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, - load_plan: HFLoadPlan, - ) -> None: - loaded_tensor = self._cat_safetensors(safetensors, load_plan) - - if ( - "fused_w1w3.weight" in load_plan.name or "fused_w2.weight" in load_plan.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) - - loaded_tensor = self._apply_load_slices(loaded_tensor, load_plan) - self._copy_loaded_tensor_to_local(loaded_tensor, local_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 e5c48db928..1e484942c1 100644 --- a/xtuner/v1/model/moe/gpt_oss.py +++ b/xtuner/v1/model/moe/gpt_oss.py @@ -12,7 +12,6 @@ from xtuner.v1.module.decoder_layer.moe_decoder_layer import MoEActFnConfig from xtuner.v1.module.rope import RopeParametersConfig from xtuner.v1.module.router.greedy import GreedyRouterConfig -from xtuner.v1.utils.load_spec import HFLoadPlan from .moe import MoE @@ -41,15 +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, - load_plan: HFLoadPlan, - ) -> None: - loaded_tensor = self._cat_safetensors(safetensors, load_plan) - - if "fused_w1w3.weight" in load_plan.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] @@ -58,20 +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 load_plan.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 load_plan.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) + loaded_tensor = loaded_tensor.transpose(1, 2).reshape(-1) - loaded_tensor = self._apply_load_slices(loaded_tensor, load_plan) - self._copy_loaded_tensor_to_local(loaded_tensor, local_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 225fac993c..028e9de28a 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 — needed by GroupedLinear for per-expert column-parallel weights + # so HF save can reconstruct the full tensor via `reconstruct_full_tensor`. + 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. @@ -1276,7 +1317,7 @@ def fully_shard( @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 @@ -1288,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 @@ -1300,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): @@ -1340,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. @@ -1391,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" @@ -1414,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 6e742bbe21..33414ef06b 100644 --- a/xtuner/v1/model/moe/qwen3_5_text.py +++ b/xtuner/v1/model/moe/qwen3_5_text.py @@ -10,11 +10,10 @@ 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 -from xtuner.v1.utils.load_spec import HFLoadPlan from .qwen3vl_text import Qwen3VLTextMoE @@ -123,27 +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, - load_plan: HFLoadPlan, - ) -> None: - loaded_tensor = self._cat_safetensors(safetensors, load_plan) - - if "fused_w1w3.weight" in load_plan.name and "mtp" not in load_plan.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 load_plan.name and "mtp" not in load_plan.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) - loaded_tensor = self._apply_load_slices(loaded_tensor, load_plan) - self._copy_loaded_tensor_to_local(loaded_tensor, local_tensor) + return loaded_tensor def param_to_safetensor( self, @@ -173,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 853c158a6b..45742d8527 100644 --- a/xtuner/v1/model/moe/qwen3vl_text.py +++ b/xtuner/v1/model/moe/qwen3vl_text.py @@ -5,7 +5,6 @@ from xtuner.v1.data_proto import SequenceContext from xtuner.v1.utils.activation_offload import async_save_on_cpu -from xtuner.v1.utils.load_spec import HFLoadPlan from .moe import MoELossContextDict, MoEModelOutputs from .qwen3 import Qwen3MoE, Qwen3MoE30BA3Config, Qwen3MoE235BA22Config @@ -36,15 +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, - load_plan: HFLoadPlan, - ) -> None: - loaded_tensor = self._cat_safetensors(safetensors, load_plan) - - if "fused_w1w3.weight" in load_plan.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] @@ -52,13 +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 load_plan.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) - loaded_tensor = self._apply_load_slices(loaded_tensor, load_plan) - self._copy_loaded_tensor_to_local(loaded_tensor, local_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 From 189dcf363cc53bb82bd8dab07f8c727cddffc6c7 Mon Sep 17 00:00:00 2001 From: zhaopenghao Date: Tue, 11 Aug 2026 00:25:27 +0000 Subject: [PATCH 5/7] [Refactor] Generalize HF load plans for expert TP --- docs/design/load_spec_refactor.md | 586 ++++++++++++------------- tests/utils/test_interleaved_shard.py | 20 +- tests/utils/test_load_spec.py | 597 ++++++++++++++++++------- xtuner/v1/model/base.py | 127 +----- xtuner/v1/model/moe/moe.py | 4 +- xtuner/v1/utils/interleaved_shard.py | 119 +---- xtuner/v1/utils/load_spec.py | 600 +++++++++++++++++--------- 7 files changed, 1159 insertions(+), 894 deletions(-) diff --git a/docs/design/load_spec_refactor.md b/docs/design/load_spec_refactor.md index 785a4bec17..e734c2d080 100644 --- a/docs/design/load_spec_refactor.md +++ b/docs/design/load_spec_refactor.md @@ -1,403 +1,355 @@ -# LoadSpec 设计 +# LoadSpec、DTensor 与 DCP 设计 -> 面向 `xtuner/v1/utils/load_spec.py` 与 `xtuner/v1/model/base.py` 的加载/保存路径。 -> TP 设计（`dense_tp.md`）依赖本文档描述的抽象。 +## 结论 -## TL;DR +HF checkpoint、DTensor runtime layout 和 DCP checkpoint 描述的是同一个逻辑 +tensor，但面向的协议不同： -LoadSpec 描述 xtuner 运行时 tensor 与 HF safetensors 之间的**纯布局映射**。对一个 -param，它回答两件事： +- `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。 -1. 这个 param 由哪些 HF key 组成？怎么拼？ — `global_hf_keys` + `fused_dim` -2. 本 rank 持有全量 tensor 的哪一块？ — `shards`（按外到内顺序施加） +~~~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"] +~~~ -加载/保存执行路径不直接读 LoadSpec，而是调用 `plan_hf_load()` / -`plan_hf_save(...)` 拿到一份**不可变的 plan**，按 plan 驱动 IO 与通信。 +核心边界是：HF plan 负责“checkpoint 格式与 runtime tensor 之间如何搬运”，DCP +planner 负责“local storage 对应哪些全局坐标”。两者复用相同的布局数学，不强行共用 +协议对象。 -**核心约束**：LoadSpec 只承担"同 dtype 下的形状/索引映射"。fp8 的量化反量化、 -padding 的 zero-fill 等 dtype 语义都住在 `base.py` 的 load/save 路径里， -LoadSpec 不感知。 +## 一、LoadSpec 与 Plan 的职责 ---- +### 1.1 为什么需要两层对象 -## 1. 设计理念 +可以把 `LoadSpec` 理解成地图，把 Plan 理解成一次具体行程：地图不会保存“今天从哪条 +路走”，而同一张地图可以生成完整保存、保留 EP 或只 gather FSDP 等不同路线。 -### 1.1 单一抽象，两条正交轴 - -原先三类映射（SAME / FUSED / SHARD）统一成一个 schema 上的两个正交维度： - -| 问题 | 表达 | -| --- | --- | -| 这个 param 对应几个 HF key？怎么拼？ | `len(global_hf_keys)`；多 key 时 `fused_dim` 指定拼接维 | -| 本 rank 持有哪一块？ | `shards`（可为空；按施加顺序排列） | - -消费方用派生属性 `is_fused` / `is_sharded` 查询，**不需要**任何枚举分支。 - -### 1.2 多维切分按顺序叠加 - -`shards` 是列表，原生支持 TP × FSDP、EP × FSDP 等多轴组合。每条 -`ShardDescriptor.start/end` 的含义是"在**前面所有** descriptor 切完之后的 -子 tensor 上的偏移"。这条规则完全对齐 DTensor `placements` 从 `mesh_dim=0` 到 -`mesh_dim=N-1` 逐步施加的语义 —— 你可以把 `shards[i]` 理解成 `placements[i]` -在"此刻本 rank 实际持有"这个问题上的等价形式。 - -### 1.3 Plan 是冻结快照 - -`plan_hf_load()` / `plan_hf_save(...)` 返回的是 Pydantic dataclass： - -- 一次性从当前 LoadSpec 状态计算出执行所需的全部信息； -- 不持有对 LoadSpec 的引用； -- 执行器（`_load_hf_param` / `unshard_tensors_for_hf_save` / `_split_hf_tensors_for_save`） - 只读 plan，**不读** LoadSpec。 - -这条边界保证"布局规划"和"IO/通信执行"解耦。未来要接入新的持久化格式（例如 -DCP），只需要替换 plan 的消费者，不牵涉 LoadSpec 内部结构。 - -### 1.4 fp8 与 LoadSpec 解耦 - -LoadSpec 是"同 dtype 下的布局描述"。fp8 涉及的两件事 —— 量化/反量化、运行时 -padding —— 归属如下： - -- **运行时 padding**：用 `LoadSpec.origin_shape` 表达 checkpoint-visible shape - （剥掉运行时 padding 之后）。今天这个字段的唯一来源是 fp8 tensor metadata； - 它只记录 shape，不记录 dtype / wrapper 类型。 -- **量化/反量化**：只在 `base.py._to_float8` / 反量化分支里现场判断（通过 - `is_float8_weight(tensor)`）。LoadSpec 不包含 `runtime_is_float8` 这类 - dtype-specific 字段。 +| 对象 | 负责 | 不负责 | +| --- | --- | --- | +| `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 | -### 1.5 Spec → Plan → Executor 的分层 +Plan 生成后，executor 不再回读 `LoadSpec`。因此 Plan 是 self-contained 的 +rank-local 程序，也可以独立测试。 -``` -┌────────────────────┐ plan_hf_load() ┌──────────────┐ -│ │ ──────────────────▶│ HFLoadPlan │──▶ _load_hf_param -│ LoadSpec │ └──────────────┘ -│ (pure layout) │ plan_hf_save(...) ┌──────────────┐ -│ │ ──────────────────▶│ HFSavePlan │──▶ unshard_tensors_for_hf_save -└────────────────────┘ └──────────────┘ │ - ▼ - _split_hf_tensors_for_save -``` +### 1.2 LoadSpec 保存哪些 shape -"Spec 是源、Plan 是派生、Executor 只依赖 Plan"。这条线保持单向。 +- `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. 数据模型 +## 二、统一的分片描述 -### 2.1 `ShardDescriptor` +### 2.1 ShardDescriptor ```python class ShardDescriptor(BaseModel): - dim: int # 被切的维 - start: int # 在"前面切完的 sub-tensor"上的起点 - end: int # 在"前面切完的 sub-tensor"上的终点 - group: dist.ProcessGroup # 产生这次切分的通信组 + dim: int + group: dist.ProcessGroup + interleave_factor: int = 1 ``` -`group` 是 load/save 双向通信域。load 时只需要知道本 rank 的范围；save 时需要沿 -`group` 做 all-gather 复原全量 tensor。 +- `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`。 -### 2.2 `LoadSpec` +even interleave 设当前维长度为 `N`，group size 为 `M`，factor 为 `F`，要求： -```python -class LoadSpec(BaseModel): - name: str # xtuner 侧 fully-qualified param name - global_hf_keys: list[str] # 对应的 HF key 列表（按 fused_dim 拼接顺序） - global_shape: tuple[int, ...] # 全量 tensor（fused 之后）的 runtime shape - # 可能包含运行时 padding（例如 fp8 的 FSDP 对齐 pad） - fused_dim: int | None = None # 多 HF key 时的拼接维；单 key 时必须为 None - shards: list[ShardDescriptor] = [] # 从外到内的切分列表 - origin_shape: tuple[int, ...] | None = None # checkpoint-visible shape after runtime padding is trimmed - # None 表示"runtime shape 就是 checkpoint shape" -``` - -派生属性： - -```python -is_fused # len(global_hf_keys) > 1 -is_sharded # bool(shards) -unpadded_global_shape # origin_shape or global_shape +```text +N % (M * F) == 0 +run_size = N / (M * F) ``` -**不变量**（`model_post_init` 强制）： - -- `is_fused` ⇔ `fused_dim is not None`； -- 每条 shard 的 `start/end` 必须落在"前面切完之后的 sub-tensor"范围内； -- 若 `origin_shape` 给定，它的秩与 `global_shape` 相同，且每维 `≤ global_shape`。 +rank `r` 持有第 `(j * M + r)` 个 run，`j = 0 ... F - 1`。当前不支持 uneven +interleave；不整除会明确报错。 -### 2.3 `HFLoadPlan` +### 2.2 DTensor placement 如何表达 Expert TP -`plan_hf_load()` 的产出： +MoE column-parallel 权重需要 EP 分 expert、TP 再在每个 expert/projection 内部分列： ```python -class HFLoadPlan(BaseModel): - name: str - hf_keys: list[str] # 本 rank 实际需要读的 HF key - fused_dim: int | None = None # 多 key 时的拼接维 - slices: list[LoadSlice] = [] # 读完拼接后，再做的 narrow 列表 - zero_fill: bool = False # 本 rank 完全落在运行时 padding 区，跳过 IO +placements = ( + Shard(0), + InterleavedShard(0, num_local_stripes=local_experts * fused_projections), +) ``` -`slices` 的 start/end 是**相对已加载 tensor 的坐标**，不是相对 `global_shape`。 -zero_fill=True 时 `hf_keys` 和 `slices` 都为空。 +`InterleavedShard` 继承 PyTorch 私有 `_StridedShard`。其 `split_factor` 不是“当前 +rank 数”，而是当前切分前已经存在的逻辑 stripe 数；在这里就是本 EP rank 的 +`expert × projection` 数。 -### 2.4 `HFSavePlan` +DTensor 构造时 placements 按 mesh 维从左到右作用；从 placements 推导 carving +order、HF save 重建时则按相反方向撤销。FSDP2 还会在 mesh 最左侧 prepend 一个 +`_StridedShard` 标签，但它实际切的是已经完成 EP/ETP 的 local parameter。因此 +`LoadSpec.from_tensor()` 在唯一的转换边界上做两件事： -`plan_hf_save(...)` 的产出，承载两类信息： +1. 真正的 `InterleavedShard` 转成 `interleave_factor > 1`。 +2. FSDP prepend placement 归一化成最后应用的 continuous descriptor。 -```python -class HFSavePlan(BaseModel): - name: str - hf_keys: list[str] # 当前 save tensor 最终要写/同步的 HF keys - global_shape: tuple[int, ...] - unpadded_global_shape: tuple[int, ...] - fused_dim: int | None = None - distributed_save: bool = False - preserves_shards: bool = False # True 表示 hf_keys 来自保留 shard 后的局部 tensor - unshard_steps: list[SaveShardStep] = [] # 所有 shard 的逆操作 + preserved 标记 -``` - -`SaveShardStep` 记录一次 shard 在"施加前的 runtime shape / checkpoint-visible -shape"两个快照 —— save 执行时倒序跑每一步、all-gather 还原、narrow 回 -checkpoint-visible shape。`preserved` 标记把某些 shard 排除在 all-gather 之外 -（见 §3.3）。`HFSavePlan.hf_keys` 始终是执行器要处理的 key 集合：普通 save 下 -它是完整 HF key list，preserved shard save 下它是当前局部 shard 覆盖的 key list。 +后续 plan 不再依赖 `_StridedShard` 类型或 `DTensorSpec.shard_order`。 ---- +### 2.3 通俗例子：两位仓库员各拿每个货架的一半 -## 3. 计划生成 +设 fused W1/W3 的 dim-0 共 16 行，业务顺序为： -### 3.1 `plan_hf_load()` - -不接受参数 —— 本 rank 的所有信息已经在 LoadSpec 里。步骤： +```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] +``` -1. 计算本 rank 最终持有的区间 `final_intervals`（顺序应用 `shards`）； -2. 用 `unpadded_global_shape` 裁剪掉运行时 padding 部分；若裁完为空，返回 - `zero_fill=True`； -3. 若 `is_fused`，按 `fused_dim` 上的区间算出需要的 HF key 下标范围（floor/ceil - 支持 mid-key shard，例如 FSDP 在 EP-local 专家 tensor 内部再切）； -4. 对每个 dim，如果"最终区间"比"加载后的 tensor 区间"窄，生成一条 `LoadSlice`。 +`EP=2, TP=2` 时，每个 EP rank 只有一个 expert，因 gate/up 融合所以 +`interleave_factor=2`： -### 3.2 `plan_hf_save(distributed_save=, preserve_process_group=, gather_process_group=)` +| `(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]` | -三个参数对应三种 save 策略，互斥使用： +这像两位仓库员分别负责每个货架的左半或右半，而不是一人搬走完整货架。若误用第二个 +`Shard(0)`，TP rank 会拿到连续的完整 projection 子集，gate/up 的列并行语义就错了。 -| 参数 | 用途 | -| --- | --- | -| `distributed_save=True` | HF save：非 fused tensor 只在 rank0 写；fused tensor 的 HF key 在 save rank 间分配 | -| `preserve_process_group=ep_group` | RL 权重同步：保留 EP 在 `fused_dim` 上的 shard，每个 EP rank 只流自己的 expert key；其他 shard 照常 all-gather | -| `gather_process_group=fsdp_group` | FSDP-only all-gather：只 gather 这个 group 的 shard，其他 shard 保留 | +这种 `(Shard, InterleavedShard)` 同维布局在部分 PyTorch 版本无法归约出 +`shard_order`，因此 `full_tensor()` / `redistribute()` 不能作为稳定实现。XTuner +训练直接使用 `DTensor.from_local()` / `to_local()`，HF 和 DCP 则使用本文的自定义 +布局编译路径。 -策略统一落到 `_preserved_shard_indices` 这一步上 —— 决定哪些 `LoadSpec.shards` -需要保留。之后 `_save_shard_steps` 给每个 shard 生成带 `preserved` 标记的 -`SaveShardStep`。若有 preserved shard，`LoadSpec` 直接从这些 shard 推导 -`HFSavePlan.hf_keys`；save plan 只暴露最终要写/同步的 HF keys，以及 -`preserves_shards` 说明这些 keys 来自局部 tensor 还是完整 tensor。 +## 三、HF Load -### 3.3 preserve vs gather 的正交性 +### 3.1 Plan 生成 -`preserve_process_group` 是"显式保留某个 group"的策略，`gather_process_group` -是"显式 gather 某个 group（其余保留）"的策略。两者不能同时使用（assert 拦截）。 -在今天的代码里： +`LoadSpec.plan_hf_load()` 从完整 canonical tensor 开始，按 descriptors 正向计算 +当前 rank 的 ownership： -- 普通 HF save：两者都不传，全部 all-gather； -- RL 权重同步：传 `preserve_process_group=ep_group`； -- `_fsdp_foreach_allgather`：传 `gather_process_group=fsdp_group`，只做 FSDP - 层的 all-gather，不动 EP / TP。 +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 大小展开逐元素索引。 -## 4. 执行 +### 3.2 固定执行顺序 -### 4.1 加载路径 +~~~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"] +~~~ -```python -def _load_hf_param(self, param, load_spec, loader): - plan = load_spec.plan_hf_load() - if plan.zero_fill: - # 本 rank 只持有运行时 padding，写 0 返回 - local_tensor.zero_() - return [] - # 按 plan.hf_keys 逐个读（fp8 走 dequant 分支，这里 base.py 现场处理） - loaded_tensors = self._load_hf_keys(plan, loader, ...) - # 拼接 + narrow 全部交给 safetensors_to_params - self.safetensors_to_params(loaded_tensors, local_tensor, plan) -``` +必须先 canonicalize 再按 ownership copy。例如 GPT-OSS 的 HF +`gate_up_proj` 先 transpose、重排 gate/up 并 flatten，之后 dim-0 才与 +`ShardDescriptor` 描述的 expert-major runtime 维一致；先 slice HF tensor 会切错 +语义维。 -`safetensors_to_params` 的签名是 `(safetensors, local_tensor, plan)`。三个 MoE -子类（`gpt_oss`、`qwen3_5_text`、`qwen3vl_text`）按 `plan.name` 做 reshape / -transpose 等模型特有变换后，调通用的 `_apply_load_slices` + `_copy_loaded_tensor_to_local`。 +### 3.3 模型差异被限制在 adapter -### 4.2 保存路径 +| 模型 | 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 | -所有 save 场景（HF save、RL 权重同步、FSDP-only gather）共用一条管道： +新增模型只需提供 key mapping 和 canonical adapter，不复制 EP/ETP/FSDP load +流程。 -```python -save_items = [HFSaveItem(tensor, load_spec.plan_hf_save(...)) for ...] -full_tensors = unshard_tensors_for_hf_save(save_items) -for full_tensor, item in zip(full_tensors, save_items): - names, tensors = self._split_hf_tensors_for_save(full_tensor, item.save_plan) -``` +## 四、HF Save、padding 与重建 -`unshard_tensors_for_hf_save` 自带**依赖感知的批量 foreach all-gather**： +### 4.1 逆分片 -- 同一个 tensor 的多个 step 必须串行（例如 "先还原 FSDP，再还原 EP"）； -- 不同 tensor 的 step 如果 `(group, dtype)` 兼容，可以 foreach 批到同一次 NCCL 调用。 +`plan_hf_save()` 为每个 descriptor 生成 `SaveShardStep`，executor 逆序执行所有未 +preserve 的 step： -每一轮由 `_build_ready_save_unshard_groups` 从每个 pending 队列取头部 step，按 -group + dtype 分桶；`_foreach_all_gather_save_shards` 跑一次批量 gather；下一轮 -再消费队列的下一层。MoE EP+FSDP 的 save 就是这样两轮跑完的。 +- 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。 -### 4.3 `HFSaveItem` +例如 `TP=2, F=2`，gather 得到： -```python -class HFSaveItem(NamedTuple): - tensor: torch.Tensor - save_plan: HFSavePlan +```text +rank0: [A0, A2] +rank1: [A1, A3] ``` -这是**跨 LoadSpec 和 BaseModel 边界**的 bundle：一边是 runtime tensor（模型侧 -概念，带 fp8 wrapper / DTensor wrapper），一边是纯布局的 `HFSavePlan`。它的 -归属地是 `base.py` ——`load_spec.py` 保持"不认识模型侧概念"。 -`unshard_tensors_for_hf_save` 的签名使用两个平行列表（`list[torch.Tensor]` + -`list[HFSavePlan]`）而不是 `list[HFSaveItem]`，避免 `load_spec.py` 反向依赖 -`base.py`。 - ---- +deinterleave 后必须为 `[A0, A1, A2, A3]`，不能像 continuous shard 一样直接按 +rank concat。 -## 5. 调用时机 +### 4.2 collective padding 与 FP8 padding -`_init_load_spec` 被定位为"从当前 DTensor 布局反推 HF 映射的纯函数"。 -调用约定：**谁改 param 布局谁负责重算，后者覆盖前者**。 +两类 padding 分层处理： -| 时机 | 调用方 | spec 代表 | +| padding | 原因 | 处理位置 | | --- | --- | --- | -| 子类 `__init__` 末尾 | 子类自己 | 构建完成时的布局（EP-only / Replicate / 其它 init-time 切分） | -| `parallelize(tp_mesh)` 结束 | `BaseModel.parallelize` | TP + 已有切分 | -| `fully_shard` 结束 | `BaseModel.fully_shard` | 叠加 FSDP（训练态） | -| `Float8Handler.pad_for_fsdp` 回调 | 回调内 | fp8 pad 后的真实 shape | +| 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` | -`from_hf` / `save_hf` 入口有 assert 兜底： +通俗例子：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 补回来。 -```python -assert "load_spec_mapping" in self.__dict__, ( - f"{type(self).__name__}.__init__ must call self._init_load_spec() at the end." -) -``` +load 侧没有 collective trim：plan 只复制 `origin_shape` 内存在的数据，并先把未写 +runtime padding 清零。 -这条约定是硬契约；子类若跳过会在第一次 load/save 时被抓。 +### 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 | -## 6. 示例 +`distributed_save` 是写出分配策略，不是部分重建策略；当前仍先正确重建完整 +canonical tensor，再决定哪些 rank 写哪些 keys。 -### 6.1 Dense, tp=2, fsdp=4, `q_proj.weight` +### 4.4 reconstruct_full_tensor -```python -LoadSpec( - name="layers.0.self_attn.q_proj.weight", - global_hf_keys=["model.layers.0.self_attn.q_proj.weight"], - global_shape=(n*d, h), - fused_dim=None, - shards=[ - ShardDescriptor(dim=0, start=tp_start, end=tp_end, group=tp_group), - ShardDescriptor(dim=0, start=fsdp_start, end=fsdp_end, group=fsdp_group), - ], -) -``` +`reconstruct_full_tensor(dt)` 仍是 public convenience API，因为 PyTorch +`DTensor.full_tensor()` 不稳定支持 `(Shard, InterleavedShard)`。它现在只是薄封装： -`fsdp_start/end` 相对于"已经被 TP 切过的 sub-tensor"而言，不是相对 -`global_shape`。 +1. 从 DTensor placements 构造 descriptor history。 +2. 生成不 preserve 任何 shard 的 full-runtime save steps。 +3. 调用与 `HFSavePlan` 相同的逆分片 executor。 +4. 返回 `global_shape` 的 runtime tensor。 -### 6.2 MoE, ep=8, fsdp=4, fused expert weight +它不做模型格式转换，也不 trim `origin_shape`；`distributed_save` 和 preserve group +属于 `HFSavePlan` 策略，不进入这个 API。生产 HF save 不再通过 +`save_requires_full_reconstruct` 走特殊旁路。 -```python -LoadSpec( - name="layers.0.experts.fused_w1w3.weight", - global_hf_keys=[f"model.layers.0.mlp.experts.{i}.gate_proj.weight" for i in range(64)] - + [f"model.layers.0.mlp.experts.{i}.up_proj.weight" for i in range(64)], - global_shape=(128 * I_padded, H), # I_padded 含 fp8 FSDP 对齐 pad - fused_dim=0, - shards=[ - ShardDescriptor(dim=0, start=ep_start, end=ep_end, group=ep_group), - ShardDescriptor(dim=0, start=fsdp_start, end=fsdp_end, group=fsdp_group), - ], - origin_shape=(128 * I, H), # 剥掉 pad 后的 checkpoint shape -) -``` +## 五、DCP 的独立协议 + +### 5.1 为什么默认 DCP 不够 -RL 权重同步调用 `plan_hf_save(preserve_process_group=ep_group)` —— EP shard 被 -标记 preserved，保存管道只做 FSDP 还原，结果留在 EP-local 坐标系；再由 -`_request_ep_sequential_update` 按 EP rank 顺序广播。 +普通 DTensor rank 可由一个 `local_shape + global_offset` 描述；interleaved rank +对应多个不连续区间。若把 `[0,1,4,5]` 错当成从 offset 0 开始的连续四行，DCP 会把 +后两行写到 global rows 2、3，checkpoint 看似成功但内容错误。 -### 6.3 embed_tokens, 纯 FSDP +`compute_runs()` 只根据 global shape、mesh coordinate 和 placements 做几何计算， +不读取 tensor value，也不发 collective。它把 local tensor 拆为若干“local 连续且 +global 也连续”的 `Run`： ```python -LoadSpec( - name="embed_tokens.weight", - global_hf_keys=["model.embed_tokens.weight"], - global_shape=(V, H), - fused_dim=None, - shards=[ShardDescriptor(dim=0, start=fsdp_start, end=fsdp_end, group=fsdp_group)], +Run( + global_offset=(global_row, 0), + sizes=(num_rows, hidden_size), + local_start=local_row, + local_size=num_rows, ) ``` ---- +### 5.2 通俗例子：给一本交错装订的书编页码 -## 7. 为什么这样设计 +沿用 `(ep=0,tp=0)` 的 local rows `[0,1,4,5]`。它在内存中是连续四行，但 +`compute_runs()` 返回两段： -几个关键取舍的归档。 - -### 7.1 为什么 `shards` 是列表而不是单轴四元组 - -旧的 `(dim, shard_start, shard_end, group)` 只表达一刀。TP × FSDP 或 EP × FSDP -是常见组合，旧 schema 只能靠"加载时临时推导第二刀"这种硬编码绕过（ -`FSDP_SHARD_DIM == 0` 就是这条路径的残留）。列表 + DTensor 施加顺序是最小的 -统一表达。 - -### 7.2 为什么删 `LoadEnum` - -`SAME/FUSED/SHARD` 给定 `global_hf_keys` 和 `shards` 后是可派生的。保留它相当于 -同一份状态的两种表达，下游分支要同步维护。直接用 `is_fused` / `is_sharded` 两个 -独立 bool 可以正交表达所有组合（包括原本需要新造 `FUSED_SHARD` 的情况）。 - -### 7.3 为什么 fp8 不进 LoadSpec - -LoadSpec 的定位是"同 dtype 下的映射"。fp8 涉及的反量化需要的是 tensor 的真实 -dtype / wrapper 类型，这些只有在 IO 路径里拿到 runtime tensor 才能判断。若把 -`runtime_is_float8` 放进 spec，一方面是状态重复（`is_float8_weight(tensor)` 已经 -是事实来源），另一方面污染 LoadSpec 的语义 —— 它不再是纯布局描述。 - -`origin_shape` 是 checkpoint-visible shape。它今天只服务 fp8 runtime padding， -但仍然只携带 shape 信息；fp8 的 dtype / wrapper 判断不进入 LoadSpec。 - -### 7.4 为什么 `unshard_tensors_for_hf_save` 住在 `load_spec.py` - -尽管它做的是分布式 all-gather，但它**只依赖 HFSavePlan + 一个通信原语**。把它 -放在 `load_spec.py` 让"spec → plan → 执行"三层都在一个文件里闭环，调用方 -（base.py）只需要准备 `(tensor, plan)` 对，不需要理解 shard 调度。 - -若将来 `unshard_tensors_for_hf_save` 进一步膨胀，可以拆到独立模块（例如 -`save_runner.py`），但当前规模尚不需要。 - -### 7.5 为什么保存不用 `_fuse_contiguous_chunks_without_alloc` - -旧代码对 `dim == 0` 的单 tensor all-gather 用过一个零拷贝 view 合并优化。这条 -优化只在"一次 gather 一个 tensor"时成立 —— 当前批量 foreach 把多个 tensor 交错 -塞进同一个扁平缓冲区，per-tensor chunks 不再连续，这条路径失效。换掉 NCCL 调用 -次数（O(num_tensors) → O(rounds)）比 dim=0 多一次 cat alloc 更划算。如果某个 -特定场景发现这次 trade-off 不值，可以单独给那条路走非批量路径，但默认策略保持 -批量。 - ---- +```text +Run(global rows [0,2), local rows [0,2)) +Run(global rows [4,6), local rows [2,4)) +``` -## 8. 测试 +就像一本小册子依次装订了原书第 0、1、4、5 页；保存时必须在目录中记录成两段， +不能声称它是原书第 0 到 3 页。 -核心测试都在 `tests/utils/test_load_spec.py`： +Save planner 为它生成： -- `TestLoadSpecSchema`：字段契约 + `shards` 顺序验证； -- `TestHFLoadPlan`：`plan_hf_load` 在 fused / non-fused / fp8 padding 下的产出； -- `TestHFSavePolicy`：`distributed_save` 的 HF key 分配规则。 +```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) +``` -行为等价性由 `tests/model/test_qwen3_dense.py::test_save_hf` 和 -`tests/model/test_qwen3_moe.py::test_save_hf` 的 safetensors bit-equal 保证。 +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/utils/test_interleaved_shard.py b/tests/utils/test_interleaved_shard.py index ba923f8edb..70843ff8c3 100644 --- a/tests/utils/test_interleaved_shard.py +++ b/tests/utils/test_interleaved_shard.py @@ -15,6 +15,7 @@ from __future__ import annotations +import importlib.util as _ilu import os import shutil import sys @@ -27,13 +28,12 @@ from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard from torch.distributed.tensor import DTensor, Shard, distribute_tensor + _HERE = os.path.dirname(os.path.abspath(__file__)) _REPO_ROOT = os.path.dirname(os.path.dirname(_HERE)) sys.path.insert(0, _REPO_ROOT) # Import the module directly to avoid pulling in xtuner package's heavy deps (loguru etc.) that # aren't required for this unit test. -import importlib.util as _ilu - _spec = _ilu.spec_from_file_location( "interleaved_shard", os.path.join(_REPO_ROOT, "xtuner", "v1", "utils", "interleaved_shard.py"), @@ -42,7 +42,6 @@ _mod = _ilu.module_from_spec(_spec) _spec.loader.exec_module(_mod) InterleavedShard = _mod.InterleavedShard -compute_runs = _mod.compute_runs has_interleaved_placement = _mod.has_interleaved_placement reconstruct_full_tensor = _mod.reconstruct_full_tensor @@ -108,7 +107,6 @@ def test_2d_layout_and_reconstruct(): def test_hf_round_trip(): """Exercise InterleavedShard through BaseModel's public HF save/load API.""" from transformers import PretrainedConfig - from xtuner.v1.model.base import BaseModel, XTunerBaseModelConfig class _ToyConfig(XTunerBaseModelConfig): @@ -205,17 +203,17 @@ def test_post_fully_shard_reconstruct(): f"reconstruct mismatch on post-FSDP layout: max_diff={(full - g).abs().max().item()}" ) - # HF load uses compute_runs to copy from the concatenated global tensor into the post-FSDP - # local tensor. This must describe FSDP's prepended shard as a contiguous cut; otherwise a - # valid HF checkpoint is loaded into the wrong local rows before training starts. + # Exercise the public load plan on the post-FSDP layout. LoadSpec converts + # InterleavedShard runs into a generic canonical source-to-local copy program. + from xtuner.v1.utils.load_spec import LoadSpec + local = model.weight._local_tensor loaded_local = torch.empty_like(local, dtype=g.dtype) - for run in compute_runs(model.weight): - loaded_slice = g.narrow(0, run.global_offset[0], run.local_size) - loaded_local.narrow(0, run.local_start, run.local_size).copy_(loaded_slice) + load_spec = LoadSpec.from_tensor(name="weight", hf_keys=["weight"], tensor=model.weight) + load_spec.plan_hf_load().load_into([g], loaded_local, lambda _, tensor: tensor) expected_local = local.to(g.dtype) assert torch.allclose(loaded_local, expected_local), ( - f"compute_runs load mismatch on post-FSDP layout: " + f"load plan mismatch on post-FSDP layout: " f"max_diff={(loaded_local - expected_local).abs().max().item()}" ) diff --git a/tests/utils/test_load_spec.py b/tests/utils/test_load_spec.py index f26ca85498..cf2c0e9cc5 100644 --- a/tests/utils/test_load_spec.py +++ b/tests/utils/test_load_spec.py @@ -1,4 +1,5 @@ import os +from collections.abc import Callable import pytest import torch @@ -16,10 +17,6 @@ @pytest.fixture(scope="module") def single_rank_group() -> dist.ProcessGroup: - # ShardDescriptor.group is typed as `dist.ProcessGroup`; Pydantic enforces - # the isinstance check even with `arbitrary_types_allowed=True`, so schema - # tests need a real (but minimal) process group. A single-rank gloo group - # is sufficient and avoids any CUDA / multi-process plumbing. if not dist.is_initialized(): os.environ.setdefault("RANK", "0") os.environ.setdefault("WORLD_SIZE", "1") @@ -31,9 +28,45 @@ def single_rank_group() -> dist.ProcessGroup: return group -class TestLoadSpecSchema: - """New-schema fields should describe layout without legacy dispatch state.""" +@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", @@ -63,7 +96,6 @@ def test_from_tensor_derives_plain_tensor_layout(self) -> None: assert spec.origin_shape == (120, 64) def test_from_tensor_derives_dtensor_shards(self, single_rank_group: dist.ProcessGroup) -> None: - assert single_rank_group is not None mesh = DeviceMesh("cpu", [0]) tensor = distribute_tensor(torch.empty(128, 64), mesh, [DTensorShard(0)]) @@ -72,154 +104,328 @@ def test_from_tensor_derives_dtensor_shards(self, single_rank_group: dist.Proces assert spec.global_hf_keys == ["gate"] assert spec.global_shape == (128, 64) assert spec.fused_dim is None - assert [(shard.dim, shard.start, shard.end) for shard in spec.shards] == [(0, 0, 128)] + assert [(shard.dim, shard.interleave_factor) for shard in spec.shards] == [(0, 1)] - def test_dtensor_shards_follow_explicit_placement_order(self, single_rank_group: dist.ProcessGroup) -> None: + 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_local_rank(self, mesh_dim: int) -> int: - return (1, 0)[mesh_dim] - def get_group(self, mesh_dim: int) -> dist.ProcessGroup: return single_rank_group class FakeDTensor: - shape = (8,) - placements = (_StridedShard(0, split_factor=2), DTensorShard(0)) + 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.start, shard.end) for shard in shards] == [(0, 0, 4), (0, 2, 4)] + 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=[ - "model.layers.0.mlp.experts.0.gate_proj.weight", - "model.layers.0.mlp.experts.0.up_proj.weight", - ], + global_hf_keys=["gate", "up"], global_shape=(256, 64), ) - def test_multi_axis_shards_preserve_order(self, single_rank_group: dist.ProcessGroup) -> None: - ep = ShardDescriptor(dim=0, start=64, end=128, group=single_rank_group) - fsdp = ShardDescriptor(dim=0, start=16, end=32, group=single_rank_group) + 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=[ - "model.layers.0.mlp.experts.0.gate_proj.weight", - "model.layers.0.mlp.experts.0.up_proj.weight", - ], + global_hf_keys=["gate", "up"], global_shape=(256, 64), fused_dim=0, - shards=[ep, fsdp], + shards=[ + ShardDescriptor(dim=0, group=single_rank_group), + ShardDescriptor(dim=0, group=single_rank_group, interleave_factor=2), + ], ) - assert [(shard.start, shard.end) for shard in spec.shards] == [(64, 128), (16, 32)] + 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_ordered_shard_bounds_are_validated(self, single_rank_group: dist.ProcessGroup) -> None: - with pytest.raises(ValidationError, match="Invalid shard descriptor"): + 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=["model.layers.0.mlp.experts.0.gate_proj.weight"], - global_shape=(128, 64), - shards=[ - ShardDescriptor(dim=0, start=64, end=128, group=single_rank_group), - ShardDescriptor(dim=0, start=65, end=80, group=single_rank_group), - ], + global_hf_keys=["gate"], + global_shape=(10, 4), + shards=[ShardDescriptor(dim=0, group=single_rank_group, interleave_factor=3)], ) - def test_zero_size_dtensor_shards_are_valid(self, single_rank_group: dist.ProcessGroup) -> None: + 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, start=1, end=1, group=single_rank_group)], + shards=[ShardDescriptor(dim=0, group=single_rank_group)], ) plan = spec.plan_hf_load() - assert plan.zero_fill is True assert plan.hf_keys == [] + target = torch.empty(0, 1, 1024) + plan.load_into([], target, lambda _, tensor: tensor) + assert target.numel() == 0 class TestHFLoadPlan: - """LoadSpec should derive HF read plans from shards only.""" - - def test_fused_slice_selects_overlapping_hf_keys(self, single_rank_group: dist.ProcessGroup) -> None: + 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, 64), + global_shape=(400, 8), fused_dim=0, - shards=[ShardDescriptor(dim=0, start=150, end=260, group=single_rank_group)], + shards=[ShardDescriptor(dim=0, group=single_rank_group, interleave_factor=2)], ) plan = spec.plan_hf_load() - assert plan.hf_keys == ["k1", "k2"] - assert plan.fused_dim == 0 - assert [(load_slice.dim, load_slice.start, load_slice.end) for load_slice in plan.slices] == [(0, 50, 160)] - assert not hasattr(plan, "loaded_shape") + 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_slice_keeps_single_hf_key(self, single_rank_group: dist.ProcessGroup) -> None: + 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, start=64, end=192, group=single_rank_group)], + 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) - assert plan.hf_keys == ["q_proj"] - assert plan.fused_dim is None - assert [(load_slice.dim, load_slice.start, load_slice.end) for load_slice in plan.slices] == [(1, 64, 192)] + torch.testing.assert_close(target, full[:, 128:]) - def test_origin_shape_clips_runtime_padding(self, single_rank_group: dist.ProcessGroup) -> None: + 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, 64), + global_shape=(480, 8), fused_dim=0, - shards=[ShardDescriptor(dim=0, start=350, end=450, group=single_rank_group)], - origin_shape=(400, 64), + shards=[ShardDescriptor(dim=0, group=single_rank_group)], + origin_shape=(400, 8), ) plan = spec.plan_hf_load() assert plan.hf_keys == ["k3"] - assert [(load_slice.dim, load_slice.start, load_slice.end) for load_slice in plan.slices] == [(0, 50, 100)] - assert plan.zero_fill is False - - def test_origin_shape_returns_zero_fill_for_pad_only_rank(self, single_rank_group: dist.ProcessGroup) -> None: + 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, 64), + global_shape=(480, 8), fused_dim=0, - shards=[ShardDescriptor(dim=0, start=420, end=480, group=single_rank_group)], - origin_shape=(400, 64), + shards=[ShardDescriptor(dim=0, group=single_rank_group)], + origin_shape=(400, 8), ) plan = spec.plan_hf_load() - assert plan.zero_fill is True assert plan.hf_keys == [] - assert plan.slices == [] + 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), :]) -class TestHFSavePolicy: - """HF save should preserve the old distributed write policy from the new schema.""" + 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 @@ -233,49 +439,56 @@ def test_fused_keys_are_split_across_save_ranks(self, monkeypatch: pytest.Monkey 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), - } + 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) -> None: + 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, start=100, end=200, group=single_rank_group)], + 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) -> None: + 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, start=50, end=150, group=single_rank_group)], + 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 TestHFSaveUnshardScheduler: - """Save unshard should batch independent work without violating per-tensor dependencies.""" - +class TestHFSaveUnshard: @staticmethod - def _patch_foreach_all_gather(monkeypatch: pytest.MonkeyPatch) -> list[dict[str, object]]: + 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( @@ -289,119 +502,191 @@ def fake_foreach_all_gather( "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_single_tensor_single_step( - self, monkeypatch: pytest.MonkeyPatch, single_rank_group: dist.ProcessGroup - ) -> None: - calls = self._patch_foreach_all_gather(monkeypatch) - spec = LoadSpec( - name="layers.0.mlp.gate.weight", - global_hf_keys=["gate"], - global_shape=(4, 2), - shards=[ShardDescriptor(dim=0, start=1, end=3, group=single_rank_group)], - ) - - output = unshard_tensors_for_hf_save( - [torch.ones(2, 2)], - [spec.plan_hf_save()], - ) - - assert [tuple(tensor.shape) for tensor in output] == [(4, 2)] - assert [call["shapes"] for call in calls] == [[(4, 2)]] - - def test_same_group_same_dtype_tensors_are_batched( - self, monkeypatch: pytest.MonkeyPatch, single_rank_group: dist.ProcessGroup + 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="layers.0.mlp.gate.weight", - global_hf_keys=["gate"], - global_shape=(4, 2), - shards=[ShardDescriptor(dim=0, start=1, end=3, group=single_rank_group)], + 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="layers.0.mlp.up.weight", - global_hf_keys=["up"], - global_shape=(6, 2), - shards=[ShardDescriptor(dim=0, start=2, end=5, group=single_rank_group)], + 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(2, 2), torch.ones(3, 2)], + [torch.ones(8, 2), torch.ones(4, 2)], [spec.plan_hf_save() for spec in specs], ) - assert [tuple(tensor.shape) for tensor in output] == [(4, 2), (6, 2)] - assert [call["shapes"] for call in calls] == [[(4, 2), (6, 2)]] + 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_same_group_different_dtype_tensors_are_split( - self, monkeypatch: pytest.MonkeyPatch, single_rank_group: dist.ProcessGroup + 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="layers.0.mlp.gate.weight", - global_hf_keys=["gate"], + name=name, + global_hf_keys=[name], global_shape=(4, 2), - shards=[ShardDescriptor(dim=0, start=1, end=3, group=single_rank_group)], - ), - LoadSpec( - name="layers.0.mlp.up.weight", - global_hf_keys=["up"], - global_shape=(4, 2), - shards=[ShardDescriptor(dim=0, start=1, end=3, group=single_rank_group)], - ), + shards=[ShardDescriptor(dim=0, group=single_rank_group)], + ) + for name in ("gate", "up") ] output = unshard_tensors_for_hf_save( - [torch.ones(2, 2, dtype=torch.float32), torch.ones(2, 2, dtype=torch.float64)], + [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_multi_step_tensor_waits_for_previous_step( - self, monkeypatch: pytest.MonkeyPatch, single_rank_group: dist.ProcessGroup + 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: - calls = self._patch_foreach_all_gather(monkeypatch) - specs = [ - LoadSpec( - name="layers.0.experts.fused_w1w3.weight", - global_hf_keys=["k0", "k1"], - global_shape=(8, 2), - fused_dim=0, - shards=[ - ShardDescriptor(dim=0, start=0, end=4, group=single_rank_group), - ShardDescriptor(dim=0, start=1, end=3, group=single_rank_group), - ], - ), - LoadSpec( - name="layers.0.mlp.gate.weight", - global_hf_keys=["gate"], - global_shape=(4, 2), - shards=[ShardDescriptor(dim=0, start=1, end=3, group=single_rank_group)], - ), - ] + 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.ones(2, 2), torch.ones(2, 2)], - [spec.plan_hf_save() for spec in specs], + [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)], ) - assert [tuple(tensor.shape) for tensor in output] == [(8, 2), (4, 2)] - assert [call["shapes"] for call in calls] == [[(4, 2), (4, 2)], [(8, 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,)] -class TestBaseModelHFSave: - """BaseModel save should preserve state semantics outside LoadSpec.""" + 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: diff --git a/xtuner/v1/model/base.py b/xtuner/v1/model/base.py index 663b7b908f..4887983492 100644 --- a/xtuner/v1/model/base.py +++ b/xtuner/v1/model/base.py @@ -48,7 +48,6 @@ 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 ( - HFLoadPlan, HFSavePlan, LoadSpec, unshard_tensors_for_hf_save, @@ -942,77 +941,16 @@ 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, - load_plan: HFLoadPlan, - ) -> None: - """Copy loaded HF tensors into a local parameter tensor. - - Args: - safetensors (list[torch.Tensor]): HF tensors loaded for ``load_plan.hf_keys``, in key order. - local_tensor (torch.Tensor): Destination local parameter or buffer tensor. - load_plan (HFLoadPlan): Plan whose ``slices`` are relative to ``safetensors`` after concatenation. - """ - loaded_tensor = self._cat_safetensors(safetensors, load_plan) - loaded_tensor = self.hf_tensor_to_canonical(load_plan.name, loaded_tensor) - loaded_tensor = self._apply_load_slices(loaded_tensor, load_plan) - self._copy_loaded_tensor_to_local(loaded_tensor, local_tensor) - def hf_tensor_to_canonical(self, name: str, loaded_tensor: torch.Tensor) -> torch.Tensor: - """Convert one loaded HF tensor to XTuner's canonical layout.""" - return loaded_tensor + """Convert one checkpoint tensor to XTuner's unsharded canonical + layout. - def _cat_safetensors(self, safetensors: list[torch.Tensor], load_plan: HFLoadPlan) -> torch.Tensor: - assert safetensors, f"Internal Error. No safetensors were loaded for {load_plan.name}" - if len(safetensors) > 1: - dim = load_plan.fused_dim - assert dim is not None, "Internal Error dim must not be None when len(safetensors) > 1" - return torch.cat(safetensors, dim=dim) - return safetensors[0] - - def _apply_load_slices(self, loaded_tensor: torch.Tensor, load_plan: HFLoadPlan) -> torch.Tensor: - for load_slice in load_plan.slices: - start = min(load_slice.start, loaded_tensor.shape[load_slice.dim]) - end = min(load_slice.end, loaded_tensor.shape[load_slice.dim]) - assert start <= end, f"Invalid load slice [{start}, {end}) for {load_plan.name}" - loaded_tensor = loaded_tensor.narrow(load_slice.dim, start, end - start) + ``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 _copy_loaded_tensor_to_local(self, loaded_tensor: torch.Tensor, local_tensor: torch.Tensor) -> None: - if loaded_tensor.shape == local_tensor.shape: - local_tensor.copy_(loaded_tensor) - return - - assert loaded_tensor.dim() == local_tensor.dim(), ( - f"Loaded tensor shape {tuple(loaded_tensor.shape)} is incompatible with local tensor shape " - f"{tuple(local_tensor.shape)}" - ) - # HF checkpoints never store FSDP padding. After applying the LoadPlan slices, only the FSDP shard dim may be - # shorter than the runtime local tensor; all other dims must match exactly. - non_pad_dim_matches = all( - loaded_tensor.shape[dim] == local_tensor.shape[dim] - for dim in range(local_tensor.dim()) - if dim != self.FSDP_SHARD_DIM - ) - assert non_pad_dim_matches, ( - f"Loaded tensor shape {tuple(loaded_tensor.shape)} is incompatible with local tensor shape " - f"{tuple(local_tensor.shape)}; padding is only expected on dim {self.FSDP_SHARD_DIM}" - ) - non_pad_len = loaded_tensor.shape[self.FSDP_SHARD_DIM] - assert non_pad_len <= local_tensor.shape[self.FSDP_SHARD_DIM], ( - f"Loaded tensor shape {tuple(loaded_tensor.shape)} is larger than local tensor shape " - f"{tuple(local_tensor.shape)}" - ) - local_tensor.narrow(self.FSDP_SHARD_DIM, 0, non_pad_len).copy_(loaded_tensor) - - if non_pad_len < local_tensor.shape[self.FSDP_SHARD_DIM]: - assert self.config.float8_cfg is not None - pad_len = local_tensor.shape[self.FSDP_SHARD_DIM] - non_pad_len - # Torch casts the scalar to the destination dtype; for fp8 this writes the canonical zero value. - local_tensor.narrow(self.FSDP_SHARD_DIM, non_pad_len, pad_len).copy_(0.0) # type: ignore - def param_to_safetensor( self, safetensor: torch.Tensor, @@ -1374,20 +1312,9 @@ def _get_hf_param( buffer_names = {self._clean_param_name(name) for name, _ in self.named_buffers()} for param, load_spec in params: - # InterleavedShard-bearing DTensors (e.g. fused MoE column-parallel weights) have - # `shard_order=None`; their layout cannot be described by per-step ShardDescriptors. - # Materialize the global tensor up-front via `reconstruct_full_tensor` and treat the - # result as already-unsharded by the rest of the save pipeline (load_spec.shards is - # empty, so `unshard_tensors_for_hf_save` becomes a no-op for these items). - if load_spec.needs_full_reconstruct: - assert isinstance(param, DTensor), ( - f"needs_full_reconstruct=True implies a DTensor param, got {type(param).__name__}" - ) - from xtuner.v1.utils.interleaved_shard import reconstruct_full_tensor - - runtime_tensor = reconstruct_full_tensor(param) - else: - runtime_tensor = param._local_tensor if isinstance(param, DTensor) else param + # 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 runtime_is_float8 = is_float8_weight(runtime_tensor) is_buffer = load_spec.name in buffer_names if runtime_tensor.is_floating_point() and not is_buffer: @@ -1882,23 +1809,14 @@ def _load_hf_param( ) -> list[str]: """Unified HF load path for a single parameter / buffer. - ``LoadSpec.plan_hf_load`` computes this rank's HF keys and loaded-tensor-relative slices from the new - schema. This method only executes that plan: load keys, dequantize fp8 when needed, then hand off to - ``safetensors_to_params`` for cat + narrow + copy. + ``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. 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() - if load_plan.zero_fill: - # No checkpoint key overlaps this rank. This can be fp8 runtime padding, or a legal zero-sized DTensor - # shard when a tiny tensor dimension is split across more ranks than it has elements. - assert load_spec.origin_shape is not None or local_tensor.numel() == 0, ( - "Empty load plan is only legal for runtime pad-only or zero-sized local tensors" - ) - local_tensor.zero_() # type: ignore - return [] missing_keys: list[str] = [] loaded_tensors: list[torch.Tensor] = [] @@ -1920,29 +1838,10 @@ def _load_hf_param( if missing_keys: return missing_keys - if load_spec.needs_full_reconstruct: - # InterleavedShard-style placements: this rank owns N contiguous "runs" of rows in - # the global tensor (one per local expert). Copy each run from the concatenated - # HF tensor to the matching slice of the local tensor. - assert isinstance(param, DTensor), ( - f"needs_full_reconstruct=True implies a DTensor param, got {type(param).__name__}" - ) - from xtuner.v1.utils.interleaved_shard import compute_runs - - loaded_tensor = self._cat_safetensors(loaded_tensors, load_plan) - # Interleaved runs use canonical global coordinates. Packed model - # formats such as GPT-OSS must be converted before applying them. - loaded_tensor = self.hf_tensor_to_canonical(load_plan.name, loaded_tensor) - local = param._local_tensor - for run in compute_runs(param): - loaded_slice = loaded_tensor.narrow(0, run.global_offset[0], run.local_size) - local.narrow(0, run.local_start, run.local_size).copy_(loaded_slice) - return [] - - self.safetensors_to_params( + load_plan.load_into( loaded_tensors, local_tensor, - load_plan, + canonicalize=self.hf_tensor_to_canonical, ) return [] diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py index 028e9de28a..50fbfa3869 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/xtuner/v1/model/moe/moe.py @@ -218,8 +218,8 @@ def __init__(self, config: MoEConfig): ) 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 — needed by GroupedLinear for per-expert column-parallel weights - # so HF save can reconstruct the full tensor via `reconstruct_full_tensor`. + # 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( diff --git a/xtuner/v1/utils/interleaved_shard.py b/xtuner/v1/utils/interleaved_shard.py index 7f3e509078..3f2d317edb 100644 --- a/xtuner/v1/utils/interleaved_shard.py +++ b/xtuner/v1/utils/interleaved_shard.py @@ -19,8 +19,8 @@ * Forward / backward read ``weight.to_local()`` so the op dispatcher is never invoked on InterleavedShard parameters. - * Save / load are routed through :func:`reconstruct_full_tensor` (this module) and the LoadSpec - machinery, neither of which depends on ``shard_order``. + * 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`. @@ -31,7 +31,6 @@ from typing import NamedTuple import torch -import torch.distributed._functional_collectives as funcol from torch.distributed.tensor import DTensor, Shard from torch.distributed.tensor.placement_types import _StridedShard @@ -49,8 +48,8 @@ class Run(NamedTuple): """One contiguous run of global indices that the current rank owns on the sharded dim. - Used by both the HF save path (build per-run WriteItems / per-run slices) and the HF load - path (per-run narrow + copy from the loaded global tensor). + 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. @@ -214,31 +213,12 @@ def _is_real_strided(placement, mesh_dim: int) -> bool: def reconstruct_full_tensor(dt: DTensor) -> torch.Tensor: - """Reconstruct the global tensor from a DTensor's local data, even when the - spec contains placements that PyTorch's ``redistribute`` cannot handle - (``shard_order=None``). + """Reconstruct a full runtime tensor through the shared SavePlan executor. - Why a custom routine: ``DTensor.full_tensor()`` goes through ``redistribute`` which asserts - ``shard_order is not None`` in torch 2.10. For our ``(Shard, InterleavedShard)`` placement - that assert fires. We bypass redistribute by emitting collectives directly. - - Algorithm: - - 1. **Phase 1 — undo FSDP-prepended _StridedShard (mesh_dim 0) as plain Shard.** FSDP2 - actually chunks the parameter contiguously (``_chunk_with_empty``) regardless of the - strided label. So the right undo is a plain ``all_gather`` along the FSDP mesh dim. - After this phase every rank holds the pre-FSDP local. - - 2. **Phase 2 — undo remaining placements in REVERSE mesh-dim order:** - - * ``InterleavedShard`` (= real strided): ``all_gather`` along the placement's mesh dim, - then scatter the gathered chunks back to their correct global positions using - ``_local_shard_size_and_offset(return_first_offset=False)``. - * Plain ``Shard``: ``all_gather`` and concatenate. - - The reverse direction is essential because ``InterleavedShard.split_factor`` is defined - relative to the size of the tensor *after* the placements to its right have already - been undone. Doing TP undo before EP undo keeps the sf math consistent. + 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 @@ -247,45 +227,17 @@ def reconstruct_full_tensor(dt: DTensor) -> torch.Tensor: if not isinstance(dt, DTensor): raise TypeError(f"reconstruct_full_tensor expects a DTensor, got {type(dt).__name__}") - mesh = dt.device_mesh - placements = list(dt.placements) - # Make sure the working buffer is contiguous so all_gather copies see a well-defined layout. - result = dt._local_tensor.contiguous() - - # Phase 1: FSDP-prepended _StridedShard at mesh_dim 0 → plain gather. - for mesh_dim, placement in enumerate(placements): - if not _is_fsdp_prepended_strided(placement, mesh_dim): - continue - result = _all_gather_plain(result, placement.dim, mesh.get_group(mesh_dim)) - - # Phase 2: remaining placements in reverse mesh-dim order. - for mesh_dim in reversed(range(len(placements))): - placement = placements[mesh_dim] - if not isinstance(placement, (Shard, _StridedShard)): - continue - if _is_fsdp_prepended_strided(placement, mesh_dim): - continue # already handled in Phase 1 - if _is_real_strided(placement, mesh_dim): - result = _undo_strided(result, placement, mesh, mesh_dim) - else: - # Plain Shard or _StridedShard with sf == 1 (degenerate). - result = _all_gather_plain(result, placement.dim, mesh.get_group(mesh_dim)) - - return result - - -# --------------------------------------------------------------------------- -# Internal collective helpers -# --------------------------------------------------------------------------- + from xtuner.v1.utils.load_spec import LoadSpec, unshard_tensors_for_hf_save - -def _all_gather_plain(local: torch.Tensor, tensor_dim: int, group) -> torch.Tensor: - """``all_gather_tensor`` along ``tensor_dim`` then materialize the async - wrapper.""" - gathered = funcol.all_gather_tensor(local, gather_dim=tensor_dim, group=group) - if isinstance(gathered, funcol.AsyncCollectiveTensor): - gathered = gathered.wait() - return gathered + 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]: @@ -298,7 +250,7 @@ def compute_runs(dt: DTensor) -> list[Run]: 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, mirroring ``reconstruct_full_tensor`` which undoes FSDP first. + 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. @@ -370,34 +322,3 @@ def compute_runs(dt: DTensor) -> list[Run]: ) ) return runs - - -def _undo_strided( - local: torch.Tensor, - placement, - mesh, - mesh_dim: int, -) -> torch.Tensor: - """``all_gather`` + scatter for a strided placement. - - Each rank in the mesh dim group holds a strided chunk per ``placement``'s spec. After - ``all_gather`` the result is the concatenation of those chunks in rank order. To recover - the original layout we re-index each rank's chunk back to its true positions using - ``_local_shard_size_and_offset(return_first_offset=False)`` which returns the global - indices the rank owned within the post-undo tensor. - """ - tensor_dim = placement.dim - mesh_size = mesh.size(mesh_dim) - group = mesh.get_group(mesh_dim) - - gathered = _all_gather_plain(local, tensor_dim, group) - current_size = gathered.shape[tensor_dim] - - all_indices: list[int] = [] - for r in range(mesh_size): - all_indices.extend(_strided_indices(placement, current_size, mesh_size, r)) - - indices_tensor = torch.tensor(all_indices, device=gathered.device, dtype=torch.long) - new_result = torch.empty_like(gathered) - new_result.index_copy_(tensor_dim, indices_tensor, gathered) - return new_result diff --git a/xtuner/v1/utils/load_spec.py b/xtuner/v1/utils/load_spec.py index 991cac7198..190e3a3579 100644 --- a/xtuner/v1/utils/load_spec.py +++ b/xtuner/v1/utils/load_spec.py @@ -1,5 +1,6 @@ import math -from typing import NamedTuple +from itertools import product +from typing import Callable, NamedTuple import torch import torch.distributed as dist @@ -19,56 +20,95 @@ def _is_same_process_group(left: dist.ProcessGroup, right: dist.ProcessGroup) -> class ShardDescriptor(BaseModel): - """A single partition applied to the fused full tensor. + """One runtime partition applied to the canonical full tensor. - The full tensor is obtained by concatenating every ``LoadSpec.global_hf_keys`` along - ``LoadSpec.fused_dim`` (or taking the sole HF tensor when ``len(global_hf_keys) == 1``). - Descriptors are applied in order; later descriptors use offsets relative to the sub-tensor produced by all - earlier descriptors, matching DTensor placement semantics. + 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. - start (int): Inclusive start offset relative to the current sub-tensor. - end (int): Exclusive end offset relative to the current sub-tensor. 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 - start: int - end: 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 local_size(self, dim_size: int) -> int: + return sum(end - start for start, end in self.local_intervals(dim_size)) -def _dtensor_shards(tensor: DTensor) -> list[ShardDescriptor]: - current_shape = list(tensor.shape) - shards: list[ShardDescriptor] = [] - for mesh_dim, placement in _ordered_dtensor_placements(tensor): - if not isinstance(placement, Shard): - continue - # DTensor placement order is not always the raw mesh-dim order. FSDP2 can represent right-to-left sharding - # with _StridedShard, and PyTorch's checkpoint offset helper first expands that into the effective shard - # order. LoadSpec must preserve the same order so its descriptor intervals match DTensor local tensors. - # - # XTuner may initialize modules while the default device is "meta". PyTorch's Shard placement helpers can - # inherit that default device for temporary shape arithmetic, so force XTuner's real runtime device before - # calling the helper. - with torch.device(get_device()): - local_size, offset = placement._local_shard_size_and_offset( # type: ignore[attr-defined] - current_shape[placement.dim], - tensor.device_mesh.size(mesh_dim), - tensor.device_mesh.get_local_rank(mesh_dim), +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=placement.dim, - start=offset, - end=offset + local_size, + 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 + ), ) ) - current_shape[placement.dim] = local_size return shards @@ -128,50 +168,211 @@ def _ordered_dtensor_placements(tensor: DTensor) -> list[tuple[int, object]]: return ordered -class LoadSlice(BaseModel): - """A narrow operation in the loaded HF tensor coordinate system. +class OwnedRegion(BaseModel): + """One contiguous region of the global tensor owned by this rank. - Args: - dim (int): Tensor dimension to narrow. - start (int): Inclusive start offset in the loaded tensor. - end (int): Exclusive end offset in the loaded tensor. + ``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") - dim: int - start: int - end: int + 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): - """Execution plan for reading HF safetensors into one local tensor. + """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. - slices (list[LoadSlice]): Narrow operations to apply after loading. Offsets are relative to the loaded - tensor, not the original ``LoadSpec.global_shape``. - zero_fill (bool): Whether this rank falls entirely in a padded region and should skip checkpoint reads. + 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 - slices: list[LoadSlice] = Field(default_factory=list) - zero_fill: bool = False + 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, + safetensors: 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(safetensors) == len(self.hf_keys), ( + f"Loaded {len(safetensors)} 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 -def _final_intervals( + assert safetensors, f"Internal Error. No safetensors were loaded for {self.name}" + if len(safetensors) == 1: + loaded_tensor = safetensors[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(safetensors, 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}" + ) + + 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) + + @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[tuple[int, int]]: - intervals = [(0, dim_size) for dim_size in global_shape] +) -> 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: - current_start, _ = intervals[shard.dim] - intervals[shard.dim] = (current_start + shard.start, current_start + shard.end) - return intervals + 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 _local_shape_for_shards( + global_shape: tuple[int, ...], + shards: list[ShardDescriptor], + *, + visible_shape: tuple[int, ...] | None = None, +) -> tuple[int, ...]: + segments_by_dim = _layout_segments(global_shape, shards) + if visible_shape is None: + return tuple(sum(size for _, size in segments) for segments in segments_by_dim) + + assert len(visible_shape) == len(global_shape) + 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 _owned_regions_for_shards( + global_shape: tuple[int, ...], + shards: list[ShardDescriptor], + *, + visible_shape: tuple[int, ...], +) -> list[OwnedRegion]: + """Compile rank ownership into rectangular global-to-local copies.""" + segments_by_dim = _layout_segments(global_shape, shards) + 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): @@ -198,7 +399,6 @@ class SaveShardStep(BaseModel): load_spec_shard_index (int): Index of ``shard`` in the original ``LoadSpec.shards`` list. shard (ShardDescriptor): Shard descriptor this save step reverses. shape_before_shard (tuple[int, ...]): Runtime tensor shape immediately before ``shard`` was applied. - unpadded_shape_before_shard (tuple[int, ...]): Checkpoint-visible shape before ``shard`` was applied. preserved (bool): Whether this shard should remain applied instead of being all-gathered. """ @@ -206,7 +406,6 @@ class SaveShardStep(BaseModel): load_spec_shard_index: int shard: ShardDescriptor shape_before_shard: tuple[int, ...] - unpadded_shape_before_shard: tuple[int, ...] preserved: bool = False @@ -218,6 +417,8 @@ class HFSavePlan(BaseModel): hf_keys (list[str]): HF keys represented by the tensor after this plan's pending unshard steps finish. global_shape (tuple[int, ...]): Runtime full tensor shape before any shard is applied. unpadded_global_shape (tuple[int, ...]): Checkpoint-visible full tensor shape after removing runtime padding. + 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 @@ -232,6 +433,8 @@ class HFSavePlan(BaseModel): hf_keys: list[str] global_shape: tuple[int, ...] unpadded_global_shape: tuple[int, ...] + runtime_output_shape: tuple[int, ...] + output_shape: tuple[int, ...] fused_dim: int | None = None distributed_save: bool = False preserves_shards: bool = False @@ -240,39 +443,6 @@ class HFSavePlan(BaseModel): def _pending_unshard_steps(self) -> list[SaveShardStep]: return [step for step in reversed(self.unshard_steps) if not step.preserved] - def _preserved_shards(self) -> list[ShardDescriptor]: - return [step.shard for step in self.unshard_steps if step.preserved] - - def _expected_unsharded_shape(self) -> tuple[int, ...]: - """Return the save tensor shape after intentionally preserved shards - remain applied. - - The save path starts from the local tensor and all-gathers every pending shard step. If no shard is preserved, - the final shape should be ``unpadded_global_shape``. If some shards are preserved, for example an EP shard - during RL weight sync, the final tensor should still be cut by those preserved shards. This helper applies - only the preserved shard descriptors to ``unpadded_global_shape`` to compute that expected partially-unsharded - shape for the final assert. - - Example: - Suppose the runtime full tensor shape is ``(16, 8)`` because fp8 padding added rows, while - ``unpadded_global_shape == (14, 8)`` is the shape that should exist in HF. If the preserved EP shard is - ``ShardDescriptor(dim=0, start=8, end=16)``, that shard owns runtime rows ``[8, 16)``. The last two rows - are padding-only in HF coordinates, so the checkpoint-visible interval is clipped to ``[8, 14)`` and the - expected preserved tensor shape is ``(6, 8)``. If a shard were ``[14, 16)``, both boundaries would clip to - ``14`` and the expected shape on that rank would be ``(0, 8)``. - - Returns: - tuple[int, ...]: Expected shape after the preserved shards are still applied. - """ - effective_shape = list(self.unpadded_global_shape) - for shard in self._preserved_shards(): - # ShardDescriptor offsets are defined against the runtime shape, which may include XTuner-only padding. - # Clip preserved shard boundaries to the currently visible unpadded shape before computing its length. - clipped_start = min(shard.start, effective_shape[shard.dim]) - clipped_end = min(shard.end, effective_shape[shard.dim]) - effective_shape[shard.dim] = max(0, clipped_end - clipped_start) - return tuple(effective_shape) - class _SaveUnshardGroup(NamedTuple): """One compatible foreach all-gather batch in the save unshard loop. @@ -344,12 +514,22 @@ def unshard_tensors_for_hf_save( for index, gathered_tensor in zip(unshard_group.tensor_indices, gathered_tensors, strict=True): tensor_list[index] = gathered_tensor - for tensor, save_plan in zip(tensor_list, save_plans, strict=True): - expected_shape = save_plan._expected_unsharded_shape() - assert tuple(tensor.shape) == expected_shape, ( - f"Saved tensor shape {tuple(tensor.shape)} is incompatible with HFSavePlan global_shape=" - f"{save_plan.global_shape} and unpadded_global_shape={save_plan.unpadded_global_shape} " - f"for {save_plan.name}" + # Collective steps reconstruct runtime shapes only. Remove checkpoint-invisible + # FP8 tail padding once, after every requested gather has completed. + for index, (tensor, save_plan) in enumerate(zip(tensor_list, save_plans, strict=True)): + 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}" + ) + for dim, output_size in enumerate(save_plan.output_shape): + assert output_size <= tensor.shape[dim] + if output_size < tensor.shape[dim]: + tensor = tensor.narrow(dim, 0, output_size) + tensor_list[index] = tensor.contiguous() + + assert tuple(tensor.shape) == save_plan.output_shape, ( + f"Saved tensor shape {tuple(tensor.shape)} is incompatible with HFSavePlan output_shape=" + f"{save_plan.output_shape} for {save_plan.name}" ) return tensor_list @@ -420,6 +600,17 @@ def _pad_tensor_for_save_shard(tensor: torch.Tensor, shard_step: SaveShardStep) 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 = shard_step.shard.local_size(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, ( @@ -440,8 +631,22 @@ def _merge_gathered_save_shard( shard_step: SaveShardStep, ) -> torch.Tensor: dim = shard_step.shard.dim - gathered_tensor = torch.cat(gathered_chunks, dim=dim) - return gathered_tensor.narrow(dim, 0, shard_step.unpadded_shape_before_shard[dim]).contiguous() + 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): @@ -458,7 +663,8 @@ class LoadSpec(BaseModel): 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. - needs_full_reconstruct (bool): Whether HF I/O must use the explicit InterleavedShard reconstruction/run path. + 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") @@ -468,12 +674,7 @@ class LoadSpec(BaseModel): fused_dim: int | None = None shards: list[ShardDescriptor] = Field(default_factory=list) origin_shape: tuple[int, ...] | None = None - # When True, this tensor's layout cannot be described by the ``shards`` list — typically an - # ``InterleavedShard``-bearing DTensor whose spec has ``shard_order=None``. The HF save path - # must call :func:`xtuner.v1.utils.interleaved_shard.reconstruct_full_tensor` on the param at - # save time to materialize the global tensor, and treat the result as already-unsharded - # (i.e. ``shards`` is empty, no per-step all-gather work needed). - needs_full_reconstruct: bool = False + local_shape: tuple[int, ...] | None = None @computed_field # type: ignore[prop-decorator] @property @@ -517,14 +718,8 @@ def from_tensor( LoadSpec: Spec derived from the runtime tensor layout. """ global_hf_keys = list(hf_keys) - shards: list[ShardDescriptor] = [] - needs_full_reconstruct = False - if isinstance(tensor, DTensor): - from xtuner.v1.utils.interleaved_shard import has_interleaved_placement - - needs_full_reconstruct = has_interleaved_placement(tensor) - if not needs_full_reconstruct: - shards = _dtensor_shards(tensor) + 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, @@ -532,7 +727,7 @@ def from_tensor( fused_dim=0 if len(global_hf_keys) > 1 else None, shards=shards, origin_shape=origin_shape, - needs_full_reconstruct=needs_full_reconstruct, + local_shape=tuple(local_tensor.shape), ) def plan_hf_load(self) -> HFLoadPlan: @@ -543,38 +738,69 @@ def plan_hf_load(self) -> HFLoadPlan: runtime layout that this rank owns. Returns: - HFLoadPlan: The selected HF keys and loaded-tensor-relative slices for this rank. + HFLoadPlan: The selected HF keys and canonical source-to-local copy program for this rank. """ - effective_intervals = self._effective_intervals_for_shards(self.shards) - if effective_intervals is None: - return HFLoadPlan(name=self.name, hf_keys=[], fused_dim=self.fused_dim, zero_fill=True) + target_shape = self._runtime_local_shape() + owned_regions = _owned_regions_for_shards( + self.global_shape, + self.shards, + 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, + ) - loaded_starts = [0 for _ in self.global_shape] - loaded_ends = list(self.unpadded_global_shape) - key_start, key_end = self._local_hf_key_indices(effective_intervals) + # Select the smallest contiguous HF-key envelope covering every owned region. Ordinary FSDP/EP produces one + # region. Interleaved Expert TP produces multiple runs, but their envelope still lets EP ranks avoid reading + # experts owned by other EP ranks when the checkpoint stores one key per expert. + envelope = [ + ( + min(region.global_offsets[dim] for region in owned_regions), + max(region.global_offsets[dim] + region.sizes[dim] for region in owned_regions), + ) + for dim in range(len(self.global_shape)) + ] + key_start, key_end = self._local_hf_key_indices(envelope) hf_keys = self.global_hf_keys[key_start:key_end] + loaded_starts = [0 for _ in self.global_shape] + loaded_ends = list(self.unpadded_global_shape) if self.is_fused: key_size = self._fused_key_size() assert self.fused_dim is not None loaded_starts[self.fused_dim] = key_start * key_size loaded_ends[self.fused_dim] = key_end * key_size - slices: list[LoadSlice] = [] - for dim, (effective_start, effective_end) in enumerate(effective_intervals): - loaded_start = loaded_starts[dim] - loaded_end = loaded_ends[dim] - if effective_start == loaded_start and effective_end == loaded_end: - continue - slices.append( - LoadSlice( - dim=dim, - start=effective_start - loaded_start, - end=effective_end - loaded_start, - ) + copy_regions = [ + LoadCopyRegion( + source_offsets=tuple( + region.global_offsets[dim] - loaded_starts[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, slices=slices) + return HFLoadPlan( + name=self.name, + hf_keys=hf_keys, + fused_dim=self.fused_dim, + canonical_source_shape=tuple(end - start for start, end in zip(loaded_starts, loaded_ends)), + target_shape=target_shape, + copy_regions=copy_regions, + zero_unwritten_target=copied_numel < target_numel, + ) def plan_hf_save( self, @@ -605,10 +831,18 @@ def plan_hf_save( ) unshard_steps = self._save_shard_steps(preserved_shard_indices) preserved_shards = [step.shard for step in unshard_steps if step.preserved] - hf_keys = ( - self._local_hf_keys_for_shards(preserved_shards, require_fused_key_aligned=True) - if preserved_shards - else list(self.global_hf_keys) + 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) + + runtime_output_shape = _local_shape_for_shards(self.global_shape, preserved_shards) + output_shape = _local_shape_for_shards( + self.global_shape, + preserved_shards, + visible_shape=self.unpadded_global_shape, ) return HFSavePlan( @@ -616,6 +850,8 @@ def plan_hf_save( hf_keys=hf_keys, global_shape=self.global_shape, unpadded_global_shape=self.unpadded_global_shape, + runtime_output_shape=runtime_output_shape, + output_shape=output_shape, fused_dim=self.fused_dim, distributed_save=distributed_save, preserves_shards=bool(preserved_shards), @@ -630,27 +866,14 @@ def model_post_init(self, _) -> None: self._validate_origin_shape() self._validate_shards() - def _effective_intervals_for_shards( - self, - shards: list[ShardDescriptor], - ) -> list[tuple[int, int]] | None: - effective_shape = self.unpadded_global_shape - assert len(effective_shape) == len(self.global_shape), ( - f"origin_shape={effective_shape} must have the same rank as global_shape={self.global_shape}" - ) - assert all(effective <= global_ for effective, global_ in zip(effective_shape, self.global_shape)), ( - f"origin_shape={effective_shape} must not exceed global_shape={self.global_shape}" - ) - - final_intervals = _final_intervals(self.global_shape, shards) - effective_intervals: list[tuple[int, int]] = [] - for dim, (start, end) in enumerate(final_intervals): - effective_start = min(start, effective_shape[dim]) - effective_end = min(end, effective_shape[dim]) - if effective_start >= effective_end: - return None - effective_intervals.append((effective_start, effective_end)) - return effective_intervals + def _runtime_local_shape(self) -> tuple[int, ...]: + derived_shape = _local_shape_for_shards(self.global_shape, self.shards) + if self.local_shape is not None: + assert self.local_shape == derived_shape, ( + f"Recorded local_shape={self.local_shape} does not match descriptor-derived shape " + f"{derived_shape} for {self.name}" + ) + return derived_shape 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" @@ -680,7 +903,7 @@ def _local_hf_key_indices( ) # 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; LoadSlice later trims load tensors to the exact local range. + # 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), ( @@ -694,11 +917,22 @@ def _local_hf_keys_for_shards( *, require_fused_key_aligned: bool = False, ) -> list[str]: - effective_intervals = self._effective_intervals_for_shards(shards) - if effective_intervals is None: + regions = _owned_regions_for_shards( + self.global_shape, + shards, + visible_shape=self.unpadded_global_shape, + ) + if not regions: return [] + envelope = [ + ( + min(region.global_offsets[dim] for region in regions), + max(region.global_offsets[dim] + region.sizes[dim] for region in regions), + ) + for dim in range(len(self.global_shape)) + ] key_start, key_end = self._local_hf_key_indices( - effective_intervals, + envelope, require_fused_key_aligned=require_fused_key_aligned, ) return self.global_hf_keys[key_start:key_end] @@ -720,11 +954,12 @@ def _validate_shards(self) -> None: assert 0 <= shard.dim < len(current_shape), ( f"Invalid shard dim {shard.dim} for global_shape={self.global_shape}" ) - current_size = current_shape[shard.dim] - assert 0 <= shard.start <= shard.end <= current_size, ( - f"Invalid shard descriptor {shard} against current_shape={tuple(current_shape)}" - ) - current_shape[shard.dim] = shard.end - shard.start + current_shape[shard.dim] = shard.local_size(current_shape[shard.dim]) + + assert self.local_shape is None or tuple(current_shape) == self.local_shape, ( + f"Recorded local_shape={self.local_shape} does not match descriptor-derived shape " + f"{tuple(current_shape)} for {self.name}" + ) def _preserved_shard_indices( self, @@ -778,29 +1013,9 @@ def _save_shard_steps(self, preserved_shard_indices: set[int]) -> list[SaveShard """Convert ``LoadSpec.shards`` into save-time reverse-unshard work items. - ``LoadSpec.shards`` is ordered in the forward partitioning direction: start from the full runtime tensor, - apply one shard after another, and end at this rank's local tensor. The returned steps keep that same - largest-to-smallest order. Each step snapshots the runtime shape and the unpadded checkpoint-visible shape - that existed immediately before its shard was applied. - - Save executes these steps in reverse. Starting from the smallest local tensor, each reverse step all-gathers - one shard and narrows the gathered tensor back to ``unpadded_shape_before_shard``. This is how the save path - reconstructs the original shape information one partition layer at a time, while still avoiding fp8 runtime - padding in the checkpoint-visible tensor. - - Example: - Suppose ``global_shape=(16, 8)``, ``unpadded_global_shape=(14, 8)``, and - ``LoadSpec.shards == [ep(dim=0, start=8, end=16), fsdp(dim=0, start=3, end=5)]``. The returned steps are - in forward order: - - * ``ep_step`` records ``shape_before_shard=(16, 8)`` and - ``unpadded_shape_before_shard=(14, 8)``. - * ``fsdp_step`` records ``shape_before_shard=(8, 8)`` and - ``unpadded_shape_before_shard=(6, 8)``. - - A local save tensor has shape ``(2, 8)``. Save runs ``[fsdp_step, ep_step]``: gather FSDP back toward - ``(6, 8)``, then gather EP back toward ``(14, 8)``. If EP is preserved, only ``fsdp_step`` remains - pending and the result stays EP-local. + ``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. @@ -809,7 +1024,6 @@ def _save_shard_steps(self, preserved_shard_indices: set[int]) -> list[SaveShard list[SaveShardStep]: Work items in the same largest-to-smallest order as ``LoadSpec.shards``. """ current_shape = list(self.global_shape) - effective_shape = list(self.unpadded_global_shape) steps: list[SaveShardStep] = [] for shard_index, shard in enumerate(self.shards): @@ -818,12 +1032,8 @@ def _save_shard_steps(self, preserved_shard_indices: set[int]) -> list[SaveShard load_spec_shard_index=shard_index, shard=shard, shape_before_shard=tuple(current_shape), - unpadded_shape_before_shard=tuple(effective_shape), preserved=shard_index in preserved_shard_indices, ) ) - effective_start = min(shard.start, effective_shape[shard.dim]) - effective_end = min(shard.end, effective_shape[shard.dim]) - effective_shape[shard.dim] = max(0, effective_end - effective_start) - current_shape[shard.dim] = shard.end - shard.start + current_shape[shard.dim] = shard.local_size(current_shape[shard.dim]) return steps From c01d12890d8b6685b7aad702cf99acb98e4641c7 Mon Sep 17 00:00:00 2001 From: zhaopenghao Date: Tue, 11 Aug 2026 14:38:09 +0000 Subject: [PATCH 6/7] [Refactor] Simplify HF load and save flows --- tests/utils/test_interleaved_shard.py | 358 +++++++++++--------------- xtuner/v1/model/base.py | 168 +++++++----- xtuner/v1/utils/load_spec.py | 275 ++++++++------------ 3 files changed, 369 insertions(+), 432 deletions(-) diff --git a/tests/utils/test_interleaved_shard.py b/tests/utils/test_interleaved_shard.py index 70843ff8c3..57ce140cac 100644 --- a/tests/utils/test_interleaved_shard.py +++ b/tests/utils/test_interleaved_shard.py @@ -1,24 +1,13 @@ -"""Unit tests for ``xtuner.v1.utils.interleaved_shard``. +"""Distributed behavior tests for ``xtuner.v1.utils.interleaved_shard``. -These tests cover the InterleavedShard placement and the ``reconstruct_full_tensor`` helper -across the layouts that XTuner actually uses: - - * Plain ``(Shard, InterleavedShard)`` on a 2D (ep, tp) mesh — the layout produced by - ``GroupedLinear`` when TP is enabled. - * The post-``fully_shard`` 3D layout with FSDP prepended on top — what HF save sees in - practice. - -Run with:: - - torchrun --nproc-per-node=8 tests/utils/test_interleaved_shard.py +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 importlib.util as _ilu -import os import shutil -import sys import tempfile import torch @@ -28,223 +17,182 @@ from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard from torch.distributed.tensor import DTensor, Shard, distribute_tensor - -_HERE = os.path.dirname(os.path.abspath(__file__)) -_REPO_ROOT = os.path.dirname(os.path.dirname(_HERE)) -sys.path.insert(0, _REPO_ROOT) -# Import the module directly to avoid pulling in xtuner package's heavy deps (loguru etc.) that -# aren't required for this unit test. -_spec = _ilu.spec_from_file_location( - "interleaved_shard", - os.path.join(_REPO_ROOT, "xtuner", "v1", "utils", "interleaved_shard.py"), +from xtuner._testing import DeterministicDDPTestCase +from xtuner.v1.utils.interleaved_shard import ( + InterleavedShard, + has_interleaved_placement, + reconstruct_full_tensor, ) -assert _spec is not None and _spec.loader is not None -_mod = _ilu.module_from_spec(_spec) -_spec.loader.exec_module(_mod) -InterleavedShard = _mod.InterleavedShard -has_interleaved_placement = _mod.has_interleaved_placement -reconstruct_full_tensor = _mod.reconstruct_full_tensor NUM_EXPERTS = 4 OUT_PER_EXPERT = 4 IN_FEATURES = 8 -GLOBAL_ROWS = NUM_EXPERTS * OUT_PER_EXPERT # 16 +GLOBAL_ROWS = NUM_EXPERTS * OUT_PER_EXPERT def _build_expected_local( - g: torch.Tensor, + global_tensor: torch.Tensor, ep_rank: int, tp_rank: int, ep_size: int, tp_size: int, ) -> torch.Tensor: - """Hand-computed per-expert column parallel slice.""" + """Return the hand-computed per-expert column-parallel slice.""" experts_per_ep = NUM_EXPERTS // ep_size - rows_per_expert = g.shape[0] // NUM_EXPERTS + 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(g[row_start : row_start + rows_per_tp_per_expert]) + chunks.append(global_tensor[row_start : row_start + rows_per_tp_per_expert]) return torch.cat(chunks, dim=0) -def test_2d_layout_and_reconstruct(): - """Build a DTensor on (ep, tp) with (Shard, InterleavedShard) and reconstruct.""" - 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") - - g = torch.arange(GLOBAL_ROWS * IN_FEATURES, device="cuda", dtype=torch.float32).reshape( - GLOBAL_ROWS, IN_FEATURES - ) - dist.broadcast(g, src=0) - - placements = (Shard(0), InterleavedShard(0, num_local_stripes=NUM_EXPERTS // 2)) - dt = distribute_tensor(g, mesh, placements) - - # Layout correctness: per-rank local matches hand-computed per-expert column parallel. - expected = _build_expected_local(g, ep_rank, tp_rank, 2, 2) - assert torch.allclose(dt.to_local(), expected), ( - f"rank {dist.get_rank()} local mismatch" - ) - - # Detection helper works on this placement. ``shard_order`` only exists on torch>=2.10; - # the implementation guards with ``getattr`` so the test must too. - assert has_interleaved_placement(dt), "shard_order should be None for this placement" - assert getattr(dt._spec, "shard_order", None) is None - - # Reconstruct gives back the global tensor. - full = reconstruct_full_tensor(dt) - assert torch.allclose(full, g), ( - f"reconstruct mismatch on 2D layout: max_diff={(full - g).abs().max().item()}" - ) - - -def test_hf_round_trip(): - """Exercise InterleavedShard through BaseModel's public HF save/load API.""" - 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] - - 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_dir = tempfile.mkdtemp() if dist.get_rank() == 0 else None - checkpoint_dirs = [checkpoint_dir] - 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) - assert torch.equal(restored.weight.to_local(), model.weight.to_local()) - finally: - dist.barrier() - if dist.get_rank() == 0: - shutil.rmtree(checkpoint_dir) +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): + def __init__(self, weight: DTensor): super().__init__() self.weight = nn.Parameter(weight) - def forward(self, x): - w = self.weight.to_local() if isinstance(self.weight, DTensor) else self.weight - return torch.nn.functional.linear(x, w) - - -def test_post_fully_shard_reconstruct(): - """Layout after FSDP wraps the (ep, tp) DTensor — the case HF save actually sees.""" - mesh = init_device_mesh("cuda", (2, 2, 2), mesh_dim_names=("fsdp", "ep", "tp")) - ep_tp = mesh["ep", "tp"] - fsdp_mesh = mesh["fsdp"] - - g = torch.arange(GLOBAL_ROWS * IN_FEATURES, device="cuda", dtype=torch.float32).reshape( - GLOBAL_ROWS, IN_FEATURES - ) - dist.broadcast(g, src=0) - - placements = (Shard(0), InterleavedShard(0, num_local_stripes=NUM_EXPERTS // 2)) - dt = distribute_tensor(g, ep_tp, placements) - - model = _ToyGroupedLinear(dt).cuda() - fully_shard( - model, - mesh=fsdp_mesh, - mp_policy=MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32), - reshard_after_forward=True, - ) - - # Sanity: a forward pass through the wrapped model still produces the right output. - x = torch.randn(6, IN_FEATURES, device="cuda", dtype=torch.bfloat16) - dist.broadcast(x, src=0) - y = model(x) - ep_rank = mesh.get_local_rank("ep") - tp_rank = mesh.get_local_rank("tp") - expected_local = _build_expected_local(g, ep_rank, tp_rank, 2, 2).to(torch.bfloat16) - expected_y = torch.nn.functional.linear(x, expected_local) - assert torch.allclose(y.detach(), expected_y, atol=1e-2, rtol=1e-2) - y.sum().backward() - - # Detection helper still recognizes the wrapped DTensor. - assert has_interleaved_placement(model.weight) - - # Reconstruct from the post-FSDP local matches the original global. - full = reconstruct_full_tensor(model.weight) - assert torch.allclose(full, g), ( - f"reconstruct mismatch on post-FSDP layout: max_diff={(full - g).abs().max().item()}" - ) - - # Exercise the public load plan on the post-FSDP layout. LoadSpec converts - # InterleavedShard runs into a generic canonical source-to-local copy program. - from xtuner.v1.utils.load_spec import LoadSpec - - local = model.weight._local_tensor - loaded_local = torch.empty_like(local, dtype=g.dtype) - load_spec = LoadSpec.from_tensor(name="weight", hf_keys=["weight"], tensor=model.weight) - load_spec.plan_hf_load().load_into([g], loaded_local, lambda _, tensor: tensor) - expected_local = local.to(g.dtype) - assert torch.allclose(loaded_local, expected_local), ( - f"load plan mismatch on post-FSDP layout: " - f"max_diff={(loaded_local - expected_local).abs().max().item()}" - ) - - -def main(): - local_rank = int(os.environ["LOCAL_RANK"]) - torch.cuda.set_device(local_rank) - dist.init_process_group(backend="nccl") - world = dist.get_world_size() - - rank = dist.get_rank() - if world == 4: - test_2d_layout_and_reconstruct() - test_hf_round_trip() - if rank == 0: - print("[2d_layout_and_hf_round_trip] PASSED", flush=True) - elif world == 8: - test_post_fully_shard_reconstruct() - if rank == 0: - print("[post_fully_shard_reconstruct] PASSED", flush=True) - else: - if rank == 0: - print( - f"World size {world} not handled (expected 4 or 8). Skipping.", flush=True - ) - dist.destroy_process_group() - sys.exit(0) - - dist.barrier() - dist.destroy_process_group() - - -if __name__ == "__main__": - main() + 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/xtuner/v1/model/base.py b/xtuner/v1/model/base.py index 4887983492..24a7959f6d 100644 --- a/xtuner/v1/model/base.py +++ b/xtuner/v1/model/base.py @@ -534,6 +534,7 @@ class _HFSaveBucketItem(NamedTuple): tensor: torch.Tensor save_plan: HFSavePlan runtime_is_float8: bool + byte_size: int class BaseModel(nn.Module): @@ -666,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 @@ -870,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( @@ -1307,42 +1313,30 @@ def _get_hf_param( if bucket_size is None: bucket_size = self.config.hf_save_cfg.bucket_size - safetensor_size = 0 + 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: - # 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 - runtime_is_float8 = is_float8_weight(runtime_tensor) - is_buffer = load_spec.name in buffer_names - if runtime_tensor.is_floating_point() and not is_buffer: - save_dtype = self._get_save_dtype(load_spec.global_hf_keys[0], torch.bfloat16) - local_tensor = runtime_tensor.to(dtype=save_dtype) - else: - # Persistent buffers, e.g. FoPE rotary coefficients, are part of HF state but are not trainable - # weights. Keep the legacy behavior and write them in their runtime dtype instead of save_dtype. - local_tensor = runtime_tensor - tensor_size = self._get_tensor_size(runtime_tensor, dtype) - - if safetensor_size + tensor_size > bucket_size and bucket: + 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, ) - safetensor_size = 0 + bucket_bytes = 0 bucket = [] - safetensor_size += tensor_size - save_plan = load_spec.plan_hf_save( - distributed_save=distributed_save, - preserve_process_group=preserved_fused_shard_group, - ) - bucket.append( - _HFSaveBucketItem(tensor=local_tensor, save_plan=save_plan, runtime_is_float8=runtime_is_float8) - ) + bucket_bytes += save_item.byte_size + bucket.append(save_item) if bucket: yield self._build_hf_param_bucket( @@ -1351,6 +1345,38 @@ def _get_hf_param( device=device, ) + 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), + ) + def _load_spec_params(self) -> list[tuple[torch.Tensor, LoadSpec]]: ret: list[tuple[torch.Tensor, LoadSpec]] = [] for name, param in self.state_dict().items(): @@ -1785,25 +1811,6 @@ 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_hf_param( self, param: torch.Tensor, load_spec: LoadSpec, checkpoint_loader: HFCheckpointLoader ) -> list[str]: @@ -1817,24 +1824,11 @@ def _load_hf_param( """ local_tensor = param._local_tensor if isinstance(param, DTensor) else param load_plan = load_spec.plan_hf_load() - - missing_keys: list[str] = [] - loaded_tensors: list[torch.Tensor] = [] - for hf_key in load_plan.hf_keys: - 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)." - ) - weight = self._load_fp8(hf_key, checkpoint_loader) - else: - weight = checkpoint_loader.load(hf_key) - if weight is None: - missing_keys.append(hf_key) - continue - loaded_tensors.append(weight.to(local_tensor.device)) - + loaded_tensors, missing_keys = self._read_hf_tensors( + load_plan.hf_keys, + checkpoint_loader, + device=local_tensor.device, + ) if missing_keys: return missing_keys @@ -1845,6 +1839,54 @@ def _load_hf_param( ) return [] + 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] = [] + for hf_key in hf_keys: + tensor = self._read_hf_tensor(hf_key, checkpoint_loader) + if tensor is None: + missing_keys.append(hf_key) + else: + loaded_tensors.append(tensor.to(device)) + return loaded_tensors, missing_keys + + 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) + + 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 + + from xtuner.v1.float8.triton_kernels import per_block_dequant_gemm + + return per_block_dequant_gemm( + loaded_tensor_fp8.to(DEVICE), + loaded_tensor_scales.to(DEVICE), + ) + def _has_meta_param(self, module: nn.Module, recurse: bool = False) -> bool: """Check if the module has meta parameters.""" for data in chain(module.parameters(recurse=recurse), module.buffers(recurse=False)): diff --git a/xtuner/v1/utils/load_spec.py b/xtuner/v1/utils/load_spec.py index 190e3a3579..f849ed8dca 100644 --- a/xtuner/v1/utils/load_spec.py +++ b/xtuner/v1/utils/load_spec.py @@ -68,9 +68,6 @@ def local_intervals(self, dim_size: int) -> list[tuple[int, int]]: for run_index in range(self.interleave_factor) ] - def local_size(self, dim_size: int) -> int: - return sum(end - start for start, end in self.local_intervals(dim_size)) - def _dtensor_shards(tensor: DTensor) -> list[ShardDescriptor]: try: @@ -168,7 +165,7 @@ def _ordered_dtensor_placements(tensor: DTensor) -> list[tuple[int, object]]: return ordered -class OwnedRegion(BaseModel): +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, @@ -218,7 +215,7 @@ class HFLoadPlan(BaseModel): @torch.no_grad() def load_into( self, - safetensors: list[torch.Tensor], + checkpoint_tensors: list[torch.Tensor], local_tensor: torch.Tensor, canonicalize: Callable[[str, torch.Tensor], torch.Tensor], ) -> None: @@ -231,8 +228,8 @@ def load_into( f"Load target shape {tuple(local_tensor.shape)} does not match planned shape " f"{self.target_shape} for {self.name}" ) - assert len(safetensors) == len(self.hf_keys), ( - f"Loaded {len(safetensors)} tensors for {len(self.hf_keys)} planned HF keys of {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: @@ -240,14 +237,27 @@ def load_into( if not self.copy_regions: return - assert safetensors, f"Internal Error. No safetensors were loaded for {self.name}" - if len(safetensors) == 1: - loaded_tensor = safetensors[0] + 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(safetensors, dim=self.fused_dim) + 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 @@ -255,11 +265,7 @@ def load_into( f"Canonical HF tensor shape {tuple(canonical_tensor.shape)} does not match planned shape " f"{self.canonical_source_shape} for {self.name}" ) - - 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) + return canonical_tensor @staticmethod def _narrow_region( @@ -310,17 +316,15 @@ def _layout_segments( return segments_by_dim -def _local_shape_for_shards( - global_shape: tuple[int, ...], - shards: list[ShardDescriptor], +def _shape_from_segments( + segments_by_dim: list[list[tuple[int, int]]], *, visible_shape: tuple[int, ...] | None = None, ) -> tuple[int, ...]: - segments_by_dim = _layout_segments(global_shape, shards) if visible_shape is None: return tuple(sum(size for _, size in segments) for segments in segments_by_dim) - assert len(visible_shape) == len(global_shape) + assert len(visible_shape) == len(segments_by_dim) return tuple( sum( max(0, min(global_start + size, visible_size) - min(global_start, visible_size)) @@ -330,14 +334,12 @@ def _local_shape_for_shards( ) -def _owned_regions_for_shards( - global_shape: tuple[int, ...], - shards: list[ShardDescriptor], +def _visible_regions_from_segments( + segments_by_dim: list[list[tuple[int, int]]], *, visible_shape: tuple[int, ...], -) -> list[OwnedRegion]: +) -> list[_OwnedRegion]: """Compile rank ownership into rectangular global-to-local copies.""" - segments_by_dim = _layout_segments(global_shape, shards) if any(not segments for segments in segments_by_dim): return [] @@ -350,7 +352,7 @@ def _owned_regions_for_shards( local_offset += size located_segments.append(current) - regions: list[OwnedRegion] = [] + regions: list[_OwnedRegion] = [] for segment_tuple in product(*located_segments): global_offsets: list[int] = [] local_offsets: list[int] = [] @@ -366,7 +368,7 @@ def _owned_regions_for_shards( sizes.append(clipped_size) else: regions.append( - OwnedRegion( + _OwnedRegion( global_offsets=tuple(global_offsets), local_offsets=tuple(local_offsets), sizes=tuple(sizes), @@ -383,12 +385,6 @@ class SaveShardStep(BaseModel): 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. - ``load_spec_shard_index`` is only needed when some original shards should stay sharded. RL weight sync preserves - the EP shard on the fused HF dimension so each EP rank streams only its local expert keys, while later shards such - as FSDP still need to be all-gathered. Because execution reverses and groups the work items, their list positions - no longer match ``LoadSpec.shards``. The original index is the stable handle used by the save plan to decide - which work items to skip and which preserved shards should define the final expected shape. - 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. @@ -396,14 +392,12 @@ class SaveShardStep(BaseModel): EP-local. Args: - load_spec_shard_index (int): Index of ``shard`` in the original ``LoadSpec.shards`` list. 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") - load_spec_shard_index: int shard: ShardDescriptor shape_before_shard: tuple[int, ...] preserved: bool = False @@ -415,33 +409,28 @@ class HFSavePlan(BaseModel): 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. - global_shape (tuple[int, ...]): Runtime full tensor shape before any shard is applied. - unpadded_global_shape (tuple[int, ...]): Checkpoint-visible full tensor shape after removing runtime padding. 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. - preserves_shards (bool): Whether the save tensor intentionally remains sharded by some original - ``LoadSpec.shards`` entries. 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] - global_shape: tuple[int, ...] - unpadded_global_shape: tuple[int, ...] runtime_output_shape: tuple[int, ...] output_shape: tuple[int, ...] fused_dim: int | None = None distributed_save: bool = False - preserves_shards: bool = False unshard_steps: list[SaveShardStep] = Field(default_factory=list) - def _pending_unshard_steps(self) -> list[SaveShardStep]: - return [step for step in reversed(self.unshard_steps) if not step.preserved] + @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): @@ -474,38 +463,13 @@ def unshard_tensors_for_hf_save( if not tensors: return [] - # Shallow-copy the list, not the tensors. Entries with no gather work can be returned as-is, while entries - # that do need all-gather are overwritten in this working list with their gathered tensor. tensor_list = list(tensors) - - # Convert each tensor's forward shard history into the save-time work queue. Save must undo shards from - # inner to outer, so the steps are reversed; preserved shards, such as an EP shard kept local for RL weight - # sync, are removed from the queue and only used later to compute the expected partially-unsharded shape. - - # Example: - # tensor A: [ep_a(index=0), fsdp_a(index=1)], preserved {0} -> pending [fsdp_a] - # tensor B: [ep_b(index=0), fsdp_b(index=1)], preserved {} -> pending [fsdp_b, ep_b] - # tensor C: [fsdp_c(index=0)], preserved {} -> pending [fsdp_c] - # tensor D: [tp_d(index=0)], preserved {} -> pending [tp_d] - # tensor E: [ep_e(index=0)], preserved {0} -> pending [] - # This produces one pending queue per tensor; the loop below consumes compatible queue heads by group. - pending_shard_steps_list = [save_plan._pending_unshard_steps() for save_plan in save_plans] - - while True: - # Build one all-gather round. For one tensor, reverse-unshard steps must run one by one: if a local - # tensor needs to undo FSDP and then EP, the EP gather must use the tensor produced by the FSDP gather. - # `_build_ready_save_unshard_groups` consumes `pending_shard_steps_list` gradually. For example, a queue - # `[fsdp_step, ep_step]` contributes `fsdp_step` in the first round; after its gathered tensor is written - # back, the next loop consumes `ep_step`. Independent tensors with compatible group/dtype can still be - # batched together in each round. - # - # With the A-E example above, round 1 consumes fsdp_a/fsdp_b/fsdp_c together if they share group/dtype, - # and consumes tp_d in another group. tensor E contributes no work. Round 2 can then consume ep_b, because - # ep_b must use tensor B after fsdp_b has been gathered and written back. - unshard_groups = _build_ready_save_unshard_groups(tensor_list, pending_shard_steps_list) - if not unshard_groups: - break - + # 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, @@ -514,27 +478,14 @@ def unshard_tensors_for_hf_save( for index, gathered_tensor in zip(unshard_group.tensor_indices, gathered_tensors, strict=True): tensor_list[index] = gathered_tensor - # Collective steps reconstruct runtime shapes only. Remove checkpoint-invisible - # FP8 tail padding once, after every requested gather has completed. - for index, (tensor, save_plan) in enumerate(zip(tensor_list, save_plans, strict=True)): - 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}" - ) - for dim, output_size in enumerate(save_plan.output_shape): - assert output_size <= tensor.shape[dim] - if output_size < tensor.shape[dim]: - tensor = tensor.narrow(dim, 0, output_size) - tensor_list[index] = tensor.contiguous() - - assert tuple(tensor.shape) == save_plan.output_shape, ( - f"Saved tensor shape {tuple(tensor.shape)} is incompatible with HFSavePlan output_shape=" - f"{save_plan.output_shape} for {save_plan.name}" - ) - return tensor_list + # 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 _build_ready_save_unshard_groups( +def _take_ready_save_unshard_groups( tensor_list: list[torch.Tensor], pending_shard_steps_list: list[list[SaveShardStep]], ) -> list[_SaveUnshardGroup]: @@ -575,6 +526,22 @@ def _build_ready_save_unshard_groups( 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], @@ -601,7 +568,7 @@ def _pad_tensor_for_save_shard(tensor: torch.Tensor, shard_step: SaveShardStep) dim = shard_step.shard.dim shard_dim_size = shard_step.shape_before_shard[dim] - expected_local_size = shard_step.shard.local_size(shard_dim_size) + 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}" @@ -740,10 +707,10 @@ def plan_hf_load(self) -> HFLoadPlan: Returns: HFLoadPlan: The selected HF keys and canonical source-to-local copy program for this rank. """ - target_shape = self._runtime_local_shape() - owned_regions = _owned_regions_for_shards( - self.global_shape, - self.shards, + 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: @@ -755,31 +722,11 @@ def plan_hf_load(self) -> HFLoadPlan: zero_unwritten_target=math.prod(target_shape) > 0, ) - # Select the smallest contiguous HF-key envelope covering every owned region. Ordinary FSDP/EP produces one - # region. Interleaved Expert TP produces multiple runs, but their envelope still lets EP ranks avoid reading - # experts owned by other EP ranks when the checkpoint stores one key per expert. - envelope = [ - ( - min(region.global_offsets[dim] for region in owned_regions), - max(region.global_offsets[dim] + region.sizes[dim] for region in owned_regions), - ) - for dim in range(len(self.global_shape)) - ] - key_start, key_end = self._local_hf_key_indices(envelope) - hf_keys = self.global_hf_keys[key_start:key_end] - - loaded_starts = [0 for _ in self.global_shape] - loaded_ends = list(self.unpadded_global_shape) - if self.is_fused: - key_size = self._fused_key_size() - assert self.fused_dim is not None - loaded_starts[self.fused_dim] = key_start * key_size - loaded_ends[self.fused_dim] = key_end * key_size - + hf_keys, source_offsets, source_shape = self._hf_load_source(owned_regions) copy_regions = [ LoadCopyRegion( source_offsets=tuple( - region.global_offsets[dim] - loaded_starts[dim] for dim in range(len(self.global_shape)) + region.global_offsets[dim] - source_offsets[dim] for dim in range(len(self.global_shape)) ), target_offsets=region.local_offsets, sizes=region.sizes, @@ -796,12 +743,33 @@ def plan_hf_load(self) -> HFLoadPlan: name=self.name, hf_keys=hf_keys, fused_dim=self.fused_dim, - canonical_source_shape=tuple(end - start for start, end in zip(loaded_starts, loaded_ends)), + 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, *, @@ -838,23 +806,20 @@ def plan_hf_save( # tensors, so its key list is informational and needs no key alignment. hf_keys = list(self.global_hf_keys) - runtime_output_shape = _local_shape_for_shards(self.global_shape, preserved_shards) - output_shape = _local_shape_for_shards( - self.global_shape, - preserved_shards, + 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, - global_shape=self.global_shape, - unpadded_global_shape=self.unpadded_global_shape, runtime_output_shape=runtime_output_shape, output_shape=output_shape, fused_dim=self.fused_dim, distributed_save=distributed_save, - preserves_shards=bool(preserved_shards), unshard_steps=unshard_steps, ) @@ -866,15 +831,6 @@ def model_post_init(self, _) -> None: self._validate_origin_shape() self._validate_shards() - def _runtime_local_shape(self) -> tuple[int, ...]: - derived_shape = _local_shape_for_shards(self.global_shape, self.shards) - if self.local_shape is not None: - assert self.local_shape == derived_shape, ( - f"Recorded local_shape={self.local_shape} does not match descriptor-derived shape " - f"{derived_shape} for {self.name}" - ) - return derived_shape - 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) @@ -884,9 +840,9 @@ def _fused_key_size(self) -> int: ) return int(key_size) - def _local_hf_key_indices( + def _hf_key_range_for_regions( self, - effective_intervals: list[tuple[int, int]], + regions: list[_OwnedRegion], *, require_fused_key_aligned: bool = False, ) -> tuple[int, int]: @@ -895,7 +851,8 @@ def _local_hf_key_indices( assert self.fused_dim is not None key_size = self._fused_key_size() - fused_start, fused_end = effective_intervals[self.fused_dim] + 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 " @@ -917,22 +874,15 @@ def _local_hf_keys_for_shards( *, require_fused_key_aligned: bool = False, ) -> list[str]: - regions = _owned_regions_for_shards( - self.global_shape, - shards, + 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 [] - envelope = [ - ( - min(region.global_offsets[dim] for region in regions), - max(region.global_offsets[dim] + region.sizes[dim] for region in regions), - ) - for dim in range(len(self.global_shape)) - ] - key_start, key_end = self._local_hf_key_indices( - envelope, + 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] @@ -949,16 +899,12 @@ def _validate_origin_shape(self) -> None: ) def _validate_shards(self) -> None: - current_shape = list(self.global_shape) - for shard in self.shards: - assert 0 <= shard.dim < len(current_shape), ( - f"Invalid shard dim {shard.dim} for global_shape={self.global_shape}" - ) - current_shape[shard.dim] = shard.local_size(current_shape[shard.dim]) + 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 tuple(current_shape) == self.local_shape, ( + 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"{tuple(current_shape)} for {self.name}" + f"{derived_shape} for {self.name}" ) def _preserved_shard_indices( @@ -1029,11 +975,12 @@ def _save_shard_steps(self, preserved_shard_indices: set[int]) -> list[SaveShard for shard_index, shard in enumerate(self.shards): steps.append( SaveShardStep( - load_spec_shard_index=shard_index, shard=shard, shape_before_shard=tuple(current_shape), preserved=shard_index in preserved_shard_indices, ) ) - current_shape[shard.dim] = shard.local_size(current_shape[shard.dim]) + current_shape[shard.dim] = sum( + end - start for start, end in shard.local_intervals(current_shape[shard.dim]) + ) return steps From 33f880871f440597c0f6fdc6c30573a15dec220f Mon Sep 17 00:00:00 2001 From: zhaopenghao Date: Thu, 13 Aug 2026 03:58:14 +0000 Subject: [PATCH 7/7] [CI] Run DSA multiprocess tests before TileLang JIT --- tests/module/attention/test_dsa_mla.py | 139 +++++++++++++------------ 1 file changed, 71 insertions(+), 68 deletions(-) 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)