diff --git a/docs/design/decouple_ep_fsdp.md b/docs/design/decouple_ep_fsdp.md new file mode 100644 index 0000000000..37603d20e7 --- /dev/null +++ b/docs/design/decouple_ep_fsdp.md @@ -0,0 +1,355 @@ +# MoE EP 与 FSDP 解耦（dp2ep 布局） + +Expert Parallelism（EP）与 FSDP 解耦，是把 EP 从"与 FSDP 正交的 mesh 维度"改成"FSDP shard 维度的子维度"（torchtitan 称之为 dp2ep）。解耦后 dense 参数在完整的 data-parallel shard 域上分片、不再沿 EP 复制；routed expert 保持按 EP 切分专家，再只沿剩余的 `efsdp` 维做 FSDP。配置入口是 `FSDPConfig.decouple_ep_fsdp`，默认关闭，关闭时旧路径不受影响。 + +```mermaid +flowchart LR + R["root mesh (replicate, efsdp, ep)"] --> EP["ep_mesh = root[ep]
token dispatch / expert 归属"] + R --> DS["fsdp_mesh = flatten(efsdp, ep) = dp_shard
dense 参数的 FSDP 分片域"] + R --> EF["expert_fsdp_mesh = root[efsdp]
routed expert 在 EP 之上的 FSDP 分片域"] + R -. "replicate > 1 (HSDP)" .-> HS["hsdp_mesh = root[replicate, dp_shard]
expert_fsdp_mesh = root[replicate, efsdp]"] +``` + +## 1. 问题：EP 把 FSDP 分片度锁死 + +旧路径的 MoE 训练 mesh 是二维 `(fsdp, ep)`，且 `fsdp = world_size / ep_size`。两类参数的 data-parallel 语义不同，却共用这一张 mesh： + +| 参数 | 旧路径 placement | 后果 | +|---|---|---| +| routed expert | `Shard(0)` on `ep`，再 FSDP over `fsdp` | 正确，且与新路径数学同构 | +| dense（attention / router / shared expert / embedding / lm_head） | `Replicate()` on `ep`（`_replicate_other_params`），再 FSDP over `fsdp` | 参数、梯度、优化器状态在 EP 维复制 `ep` 份；每步一次额外的 coalesced all-reduce 补梯度 | +| HSDP | 要求 `ep_size == 1`（`FSDPConfig` 校验） | HSDP 与 EP 不能共存 | + +EP 越大，dense 侧越亏。以 GLM-5.2-30B 单机 8 卡为例，EP=8 时每卡保存一份完整的 dense 参数与 fp32 优化器状态，legacy 路径 EP=8 的 reserved 显存顶到约 132 GiB 的 allocator 上限（allocated 峰值 120 GiB），每步触发 allocator 释放重试，步时退化到 EP=4 的 7 倍（见 §5）。 + +## 2. 目标布局 + +```text +dp_shard = hsdp_sharding_size or world_size # dense 参数的 FSDP 分片数 +replicate = world_size / dp_shard # HSDP 副本数，无 HSDP 时为 1 +efsdp = dp_shard / ep # routed expert 在 EP 之上的 FSDP 分片数 + +root mesh = (replicate, efsdp, ep) # ep 保持最内维：EP group 仍是连续 rank +dense : FSDP over flatten(efsdp, ep) = dp_shard (+ replicate 上 HSDP) +expert : Shard(0) over ep + FSDP over efsdp (+ replicate 上 HSDP) +``` + +约束从 `fsdp = world / ep`（死）变成 `ep | dp_shard`（活）。routed expert 的数学与旧路径完全相同：旧路径在 `world/ep` 个 rank 上 FSDP，新路径在 `efsdp = dp_shard/ep` 个 rank 上 FSDP，无 HSDP 时二者就是同一组 rank。新旧差异全部集中在 dense 侧。 + +单机 8 卡的三个典型拓扑： + +| 配置 | root mesh | dense 每卡持有 | routed expert 每卡持有 | +|---|---|---|---| +| EP=8（旧） | `(fsdp=1, ep=8)` | 100%（8 份复制） | 1/8 的 expert | +| EP=8，解耦 | `(1, efsdp=1, ep=8)` | 1/8 | 1/8 的 expert（与旧路径相同） | +| EP=4（旧） | `(fsdp=2, ep=4)` | 1/2（4 份复制） | 1/4 的 expert，再在 rank r 与 r+4 之间 FSDP 对半 | +| EP=4，解耦 | `(1, efsdp=2, ep=4)` | 1/8 | 同上（rank r 与 r+4 构成 efsdp group） | +| 16 卡，EP=8，`hsdp_sharding_size=8` | `(replicate=2, efsdp=1, ep=8)` | 节点内 1/8，跨节点复制 | 节点内 EP=8，跨节点复制并由 HSDP 归约 | + +"EP=8 且 FSDP=8"不再是矛盾配置：对 dense 而言 FSDP shard 数是 8；对 expert 而言 `efsdp = 1`，因为 expert 已被 EP 切成 8 份。 + +## 3. 实现 + +各小节的"位置"行给出文件与行号，行号对应分支 `feat/decouple-ep-fsdp` 的代码提交 `0a7c6523`（2026-09-17，基于 upstream/main `6c06f96e`）。 + +### 3.1 配置与校验 + +位置：`xtuner/v1/config/fsdp.py` 44–73，`xtuner/v1/model/moe/moe.py` 1764–1771。 + +```python +# xtuner/v1/config/fsdp.py +decouple_ep_fsdp: bool = False # 默认关闭，旧路径不受影响 + +@model_validator(mode="after") +def _validate_ep_fsdp_topology(self): + # ValueError（pydantic 包装为 ValidationError），不是 assert：python -O 下同样生效 + if self.ep_size < 1 or (self.hsdp_sharding_size is not None and self.hsdp_sharding_size < 1): + raise ValueError(...) # 必须是正整数 + if self.hsdp_sharding_size is not None: + if self.decouple_ep_fsdp: + if self.hsdp_sharding_size % self.ep_size != 0: # 必须存在整数 efsdp + raise ValueError(...) + elif self.ep_size != 1: # 旧约束只在旧路径保留 + raise ValueError(...) +``` + +`MoE._init_decoupled_device_mesh` 再次校验 `world_size % dp_shard == 0` 与 `dp_shard % ep_size == 0`，同样抛 `ValueError`。ExpertTP（`expert_tp_size > 1`）与解耦同时开启时抛 `NotImplementedError`，ETP 的 `(Shard, InterleavedShard)` 二维 placement 需要单独设计 mesh 维。 + +### 3.2 单一 root mesh + +位置：`xtuner/v1/model/moe/moe.py` 1743–1808（`_init_decoupled_device_mesh`），1576–1581（分流）。 + +```python +# xtuner/v1/model/moe/moe.py::_init_decoupled_device_mesh（节选） +world_size = dist.get_world_size() +dp_shard = fsdp_config.hsdp_sharding_size or world_size +replicate_size = world_size // dp_shard +efsdp_size = dp_shard // ep_size + +root_mesh = init_device_mesh( + device, (replicate_size, efsdp_size, ep_size), + mesh_dim_names=(f"{prefix}.replicate", f"{prefix}.efsdp", f"{prefix}.ep"), +) +self._world_mesh = root_mesh +self.ep_mesh = root_mesh[ep_name] # dispatcher 用，接口不变 +self.fsdp_mesh = root_mesh[efsdp_name, ep_name]._flatten(dp_shard_name) # dense 的一维 dp_shard +if replicate_size > 1: # HSDP + self.hsdp_mesh = root_mesh[replicate_name, dp_shard_name] + self.expert_fsdp_mesh = root_mesh[replicate_name, efsdp_name] +else: + self.hsdp_mesh = None + self.expert_fsdp_mesh = root_mesh[efsdp_name] +``` + +所有子 mesh 从同一个 root 派生： + +- `ep_mesh = root[ep]`：dispatcher / DeepEP 拿到的 EP group 不变，仍是节点内连续 rank。`MoE.__init__` 里先建的 `ep_mesh` 通过 PyTorch 的 mesh 相等性 hash 被重新挂到这个 root 上，因此 `ep` 维必须沿用旧名字 `{prefix}.ep`； +- `fsdp_mesh = root[efsdp, ep]._flatten("dp_shard")`：语义仍是"dense 参数的一维 shard group"，`_fsdp_foreach_allgather`、`Float8Handler.build_reduce_mesh` 等既有消费者无需改动； +- `hsdp_mesh = root[replicate, dp_shard]`（仅 `replicate > 1`）：dense 的 HSDP mesh，与旧路径契约一致； +- `expert_fsdp_mesh = root[efsdp]` 或 `root[replicate, efsdp]`：新增属性，只在解耦路径上非空，dense 模型恒为 `None`。 + +同源很重要：FSDP2 要求已是 DTensor 的参数与 FSDP mesh 共享同一个 root，expert 参数最终要同时带 EP shard 和 `efsdp` 上的 `_StridedShard`。`efsdp = 1` 时该维保留而不是省略——对 expert 做的那次 `fully_shard` 是 mixed precision policy 生效的地方（torchtitan #1324 的同一注释）。 + +### 3.3 两级 `fully_shard` + +位置：`xtuner/v1/model/moe/moe.py` 1310–1342（decoder layer），1370–1380（prefetch），1419–1451（MTP），1723–1741（helper）。 + +```python +# xtuner/v1/model/moe/moe.py::fully_shard（节选） +decoupled = self.fsdp_config.decouple_ep_fsdp +if not decoupled and (self.ep_mesh.size() > 1 or tp_enabled): + self._replicate_other_params(self) # 旧路径：dense 在 ep 维 Replicate；解耦路径跳过 + +for layer_idx, layer in self.layers.items(): + if decoupled: # 1) 先包 MoEBlock -> expert_fsdp_mesh + self._fully_shard_expert_blocks(layer, mp_policy=mp_policy, reshard_after_forward=...) + ... + self._fully_shard( # 2) 再包整层 -> dense 的 fsdp_mesh / hsdp_mesh + mesh=self.fsdp_mesh if self.hsdp_mesh is None else self.hsdp_mesh, module=layer, ... + ) + +for layer_cur, layer_next in zip(layers[:-1], layers[1:]): + if decoupled: # expert 的 all-gather 与下一层 dense 一起预取 + layer_cur.set_modules_to_forward_prefetch([*self._expert_blocks(layer_cur), layer_next]) + +@staticmethod +def _expert_blocks(module): + return [m for m in module.modules() if isinstance(m, MoEBlock)] + +def _fully_shard_expert_blocks(self, module, mp_policy, reshard_after_forward): + for expert_block in self._expert_blocks(module): + self._fully_shard(mesh=self.expert_fsdp_mesh, module=expert_block, mp_policy=mp_policy, ...) +``` + +解耦路径跳过 `_replicate_other_params`：dense 参数进入 `fully_shard` 时是普通 tensor，由 FSDP 在完整 `dp_shard` 上分片。每个 decoder layer 先把所有 `MoEBlock`（只含 routed expert 的 grouped linear 与激活）在 `expert_fsdp_mesh` 上 `fully_shard`，再在 dense mesh 上包装整层；FSDP2 会保留内层的 expert wrapper，把 expert 参数从外层 param group 中排除。MTP layer 使用同样的两级包装。forward prefetch 列表同时包含下一层与其 expert block，让 expert 的 all-gather 与 attention 重叠，而不是串行等在 MoE block 前。 + +### 3.4 梯度归约 + +位置：`xtuner/v1/model/moe/moe.py` 1482–1486（分流），1673–1722（`_scale_and_reduce_grad_decoupled`）。 + +不变式：每个参数的最终梯度等于全 data-parallel world 上的均值，与 `ep` / `efsdp` 取值无关。 + +```python +# xtuner/v1/model/moe/moe.py::_scale_and_reduce_grad_decoupled（节选） +for name, param in self.trainable_parameters(): + if param.grad is None: + continue + if ep_size > 1 and ".experts" in name: + param.grad.div_(ep_size) # expert：FSDP 已在 efsdp 上求均值，剩余因子是 ep + continue + if not isinstance(param, DTensor): + continue + if any(isinstance(p, Shard) for p in param.placements): + continue # FSDP 管理的参数（含 _StridedShard）：FSDP 已归约 + replicate_dims = [names[i] for i, p in enumerate(param.placements) if isinstance(p, Replicate)] + flat_mesh = param.device_mesh[replicate_dims]._flatten() if len(replicate_dims) > 1 else param.device_mesh[replicate_dims[0]] + grad = param.grad.to_local() + grad.div_(flat_mesh.size()) # 全 Replicate 的 ignored fp32 参数：手工求均值 + grads_by_group.setdefault(flat_mesh.get_group(), []).append(grad) + +for group, grads in grads_by_group.items(): # 每个 group 一次 coalesced all-reduce + with dist._coalescing_manager(group=group): + for grad in grads: + dist.all_reduce(grad, ReduceOp.SUM, group=group) +``` + +| 参数类 | 谁完成归约 | 解耦路径的额外动作 | +|---|---|---| +| dense FSDP 参数（任一 `Shard` placement） | FSDP 在 `dp_shard` 上 reduce-scatter；HSDP 再跨 `replicate` all-reduce | 无（旧路径的手工跨 EP all-reduce 删除） | +| routed expert | FSDP 在 `efsdp` 上 reduce-scatter，但每个 expert 已看过整个 EP group 路由来的 token | `grad.div_(ep)`，与 NeMo AutoModel 的 `ep_ratio = dp_shard / efsdp` 相同 | +| FSDP ignored、各维全 `Replicate` 的 fp32 参数（`fp32_keys_pattern` 命中的 dense 参数） | 无人归约 | 沿所有 `Replicate` 维 flatten 后做一次 coalesced all-reduce 求均值 | + +这也是解耦路径不能复用旧梯度逻辑的原因：旧 dense 是 EP replicate，需要手工跨 EP 归约；新 dense 已由 FSDP 在完整 `dp_shard` 上归约，再做一次就是重复归约。注意第三行与 HSDP 无关：无 HSDP 时 world mesh 仍是三维 `(1, efsdp, ep)`，ignored 参数在三个维度上都是 `Replicate`，仍需这次 all-reduce。 + +### 3.5 FP8 + +位置：`xtuner/v1/float8/float8_handler.py` 121–164、166–238、316–335；`xtuner/v1/float8/fsdp_utils.py` 127–137；`xtuner/v1/model/moe/moe.py` 1281–1292；`xtuner/v1/model/base.py` 1015。 + +tile-wise FP8 需要知道每个本地权重会被多少个 FSDP rank 切分（决定 padding），以及 scale 的 reduce-max 应跨哪些 rank。解耦后 dense 与 expert 的答案不同： + +```python +# xtuner/v1/float8/float8_handler.py::pad_for_fsdp（节选） +num_fsdp_chunks = fsdp_mesh.size(-1) # dense：dp_shard 份 +if expert_fsdp_mesh is not None and isinstance(module, TileWiseFloat8GroupedLinear): + num_fsdp_chunks = expert_fsdp_mesh.size(-1) # routed expert：efsdp 份 +padded_out_features = Float8Handler.get_num_features_after_pad(tensor_size, 0, num_fsdp_chunks) + +# ::_build_decoupled_reduce_meshes（节选） +dense_shard_size = fsdp_mesh.size(-1) # 连续 rank，stride 1 +expert_shard_size = expert_fsdp_mesh.size(-1) +expert_stride = world_size // expert_fsdp_mesh.size() # == ep_size +self.tilewise_reduce_mesh_mapping = self._build_strided_reduce_mesh_mapping( + model, (TileWiseFloat8Linear,), dense_shard_size, 1) +self.expert_tilewise_reduce_mesh_mapping = self._build_strided_reduce_mesh_mapping( + model, (TileWiseFloat8GroupedLinear,), expert_shard_size, expert_stride) +``` + +- dense linear 按 `dp_shard` 的 chunk 数 padding，reduce mesh 的 rank stride 为 1； +- grouped-expert linear 按 `efsdp` 的 chunk 数 padding，reduce mesh 的 rank stride 为 `ep`； +- scale 预计算（`precompute_tilewise_float8_scale_for_fsdp`）新增 `module_types` 参数，按 dense / expert 两类各执行一次。 + +旧路径 `expert_fsdp_mesh is None`，走原有的单 mesh 代码。tensor-wise FP8 只作用于 dense linear，shard group 仍是 `fsdp_mesh`，无需改动。 + +### 3.6 HF 权重与 DCP + +位置：无代码改动。依赖 `xtuner/v1/utils/load_spec.py` 的 `LoadSpec.from_tensor`、`RuntimeLayout.from_dtensor`、`plan_hf_save`；验证用例 `tests/engine/test_decoupled_ep_fsdp_train_engine.py::TestDecoupledEpFsdpCheckpoint`，实验脚本 `tests/model/run_decoupled_ep_fsdp_ckpt.py`。 + +`LoadSpec.from_tensor` 直接从 DTensor placement 记录 shard 历史，新 placement 无需特殊处理： + +```text +routed expert（解耦 + HSDP） + DTensor placement : (Replicate, _StridedShard(0), Shard(0)) on (replicate, efsdp, ep) + LoadSpec.shards : [ShardDescriptor(dim=0, group=ep), ShardDescriptor(dim=0, group=efsdp)] # Replicate 不产生 shard +dense（解耦 + HSDP） + DTensor placement : (Replicate, Shard(0)) on (replicate, dp_shard) + LoadSpec.shards : [ShardDescriptor(dim=0, group=dp_shard)] +``` + +同步 HF save 按 shard 描述逆向 unshard 后再走模型已有的 HF key / fused-expert 映射；DCP 按 DTensor placement 保存并在载入时 reshard。因此没有为任何模型重写 `state_dict_adapter`，工作量在于让新 mesh 产出正确的 `LoadSpec`。 + +### 3.7 RL 权重同步 + +位置：`xtuner/v1/model/base.py` 1916–1947；`xtuner/v1/rl/weight_update/weight_iterator.py` 142–151、198、207–230、240–249；测试 `tests/rl/test_weight_iterator.py` 237–427。 + +Turbomind 的 layer-wise 更新只 gather FSDP shard、保留 EP-local 的 expert 切片。`BaseModel._fsdp_foreach_allgather` 按每个 `LoadSpec` 选择 gather group： + +```python +# xtuner/v1/model/base.py::_fsdp_gather_group +def _fsdp_gather_group(self, load_spec, fsdp_group, expert_fsdp_group): + if expert_fsdp_group is None: + return fsdp_group + if any(self._is_same_process_group(shard.group, expert_fsdp_group) for shard in load_spec.shards): + return expert_fsdp_group # spec 带 efsdp shard：只 gather efsdp，EP shard 保留 + return fsdp_group # dense：gather dp_shard +``` + +compose 模型（如 Qwen3.5-VL MoE）曾有一个 owner 选择错误：`WeightIterator.iter_layer_batches` 从 language tower 取参数与 `LoadSpec`，却用外层 compose 模型的 mesh 做 gather。compose 模型本身在 world mesh 上包装、没有 `expert_fsdp_mesh`，也不维护自己的 `load_spec_mapping`，于是 language tower 的 `efsdp` shard（HSDP 时连 `dp_shard` shard）被当作"应保留"的 shard 而不 gather，Turbomind 收到的是 rank-local 碎片。现在 gather 由参数所属的子模块执行，并从该子模块的 `load_spec_mapping` 取 spec： + +```python +# xtuner/v1/rl/weight_update/weight_iterator.py::_param_owner +def _param_owner(model, name): + if isinstance(model.config, BaseComposeConfig): + for submodule in ("language_model", "vision_tower", "multi_modal_projector"): + if name.startswith(f"{submodule}."): + return getattr(model, submodule), name[len(submodule) + 1 :] + return model, name # plain 模型：owner 就是自己，行为不变 +``` + +`tests/rl/test_weight_iterator.py::TestLayerBatchesGatherWithParamOwner` 用假 process group 固定了 `efsdp = 2` 与 HSDP 两种拓扑下每个 tensor 的完整形状与取值。 + +## 4. 与参考实现的对照 + +| 项目 | mesh | expert 的 FSDP 维 | 梯度修正 | HSDP + EP | +|---|---|---|---|---| +| torchtitan（#1324 起） | `(pp, dp_replicate, dp_shard_mod_ep, dp_shard_in_ep, cp, tp)` | `dp_shard_mod_ep = dp_shard·cp / ep`，派生 | FSDP gradient divide factor（#1551） | 首日支持 | +| NeMo AutoModel | `device_mesh(dp_replicate, dp_shard_cp)` + `moe_mesh(ep, ep_shard)` | `ep_shard`，派生 | 显式 `grad.div_(ep_ratio)` | 支持 | +| Megatron-Core（Megatron-FSDP） | expert 参数在 expert data-parallel group 上分片 | `dp·cp / ep`，派生 | 框架内部 | 支持 | +| XTuner 本方案 | `(replicate, efsdp, ep)` | `efsdp = dp_shard / ep`，派生 | 显式 `grad.div_(ep)` | mesh 与数学路径支持，多机未实测 | + +三个参考实现都把 expert 的 FSDP 维当作派生量而不是独立配置，本方案与之一致。实现上采用 torchtitan #1324 的"两次 `fully_shard`"法，对 torch 版本要求最低；`shard_placement_fn` 单次包装、expert dim-1 分片（#1561）留作后续优化。 + +## 5. 验证结论与边界 + +| 项目 | 证据 | 边界 | 链接 | +|---|---|---|---| +| mesh / placement / 两级 FSDP（L0） | fake-PG 下 8/16/64 rank 共 31 例（含配置校验与运行时校验的回归）；旧布局在开关关闭时逐项固定 | 需要一张空闲 GPU 建 CUDA context，不通信 | [test_decoupled_ep_fsdp_mesh.py](../../tests/model/test_decoupled_ep_fsdp_mesh.py)、[baseline.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/baseline.md) | +| BF16 训练语义（L1） | tiny Qwen3-MoE 50 步：解耦 EP8 vs 旧 EP8 loss 最大相对差 `1.9e-5`（旧 EP8 vs EP1 为 `2.5e-5`）；3.4B 模型每卡 dense 参数 836 → 105 MiB，峰值 allocated 11861 → 8203 MiB，步时 170 → 158 ms | 单机；reduction order 不同，不是逐 bit 相同 | [test_decoupled_ep_fsdp_train_engine.py](../../tests/engine/test_decoupled_ep_fsdp_train_engine.py)（8 卡 gate：loss 相对差 ≤ 1e-3、总/逐参数梯度范数 ≤ 1e-2、每卡参数显存 ±25%）、[L1.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/L1.md)、[run_decoupled_ep_fsdp_numerics.py](../../tests/model/run_decoupled_ep_fsdp_numerics.py) | +| `efsdp > 1`、HSDP + EP（L2） | 单机缩比拓扑 `(1,2,4)`、`(2,1,4)`、`(2,2,2)` 均在噪声内；16/64 rank 形状由 L0 固定 | 真正的双节点 collective 未跑 | [test_decoupled_ep_fsdp_train_engine.py](../../tests/engine/test_decoupled_ep_fsdp_train_engine.py)（同上阈值，拓扑 `(1,2,4)`、`(2,1,4)`、`(2,2,2)`）、[L2.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/L2.md)、[decisions.md D10](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/decisions.md) | +| HF / DCP（L3） | 刚 `from_hf` 后导出 807 个 key 全部 bit-exact；DCP 同布局 resume 与跨布局 reshard 后 loss 连续 | 未覆盖 async save、FP8 checkpoint、GLM/MTP 导出、world size 变化 | [test_decoupled_ep_fsdp_train_engine.py](../../tests/engine/test_decoupled_ep_fsdp_train_engine.py)（8 卡 gate：step-0 导出逐 bit 相同、resume 与跨布局 loss 相对差 ≤ 1e-3、导出差异 ≤ 数个 bf16 ulp）、[L3.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/L3.md)、[run_decoupled_ep_fsdp_ckpt.py](../../tests/model/run_decoupled_ep_fsdp_ckpt.py) | +| FP8 | 3.4B 与 GLM-5.2 tile-wise FP8 训练 20 步稳定，差异量级接近同配置重复噪声 | FP8 + HSDP、FP8 + checkpoint 未测 | [L3.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/L3.md)、[GLM52.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/GLM52.md) | +| GLM-5.2-30B，8×H200 目标配方 | EP4：1.688 s / 99.5 GiB → 1.658 s / 83.5 GiB；EP8：11.8 s / 120.0 GiB → 1.64 s / 76.3 GiB | EP8 的加速来自离开 allocator 上限（legacy 每步 2–3 次 alloc retry，解耦为 0），不是 collective 本身快 7 倍 | [GLM52.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/GLM52.md) | +| RL 权重同步 | plain MoE 的 per-spec gather 逻辑成立；compose 的 owner 错误已修复并有单测 | 未做真实 Turbomind 端到端对拍 | [test_weight_iterator.py](../../tests/rl/test_weight_iterator.py#L237-L427) | +| legacy 兼容 | 开关默认关闭，旧分支执行逻辑未改动，L0 旧布局回归通过 | 未做 base-vs-branch 端到端逐 tensor 对照 | [baseline.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/baseline.md)、[decisions.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/decisions.md) | + +`reports/` 下的报告不随代码 PR 合入，链接指向 fork 分支 `feat/decouple-ep-fsdp` 上的固定地址；测试与脚本是仓库内相对链接。 + +## 6. 已知限制与后续工作 + +### 6.1 `fp32_keys_pattern` 命中 routed expert 时梯度不归约 + +`HFSaveCfg.fp32_keys_pattern` 用 HF key 正则把少量参数保留在 fp32 并排除在 FSDP 之外，内置用途是 Qwen3.5 linear-attention 的 `A_log`、norm 这类 dense 小参数。`BaseModel._fully_shard` 对命中的参数分两种处理：普通 tensor 先 `distribute_tensor(Replicate on world_mesh)` 再放进 `ignored_params`；已经是 DTensor 的参数原样放进 `ignored_params`。routed expert 在 `build_grouped_linear` 里就已经是 `Shard(0)` on `ep` 的 DTensor，走的是第二种。 + +于是一旦某条 pattern 命中 expert，会同时发生两件事： + +1. expert 不再被 `_fully_shard_expert_blocks` 沿 `efsdp` 分片，每个 `efsdp` rank（HSDP 时每个 replica）各持一份完整的 EP-local expert； +2. `_scale_and_reduce_grad_decoupled` 看到参数名含 `.experts`，只做 `/ep` 就 `continue`，不会走第三行的 Replicate 维 all-reduce。 + +每个 `efsdp` rank 的 expert 梯度只含自己 EP group 的 token 贡献，optimizer 一步之后同一个 expert 在不同 `efsdp` rank 上的权重就不再相同，训练静默出错。`efsdp = replicate = 1`（如单机 EP=8）时没有第二份副本，问题不显现；当前内置 pattern 也都只匹配 dense，GLM-5.2 没有 pattern，所以现有实验没有触发。旧路径同样存在这个边界：ignored expert 不会沿 `fsdp` 分片，`fsdp > 1` 时一样分叉。 + +修法：短期在解耦路径检测到 pattern 命中 `MoEBlock` 参数时直接报错；完整方案是把这类 expert 表达成 `Replicate(replicate), Replicate(efsdp), Shard(0)(ep)`，先走 §3.4 第三行的 Replicate 维 all-reduce 求均值，再 `/ep`。这与 §7.2 的 `expert_fsdp = False` 模式是同一套实现。 + +### 6.2 "什么是 expert 参数"有两套判定标准 + +| 环节 | 判定依据 | 位置 | +|---|---|---| +| FSDP 包装 | 模块类型 `isinstance(m, MoEBlock)` | `MoE._expert_blocks` | +| 梯度缩放 | 参数名 `".experts" in name` | `MoE._scale_and_reduce_grad_decoupled`（旧路径 `scale_and_reduce_grad` 同样） | +| FP8 padding / reduce mesh | 模块类型 `TileWiseFloat8GroupedLinear` | `Float8Handler.pad_for_fsdp` / `_build_decoupled_reduce_meshes` | +| RL gather | `LoadSpec.shards` 中是否有 `efsdp` group 的 shard | `BaseModel._fsdp_gather_group` | + +今天它们恰好一致，因为 `MoEDecoderLayer` 把 `MoEBlock` 挂在 `experts` 属性下（`shared_experts` 不含 `.experts` 子串，不会误伤）。两个方向都可能失效： + +- 新模型把 `MoEBlock` 挂到别的属性名（如 `moe`）：expert 会被 `expert_fsdp_mesh` 分片，reduce-scatter 只在 `efsdp` 上求均值，但名字里没有 `.experts`，不做 `/ep`，梯度被放大 `ep` 倍，等价于 expert 的学习率乘以 `ep`，没有任何报错； +- 某个 dense 参数名恰好含 `.experts`（如 `experts_router`）：会被 `/ep`，梯度缩小 `ep` 倍。 + +修法：在 `_fully_shard_expert_blocks` 包装时记录 expert 参数的 id / FQN 集合，梯度缩放、FP8、save / RL 都查这一份，删除按名字匹配的分支。 + +### 6.3 其它 + +- 两次不同 mesh 的 `fully_shard` 与 `DeviceMesh._flatten(name)` 只在 torch 2.8 / 2.9 验证；`pyproject.toml` 声明 `torch>=2.6`，需要补 CI 或为该开关声明更高的最低版本。 +- ExpertTP 与解耦互斥；expert dim-1 分片、`shard_placement_fn` 单次包装留作优化。 + +## 7. 开放问题（请架构师决策） + +### 7.1 是否保留 HSDP 与 EP 的共存 + +保留的成本很低。`replicate` 维由同一个 `init_device_mesh` 调用产生，HSDP 专属代码只有 `hsdp_mesh` / `expert_fsdp_mesh` 的二维选择；`replicate` 上的 all-reduce 由 FSDP2 完成。梯度路径没有任何 HSDP 分支：§3.4 三行规则中，前两行与 HSDP 无关，第三行在无 HSDP 时同样需要（world mesh 仍是三维）。去掉 HSDP 能省的只是几行 mesh 选择和 L2 的两个验证拓扑，梯度处理不会因此变简单。 + +保留的收益是多机拓扑：EP 留在节点内时，`hsdp_sharding_size = 8` 让 dense 的 all-gather / reduce-scatter 也留在 NVLink 域内，跨机只剩一次梯度 all-reduce；这是 torchtitan、AutoModel 都支持的组合，也是 DESIGN 中 64 卡场景的前提。目前 XTuner 对 HSDP 使用确实不多（compose 模型的 `init_world_mesh` 还标着 TODO），且本方案的 HSDP + EP 只在单机缩比拓扑验证过。 + +建议：保留 mesh 与数学路径，在文档中标为"experimental，多机未验证"；若维护者希望收窄支持面，只需在 `FSDPConfig` 中对 `decouple_ep_fsdp and hsdp_sharding_size is not None` 抛错，待双节点验证后再放开，代码不必删除。 + +### 7.2 `efsdp` 是否作为独立配置暴露 + +三个参考实现都不暴露：expert 的 data-parallel 语义要求 `ep × efsdp = dp_shard`，即每个 token 在整个 world 上恰好被计算一次，`efsdp` 因此只能是派生量。Megatron-Core 只暴露 `--expert-model-parallel-size`，expert 的 FSDP 分片域是 `dp·cp / ep`；torchtitan、AutoModel 同理。 + +接入 MoonEP / UltraEP 这类动态冗余 expert 的通信库时，需要的不是另一个 `efsdp` 数值，而是一种新模式："expert 不做 FSDP"。从公开资料看，这两个库在每层 MoE 前按实时负载把热 expert 复制到其它 EP rank 并预取权重，要求 EP-local 的 expert 权重在 MoE 计算窗口之外也是完整可寻址的；而 FSDP 分片的 expert 只在 all-gather 窗口内完整，且 FSDP 单元是整个 `MoEBlock`。当 `ep == dp_shard` 时（如单机 EP=8）`efsdp = 1` 已自动满足；当 `ep < dp_shard` 时，"不做 FSDP"意味着 expert 在 `efsdp` 维上复制，这是 placement 的改变而不是分片数的改变： + +```text +expert_fsdp = True （现状）: Shard(0) on ep + _StridedShard on efsdp FSDP reduce-scatter，再 /ep +expert_fsdp = False（新增）: Shard(0) on ep + Replicate on (replicate, efsdp) 手工 all-reduce 求均值，再 /ep +``` + +后者的梯度处理正是 §3.4 第三行已经存在的"沿 Replicate 维 all-reduce"分支，也正是修复 §6.1（fp32 pattern 命中 expert）所需的实现，两件事可以一次做完。代价是 expert 参数与优化器状态在 `efsdp` 维复制、失去 FSDP wrapper 提供的 mixed precision（需手工 cast）。 + +建议：`efsdp` 保持派生；不新增整数配置。等 MoonEP / UltraEP 接入启动时，在 `FSDPConfig` 增加 `expert_fsdp: bool = True`（或 `expert_shard_mode: "fsdp" | "replicate"`）。root mesh 不变，`ep_mesh` 仍是最内层连续 rank 的 group，dispatcher 的替换与本方案正交。 + +## 8. 合入建议 + +PR 的提交按 torchtitan 的演进顺序排列，维护者若希望拆分，可按同一顺序切分： + +1. 基础解耦：配置开关、root mesh、两级 `fully_shard`、梯度归约、L0 fake-PG 测试； +2. 生态：FP8 per-class padding / reduce mesh、`LoadSpec` 驱动的 HF / DCP、RL per-spec gather 与 compose owner 修复； +3. 文档与后续：本文、§6 的限制项。 + +`decouple_ep_fsdp` 默认 `False`，所有代码提交都不改动旧路径的执行逻辑，出问题时一键回退。 diff --git a/examples/v1/config/sft_glm5p2.py b/examples/v1/config/sft_glm5p2.py index 2b991b92c0..947f21a342 100644 --- a/examples/v1/config/sft_glm5p2.py +++ b/examples/v1/config/sft_glm5p2.py @@ -108,10 +108,15 @@ def _get_float8_config() -> Float8Config | None: else: raise ValueError(f"Unsupported OPTIMIZER={optimizer!r}. Use adamw or muon.") lr_cfg = LRConfig(lr_type=os.environ.get("LR_TYPE", "cosine"), warmup_ratio=float(os.environ.get("WARMUP_RATIO", "0"))) +hsdp_sharding_size = os.environ.get("HSDP_SHARDING_SIZE") fsdp_cfg = FSDPConfig( cpu_offload=_get_bool_env("CPU_OFFLOAD", False), ep_size=ep_size, torch_compile=_get_bool_env("TORCH_COMPILE", False), + # DECOUPLE_EP_FSDP=1: dp2ep layout — dense params sharded over the full FSDP mesh instead of + # being replicated `ep_size` times; experts sharded `dp_shard / ep_size` ways on top of EP. + decouple_ep_fsdp=_get_bool_env("DECOUPLE_EP_FSDP", False), + hsdp_sharding_size=int(hsdp_sharding_size) if hsdp_sharding_size else None, ) trainer = TrainerConfig( diff --git a/tests/engine/test_decoupled_ep_fsdp_train_engine.py b/tests/engine/test_decoupled_ep_fsdp_train_engine.py new file mode 100644 index 0000000000..f3cbe177c3 --- /dev/null +++ b/tests/engine/test_decoupled_ep_fsdp_train_engine.py @@ -0,0 +1,324 @@ +"""GPU gates for the decoupled EP/FSDP layout (``FSDPConfig.decouple_ep_fsdp``). + +Every test trains the same tiny random Qwen3-MoE on the same token stream under several EP/FSDP +layouts and fails when the decoupled layout departs from the legacy one: + +* ``TestDecoupledEpFsdpNumerics`` (L1 / L2): loss curves, total gradient norms, per-parameter + gradient norms at step 0 (identical weights and data, so they catch a wrong expert-gradient scale + that AdamW would hide from the loss curve; after the first update bf16 weight differences flip + individual top-k routing decisions and router gradients legitimately differ by percents) and + per-rank parameter memory, for ``efsdp == 1``, ``efsdp > 1`` and HSDP + EP. +* ``TestDecoupledEpFsdpCheckpoint`` (L3): bit-exact HF export right after ``from_hf``, DCP resume, + HF export after resume and cross-layout DCP resharding. + +The observed noise floor between layouts is ~1e-5 on losses and ≤ 2e-3 on per-parameter grad +norms (see ``docs/design/decouple_ep_fsdp.md`` §5); the gates sit 1-2 orders of magnitude above +it and orders of magnitude below any layout bug seen so far. +""" + +import shutil +import tempfile +import unittest +from pathlib import Path +from typing import Any + +import pytest +import torch +import torch.distributed as dist +from safetensors.torch import save_file + +from xtuner._testing import DeterministicDDPTestCase +from xtuner._testing.decoupled_ep_fsdp import ( + LayoutMode, + assert_hf_dirs_close, + build_engine, + build_hf_checkpoint, + load_hf_dir, + max_rel_diff, + parse_mode, + release, + run_mode, + train, +) +from xtuner.v1.model.moe.qwen3 import Qwen3MoEConfig + + +SEQ_LEN = 2048 +LR = 1e-4 +# Loss curves of two layouts: max relative difference over all steps (noise ~2e-5). +LOSS_RTOL = 1e-3 +# Total grad norms per step and per-parameter grad norms at step 0 (noise ≤ 2e-3 dense, ≤ 2e-5 expert). +GRAD_NORM_RTOL = 1e-2 +# Per-rank parameter memory vs. the value implied by the layout (FSDP padding slack). +MEMORY_RTOL = 0.25 +# HF exports: 1 bf16 ulp is 2**-7 relative; resumed vs. continuous differs by single ulps on a +# subset of elements, exports of different layouts by a few ulps. +HF_RESUME_RTOL, HF_RESUME_ATOL = 2**-6, 1e-3 +HF_LAYOUT_RTOL, HF_LAYOUT_ATOL = 2**-6, 2e-3 + + +class _SharedTmpDirMixin: + """Temp directory created by rank 0 and shared with every rank of the process group.""" + + def _make_shared_tmpdir(self) -> Path: + holder = [tempfile.mkdtemp(prefix="decoupled_ep_fsdp_") if dist.get_rank() == 0 else None] + dist.broadcast_object_list(holder, src=0) + assert holder[0] is not None + return Path(holder[0]) + + @staticmethod + def _cleanup_shared_tmpdir(path: Path) -> None: + dist.barrier() + if dist.get_rank() == 0: + shutil.rmtree(path, ignore_errors=True) + + +def _build_tiny_checkpoint(root: Path) -> Path: + hf_dir = root / "tiny_qwen3_moe_hf" + if dist.get_rank() == 0: + build_hf_checkpoint(hf_dir, seed=0, model_size="tiny") + dist.barrier() + return hf_dir + + +@unittest.skipUnless(torch.cuda.device_count() >= 8, "requires 8 CUDA devices") +class TestDecoupledEpFsdpNumerics(_SharedTmpDirMixin, DeterministicDDPTestCase): + def test_ep8_decoupled_matches_legacy_ep8_and_ep1(self) -> None: + # L1: `efsdp == 1`; dense params go from replicated over EP to sharded over all 8 ranks. + results = self._run_layouts(("A:ep=1", "B:ep=8", "C:ep=8,decouple=1"), steps=20) + self._assert_training_matches(results["C"], results["B"]) + self._assert_training_matches(results["C"], results["A"]) + # 中文注释：每卡参数显存断言。记 D / E 为 dense / expert 参数总量： + # B（legacy，ep=8）：mesh 为 (fsdp=1, ep=8)，dense 在 ep 维整份复制 → 每卡 D；expert 按 ep 分 8 份 → E/8。 + # C（解耦，ep=8）：dense 在整个 dp_shard=8 上分片 → D/8；expert 按 ep×efsdp=8×1 分片 → E/8。 + # 所以 C 的 dense 应是 B 的 1/8，expert 相同。loss / 梯度对比区分不了"dense 真的分片了"和 + # "dense 仍在 EP 维复制"（两者数值完全一样），只有这条显存断言能证明解耦布局确实生效。 + self._assert_param_memory(results["C"], results["B"], dense_ratio=1 / 8, expert_ratio=1.0) + + def test_efsdp_and_hsdp_layouts_match_legacy_ep4(self) -> None: + # L2 on 8 ranks: `efsdp == 2` (C4), HSDP + EP with `efsdp == 1` (H41) and `efsdp == 2` (H22). + results = self._run_layouts( + ("B4:ep=4", "C4:ep=4,decouple=1", "H41:ep=4,decouple=1,hsdp=4", "H22:ep=2,decouple=1,hsdp=4"), + steps=10, + ) + for name in ("C4", "H41", "H22"): + self._assert_training_matches(results[name], results["B4"]) + # Per-rank shards. Legacy ep=4 on 8 ranks: dense over fsdp=2 (D/2), experts over ep=4 then + # fsdp=2 (E/8). Decoupled: dense over dp_shard (D/dp_shard), experts over ep * efsdp == + # dp_shard (E/dp_shard); dp_shard is 8 for C4 and 4 for the HSDP layouts H41 / H22. + # + # 中文注释：ratio = 被测布局的每卡参数量 / 参考布局 B4 的每卡参数量，按各布局的分片数推导： + # B4 （legacy，ep=4） ：mesh (fsdp=2, ep=4)；dense 在 ep 维复制、在 fsdp 维分 2 份 → D/2； + # expert 先按 ep 分 4 份、再按 fsdp 分 2 份 → E/8。 + # C4 （解耦，ep=4） ：dp_shard=8，efsdp=8/4=2；dense → D/8；expert 按 ep×efsdp=8 → E/8。 + # H41（解耦，ep=4，hsdp=4）：dp_shard=4，replicate=2，efsdp=1；dense → D/4；expert 按 4×1 → E/4。 + # H22（解耦，ep=2，hsdp=4）：dp_shard=4，replicate=2，efsdp=2；dense → D/4；expert 按 2×2 → E/4。 + # 于是 C4 vs B4：dense (D/8)/(D/2)=1/4，expert (E/8)/(E/8)=1； + # H41、H22 vs B4：dense (D/4)/(D/2)=1/2，expert (E/4)/(E/8)=2。 + # 写成 (1/8)/(1/2) 而不是 0.25，是为了让"每卡份额 / 参考份额"的推导过程留在代码里。 + self._assert_param_memory(results["C4"], results["B4"], dense_ratio=(1 / 8) / (1 / 2), expert_ratio=1.0) + self._assert_param_memory(results["H41"], results["B4"], dense_ratio=(1 / 4) / (1 / 2), expert_ratio=2.0) + self._assert_param_memory(results["H22"], results["B4"], dense_ratio=(1 / 4) / (1 / 2), expert_ratio=2.0) + + @property + def world_size(self) -> int: + return 8 + + def _run_layouts(self, specs: tuple[str, ...], steps: int) -> dict[str, dict[str, Any]]: + self.create_pg("cuda") + root = self._make_shared_tmpdir() + try: + hf_dir = _build_tiny_checkpoint(root) + results: dict[str, dict[str, Any]] = {} + for spec in specs: + mode = parse_mode(spec) + results[mode["name"]] = run_mode( + mode, + hf_dir, + steps=steps, + seq_len=SEQ_LEN, + lr=LR, + dispatcher="all2all", + fp8=False, + grad_norm_steps=(0,), + tag=self._testMethodName, + ) + dist.barrier() + return results + finally: + self._cleanup_shared_tmpdir(root) + + def _assert_training_matches(self, result: dict[str, Any], reference: dict[str, Any]) -> None: + name, ref_name = result["mode"]["name"], reference["mode"]["name"] + loss_diff = max_rel_diff(result["losses"], reference["losses"]) + self.assertLessEqual(loss_diff, LOSS_RTOL, f"{name} vs {ref_name}: loss curves differ by {loss_diff:.2e}") + grad_diff = max_rel_diff(result["grad_norms"], reference["grad_norms"]) + self.assertLessEqual(grad_diff, GRAD_NORM_RTOL, f"{name} vs {ref_name}: grad norms differ by {grad_diff:.2e}") + for step, ref_norms in reference["param_grad_norms"].items(): + norms = result["param_grad_norms"][step] + self.assertEqual(set(norms), set(ref_norms), f"{name} vs {ref_name}: parameter sets differ at step {step}") + worst = max(ref_norms, key=lambda key: abs(norms[key] - ref_norms[key]) / max(abs(ref_norms[key]), 1e-12)) + diff = abs(norms[worst] - ref_norms[worst]) / max(abs(ref_norms[worst]), 1e-12) + self.assertLessEqual( + diff, + GRAD_NORM_RTOL, + f"{name} vs {ref_name}: grad norm of {worst} differs by {diff:.2e} at step {step} " + f"({norms[worst]:.6g} vs {ref_norms[worst]:.6g})", + ) + + def _assert_param_memory( + self, result: dict[str, Any], reference: dict[str, Any], dense_ratio: float, expert_ratio: float + ) -> None: + # 中文注释：比较两个布局训练时记录的"每卡本地参数字节数"（`param_memory`：对每个参数取 + # DTensor 的本地分片求和，按名字里是否含 `.experts` 分成 expert / dense 两类）。 + # 断言 result 的值 ≈ reference 的值 × ratio，容差 MEMORY_RTOL（±25%）留给 FSDP 的 padding。 + # 期望比值由调用方按布局的分片数给出，见各调用点的推导。 + name, ref_name = result["mode"]["name"], reference["mode"]["name"] + for key, ratio in (("dense_param_mib", dense_ratio), ("expert_param_mib", expert_ratio)): + expected = reference["memory"][key] * ratio + actual = result["memory"][key] + self.assertAlmostEqual( + actual, + expected, + delta=expected * MEMORY_RTOL, + msg=f"{name} vs {ref_name}: {key} {actual:.1f} MiB, expected {expected:.1f} MiB per rank", + ) + + +@unittest.skipUnless(torch.cuda.device_count() >= 8, "requires 8 CUDA devices") +class TestDecoupledEpFsdpCheckpoint(_SharedTmpDirMixin, DeterministicDDPTestCase): + def test_hf_export_dcp_resume_and_cross_layout_reshard(self) -> None: + self.create_pg("cuda") + root = self._make_shared_tmpdir() + try: + hf_dir = _build_tiny_checkpoint(root) + vocab_size = Qwen3MoEConfig.from_hf(hf_dir).vocab_size + src_bf16 = root / "source_bf16" + if dist.get_rank() == 0: + # `save_hf` writes bf16; compare against the bf16-cast source. + src_bf16.mkdir() + tensors = {k: v.to(torch.bfloat16).contiguous() for k, v in load_hf_dir(hf_dir).items()} + save_file(tensors, str(src_bf16 / "model.safetensors")) + dist.barrier() + + modes = [ + parse_mode(s) + for s in ("A:ep=1", "C:ep=8,decouple=1", "C4:ep=4,decouple=1", "H41:ep=4,decouple=1,hsdp=4") + ] + continuous: dict[str, list[float]] = {} + for mode in modes: + continuous[mode["name"]] = self._check_round_trip( + mode, hf_dir, src_bf16, root / mode["name"], vocab_size + ) + + # Step-10 exports of every layout vs. the ep=1 baseline (same data, same optimizer). + for mode in modes[1:]: + if dist.get_rank() == 0: + assert_hf_dirs_close( + root / "A" / "hf_step10", + root / mode["name"] / "hf_step10", + rtol=HF_LAYOUT_RTOL, + atol=HF_LAYOUT_ATOL, + label=f"hf_step10 {mode['name']} vs A", + ) + dist.barrier() + + # DCP resharding across layouts: the step-5 checkpoint of `src` continues under `dst`. + for src, dst in (("C", "A"), ("A", "C"), ("C4", "H41")): + dst_mode = next(m for m in modes if m["name"] == dst) + engine = build_engine(dst_mode, hf_dir, tag=f"cross_{src}") + engine.load_dcp(root / src / "dcp_step5") + losses = train(engine, range(5, 10), vocab_size, SEQ_LEN) + release(engine) + diff = max_rel_diff(losses, continuous[src]) + self.assertLessEqual(diff, LOSS_RTOL, f"cross load {src} -> {dst}: losses differ by {diff:.2e}") + dist.barrier() + finally: + self._cleanup_shared_tmpdir(root) + + @property + def world_size(self) -> int: + return 8 + + def _check_round_trip( + self, mode: LayoutMode, hf_dir: Path, src_bf16: Path, mode_dir: Path, vocab_size: int + ) -> list[float]: + name = mode["name"] + engine = build_engine(mode, hf_dir, tag="main") + engine.from_hf(hf_path=hf_dir, strict=True) + engine.save_hf(str(mode_dir / "hf_step0")) + dist.barrier() + if dist.get_rank() == 0: + assert_hf_dirs_close( + src_bf16, mode_dir / "hf_step0", rtol=0.0, atol=0.0, label=f"{name} hf_step0 vs source" + ) + + train(engine, range(0, 5), vocab_size, SEQ_LEN) + engine.save_dcp(mode_dir / "dcp_step5") + dist.barrier() + losses_continuous = train(engine, range(5, 10), vocab_size, SEQ_LEN) + engine.save_hf(str(mode_dir / "hf_step10")) + dist.barrier() + release(engine) + + resumed = build_engine(mode, hf_dir, tag="resume") + resumed.load_dcp(mode_dir / "dcp_step5") + losses_resumed = train(resumed, range(5, 10), vocab_size, SEQ_LEN) + resumed.save_hf(str(mode_dir / "hf_step10_resumed")) + dist.barrier() + release(resumed) + + diff = max_rel_diff(losses_resumed, losses_continuous) + self.assertLessEqual(diff, LOSS_RTOL, f"{name}: resumed losses differ from the continuous run by {diff:.2e}") + if dist.get_rank() == 0: + assert_hf_dirs_close( + mode_dir / "hf_step10", + mode_dir / "hf_step10_resumed", + rtol=HF_RESUME_RTOL, + atol=HF_RESUME_ATOL, + label=f"{name} hf_step10 resumed vs continuous", + ) + dist.barrier() + return losses_continuous + + +class TestHfDirComparison: + """CPU checks of the checkpoint comparison the GPU gates rely on.""" + + @staticmethod + def _write(path: Path, tensors: dict[str, torch.Tensor]) -> Path: + path.mkdir(parents=True, exist_ok=True) + save_file(tensors, str(path / "model.safetensors")) + return path + + def test_bit_exact_passes_and_one_ulp_fails(self, tmp_path: Path) -> None: + base = {"w": torch.linspace(-1, 1, 64, dtype=torch.bfloat16)} + lhs = self._write(tmp_path / "lhs", base) + rhs = self._write(tmp_path / "rhs", {"w": base["w"].clone()}) + assert_hf_dirs_close(lhs, rhs, rtol=0.0, atol=0.0, label="same") + + nudged = base["w"].clone() + # One bf16 ulp: bump the bit pattern (a float32 `nextafter` would round back to the same bf16). + nudged[3] = (nudged[3].view(torch.int16) + 1).view(torch.bfloat16) + rhs_ulp = self._write(tmp_path / "rhs_ulp", {"w": nudged}) + with pytest.raises(AssertionError, match="one ulp: w"): + assert_hf_dirs_close(lhs, rhs_ulp, rtol=0.0, atol=0.0, label="one ulp") + # The layout tolerance admits a few ulps. + assert_hf_dirs_close(lhs, rhs_ulp, rtol=HF_LAYOUT_RTOL, atol=HF_LAYOUT_ATOL, label="tolerant") + + def test_key_and_shape_mismatches_fail(self, tmp_path: Path) -> None: + lhs = self._write(tmp_path / "lhs", {"w": torch.ones(4, dtype=torch.bfloat16)}) + with pytest.raises(AssertionError, match="key sets differ"): + assert_hf_dirs_close( + lhs, self._write(tmp_path / "keys", {"v": torch.ones(4, dtype=torch.bfloat16)}), 0, 0, "keys" + ) + with pytest.raises(AssertionError, match="expected \\(4,\\)"): + assert_hf_dirs_close( + lhs, self._write(tmp_path / "shape", {"w": torch.ones(8, dtype=torch.bfloat16)}), 0, 0, "shape" + ) + + def test_max_rel_diff(self) -> None: + assert max_rel_diff([1.0, 2.0], [1.0, 2.0]) == 0.0 + assert max_rel_diff([1.0, 2.2], [1.0, 2.0]) == pytest.approx(0.1) + with pytest.raises(ValueError, match="length mismatch"): + max_rel_diff([1.0], [1.0, 2.0]) diff --git a/tests/model/run_decoupled_ep_fsdp_ckpt.py b/tests/model/run_decoupled_ep_fsdp_ckpt.py new file mode 100644 index 0000000000..50e6ffafe4 --- /dev/null +++ b/tests/model/run_decoupled_ep_fsdp_ckpt.py @@ -0,0 +1,151 @@ +"""DCP round trip and HF export checks for the decoupled EP/FSDP layout (DESIGN.md §6 L3). + +For every mode: + +1. ``from_hf`` the tiny checkpoint and immediately ``save_hf`` → must be bit-identical to the source; +2. train 5 steps, ``save_dcp``, train 5 more steps, ``save_hf`` (the "continuous" run); +3. build a fresh engine, ``load_dcp`` the step-5 checkpoint, train steps 5-9 → losses must continue + the continuous curve, and its ``save_hf`` must match the continuous export; +4. the step-10 HF exports of all modes are compared against the ``ep=1`` baseline export. + +Optionally the step-5 DCP checkpoint of one layout is loaded into another layout +(``--cross-load SRC:DST``) to check DCP resharding across the switch. + +The pytest gate for these checks lives in ``tests/engine/test_decoupled_ep_fsdp_train_engine.py``; +this script writes the JSON behind the markdown report. + + torchrun --nproc-per-node 8 tests/model/run_decoupled_ep_fsdp_ckpt.py \\ + --modes A:ep=1 C:ep=8,decouple=1 C4:ep=4,decouple=1 H41:ep=4,decouple=1,hsdp=4 --out /tmp/l3 +""" + +import argparse +import json +import os +import shutil +from pathlib import Path +from typing import Any + +import torch +import torch.distributed as dist +from safetensors.torch import save_file + +from xtuner._testing.decoupled_ep_fsdp import ( + build_engine, + build_hf_checkpoint, + compare_hf, + load_hf_dir, + max_rel_diff, + parse_mode, + release, + train, +) +from xtuner.v1.model.moe.qwen3 import Qwen3MoEConfig + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--modes", nargs="+", required=True) + parser.add_argument("--cross-load", nargs="*", default=[], help="SRC:DST mode-name pairs") + parser.add_argument("--seq-len", type=int, default=2048) + parser.add_argument("--hf-dir", type=str, default=None) + parser.add_argument("--out", type=str, required=True) + args = parser.parse_args() + + dist.init_process_group("nccl") + rank = dist.get_rank() + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + out = Path(args.out) + hf_dir = Path(args.hf_dir) if args.hf_dir else out / "tiny_qwen3_moe_hf" + if rank == 0: + out.mkdir(parents=True, exist_ok=True) + if not (hf_dir / "config.json").exists(): + build_hf_checkpoint(hf_dir, 0, "tiny") + dist.barrier() + vocab_size = Qwen3MoEConfig.from_hf(hf_dir).vocab_size + + # Source checkpoint in bf16, the dtype `save_hf` writes. + src_bf16 = out / "source_bf16" + if rank == 0: + src_bf16.mkdir(exist_ok=True) + tensors = {k: v.to(torch.bfloat16).contiguous() for k, v in load_hf_dir(hf_dir).items()} + save_file(tensors, str(src_bf16 / "model.safetensors")) + dist.barrier() + + report: dict[str, Any] = {"modes": {}, "cross_hf_vs_baseline": {}, "cross_load": {}} + modes = [parse_mode(spec) for spec in args.modes] + for mode in modes: + name = mode["name"] + mode_dir = out / name + if rank == 0: + shutil.rmtree(mode_dir, ignore_errors=True) + mode_dir.mkdir(parents=True) + dist.barrier() + if rank == 0: + print(f"===== {mode}", flush=True) + + engine = build_engine(mode, hf_dir, "ckpt_main") + engine.from_hf(hf_path=hf_dir, strict=True) + engine.save_hf(str(mode_dir / "hf_step0")) + dist.barrier() + entry: dict[str, Any] = {"mode": mode} + if rank == 0: + entry["hf_step0_vs_source"] = compare_hf(src_bf16, mode_dir / "hf_step0") + + losses_first = train(engine, range(0, 5), vocab_size, args.seq_len) + engine.save_dcp(mode_dir / "dcp_step5") + dist.barrier() + losses_cont = train(engine, range(5, 10), vocab_size, args.seq_len) + engine.save_hf(str(mode_dir / "hf_step10")) + dist.barrier() + release(engine) + + resumed = build_engine(mode, hf_dir, "ckpt_resume") + resumed.load_dcp(mode_dir / "dcp_step5") + losses_resumed = train(resumed, range(5, 10), vocab_size, args.seq_len) + resumed.save_hf(str(mode_dir / "hf_step10_resumed")) + dist.barrier() + release(resumed) + + entry["losses_steps_0_4"] = losses_first + entry["losses_steps_5_9_continuous"] = losses_cont + entry["losses_steps_5_9_resumed"] = losses_resumed + entry["resume_max_rel_loss_diff"] = max_rel_diff(losses_resumed, losses_cont) + if rank == 0: + entry["hf_step10_resumed_vs_continuous"] = compare_hf( + mode_dir / "hf_step10", mode_dir / "hf_step10_resumed" + ) + print(json.dumps({k: v for k, v in entry.items() if k != "mode"}, indent=1), flush=True) + report["modes"][name] = entry + + if rank == 0: + baseline = modes[0]["name"] + for mode in modes[1:]: + report["cross_hf_vs_baseline"][mode["name"]] = compare_hf( + out / baseline / "hf_step10", out / mode["name"] / "hf_step10" + ) + + for pair in args.cross_load: + src, dst = pair.split(":") + dst_mode = next(m for m in modes if m["name"] == dst) + if rank == 0: + print(f"===== cross load {src} -> {dst}", flush=True) + engine = build_engine(dst_mode, hf_dir, f"ckpt_cross_{src}") + engine.load_dcp(out / src / "dcp_step5") + losses = train(engine, range(5, 10), vocab_size, args.seq_len) + release(engine) + ref = report["modes"][src]["losses_steps_5_9_continuous"] + report["cross_load"][pair] = { + "losses_steps_5_9": losses, + "max_rel_loss_diff_vs_src_continuous": max_rel_diff(losses, ref), + } + + if rank == 0: + (out / "l3.json").write_text(json.dumps(report, indent=2)) + print(json.dumps({k: report[k] for k in ("cross_hf_vs_baseline", "cross_load")}, indent=1), flush=True) + print(f"wrote {out / 'l3.json'}") + dist.barrier() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/model/run_decoupled_ep_fsdp_numerics.py b/tests/model/run_decoupled_ep_fsdp_numerics.py new file mode 100644 index 0000000000..d8dc0d7494 --- /dev/null +++ b/tests/model/run_decoupled_ep_fsdp_numerics.py @@ -0,0 +1,82 @@ +"""Multi-GPU numerical-equivalence runs for the EP/FSDP decoupling (DESIGN.md §6 L1 / L2). + +Every mode trains the same tiny Qwen3-MoE (random HF checkpoint built once from a fixed seed) +on the same token stream, so loss / grad-norm curves of different layouts can be compared. The +pytest gate for the same comparison lives in ``tests/engine/test_decoupled_ep_fsdp_train_engine.py``; +this script exists for the larger ``--model-size medium`` memory / step-time runs, ``--fp8`` and +``--dispatcher deepep``, and for producing the JSON behind the markdown reports +(``summarize_decoupled_ep_fsdp_numerics.py``). + +Example (L1): + + torchrun --nproc-per-node 8 tests/model/run_decoupled_ep_fsdp_numerics.py \\ + --modes A:ep=1 B:ep=8 C:ep=8,decouple=1 --steps 50 --out reports/l1.json + +Mode spec: ``:key=value[,key=value...]`` with keys ``ep`` (int), ``decouple`` (0/1) and +``hsdp`` (hsdp_sharding_size, int). +""" + +import argparse +import json +import os +from pathlib import Path + +import torch +import torch.distributed as dist + +from xtuner._testing.decoupled_ep_fsdp import MODEL_SIZES, build_hf_checkpoint, parse_mode, run_mode + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--modes", nargs="+", required=True) + parser.add_argument("--steps", type=int, default=50) + parser.add_argument("--seq-len", type=int, default=2048) + parser.add_argument("--lr", type=float, default=1e-4) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--model-size", choices=tuple(MODEL_SIZES), default="tiny") + parser.add_argument("--dispatcher", choices=("all2all", "deepep"), default="all2all") + parser.add_argument("--fp8", action="store_true", help="tile-wise float8 for linear and grouped linear") + parser.add_argument("--hf-dir", type=str, default=None, help="where to build / reuse the tiny HF checkpoint") + parser.add_argument("--out", type=str, required=True) + args = parser.parse_args() + + dist.init_process_group("nccl") + rank = dist.get_rank() + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + + hf_dir = Path(args.hf_dir or Path(args.out).with_suffix("")).with_name(f"{args.model_size}_qwen3_moe_hf") + if rank == 0 and not (hf_dir / "config.json").exists(): + hf_dir.mkdir(parents=True, exist_ok=True) + build_hf_checkpoint(hf_dir, args.seed, args.model_size) + dist.barrier() + + results = [] + for spec in args.modes: + mode = parse_mode(spec) + if rank == 0: + print(f"===== running mode {mode}", flush=True) + results.append(run_mode(mode, hf_dir, args.steps, args.seq_len, args.lr, args.dispatcher, args.fp8)) + dist.barrier() + if rank == 0: + last = results[-1] + print( + f"[{mode['name']}] loss[0]={last['losses'][0]:.6f} loss[-1]={last['losses'][-1]:.6f} " + f"grad_norm[0]={last['grad_norms'][0]:.6f} step_time={last['step_time_s'] * 1000:.1f}ms " + f"mem={last['memory']}", + flush=True, + ) + + if rank == 0: + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + out.write_text( + json.dumps({"args": vars(args), "world_size": dist.get_world_size(), "results": results}, indent=2) + ) + print(f"wrote {out}") + dist.barrier() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/model/summarize_decoupled_ep_fsdp_numerics.py b/tests/model/summarize_decoupled_ep_fsdp_numerics.py new file mode 100644 index 0000000000..516f5cf97c --- /dev/null +++ b/tests/model/summarize_decoupled_ep_fsdp_numerics.py @@ -0,0 +1,90 @@ +"""Turn the JSON written by ``run_decoupled_ep_fsdp_numerics.py`` into markdown tables.""" + +import argparse +import json +from pathlib import Path + + +def rel(a: float, b: float) -> float: + return abs(a - b) / max(abs(b), 1e-12) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("json_path") + parser.add_argument("--ref", default=None, help="reference mode name (default: first mode)") + args = parser.parse_args() + data = json.loads(Path(args.json_path).read_text()) + results = {r["mode"]["name"]: r for r in data["results"]} + names = list(results) + ref = args.ref or names[0] + steps = len(results[ref]["losses"]) + + print("## Modes\n") + print("| mode | ep | decouple | hsdp_sharding_size |") + print("|---|---|---|---|") + for n in names: + m = results[n]["mode"] + print(f"| {n} | {m['ep']} | {m['decouple']} | {m['hsdp']} |") + + print("\n## Loss curve (reduced_llm_loss)\n") + print("| step | " + " | ".join(names) + " | " + " | ".join(f"rel({n} vs {ref})" for n in names if n != ref) + " |") + print("|---|" + "---|" * (2 * len(names) - 1)) + for s in list(range(0, steps, 5)) + [steps - 1]: + row = [f"{results[n]['losses'][s]:.6f}" for n in names] + diffs = [f"{rel(results[n]['losses'][s], results[ref]['losses'][s]):.2e}" for n in names if n != ref] + print(f"| {s} | " + " | ".join(row) + " | " + " | ".join(diffs) + " |") + print("\n| pair | max rel loss diff over all steps | mean rel loss diff |") + print("|---|---|---|") + for n in names: + if n == ref: + continue + d = [rel(a, b) for a, b in zip(results[n]["losses"], results[ref]["losses"])] + print(f"| {n} vs {ref} | {max(d):.2e} | {sum(d) / len(d):.2e} |") + + print("\n## Total grad norm (pre-clip)\n") + print("| step | " + " | ".join(names) + " | " + " | ".join(f"rel({n} vs {ref})" for n in names if n != ref) + " |") + print("|---|" + "---|" * (2 * len(names) - 1)) + for s in list(range(0, steps, 5)) + [steps - 1]: + row = [f"{results[n]['grad_norms'][s]:.6f}" for n in names] + diffs = [f"{rel(results[n]['grad_norms'][s], results[ref]['grad_norms'][s]):.2e}" for n in names if n != ref] + print(f"| {s} | " + " | ".join(row) + " | " + " | ".join(diffs) + " |") + + print("\n## Per-parameter grad norms (max / mean relative diff vs reference)\n") + print("| step | pair | dense params | expert params | worst param |") + print("|---|---|---|---|---|") + for s in results[ref]["param_grad_norms"]: + ref_norms = results[ref]["param_grad_norms"][s] + for n in names: + if n == ref: + continue + norms = results[n]["param_grad_norms"][s] + dense, expert, worst = [], [], (0.0, "") + for k, v in ref_norms.items(): + d = rel(norms[k], v) + (expert if ".experts" in k else dense).append(d) + if d > worst[0]: + worst = (d, k) + print( + f"| {s} | {n} vs {ref} | max {max(dense):.2e} / mean {sum(dense) / len(dense):.2e} " + f"| max {max(expert):.2e} / mean {sum(expert) / len(expert):.2e} | {worst[1]} ({worst[0]:.2e}) |" + ) + + print("\n## Per-rank memory (MiB) and step time\n") + keys = [ + "dense_param_mib", + "expert_param_mib", + "allocated_after_load_mib", + "allocated_after_step0_mib", + "peak_allocated_mib", + "peak_reserved_mib", + ] + print("| mode | " + " | ".join(keys) + " | step time (ms) |") + print("|---|" + "---|" * (len(keys) + 1)) + for n in names: + mem = results[n]["memory"] + print(f"| {n} | " + " | ".join(f"{mem[k]:.1f}" for k in keys) + f" | {results[n]['step_time_s'] * 1000:.1f} |") + + +if __name__ == "__main__": + main() diff --git a/tests/model/test_decoupled_ep_fsdp_mesh.py b/tests/model/test_decoupled_ep_fsdp_mesh.py new file mode 100644 index 0000000000..ae0c1d0480 --- /dev/null +++ b/tests/model/test_decoupled_ep_fsdp_mesh.py @@ -0,0 +1,471 @@ +"""L0 placement tests for the MoE EP/FSDP mesh layouts. + +These tests run in a single process on top of a fake ``ProcessGroup`` so that +8/16/64-rank layouts can be asserted without any multi-process launch. The +model stays on the ``meta`` device; only the mesh bookkeeping and the DTensor +placements produced by ``MoE.fully_shard`` are inspected. +""" + +from collections.abc import Iterator +from contextlib import contextmanager + +import pytest +import torch +import torch.distributed as dist +from pydantic import ValidationError +from torch.distributed.device_mesh import DeviceMesh, _mesh_resources +from torch.distributed.tensor import DTensor, Replicate, Shard +from torch.distributed.tensor.placement_types import _StridedShard +from torch.testing._internal.distributed.fake_pg import FakeStore + +from xtuner.v1.config import FSDPConfig +from xtuner.v1.model.base import BaseModel +from xtuner.v1.model.moe.moe import MoE, MoEConfig +from xtuner.v1.module.attention import MHAConfig +from xtuner.v1.module.router import NoAuxRouterConfig + + +PREFIX = "l0mesh" +HIDDEN = 128 +N_EXPERTS = 64 +MOE_INTER = 128 + +DENSE_PARAMS = ( + "embed_tokens.weight", + "layers.0.mlp.gate_proj.weight", + "layers.1.self_attn.q_proj.weight", + "layers.1.gate.weight", + "layers.1.shared_experts.gate_proj.weight", + "layers.1.input_layernorm.weight", + "norm.weight", + "lm_head.weight", +) +EXPERT_PARAMS = ( + "layers.1.experts.fused_w1w3.weight", + "layers.1.experts.fused_w2.weight", +) + + +def _reset_mesh_resources() -> None: + _mesh_resources.mesh_stack.clear() + _mesh_resources.child_to_root_mapping.clear() + _mesh_resources.root_to_flatten_mapping.clear() + _mesh_resources.flatten_name_to_root_dims.clear() + + +@contextmanager +def _fake_world(world_size: int, rank: int) -> Iterator[None]: + dist.init_process_group("fake", store=FakeStore(), rank=rank, world_size=world_size) + try: + yield + finally: + dist.destroy_process_group() + _reset_mesh_resources() + + +# `MoE.__init__` unconditionally creates CUDA streams and dispatcher buffers, so a +# CUDA context is required even though the fake ProcessGroup never communicates. +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="MoE construction requires a CUDA device") + + +@pytest.fixture(autouse=True) +def _keep_model_on_meta(monkeypatch: pytest.MonkeyPatch) -> None: + # Placements are fully determined before materialization; skipping + # `_to_empty_meta` keeps the test free of device allocations. + monkeypatch.setattr(BaseModel, "_to_empty_meta", lambda self: None) + + +def _build_model(ep_size: int) -> MoE: + config = MoEConfig( + vocab_size=1024, + max_position_embeddings=256, + pad_token_id=0, + eos_token_id=0, + num_hidden_layers=2, + hidden_size=HIDDEN, + intermediate_size=256, + rms_norm_eps=1e-6, + rope_theta=1e6, + hidden_act="silu", + attention=MHAConfig(num_attention_heads=4, num_key_value_heads=4, head_dim=32), + tie_word_embeddings=False, + n_routed_experts=N_EXPERTS, + n_shared_experts=1, + num_experts_per_tok=2, + first_k_dense_replace=1, + hidden_factor=1.0, + moe_intermediate_size=MOE_INTER, + router=NoAuxRouterConfig( + scoring_func="sigmoid", + router_scaling_factor=1.0, + n_group=1, + topk_group=1, + norm_topk_prob=True, + ), + compile_cfg=False, + ep_size=ep_size, + dispatcher="all2all", + mesh_prefix=PREFIX, + ) + with torch.device("meta"): + return MoE(config) + + +def _shard_model(ep_size: int, **fsdp_kwargs) -> MoE: + model = _build_model(ep_size) + fsdp_config = FSDPConfig(ep_size=ep_size, torch_compile=False, **fsdp_kwargs) + return model.fully_shard(fsdp_config) + + +def _named_params(model: MoE) -> dict[str, DTensor]: + params: dict[str, DTensor] = {} + for name, param in model.named_parameters(): + params[name.replace("_checkpoint_wrapped_module.", "")] = param + return params + + +def _ranks(mesh: DeviceMesh) -> list: + return mesh.mesh.tolist() + + +def _legacy_ep_ranks(world_size: int, ep_size: int, rank: int) -> list[int]: + start = (rank // ep_size) * ep_size + return list(range(start, start + ep_size)) + + +def _legacy_fsdp_ranks(world_size: int, ep_size: int, rank: int) -> list[int]: + return [rank % ep_size + k * ep_size for k in range(world_size // ep_size)] + + +def _assert_placements(param: DTensor, mesh_dim_names: tuple[str, ...], placements: tuple) -> None: + assert isinstance(param, DTensor) + assert param.device_mesh.mesh_dim_names == mesh_dim_names + assert len(param.placements) == len(placements) + for actual, expected in zip(param.placements, placements): + assert type(actual) is type(expected), (actual, expected) + if isinstance(expected, _StridedShard): + assert actual.dim == expected.dim and actual.split_factor == expected.split_factor + elif isinstance(expected, Shard): + assert actual.dim == expected.dim + + +class TestLegacyMoEMeshLayout: + """Layout produced by the original EP-orthogonal-to-FSDP path. + + ``decouple_ep_fsdp`` is off by default, so these expectations are the + regression contract for the legacy path: ``root = (fsdp = world / ep, ep)`` + with ``ep`` on the innermost (contiguous-rank) dimension. + """ + + @pytest.mark.parametrize( + "world_size,ep_size,rank", + [(8, 8, 0), (8, 8, 5), (8, 4, 6), (8, 2, 3), (16, 8, 9), (64, 8, 37)], + ) + def test_mesh_ranks(self, world_size: int, ep_size: int, rank: int) -> None: + with _fake_world(world_size, rank): + model = _shard_model(ep_size) + + assert model.hsdp_mesh is None + assert model._world_mesh is not None + assert tuple(model._world_mesh.shape) == (world_size // ep_size, ep_size) + assert model._world_mesh.mesh_dim_names == (f"{PREFIX}.fsdp", f"{PREFIX}.ep") + + assert model.fsdp_mesh is not None and model.ep_mesh is not None + assert model.fsdp_mesh.mesh_dim_names == (f"{PREFIX}.fsdp",) + assert _ranks(model.fsdp_mesh) == _legacy_fsdp_ranks(world_size, ep_size, rank) + assert model.ep_mesh.mesh_dim_names == (f"{PREFIX}.ep",) + assert _ranks(model.ep_mesh) == _legacy_ep_ranks(world_size, ep_size, rank) + assert dist.get_process_group_ranks(model.ep_mesh.get_group()) == _legacy_ep_ranks( + world_size, ep_size, rank + ) + + @pytest.mark.parametrize("world_size,ep_size,rank", [(8, 8, 0), (16, 8, 9), (64, 8, 37)]) + def test_param_placements(self, world_size: int, ep_size: int, rank: int) -> None: + fsdp_size = world_size // ep_size + with _fake_world(world_size, rank): + model = _shard_model(ep_size) + params = _named_params(model) + mesh_dim_names = (f"{PREFIX}.fsdp", f"{PREFIX}.ep") + + for name in DENSE_PARAMS: + param = params[name] + # Dense params are replicated over ep and only sharded world/ep ways. + _assert_placements(param, mesh_dim_names, (Shard(0), Replicate())) + assert param.to_local().shape[0] == param.shape[0] // fsdp_size, name + + for name in EXPERT_PARAMS: + param = params[name] + _assert_placements(param, mesh_dim_names, (_StridedShard(0, split_factor=ep_size), Shard(0))) + assert param.to_local().shape[0] == param.shape[0] // world_size, name + + def test_ep1_is_plain_fsdp(self) -> None: + world_size, rank = 8, 3 + with _fake_world(world_size, rank): + model = _shard_model(ep_size=1) + params = _named_params(model) + + assert model.ep_mesh is not None and model.ep_mesh.size() == 1 + assert model.fsdp_mesh is not None and _ranks(model.fsdp_mesh) == list(range(world_size)) + for name in DENSE_PARAMS + EXPERT_PARAMS: + param = params[name] + _assert_placements(param, (f"{PREFIX}.fsdp",), (Shard(0),)) + assert param.to_local().shape[0] == param.shape[0] // world_size, name + + def test_hsdp_ep1_layout(self) -> None: + world_size, shard_size, rank = 16, 8, 11 + with _fake_world(world_size, rank): + model = _shard_model(ep_size=1, hsdp_sharding_size=shard_size) + params = _named_params(model) + + assert model.hsdp_mesh is not None + assert tuple(model.hsdp_mesh.shape) == (world_size // shard_size, shard_size) + assert model.hsdp_mesh.mesh_dim_names == (f"{PREFIX}.hsdp_replicate", f"{PREFIX}.hsdp_shard") + assert model.fsdp_mesh is not None + assert _ranks(model.fsdp_mesh) == list( + range((rank // shard_size) * shard_size, (rank // shard_size + 1) * shard_size) + ) + for name in DENSE_PARAMS + EXPERT_PARAMS: + param = params[name] + _assert_placements(param, model.hsdp_mesh.mesh_dim_names, (Replicate(), Shard(0))) + assert param.to_local().shape[0] == param.shape[0] // shard_size, name + + def test_hsdp_with_ep_is_rejected(self) -> None: + with pytest.raises(ValidationError, match="HSDP requires expert parallel size to be 1"): + FSDPConfig(ep_size=8, hsdp_sharding_size=8) + + +def _decoupled_ranks(world_size: int, ep_size: int, dp_shard: int, rank: int) -> dict[str, list]: + replicate = world_size // dp_shard + efsdp = dp_shard // ep_size + block_start = (rank // dp_shard) * dp_shard + ep_start = (rank // ep_size) * ep_size + efsdp_in_block = [(rank % ep_size) + k * ep_size for k in range(efsdp)] + return { + "ep": list(range(ep_start, ep_start + ep_size)), + "efsdp": [block_start + r for r in efsdp_in_block], + "dp_shard": list(range(block_start, block_start + dp_shard)), + "hsdp": [[r * dp_shard + j for j in range(dp_shard)] for r in range(replicate)], + "expert_hsdp": [[r * dp_shard + j for j in efsdp_in_block] for r in range(replicate)], + } + + +class TestDecoupledMoEMeshShapes: + """Mesh bookkeeping of the decoupled (dp2ep) path: ``root = (replicate, efsdp, ep)``.""" + + @pytest.mark.parametrize( + "world_size,ep_size,rank", + [(8, 8, 5), (8, 4, 6), (16, 8, 9), (64, 8, 37), (64, 32, 50)], + ) + def test_mesh_without_hsdp(self, world_size: int, ep_size: int, rank: int) -> None: + expected = _decoupled_ranks(world_size, ep_size, world_size, rank) + with _fake_world(world_size, rank): + model = _build_model(ep_size) + model._init_device_mesh(FSDPConfig(ep_size=ep_size, torch_compile=False, decouple_ep_fsdp=True)) + + root = model._world_mesh + assert root is not None + assert tuple(root.shape) == (1, world_size // ep_size, ep_size) + assert root.mesh_dim_names == (f"{PREFIX}.replicate", f"{PREFIX}.efsdp", f"{PREFIX}.ep") + assert _mesh_resources.get_root_mesh(model.ep_mesh) is root + assert _mesh_resources.get_root_mesh(model.fsdp_mesh) is root + assert _mesh_resources.get_root_mesh(model.expert_fsdp_mesh) is root + + assert model.hsdp_mesh is None + assert model.ep_mesh.mesh_dim_names == (f"{PREFIX}.ep",) + assert _ranks(model.ep_mesh) == expected["ep"] + assert dist.get_process_group_ranks(model.ep_mesh.get_group()) == expected["ep"] + assert model.expert_fsdp_mesh.mesh_dim_names == (f"{PREFIX}.efsdp",) + assert _ranks(model.expert_fsdp_mesh) == expected["efsdp"] + assert model.fsdp_mesh.mesh_dim_names == (f"{PREFIX}.dp_shard",) + assert _ranks(model.fsdp_mesh) == expected["dp_shard"] + assert dist.get_process_group_ranks(model.fsdp_mesh.get_group()) == expected["dp_shard"] + + @pytest.mark.parametrize( + "world_size,ep_size,dp_shard,rank", + [(16, 8, 8, 11), (64, 8, 32, 37), (64, 8, 8, 3)], + ) + def test_mesh_with_hsdp(self, world_size: int, ep_size: int, dp_shard: int, rank: int) -> None: + expected = _decoupled_ranks(world_size, ep_size, dp_shard, rank) + replicate = world_size // dp_shard + with _fake_world(world_size, rank): + model = _build_model(ep_size) + model._init_device_mesh( + FSDPConfig(ep_size=ep_size, torch_compile=False, decouple_ep_fsdp=True, hsdp_sharding_size=dp_shard) + ) + + root = model._world_mesh + assert root is not None + assert tuple(root.shape) == (replicate, dp_shard // ep_size, ep_size) + assert _ranks(model.ep_mesh) == expected["ep"] + assert _ranks(model.fsdp_mesh) == expected["dp_shard"] + + assert model.hsdp_mesh is not None + assert model.hsdp_mesh.mesh_dim_names == (f"{PREFIX}.replicate", f"{PREFIX}.dp_shard") + assert _ranks(model.hsdp_mesh) == expected["hsdp"] + assert model.expert_fsdp_mesh.mesh_dim_names == (f"{PREFIX}.replicate", f"{PREFIX}.efsdp") + assert _ranks(model.expert_fsdp_mesh) == expected["expert_hsdp"] + assert _mesh_resources.get_root_mesh(model.hsdp_mesh) is root + assert _mesh_resources.get_root_mesh(model.expert_fsdp_mesh) is root + + def test_config_requires_ep_to_divide_shard_size(self) -> None: + FSDPConfig(ep_size=8, hsdp_sharding_size=8, decouple_ep_fsdp=True) + FSDPConfig(ep_size=4, hsdp_sharding_size=8, decouple_ep_fsdp=True) + with pytest.raises(ValidationError, match="divisible by `ep_size`"): + FSDPConfig(ep_size=8, hsdp_sharding_size=4, decouple_ep_fsdp=True) + with pytest.raises(ValidationError, match="divisible by `ep_size`"): + FSDPConfig(ep_size=3, hsdp_sharding_size=8, decouple_ep_fsdp=True) + + def test_config_rejects_non_positive_sizes(self) -> None: + # The checks are pydantic validators raising `ValueError`, not `assert`s that `python -O` + # would drop; `ValidationError` is the `ValueError` subclass pydantic wraps them in. + with pytest.raises(ValidationError, match="`ep_size` must be a positive integer"): + FSDPConfig(ep_size=0) + with pytest.raises(ValidationError, match="`hsdp_sharding_size` must be a positive integer"): + FSDPConfig(hsdp_sharding_size=0, decouple_ep_fsdp=True) + assert issubclass(ValidationError, ValueError) + + def test_runtime_rejects_shard_size_not_dividing_world_size(self) -> None: + # `hsdp_sharding_size=3` passes the config validator (3 % ep_size == 0) but cannot tile + # an 8-rank world; the mesh construction must fail loudly instead of asserting. + with _fake_world(8, 0): + with pytest.raises(ValueError, match=r"world_size \(8\) must be divisible by hsdp_sharding_size \(3\)"): + _shard_model(1, hsdp_sharding_size=3, decouple_ep_fsdp=True) + + +class TestDecoupledMoEParamPlacements: + """DTensor placements after the two-level ``fully_shard`` of the decoupled path.""" + + @pytest.mark.parametrize("world_size,ep_size,rank", [(8, 8, 0), (16, 8, 9), (64, 8, 37)]) + def test_param_placements(self, world_size: int, ep_size: int, rank: int) -> None: + efsdp = world_size // ep_size + with _fake_world(world_size, rank): + model = _shard_model(ep_size, decouple_ep_fsdp=True) + params = _named_params(model) + + for name in DENSE_PARAMS: + param = params[name] + # Dense params: plain FSDP over the full dp_shard mesh, no EP replication. + _assert_placements(param, (f"{PREFIX}.dp_shard",), (Shard(0),)) + assert param.to_local().shape[0] == param.shape[0] // world_size, name + + for name in EXPERT_PARAMS: + param = params[name] + _assert_placements( + param, + (f"{PREFIX}.efsdp", f"{PREFIX}.ep"), + (_StridedShard(0, split_factor=ep_size), Shard(0)), + ) + assert param.device_mesh.size(0) == efsdp + assert param.to_local().shape[0] == param.shape[0] // world_size, name + + from torch.distributed.fsdp import FSDPModule + + layer = model.layers["1"] + assert isinstance(layer.experts, FSDPModule) + assert isinstance(layer, FSDPModule) + assert not isinstance(model.layers["0"].mlp, FSDPModule) + + @pytest.mark.parametrize("world_size,ep_size,dp_shard,rank", [(16, 8, 8, 11), (64, 8, 32, 37)]) + def test_param_placements_with_hsdp(self, world_size: int, ep_size: int, dp_shard: int, rank: int) -> None: + efsdp = dp_shard // ep_size + with _fake_world(world_size, rank): + model = _shard_model(ep_size, decouple_ep_fsdp=True, hsdp_sharding_size=dp_shard) + params = _named_params(model) + + for name in DENSE_PARAMS: + param = params[name] + _assert_placements(param, (f"{PREFIX}.replicate", f"{PREFIX}.dp_shard"), (Replicate(), Shard(0))) + assert param.to_local().shape[0] == param.shape[0] // dp_shard, name + + for name in EXPERT_PARAMS: + param = params[name] + _assert_placements( + param, + (f"{PREFIX}.replicate", f"{PREFIX}.efsdp", f"{PREFIX}.ep"), + (Replicate(), _StridedShard(0, split_factor=ep_size), Shard(0)), + ) + assert param.device_mesh.size(1) == efsdp + assert param.to_local().shape[0] == param.shape[0] // dp_shard, name + + def test_legacy_layout_is_untouched_by_default(self) -> None: + # Same construction as the decoupled tests, but with the flag left at its default. + world_size, ep_size, rank = 16, 8, 9 + with _fake_world(world_size, rank): + model = _shard_model(ep_size) + params = _named_params(model) + assert model.expert_fsdp_mesh is None + _assert_placements( + params["layers.1.self_attn.q_proj.weight"], (f"{PREFIX}.fsdp", f"{PREFIX}.ep"), (Shard(0), Replicate()) + ) + + +class TestDecoupledFloat8Meshes: + """Tile-wise fp8 bookkeeping on the decoupled path: per-class FSDP chunk counts and reduce meshes.""" + + def test_reduce_meshes_follow_each_class_shard_stride(self) -> None: + from xtuner.v1.float8.config import Float8Config, ScalingGranularity + + world_size, ep_size, rank = 64, 8, 37 + with _fake_world(world_size, rank): + config = _build_model(ep_size).config.model_copy( + update={ + "hidden_size": 256, + "n_routed_experts": 16, + "moe_intermediate_size": 128, + "attention": MHAConfig(num_attention_heads=4, num_key_value_heads=4, head_dim=64), + "float8_cfg": Float8Config( + scaling_granularity_gemm=ScalingGranularity.TILEWISE, + scaling_granularity_grouped_gemm=ScalingGranularity.TILEWISE, + ), + } + ) + with torch.device("meta"): + model = MoE(config) + model.fully_shard(FSDPConfig(ep_size=ep_size, torch_compile=False, decouple_ep_fsdp=True)) + params = _named_params(model) + handler = model.float8_handler + assert handler is not None + + # dense q_proj (256, 256) is sharded 64 ways -> 4 local rows -> absmax reduced over 32 contiguous ranks + dense_local = tuple(params["layers.1.self_attn.q_proj.weight"].to_local().shape) + assert dense_local == (4, 256) + assert _ranks(handler.tilewise_reduce_mesh_mapping[dense_local]) == list(range(32, 64)) + # experts: 16 experts x 256 rows = 4096 rows sharded (efsdp=8) x (ep=8) ways -> 64 local rows + # -> absmax reduced over 2 ranks that are `ep` apart (neighbours on the efsdp dim) + expert_local = tuple(params["layers.1.experts.fused_w1w3.weight"].to_local().shape) + assert expert_local == (64, 256) + assert handler.expert_tilewise_reduce_mesh_mapping is not None + assert _ranks(handler.expert_tilewise_reduce_mesh_mapping[expert_local]) == [37, 45] + assert dense_local not in handler.expert_tilewise_reduce_mesh_mapping + assert expert_local not in handler.tilewise_reduce_mesh_mapping + + # "n*128+64" straddling blocks: dense neighbours are adjacent ranks, expert neighbours are `ep` apart + assert _ranks(handler.tilewise_reduce_mesh_devided_64) == [36, 37] + assert _ranks(handler.expert_tilewise_reduce_mesh_devided_64) == [37, 45] + + def test_legacy_float8_meshes_unchanged(self) -> None: + from xtuner.v1.float8.config import Float8Config, ScalingGranularity + + world_size, ep_size, rank = 64, 8, 37 + with _fake_world(world_size, rank): + config = _build_model(ep_size).config.model_copy( + update={ + "hidden_size": 256, + "n_routed_experts": 16, + "moe_intermediate_size": 128, + "attention": MHAConfig(num_attention_heads=4, num_key_value_heads=4, head_dim=64), + "float8_cfg": Float8Config( + scaling_granularity_gemm=ScalingGranularity.TILEWISE, + scaling_granularity_grouped_gemm=ScalingGranularity.TILEWISE, + ), + } + ) + with torch.device("meta"): + model = MoE(config) + model.fully_shard(FSDPConfig(ep_size=ep_size, torch_compile=False)) + handler = model.float8_handler + assert handler is not None + assert handler.expert_tilewise_reduce_mesh_mapping is None + assert handler.expert_tilewise_reduce_mesh_devided_64 is None + # legacy: every class is sharded on the same (fsdp=8, stride ep=8) mesh + assert _ranks(handler.tilewise_reduce_mesh_devided_64) == [37, 45] diff --git a/tests/rl/test_weight_iterator.py b/tests/rl/test_weight_iterator.py index ec7a3f7884..83bdb19fe5 100644 --- a/tests/rl/test_weight_iterator.py +++ b/tests/rl/test_weight_iterator.py @@ -10,6 +10,7 @@ from xtuner.v1.float8.fsdp_utils import WeightWithDynamicTilewiseFloat8CastTensor from xtuner.v1.float8.triton_kernels.per_block_quant_gemm import per_block_quant_torch from xtuner.v1.model.base import BaseModel, HFSaveCfg, XTunerBaseModelConfig +from xtuner.v1.model.compose.base import BaseComposeConfig from xtuner.v1.rl.weight_update.data import RolloutBackend, RolloutWeightUpdateInfo, RolloutWeightUpdateTarget from xtuner.v1.rl.weight_update.weight_iterator import WeightIterator from xtuner.v1.utils import get_device @@ -231,3 +232,196 @@ def gather_train_ep_shards( state_dict[f"expert_{index}"].cpu(), torch.tensor([index], dtype=torch.bfloat16), ) + + +def _bf16_param(values: list[int]) -> nn.Parameter: + return nn.Parameter(torch.tensor(values, device=get_device(), dtype=torch.bfloat16)) + + +class _DecoupledLanguageTower(BaseModel): + """Language tower on the decoupled EP/FSDP layout as seen from rank 0. + + ``layers.0.experts.weight`` is EP-sharded and then FSDP-sharded on ``efsdp``; the dense + parameters are FSDP-sharded on ``dp_shard``. Every local block holds ``[0, 1]``. + """ + + def __init__(self, groups: dict[str, dist.ProcessGroup], expert_fsdp_ndim: int) -> None: + super().__init__(XTunerBaseModelConfig()) + layer = nn.Module() + layer.attn = nn.Module() + layer.attn.weight = _bf16_param([0, 1]) + layer.experts = nn.Module() + layer.experts.weight = _bf16_param([0, 1]) + self.layers = nn.ModuleDict({"0": layer}) + self.norm = nn.Module() + self.norm.weight = _bf16_param([0, 1]) + self.fsdp_mesh = SimpleNamespace(get_group=lambda: groups["dp_shard"]) + self.expert_fsdp_mesh = SimpleNamespace(ndim=expert_fsdp_ndim, get_group=lambda dim=None: groups["efsdp"]) + dense_shards = [ShardDescriptor(dim=0, group=groups["dp_shard"])] + self.load_spec_mapping = { + "layers.0.attn.weight": LoadSpec( + name="layers.0.attn.weight", + global_hf_keys=["attn"], + global_shape=(8,), + shards=dense_shards, + local_shape=(2,), + ), + "layers.0.experts.weight": LoadSpec( + name="layers.0.experts.weight", + global_hf_keys=[f"expert_{index}" for index in range(8)], + global_shape=(8,), + fused_dim=0, + shards=[ShardDescriptor(dim=0, group=groups["ep"]), ShardDescriptor(dim=0, group=groups["efsdp"])], + local_shape=(2,), + ), + "norm.weight": LoadSpec( + name="norm.weight", + global_hf_keys=["norm"], + global_shape=(8,), + shards=dense_shards, + local_shape=(2,), + ), + } + + def to_hf_key_list(self, key: str) -> list[str]: + return [key] + + +class _WorldShardedTower(BaseModel): + """Vision tower / projector stand-in: one parameter FSDP-sharded on the world mesh.""" + + def __init__(self, block_name: str, world_group: dist.ProcessGroup, world_size: int) -> None: + super().__init__(XTunerBaseModelConfig()) + block = nn.Module() + block.weight = _bf16_param(list(range(8 // world_size))) + setattr(self, block_name, block) + self.fsdp_mesh = SimpleNamespace(get_group=lambda: world_group) + self.load_spec_mapping = { + f"{block_name}.weight": LoadSpec( + name=f"{block_name}.weight", + global_hf_keys=[block_name], + global_shape=(8,), + shards=[ShardDescriptor(dim=0, group=world_group)], + local_shape=(8 // world_size,), + ) + } + + def to_hf_key_list(self, key: str) -> list[str]: + return [key] + + +class _ComposeModel(BaseModel): + """Compose model stand-in: like ``BaseComposeModel.fully_shard`` it is wrapped on the world + mesh and has neither an ``expert_fsdp_mesh`` nor a ``load_spec_mapping`` of its own.""" + + def __init__( + self, + language_model: BaseModel, + vision_tower: BaseModel, + projector: BaseModel, + world_group: dist.ProcessGroup, + ) -> None: + super().__init__(XTunerBaseModelConfig()) + # `iter_layer_batches` only inspects `isinstance(model.config, BaseComposeConfig)`. + self.config = BaseComposeConfig.model_construct() + self.language_model = language_model + self.vision_tower = vision_tower + self.multi_modal_projector = projector + self.fsdp_mesh = SimpleNamespace(get_group=lambda: world_group) + + +class TestLayerBatchesGatherWithParamOwner: + """``iter_layer_batches`` must gather every parameter with the module that owns it. + + Regression test for compose MoE models on the decoupled EP/FSDP layout with the IPC/Turbomind + transport: the outer compose model is wrapped on the world mesh, so gathering with it kept the + language tower's ``efsdp`` (and HSDP ``dp_shard``) shards local and streamed rank-local + fragments instead of complete weights. + """ + + @pytest.mark.parametrize( + ("world_size", "expert_fsdp_ndim"), + [ + pytest.param(4, 1, id="ep2-efsdp2"), + pytest.param(8, 2, id="hsdp-replicate2-ep2-efsdp2"), + ], + ) + def test_compose_layer_batches_are_gathered_per_owner( + self, + monkeypatch: pytest.MonkeyPatch, + world_size: int, + expert_fsdp_ndim: int, + ) -> None: + groups = {name: dist.ProcessGroup(dist.HashStore(), 0, 1) for name in ("world", "dp_shard", "ep", "efsdp")} + group_ranks = { + groups["world"]: list(range(world_size)), + groups["dp_shard"]: [0, 1, 2, 3], + groups["ep"]: [0, 1], + groups["efsdp"]: [0, 2], + } + monkeypatch.setattr( + dist, + "get_world_size", + lambda group=None: world_size if group is None else len(group_ranks[group]), + ) + monkeypatch.setattr(dist, "get_rank", lambda group=None: 0) + monkeypatch.setattr(dist, "get_process_group_ranks", lambda group: group_ranks[group]) + + gathered_groups: list[dist.ProcessGroup] = [] + + def gather_contiguous_blocks( + tensor_list: list[torch.Tensor], + group: dist.ProcessGroup, + ) -> list[list[torch.Tensor]]: + # Rank ``r`` of ``group`` holds the r-th contiguous block, each block the size of ours. + gathered_groups.append(group) + return [ + [tensor + rank * tensor.numel() for rank in range(len(group_ranks[group]))] for tensor in tensor_list + ] + + monkeypatch.setattr(load_spec_module, "foreach_all_gather", gather_contiguous_blocks) + + language_model = _DecoupledLanguageTower(groups, expert_fsdp_ndim) + vision_tower = _WorldShardedTower("patch", groups["world"], world_size) + projector = _WorldShardedTower("proj", groups["world"], world_size) + model = _ComposeModel(language_model, vision_tower, projector, groups["world"]) + rollout_info = RolloutWeightUpdateInfo( + rollout_config=cast(Any, SimpleNamespace()), + weight_update_targets=( + RolloutWeightUpdateTarget( + endpoint_rank=0, + update_ranks=(0,), + inference_engine_ranks=(0,), + server_url="http://rollout", + lifecycle_state="active", + ), + ), + train_rank=0, + transport_type="ipc", + backend="turbomind", + ) + iterator = WeightIterator( + config=SimpleNamespace(model_cfg=model.config), + engine=SimpleNamespace(model=model), + rollout_info=rollout_info, + global_hf_keys_mapping_cache={}, + ) + + state_dict = { + name: tensor.cpu() for batch in iterator.iter_layer_batches() for name, tensor in batch.state_dict.items() + } + + full = torch.arange(8, dtype=torch.bfloat16) + expected = { + # The EP shard stays local (experts 0-3 of this EP rank); only the `efsdp` shard is gathered. + "model.language_model.layers.0.mlp.experts.weight": full[:4], + "model.language_model.layers.0.attn.weight": full, + "model.language_model.norm.weight": full, + "model.vision_tower.patch.weight": full, + "model.multi_modal_projector.proj.weight": full, + } + assert set(state_dict) == set(expected) + for name, tensor in expected.items(): + torch.testing.assert_close(state_dict[name], tensor, rtol=0, atol=0) + assert {groups["efsdp"], groups["dp_shard"], groups["world"]} <= set(gathered_groups) + assert groups["ep"] not in gathered_groups diff --git a/xtuner/_testing/decoupled_ep_fsdp.py b/xtuner/_testing/decoupled_ep_fsdp.py new file mode 100644 index 0000000000..fe1611dbd7 --- /dev/null +++ b/xtuner/_testing/decoupled_ep_fsdp.py @@ -0,0 +1,461 @@ +"""Shared helpers for the decoupled EP/FSDP ("dp2ep") numerics and checkpoint checks. + +The GPU pytest gates in ``tests/engine/test_decoupled_ep_fsdp_train_engine.py`` and the manual +experiment scripts under ``tests/model/`` (``run_decoupled_ep_fsdp_numerics.py`` and +``run_decoupled_ep_fsdp_ckpt.py``) build the same tiny random Qwen3-MoE, feed it the same token +stream and compare layouts with the same functions, so a number in a report and a gate threshold +mean the same thing. +""" + +from __future__ import annotations + +import gc +import time +from collections.abc import Iterable, Sequence +from pathlib import Path +from typing import Any, TypedDict + +import torch +import torch.distributed as dist +from safetensors.torch import load_file +from torch.distributed.tensor import DTensor + +from xtuner.v1.config import AdamWConfig, FSDPConfig +from xtuner.v1.engine.train_engine import TrainEngine +from xtuner.v1.float8.config import Float8Config, ScalingGranularity +from xtuner.v1.loss.ce_loss import CELossConfig +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 +from xtuner.v1.utils.dtensor import cal_total_norm + + +DEVICE = get_device() + +MODEL_SIZES: dict[str, dict[str, int]] = { + # ~108M params: numerics (fast, 50 steps in < 10 s) + "tiny": { + "vocab_size": 4096, + "hidden_size": 512, + "intermediate_size": 1024, + "moe_intermediate_size": 256, + "num_hidden_layers": 4, + "num_attention_heads": 8, + "num_key_value_heads": 4, + "head_dim": 64, + "num_experts": 64, + "num_experts_per_tok": 4, + }, + # ~3.4B params (3.2B in routed experts): step-time / memory comparison + "medium": { + "vocab_size": 32768, + "hidden_size": 2048, + "intermediate_size": 4096, + "moe_intermediate_size": 1024, + "num_hidden_layers": 8, + "num_attention_heads": 16, + "num_key_value_heads": 4, + "head_dim": 128, + "num_experts": 64, + "num_experts_per_tok": 4, + }, +} + + +class LayoutMode(TypedDict): + """One EP/FSDP layout under test. + + ``name`` labels the layout in results, ``ep`` is the expert-parallel size, ``decouple`` selects + ``FSDPConfig.decouple_ep_fsdp`` and ``hsdp`` is ``FSDPConfig.hsdp_sharding_size`` (``None`` for + plain FSDP over the world). + """ + + name: str + ep: int + decouple: bool + hsdp: int | None + + +def parse_mode(spec: str) -> LayoutMode: + """Parse a ``:key=value[,key=value...]`` layout spec. + + Args: + spec (str): Mode spec with keys ``ep`` (int), ``decouple`` (0/1) and ``hsdp`` (int). + + Returns: + LayoutMode: The parsed layout; unspecified keys default to ``ep=1, decouple=0, hsdp=None``. + """ + name, _, kv = spec.partition(":") + mode: LayoutMode = {"name": name, "ep": 1, "decouple": False, "hsdp": None} + for item in filter(None, kv.split(",")): + key, value = item.split("=") + if key == "ep": + mode["ep"] = int(value) + elif key == "decouple": + mode["decouple"] = bool(int(value)) + elif key == "hsdp": + mode["hsdp"] = int(value) + else: + raise ValueError(f"unknown mode key {key}") + return mode + + +def build_hf_checkpoint(path: Path, seed: int, model_size: str) -> None: + """Write a randomly initialised Qwen3-MoE HF checkpoint to ``path``. + + Args: + path (Path): Output directory (created by ``save_pretrained``). + seed (int): Seed for the random initialisation. + model_size (str): A key of :data:`MODEL_SIZES`. + """ + from transformers import Qwen3MoeConfig, Qwen3MoeForCausalLM + + config = Qwen3MoeConfig( + **MODEL_SIZES[model_size], + max_position_embeddings=8192, + norm_topk_prob=True, + tie_word_embeddings=False, + bos_token_id=1, + eos_token_id=2, + rope_theta=1000000.0, + use_sliding_window=False, + max_window_layers=MODEL_SIZES[model_size]["num_hidden_layers"], + ) + torch.manual_seed(seed) + model = Qwen3MoeForCausalLM(config) + model.save_pretrained(path, safe_serialization=True) + + +def make_batch(step: int, rank: int, vocab_size: int, seq_len: int) -> ModelItem: + """Build the deterministic random token batch of ``(step, rank)``. + + Args: + step (int): Training step; together with ``rank`` it seeds the batch. + rank (int): Data-parallel rank. + vocab_size (int): Vocabulary size to sample tokens from. + seq_len (int): Number of input tokens. + + Returns: + ModelItem: Sequence context and CE loss context on :data:`DEVICE`. + """ + generator = torch.Generator().manual_seed(100_000 + step * 1024 + rank) + tokens = torch.randint(0, vocab_size, (1, seq_len + 1), generator=generator) + input_ids = tokens[:, :-1] + labels = tokens[:, 1:].to(DEVICE) + seq_ctx = SequenceContext.from_input_ids((input_ids,), device=DEVICE) + loss_cfg = CELossConfig() + loss_ctx = loss_cfg.build(data={"shifted_labels": labels}, sp_mesh=None) + loss_ctx = loss_cfg.loss_ctx_cls.build_batches([loss_ctx])[0] + return ModelItem(seq_ctx=seq_ctx, loss_ctx={"lm": loss_ctx}) + + +def build_engine( + mode: LayoutMode, + hf_path: Path, + tag: str, + dispatcher: str = "all2all", + fp8: bool = False, + lr: float = 1e-4, +) -> TrainEngine: + """Build a ``TrainEngine`` for ``mode`` without loading weights. + + Args: + mode (LayoutMode): Layout under test. + hf_path (Path): HF checkpoint the model config is read from. + tag (str): Unique tag for the mesh prefix; every engine alive in a process needs its own. + dispatcher (str): MoE dispatcher (``"all2all"`` or ``"deepep"``). + fp8 (bool): Enable tile-wise float8 for linear and grouped linear layers. + lr (float): AdamW learning rate. + + Returns: + TrainEngine: Sharded engine with freshly built (uninitialised) parameters. + """ + model_cfg = Qwen3MoEConfig.from_hf(hf_path) + model_cfg.ep_size = mode["ep"] + model_cfg.dispatcher = dispatcher + model_cfg.compile_cfg = False + model_cfg.mesh_prefix = f"{tag}_{mode['name']}" + if fp8: + model_cfg.float8_cfg = Float8Config( + scaling_granularity_gemm=ScalingGranularity.TILEWISE, + scaling_granularity_grouped_gemm=ScalingGranularity.TILEWISE, + ) + fsdp_cfg = FSDPConfig( + ep_size=mode["ep"], + decouple_ep_fsdp=mode["decouple"], + hsdp_sharding_size=mode["hsdp"], + torch_compile=False, + ) + return TrainEngine(model_cfg=model_cfg, optim_cfg=AdamWConfig(lr=lr), fsdp_cfg=fsdp_cfg) + + +def train(engine: TrainEngine, steps: Iterable[int], vocab_size: int, seq_len: int) -> list[float]: + """Run ``steps`` optimizer steps on the shared token stream. + + Args: + engine (TrainEngine): Engine to train. + steps (Iterable[int]): Step indices; they select the batches via :func:`make_batch`. + vocab_size (int): Vocabulary size of the model. + seq_len (int): Tokens per batch. + + Returns: + list[float]: ``reduced_llm_loss`` of every step. + """ + losses = [] + for step in steps: + info = engine.train_step([make_batch(step, dist.get_rank(), vocab_size, seq_len)]) + grad_norm = engine.clip_grad_norm() + engine.step_optimizer(grad_norm) + losses.append(float(info["logs_info"]["reduced_llm_loss"])) + return losses + + +def run_mode( + mode: LayoutMode, + hf_path: Path, + steps: int, + seq_len: int, + lr: float, + dispatcher: str, + fp8: bool, + grad_norm_steps: Sequence[int] = (0, 1, 25), + tag: str = "numerics", +) -> dict[str, Any]: + """Train ``mode`` from the HF checkpoint and record losses, grad norms, memory and step time. + + Args: + mode (LayoutMode): Layout under test. + hf_path (Path): HF checkpoint to load. + steps (int): Number of optimizer steps. + seq_len (int): Tokens per batch. + lr (float): AdamW learning rate. + dispatcher (str): MoE dispatcher. + fp8 (bool): Enable tile-wise float8. + grad_norm_steps (Sequence[int]): Steps whose per-parameter grad norms are recorded. + tag (str): Mesh-prefix tag passed to :func:`build_engine`. + + Returns: + dict[str, Any]: ``losses``, ``grad_norms`` (per step), ``param_grad_norms`` (per recorded + step, keyed by parameter name), ``memory`` (per-rank MiB) and ``step_time_s``. + """ + rank = dist.get_rank() + torch.cuda.reset_peak_memory_stats() + torch.cuda.synchronize() + base_allocated = torch.cuda.memory_allocated() + + engine = build_engine(mode, hf_path, tag, dispatcher=dispatcher, fp8=fp8, lr=lr) + engine.from_hf(hf_path=hf_path, strict=True) + vocab_size = engine.model.config.vocab_size + torch.cuda.synchronize() + result: dict[str, Any] = { + "mode": mode, + "memory": param_memory(engine.model), + "losses": [], + "grad_norms": [], + "param_grad_norms": {}, + "step_time_s": None, + } + result["memory"]["allocated_after_load_mib"] = (torch.cuda.memory_allocated() - base_allocated) / 2**20 + + step_times = [] + for step in range(steps): + torch.cuda.synchronize() + t0 = time.perf_counter() + batch = make_batch(step, rank, vocab_size, seq_len) + info = engine.train_step([batch]) + grad_norm = engine.clip_grad_norm() + if step in grad_norm_steps: + result["param_grad_norms"][str(step)] = per_param_grad_norms(engine.model) + engine.step_optimizer(grad_norm) + torch.cuda.synchronize() + step_times.append(time.perf_counter() - t0) + result["losses"].append(float(info["logs_info"]["reduced_llm_loss"])) + result["grad_norms"].append(float(grad_norm)) + if step == 0: + result["memory"]["allocated_after_step0_mib"] = (torch.cuda.memory_allocated() - base_allocated) / 2**20 + + warm = step_times[5:] if len(step_times) > 10 else step_times + result["step_time_s"] = sum(warm) / len(warm) + result["step_times_s"] = step_times + result["memory"]["peak_allocated_mib"] = (torch.cuda.max_memory_allocated() - base_allocated) / 2**20 + result["memory"]["peak_reserved_mib"] = torch.cuda.max_memory_reserved() / 2**20 + + release(engine) + return result + + +def param_memory(model: torch.nn.Module) -> dict[str, float]: + """Per-rank parameter bytes of the local shards, split into routed experts and everything else. + + Args: + model (torch.nn.Module): Sharded model. + + Returns: + dict[str, float]: ``expert_param_mib`` and ``dense_param_mib``. + """ + expert_bytes = 0 + dense_bytes = 0 + for name, param in model.named_parameters(): + local = param.to_local() if isinstance(param, DTensor) else param + nbytes = local.numel() * local.element_size() + if ".experts" in name: + expert_bytes += nbytes + else: + dense_bytes += nbytes + return {"expert_param_mib": expert_bytes / 2**20, "dense_param_mib": dense_bytes / 2**20} + + +def per_param_grad_norms(model: torch.nn.Module) -> dict[str, float]: + """Global L2 norm of every parameter's gradient, keyed by the unwrapped parameter name. + + Args: + model (torch.nn.Module): Sharded model after ``backward``. + + Returns: + dict[str, float]: Parameter name (without ``_checkpoint_wrapped_module.``) to grad norm. + """ + norms: dict[str, float] = {} + for name, param in model.named_parameters(): + if param.grad is None: + continue + if isinstance(param.grad, DTensor): + norm = cal_total_norm([param.grad], norm_type=2.0, foreach=True, dtype=torch.float32) + else: + norm = torch.linalg.vector_norm(param.grad, 2.0, dtype=torch.float32) + norms[clean_name(name)] = float(norm) + return norms + + +def clean_name(name: str) -> str: + """Strip the activation-checkpoint wrapper prefix from a parameter name. + + Args: + name (str): Name from ``named_parameters``. + + Returns: + str: The name without ``_checkpoint_wrapped_module.``. + """ + return name.replace("_checkpoint_wrapped_module.", "") + + +def rel_diff(a: float, b: float) -> float: + """Relative difference ``|a - b| / |b|`` (``b`` is the reference). + + Args: + a (float): Value under test. + b (float): Reference value. + + Returns: + float: The relative difference; the denominator is clamped to ``1e-12``. + """ + return abs(a - b) / max(abs(b), 1e-12) + + +def max_rel_diff(values: Sequence[float], reference: Sequence[float]) -> float: + """Largest element-wise :func:`rel_diff` of two equally long sequences. + + Args: + values (Sequence[float]): Values under test. + reference (Sequence[float]): Reference values. + + Returns: + float: ``max(rel_diff(v, r))`` over the zipped sequences. + """ + if len(values) != len(reference): + raise ValueError(f"length mismatch: {len(values)} vs {len(reference)}") + return max(rel_diff(a, b) for a, b in zip(values, reference)) + + +def load_hf_dir(path: Path) -> dict[str, torch.Tensor]: + """Load every ``*.safetensors`` file of an HF checkpoint directory. + + Args: + path (Path): Checkpoint directory. + + Returns: + dict[str, torch.Tensor]: All tensors keyed by their checkpoint name. + """ + tensors: dict[str, torch.Tensor] = {} + for file in sorted(path.glob("*.safetensors")): + tensors.update(load_file(str(file))) + return tensors + + +def compare_hf(lhs: Path, rhs: Path) -> dict[str, Any]: + """Compare two HF checkpoint directories tensor by tensor. + + Args: + lhs (Path): Reference checkpoint. + rhs (Path): Checkpoint under test. + + Returns: + dict[str, Any]: Key counts, missing / extra keys, max absolute and relative differences, + the number of mismatched elements and the worst key; ``{"error": ...}`` on a shape mismatch. + """ + a = load_hf_dir(lhs) + b = load_hf_dir(rhs) + missing = sorted(set(a) - set(b)) + extra = sorted(set(b) - set(a)) + max_abs = 0.0 + max_rel = 0.0 + mismatched = 0 + total = 0 + worst = "" + for key in sorted(set(a) & set(b)): + x = a[key].to(torch.float32) + y = b[key].to(torch.float32) + if x.shape != y.shape: + return {"error": f"shape mismatch for {key}: {tuple(x.shape)} vs {tuple(y.shape)}"} + diff = (x - y).abs() + cur = float(diff.max()) + if cur > max_abs: + max_abs, worst = cur, key + max_rel = max(max_rel, float((diff / x.abs().clamp_min(1e-6)).max())) + mismatched += int((diff > 0).sum()) + total += diff.numel() + return { + "keys": len(set(a) & set(b)), + "missing_in_rhs": missing[:5], + "extra_in_rhs": extra[:5], + "max_abs_diff": max_abs, + "max_rel_diff": max_rel, + "mismatched_elements": mismatched, + "total_elements": total, + "worst_key": worst, + } + + +def assert_hf_dirs_close(lhs: Path, rhs: Path, rtol: float, atol: float, label: str) -> None: + """Fail unless two HF checkpoint directories hold the same keys, shapes and close values. + + Args: + lhs (Path): Reference checkpoint. + rhs (Path): Checkpoint under test. + rtol (float): Relative tolerance passed to ``torch.testing.assert_close``. + atol (float): Absolute tolerance passed to ``torch.testing.assert_close``. + label (str): Prefix for the failure message. + """ + a = load_hf_dir(lhs) + b = load_hf_dir(rhs) + if set(a) != set(b): + raise AssertionError( + f"{label}: key sets differ; missing in rhs {sorted(set(a) - set(b))[:5]}, " + f"extra in rhs {sorted(set(b) - set(a))[:5]}" + ) + for key in sorted(a): + x, y = a[key], b[key] + if x.shape != y.shape or x.dtype != y.dtype: + raise AssertionError(f"{label}: {key} has {tuple(y.shape)}/{y.dtype}, expected {tuple(x.shape)}/{x.dtype}") + torch.testing.assert_close(y.float(), x.float(), rtol=rtol, atol=atol, msg=lambda m: f"{label}: {key}: {m}") + + +def release(engine: TrainEngine) -> None: + """Drop an engine and return its device memory to the allocator. + + Args: + engine (TrainEngine): Engine to release; the caller must not use it afterwards. + """ + del engine + gc.collect() + torch.cuda.empty_cache() diff --git a/xtuner/v1/config/fsdp.py b/xtuner/v1/config/fsdp.py index 7335d6d7a5..f0d023c057 100644 --- a/xtuner/v1/config/fsdp.py +++ b/xtuner/v1/config/fsdp.py @@ -1,8 +1,8 @@ -from typing import Any, Optional +from typing import Optional import torch from cyclopts import Parameter -from pydantic import BaseModel, ConfigDict, field_serializer, field_validator +from pydantic import BaseModel, ConfigDict, field_serializer, field_validator, model_validator from typing_extensions import Annotated @@ -41,10 +41,34 @@ class FSDPConfig(BaseModel): hsdp_sharding_size: Annotated[ Optional[int], Parameter(help="Sharding size for HSDP (Hybrid Sharding Data Parallel)") ] = None + # Decoupled EP/FSDP ("dp2ep") layout. When enabled, `ep_size` is a sub-dimension of the + # FSDP shard dimension instead of being orthogonal to it: routed experts are sharded + # `dp_shard / ep_size` ways on top of EP, while every other parameter is sharded over the + # full `dp_shard` (= `hsdp_sharding_size` or world size) without being replicated across EP + # ranks. Requires `dp_shard % ep_size == 0`. When disabled, the legacy layout is untouched. + decouple_ep_fsdp: Annotated[ + bool, Parameter(help="Decouple expert parallel from FSDP: shard dense params over the full FSDP mesh") + ] = False - def model_post_init(self, __context: Any) -> None: - if self.hsdp_sharding_size is not None: - assert self.ep_size == 1, "Currently, HSDP requires expert parallel size to be 1" + @model_validator(mode="after") + def _validate_ep_fsdp_topology(self) -> "FSDPConfig": + # Explicit `ValueError`s instead of `assert`: the topology checks must survive `python -O`, + # and pydantic reports them as a `ValidationError` when the config is built. + if self.ep_size < 1: + raise ValueError(f"`ep_size` must be a positive integer, got {self.ep_size}") + if self.hsdp_sharding_size is None: + return self + if self.hsdp_sharding_size < 1: + raise ValueError(f"`hsdp_sharding_size` must be a positive integer, got {self.hsdp_sharding_size}") + if self.decouple_ep_fsdp: + if self.hsdp_sharding_size % self.ep_size != 0: + raise ValueError( + "`decouple_ep_fsdp` requires `hsdp_sharding_size` to be divisible by `ep_size`, " + f"got hsdp_sharding_size={self.hsdp_sharding_size}, ep_size={self.ep_size}" + ) + elif self.ep_size != 1: + raise ValueError("Currently, HSDP requires expert parallel size to be 1") + return self @field_serializer("param_dtype", "reduce_dtype") def serialize_param_dtype(self, value: torch.dtype) -> str: diff --git a/xtuner/v1/float8/float8_handler.py b/xtuner/v1/float8/float8_handler.py index c6ecd2c66a..510483995a 100644 --- a/xtuner/v1/float8/float8_handler.py +++ b/xtuner/v1/float8/float8_handler.py @@ -42,6 +42,11 @@ class Float8Handler: fsdp_mesh: Optional[DeviceMesh] = None tilewise_reduce_mesh_devided_64: Optional[DeviceMesh] = None tilewise_reduce_mesh_mapping: Dict[Tuple[int, int], DeviceMesh] = {} + # Decoupled EP/FSDP layout only: routed experts are FSDP-sharded on `expert_fsdp_mesh` + # (efsdp, strided by ep) instead of `fsdp_mesh`, so they need their own reduce meshes. + expert_fsdp_mesh: Optional[DeviceMesh] = None + expert_tilewise_reduce_mesh_devided_64: Optional[DeviceMesh] = None + expert_tilewise_reduce_mesh_mapping: Optional[Dict[Tuple[int, int], DeviceMesh]] = None def __init__( self, @@ -113,7 +118,12 @@ def get_shard_size_on_dim(tensor: torch.Tensor | DTensor, dim: int) -> int: return RuntimeLayout.from_dtensor(tensor).shard_size(dim) @staticmethod - def pad_for_fsdp(model: nn.Module, fsdp_mesh: DeviceMesh, callback_after_pad: Callable | None = None): + def pad_for_fsdp( + model: nn.Module, + fsdp_mesh: DeviceMesh, + callback_after_pad: Callable | None = None, + expert_fsdp_mesh: DeviceMesh | None = None, + ): from xtuner.v1.float8.float8_gmm_tile_wise import TileWiseFloat8GroupedLinear from xtuner.v1.float8.float8_linear_tensor_wise import TensorWiseFloat8Linear from xtuner.v1.float8.float8_linear_tile_wise import TileWiseFloat8Linear @@ -131,20 +141,102 @@ def pad_for_fsdp(model: nn.Module, fsdp_mesh: DeviceMesh, callback_after_pad: Ca else: tensor_size = module.weight.size() parallel_size = 1 - padded_out_features = Float8Handler.get_num_features_after_pad(tensor_size, 0, fsdp_mesh.size(-1)) + num_fsdp_chunks = fsdp_mesh.size(-1) + if expert_fsdp_mesh is not None and isinstance(module, TileWiseFloat8GroupedLinear): + # Decoupled EP/FSDP: routed experts are sharded `efsdp` ways, not `dp_shard` ways. + num_fsdp_chunks = expert_fsdp_mesh.size(-1) + padded_out_features = Float8Handler.get_num_features_after_pad(tensor_size, 0, num_fsdp_chunks) padded_out_features *= parallel_size module.pad_for_fsdp(padded_out_features=padded_out_features) if callback_after_pad is not None: callback_after_pad() - def build_reduce_mesh(self, model: nn.Module, fsdp_mesh: DeviceMesh): + def build_reduce_mesh(self, model: nn.Module, fsdp_mesh: DeviceMesh, expert_fsdp_mesh: DeviceMesh | None = None): self.fsdp_mesh = fsdp_mesh + self.expert_fsdp_mesh = expert_fsdp_mesh if self.is_tilewise_fp8: + if expert_fsdp_mesh is not None: + self._build_decoupled_reduce_meshes(model, fsdp_mesh, expert_fsdp_mesh) + return if fsdp_mesh.size(-1) >= 2: self._build_reduce_mesh_devided_64(fsdp_mesh) self._build_reduce_mesh_mapping(model, fsdp_mesh) + def _build_decoupled_reduce_meshes( + self, model: nn.Module, fsdp_mesh: DeviceMesh, expert_fsdp_mesh: DeviceMesh + ) -> None: + # Decoupled EP/FSDP layout (root mesh `(replicate, efsdp, ep)`): + # * dense fp8 weights are sharded over `dp_shard` = flatten(efsdp, ep): consecutive FSDP + # shards live on consecutive ranks (stride 1); + # * routed-expert fp8 weights are sharded over `efsdp`, whose ranks are `ep` apart. + # Each class gets its own "reduce max" meshes, built with the matching rank stride. + from xtuner.v1.float8.float8_gmm_tile_wise import TileWiseFloat8GroupedLinear + from xtuner.v1.float8.float8_linear_tile_wise import TileWiseFloat8Linear + + world_size = dist.get_world_size() + dense_shard_size = fsdp_mesh.size(-1) + expert_shard_size = expert_fsdp_mesh.size(-1) + expert_stride = world_size // expert_fsdp_mesh.size() # == ep_size + + self.tilewise_reduce_mesh_devided_64 = ( + self._build_strided_reduce_mesh(2, 1) if dense_shard_size >= 2 and dense_shard_size % 2 == 0 else None + ) + self.expert_tilewise_reduce_mesh_devided_64 = ( + self._build_strided_reduce_mesh(2, expert_stride) + if expert_shard_size >= 2 and expert_shard_size % 2 == 0 + else None + ) + self.tilewise_reduce_mesh_mapping = self._build_strided_reduce_mesh_mapping( + model, (TileWiseFloat8Linear,), dense_shard_size, 1 + ) + self.expert_tilewise_reduce_mesh_mapping = self._build_strided_reduce_mesh_mapping( + model, (TileWiseFloat8GroupedLinear,), expert_shard_size, expert_stride + ) + + @staticmethod + def _build_strided_reduce_mesh(num_ranks: int, stride: int) -> DeviceMesh: + # Groups of `num_ranks` ranks spaced `stride` apart, e.g. (r, r + stride, ...). + world_size = dist.get_world_size() + assert world_size % (num_ranks * stride) == 0, (world_size, num_ranks, stride) + return init_device_mesh( + "cuda", + (world_size // (num_ranks * stride), num_ranks, stride), + mesh_dim_names=("_", "tilewise_reduce", "ep_or_tp"), + )["tilewise_reduce"] + + def _build_strided_reduce_mesh_mapping( + self, model: nn.Module, module_types: Tuple[type, ...], shard_size: int, stride: int + ) -> Dict[Tuple[int, int], DeviceMesh]: + SHARD_DIM = 0 + mapping: Dict[Tuple[int, int], DeviceMesh] = {} + for module in model.modules(): + if not isinstance(module, module_types): + continue + assert isinstance(module.weight, DTensor), ( + "`build_reduce_mesh` should be called after apply fully_shard to the model." + ) + local_shape = module.weight._local_tensor.shape + if local_shape[SHARD_DIM] >= 128: + assert local_shape[SHARD_DIM] % 128 in (0, 64), ( + f"Currently only local_shape[SHARD_DIM] % 128 == 0 or " + f"local_shape[SHARD_DIM] % 128 == 64 is supported, got {local_shape}. Please contact us." + ) + continue + assert 128 % local_shape[SHARD_DIM] == 0, ( + f"Currently only local_shape[SHARD_DIM] % 128 == 0 is supported, got {local_shape}. Please contact us." + ) + reduce_world_size = 128 // local_shape[SHARD_DIM] + if local_shape in mapping: + assert mapping[local_shape].size() == reduce_world_size + continue + assert shard_size >= reduce_world_size and shard_size % reduce_world_size == 0, ( + f"Expect FSDP shard size >= reduce_world_size and shard size % reduce_world_size == 0, " + f"got shard size = {shard_size}, reduce_world_size = {reduce_world_size}. Please contact us." + ) + mapping[local_shape] = self._build_strided_reduce_mesh(reduce_world_size, stride) + return mapping + def _build_reduce_mesh_devided_64(self, fsdp_mesh: DeviceMesh): # 为了支持 moe 参数被 fsdp 和 ep 切成 dout = n * 128 + 64 (n >= 1) 的情况 # fsdp rank 0 的后 64 个 dim 要跟 fsdp rank 1 的前 64 个 dim 共同组成一个 block @@ -221,9 +313,26 @@ def precompute_float8_dynamic_scale_for_fsdp(self, model: Union[nn.Module, List[ for m in models: if self.is_tilewise_fp8: - precompute_tilewise_float8_scale_for_fsdp( - m, self.tilewise_reduce_mesh_mapping, self.tilewise_reduce_mesh_devided_64 - ) + if self.expert_tilewise_reduce_mesh_mapping is None: + precompute_tilewise_float8_scale_for_fsdp( + m, self.tilewise_reduce_mesh_mapping, self.tilewise_reduce_mesh_devided_64 + ) + else: + from xtuner.v1.float8.float8_gmm_tile_wise import TileWiseFloat8GroupedLinear + from xtuner.v1.float8.float8_linear_tile_wise import TileWiseFloat8Linear + + precompute_tilewise_float8_scale_for_fsdp( + m, + self.tilewise_reduce_mesh_mapping, + self.tilewise_reduce_mesh_devided_64, + module_types=(TileWiseFloat8Linear,), + ) + precompute_tilewise_float8_scale_for_fsdp( + m, + self.expert_tilewise_reduce_mesh_mapping, + self.expert_tilewise_reduce_mesh_devided_64, + module_types=(TileWiseFloat8GroupedLinear,), + ) if self.is_tensorwise_fp8: assert self.fsdp_mesh is not None, "FSDP mesh must be set for tensorwise float8 training." precompute_tensorwise_float8_scale_for_fsdp(m, self.fsdp_mesh) diff --git a/xtuner/v1/float8/fsdp_utils.py b/xtuner/v1/float8/fsdp_utils.py index 8444672fcb..75b8b8372c 100644 --- a/xtuner/v1/float8/fsdp_utils.py +++ b/xtuner/v1/float8/fsdp_utils.py @@ -124,13 +124,17 @@ def precompute_tilewise_float8_scale_for_fsdp( module: nn.Module, reduce_mesh_mapping: Dict[Tuple[int, int], DeviceMesh], # absmax need to be reduced in this group reduce_mesh_devided_64: Optional[DeviceMesh] = None, # All params share the same reduce mesh + module_types: Optional[Tuple[type, ...]] = None, # restrict to these fp8 module classes (default: all) ) -> None: from xtuner.v1.float8 import TileWiseFloat8GroupedLinear, TileWiseFloat8Linear + if module_types is None: + module_types = (TileWiseFloat8Linear, TileWiseFloat8GroupedLinear) + weights: List[WeightWithDynamicTilewiseFloat8CastTensor] = [] for m in module.modules(): if ( - isinstance(m, (TileWiseFloat8Linear, TileWiseFloat8GroupedLinear)) + isinstance(m, module_types) and isinstance(m.weight, DTensor) and isinstance(m.weight._local_tensor, WeightWithDynamicTilewiseFloat8CastTensor) ): diff --git a/xtuner/v1/model/base.py b/xtuner/v1/model/base.py index 50e4ee59dc..39352cbec9 100644 --- a/xtuner/v1/model/base.py +++ b/xtuner/v1/model/base.py @@ -564,6 +564,9 @@ class BaseModel(nn.Module): load_spec_mapping: dict[str, LoadSpec] = {} fsdp_mesh: DeviceMesh | None = None hsdp_mesh: DeviceMesh | None = None + # Only set by MoE models on the decoupled EP/FSDP path: the mesh routed experts are + # `fully_shard`ed on, `(efsdp,)` or `(replicate, efsdp)` with `efsdp = dp_shard / ep_size`. + expert_fsdp_mesh: DeviceMesh | None = None fsdp_config: FSDPConfig | None = None config: XTunerBaseModelConfig @@ -1009,7 +1012,7 @@ def float8_handler(self): self._float8_handler = self.config.float8_cfg.build() if self.fsdp_mesh is not None: - self._float8_handler.build_reduce_mesh(self, self.fsdp_mesh) + self._float8_handler.build_reduce_mesh(self, self.fsdp_mesh, expert_fsdp_mesh=self.expert_fsdp_mesh) return self._float8_handler @torch.no_grad() @@ -1917,9 +1920,31 @@ def _fsdp_foreach_allgather( return 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] + expert_fsdp_group: dist.ProcessGroup | None = None + if self.expert_fsdp_mesh is not None: + # Decoupled EP/FSDP: routed experts are FSDP-sharded on the innermost (`efsdp`) dim of + # `expert_fsdp_mesh`, so their FSDP-only gather must use that group instead. + expert_fsdp_group = self.expert_fsdp_mesh.get_group(self.expert_fsdp_mesh.ndim - 1) + save_plan_list = [ + load_spec.plan_hf_save( + gather_process_group=self._fsdp_gather_group(load_spec, fsdp_group, expert_fsdp_group) + ) + for load_spec in load_spec_list + ] return unshard_tensors_for_hf_save(list(tensor_list), save_plan_list) + def _fsdp_gather_group( + self, + load_spec: LoadSpec, + fsdp_group: dist.ProcessGroup, + expert_fsdp_group: dist.ProcessGroup | None, + ) -> dist.ProcessGroup: + if expert_fsdp_group is None: + return fsdp_group + if any(self._is_same_process_group(shard.group, expert_fsdp_group) for shard in load_spec.shards): + return expert_fsdp_group + return fsdp_group + @staticmethod def _is_same_process_group(left: dist.ProcessGroup, right: dist.ProcessGroup) -> bool: if left is right: diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py index 15dc0d628e..31ad44d967 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/xtuner/v1/model/moe/moe.py @@ -1278,7 +1278,18 @@ def fully_shard( if self.config.float8_cfg is not None: # As we modify the shape of the model's parameters, # we need to reinitialize the load spec mapping. - Float8Handler.pad_for_fsdp(self, cast(DeviceMesh, self.fsdp_mesh), callback_after_pad=self._init_load_spec) + if self.fsdp_config.decouple_ep_fsdp: + # Routed experts are padded for `efsdp` FSDP chunks, everything else for `dp_shard`. + Float8Handler.pad_for_fsdp( + self, + cast(DeviceMesh, self.fsdp_mesh), + callback_after_pad=self._init_load_spec, + expert_fsdp_mesh=self.expert_fsdp_mesh, + ) + else: + Float8Handler.pad_for_fsdp( + self, cast(DeviceMesh, self.fsdp_mesh), callback_after_pad=self._init_load_spec + ) # Just for narrowing the type of self.fsdp_mesh and self.ep_mesh assert self.fsdp_mesh is not None @@ -1296,9 +1307,13 @@ def fully_shard( param.requires_grad = False 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: + decoupled = self.fsdp_config.decouple_ep_fsdp + if not decoupled and (self.ep_mesh.size() > 1 or tp_enabled): # 中文注释：不开 EP 但开启 expert TP 时，非 expert 参数仍是 TP rank 间的逻辑副本， # 需要显式放到 Replicate DTensor 上，后续梯度才会跨 expert TP 平均。 + # + # On the decoupled path dense params are NOT replicated over EP: they stay plain + # tensors and get sharded over the full `dp_shard` mesh by FSDP below. self._replicate_other_params(self) # Although rotary_emb was already constructed in __init__, it was built on the meta device. @@ -1313,6 +1328,18 @@ def fully_shard( for layer_idx, layer in tqdm(self.layers.items(), desc="[FSDP Sharding]"): layer_idx = int(layer_idx) + if decoupled: + # Routed experts get their own FSDP group on `expert_fsdp_mesh` (efsdp) before the + # layer is wrapped on the dense mesh; FSDP then excludes them from the layer group. + self._fully_shard_expert_blocks( + layer, + mp_policy=mp_policy, + reshard_after_forward=( + False + if layer_idx >= len(self.layers) - 1 and self.mtp_block is None + else self.fsdp_config.reshard_after_forward + ), + ) if self._should_recompute( layer_idx=layer_idx, mtp_idx=None, @@ -1340,7 +1367,17 @@ def fully_shard( list(self.layers.values())[:-1], list(self.layers.values())[1:], ): - layer_cur.set_modules_to_forward_prefetch([layer_next]) # type: ignore + if decoupled: + # Issue the expert all-gather of the current layer together with the next layer's + # dense all-gather so it overlaps with attention instead of stalling the MoE block. + layer_cur.set_modules_to_forward_prefetch( # type: ignore + [*self._expert_blocks(layer_cur), layer_next] + ) + else: + layer_cur.set_modules_to_forward_prefetch([layer_next]) # type: ignore + if decoupled: + last_layer = list(self.layers.values())[-1] + last_layer.set_modules_to_forward_prefetch(self._expert_blocks(last_layer)) # type: ignore self._fully_shard( mesh=self.fsdp_mesh if self.hsdp_mesh is None else self.hsdp_mesh, @@ -1379,6 +1416,10 @@ def fully_shard( self.mtp_block.layers[mtp_idx] = mtp_layer reshard_after_forward = mtp_idx != len(self.mtp_block.layers) - 1 + if decoupled: + self._fully_shard_expert_blocks( + mtp_layer, mp_policy=mp_policy, reshard_after_forward=reshard_after_forward + ) self._fully_shard( mesh=self.fsdp_mesh if self.hsdp_mesh is None else self.hsdp_mesh, mp_policy=mp_policy, @@ -1387,14 +1428,27 @@ def fully_shard( module=mtp_layer, ) if mtp_idx == 0: - layer_next.set_modules_to_forward_prefetch([mtp_layer]) # type: ignore + if decoupled: + layer_next.set_modules_to_forward_prefetch( # type: ignore + [*self._expert_blocks(layer_next), mtp_layer] + ) + else: + layer_next.set_modules_to_forward_prefetch([mtp_layer]) # type: ignore if self.config.mtp_config is not None and self.config.mtp_config.num_layers > 0: for prev_mtp_layer, next_mtp_layer in zip( list(self.mtp_block.layers)[:-1], list(self.mtp_block.layers)[1:], ): - prev_mtp_layer.set_modules_to_forward_prefetch([next_mtp_layer]) # type: ignore + if decoupled: + prev_mtp_layer.set_modules_to_forward_prefetch( # type: ignore + [*self._expert_blocks(prev_mtp_layer), next_mtp_layer] + ) + else: + prev_mtp_layer.set_modules_to_forward_prefetch([next_mtp_layer]) # type: ignore + if decoupled: + last_mtp_layer = list(self.mtp_block.layers)[-1] + last_mtp_layer.set_modules_to_forward_prefetch(self._expert_blocks(last_mtp_layer)) # type: ignore self._fully_shard( mesh=self.fsdp_mesh if self.hsdp_mesh is None else self.hsdp_mesh, @@ -1426,6 +1480,10 @@ def need_update_bias(self) -> bool: @torch.no_grad # type: ignore def scale_and_reduce_grad(self): + if self.fsdp_config is not None and self.fsdp_config.decouple_ep_fsdp: + self._scale_and_reduce_grad_decoupled() + return + # 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 @@ -1518,6 +1576,10 @@ def cal_grad_norm(self, grads: list[DTensor], dtype=torch.float32): def _init_device_mesh(self, fsdp_config: FSDPConfig): self.fsdp_config = fsdp_config + if self.fsdp_config.decouple_ep_fsdp: + self._init_decoupled_device_mesh(fsdp_config) + return + device = DEVICE world_size = dist.get_world_size() expert_tp_size = self.config.expert_tp_size if self.config.expert_tp_size > 1 else 1 @@ -1608,6 +1670,143 @@ def _init_device_mesh(self, fsdp_config: FSDPConfig): ) self.fsdp_mesh = self.hsdp_mesh[f"{self.config.mesh_prefix}.hsdp_shard"] + def _scale_and_reduce_grad_decoupled(self) -> None: + # Invariant (DESIGN §4.4): every gradient ends up as the mean over the full data-parallel + # world, independent of ep / efsdp. + # * dense params are FSDP-sharded over `dp_shard` (+ `replicate` with HSDP); FSDP's + # reduce-scatter / all-reduce already average them -> nothing to do; + # * routed experts are reduce-scattered over `efsdp = dp_shard / ep` only, while each + # expert already saw the tokens of its whole EP group -> divide by `ep`; + # * DTensors replicated on every mesh dim (fp32 params kept out of FSDP via + # `fp32_keys_pattern`) are averaged manually, as on the legacy path. + ep_size = self.ep_mesh.size() if self.ep_mesh is not None else 1 + grads_by_group: dict[dist.ProcessGroup, list[torch.Tensor]] = {} + + for name, param in self.trainable_parameters(): + if param.grad is None: + continue + + if ep_size > 1 and ".experts" in name: + param.grad.div_(ep_size) # type: ignore + continue + + if not isinstance(param, DTensor): + continue + + if any(isinstance(placement, Shard) for placement in param.placements): + # FSDP-managed (`_StridedShard` is a `Shard` subclass). + continue + + mesh_dim_names = param.device_mesh.mesh_dim_names + assert mesh_dim_names is not None # every mesh of this model is built with named dims + replicate_dim_names = tuple( + mesh_dim_names[i] for i, p in enumerate(param.placements) if isinstance(p, Replicate) + ) + if not replicate_dim_names: + continue + + if len(replicate_dim_names) > 1: + flat_mesh = param.device_mesh[replicate_dim_names]._flatten() + else: + flat_mesh = param.device_mesh[replicate_dim_names[0]] + + grad = param.grad.to_local() if isinstance(param.grad, DTensor) else param.grad + grad.div_(flat_mesh.size()) # type: ignore + grads_by_group.setdefault(flat_mesh.get_group(), []).append(grad) # type: ignore + + for group, grads in grads_by_group.items(): + with dist._coalescing_manager(group=group): + for grad in grads: + dist.all_reduce(grad, ReduceOp.SUM, group=group) + + @staticmethod + def _expert_blocks(module: nn.Module) -> list[nn.Module]: + return [submodule for submodule in module.modules() if isinstance(submodule, MoEBlock)] + + def _fully_shard_expert_blocks( + self, + module: nn.Module, + mp_policy: MixedPrecisionPolicy, + reshard_after_forward: bool, + ) -> None: + assert self.expert_fsdp_mesh is not None + assert self.fsdp_config is not None + for expert_block in self._expert_blocks(module): + self._fully_shard( + mesh=self.expert_fsdp_mesh, + mp_policy=mp_policy, + reshard_after_forward=reshard_after_forward, + offload_policy=CPUOffloadPolicy() if self.fsdp_config.cpu_offload else None, + module=expert_block, + ) + + def _init_decoupled_device_mesh(self, fsdp_config: FSDPConfig) -> None: + # Decoupled ("dp2ep") layout: EP is a sub-dimension of the FSDP shard dimension. + # + # root = (replicate, efsdp, ep) replicate = world / dp_shard (1 without HSDP) + # efsdp = dp_shard / ep + # dense params: FSDP over flatten(efsdp, ep) = dp_shard (+ replicate for HSDP) + # expert params: Shard(0) over ep (built in GroupedLinear) + FSDP over efsdp + # (+ replicate for HSDP) + # + # `ep` stays the innermost dimension so the EP group is still made of contiguous + # (intra-node) ranks, exactly like the legacy `(fsdp, ep)` root mesh. Every sub-mesh is + # derived from this single root so FSDP accepts the EP-sharded DTensor expert params. + # The `efsdp` dimension is kept even when its size is 1: the FSDP wrapping of the + # experts is what applies the mixed-precision policy to them. + if self.config.expert_tp_size > 1: + raise NotImplementedError("`decouple_ep_fsdp` with ExpertTP is not supported") + + device = DEVICE + world_size = dist.get_world_size() + ep_size = fsdp_config.ep_size + dp_shard = fsdp_config.hsdp_sharding_size if fsdp_config.hsdp_sharding_size is not None else world_size + # Explicit exceptions (not `assert`): the mesh shape below is only valid with these. + if world_size % dp_shard != 0: + raise ValueError(f"world_size ({world_size}) must be divisible by hsdp_sharding_size ({dp_shard})") + if dp_shard % ep_size != 0: + raise ValueError( + f"`decouple_ep_fsdp` requires the FSDP shard size ({dp_shard}) to be divisible by ep_size ({ep_size})" + ) + replicate_size = world_size // dp_shard + efsdp_size = dp_shard // ep_size + + prefix = self.config.mesh_prefix + replicate_name = f"{prefix}.replicate" + efsdp_name = f"{prefix}.efsdp" + ep_name = f"{prefix}.ep" + dp_shard_name = f"{prefix}.dp_shard" + + root_mesh = init_device_mesh( + device, + (replicate_size, efsdp_size, ep_size), + mesh_dim_names=(replicate_name, efsdp_name, ep_name), + ) + self._world_mesh = root_mesh + + # Same requirement as the legacy path: the `ep_mesh` created in `__init__` must map to the + # `ep` dimension of this root mesh (see the comment in `_init_device_mesh`). + new_ep_mesh = root_mesh[ep_name] + if self.ep_mesh is not None: + assert new_ep_mesh.mesh_dim_names == self.ep_mesh.mesh_dim_names, ( + f"FSDP enabled, it requires the name of new created `ep_mesh`: {new_ep_mesh.mesh_dim_names}" + f"equals to the origin one: {self.ep_mesh.mesh_dim_names}" + ) + assert torch.equal(self.ep_mesh.mesh, new_ep_mesh.mesh), ( + "FSDP enabled, it requires the `ep_size` of model config equals to the `ep_size` of FSDPConfig." + ) + else: + self.ep_mesh = new_ep_mesh + + # `fsdp_mesh` keeps its meaning of "the 1D shard group of dense parameters". + self.fsdp_mesh = root_mesh[efsdp_name, ep_name]._flatten(dp_shard_name) + if replicate_size > 1: + self.hsdp_mesh = root_mesh[replicate_name, dp_shard_name] + self.expert_fsdp_mesh = root_mesh[replicate_name, efsdp_name] + else: + self.hsdp_mesh = None + self.expert_fsdp_mesh = root_mesh[efsdp_name] + def _replicate_other_params(self, model: nn.Module): def traverse(module: nn.Module) -> None: if isinstance(module, MoEBlock): diff --git a/xtuner/v1/rl/weight_update/weight_iterator.py b/xtuner/v1/rl/weight_update/weight_iterator.py index 02da41bc47..c908d761d7 100644 --- a/xtuner/v1/rl/weight_update/weight_iterator.py +++ b/xtuner/v1/rl/weight_update/weight_iterator.py @@ -139,12 +139,16 @@ def iter_layer_batches(self): else: dtype = torch.bfloat16 - def get_params(tensor_list, name_list, save_dtype): + def get_params(owner, tensor_list, name_list, save_dtype): + # Gather with the module that owns the parameters, not with the outer model. A + # compose model is wrapped on the world mesh, while its language tower may be + # sharded on its own meshes (EP, `dp_shard`, `efsdp`, HSDP shard); only the owner + # knows which shards are FSDP shards to gather and which EP shards to keep local. _tensor_list, _spec_list = list(zip(*tensor_list)) - fsdp_unshard_tensor_list = model._fsdp_foreach_allgather(_tensor_list, _spec_list) + fsdp_unshard_tensor_list = owner._fsdp_foreach_allgather(_tensor_list, _spec_list) if save_dtype == torch.float8_e4m3fn: runtime_is_float8_list = [is_float8_weight(tensor) for tensor in _tensor_list] - fsdp_unshard_tensor_list, name_list = model._to_float8( + fsdp_unshard_tensor_list, name_list = owner._to_float8( fsdp_unshard_tensor_list, name_list, runtime_is_float8_list, @@ -191,7 +195,7 @@ def get_params(tensor_list, name_list, save_dtype): name = name.replace(".gate.", ".mlp.gate.") name_list.append(name) tensor_list.append((local_tensor, load_spec)) - fsdp_unshard_tensor_list, name_list = get_params(tensor_list, name_list, dtype) + fsdp_unshard_tensor_list, name_list = get_params(language_model, tensor_list, name_list, dtype) state_dict = dict(zip(name_list, fsdp_unshard_tensor_list)) yield WeightUpdateBatch(state_dict) @@ -200,7 +204,10 @@ def get_params(tensor_list, name_list, save_dtype): continue local_tensor = param._local_tensor if isinstance(param, DTensor) else param local_tensor = local_tensor.bfloat16() - load_spec = model.load_spec_mapping.get(name) + owner, owner_name = self._param_owner(model, name) + load_spec = owner.load_spec_mapping.get(owner._clean_param_name(owner_name)) + if load_spec is None: + raise ValueError(f"Internal Error. Parameter {name} not found in load_spec_mapping.") if isinstance(model.config, BaseComposeConfig): if "vision_tower." in name: @@ -220,7 +227,7 @@ def get_params(tensor_list, name_list, save_dtype): name = "model.embed_tokens.weight" tensor_list = [(local_tensor, load_spec)] name_list = [name] - fsdp_unshard_tensor_list, name_list = get_params(tensor_list, name_list, dtype) + fsdp_unshard_tensor_list, name_list = get_params(owner, tensor_list, name_list, dtype) state_dict = dict(zip(name_list, fsdp_unshard_tensor_list)) yield WeightUpdateBatch(state_dict) @@ -228,3 +235,15 @@ def get_params(tensor_list, name_list, save_dtype): yield WeightUpdateBatch({}, finished=True) DEVICE_MODULE.empty_cache() + + @staticmethod + def _param_owner(model: Any, name: str) -> tuple[Any, str]: + # Map a `model.state_dict()` key to the submodule that owns the parameter and to the + # key relative to that submodule, which is how each owner's `load_spec_mapping` is keyed. + # Compose models never build a `load_spec_mapping` of their own; their children do. + if isinstance(model.config, BaseComposeConfig): + for submodule in ("language_model", "vision_tower", "multi_modal_projector"): + prefix = f"{submodule}." + if name.startswith(prefix): + return getattr(model, submodule), name[len(prefix) :] + return model, name