已开启
[Bug]: dualmode Trainer FSDP optimizer layout loses TP global shape #330
MengXY107创建于 8月13日
MengXY107
8月13日 评论：
8月13日 评论：
本次更改说明
修改背景：dualmode Trainer 的生产模式在 TP/EP 后把普通 local Parameter 交给 FSDP，参数本身不再包含 DTensor global shape。不能默认把传入 param.size() 当作最终 global shape，否则 optimizer layout 的 tensor metadata 与 FSDP+TP/EP placements 不一致，Muon reshape 会读取错误的预期形状。
本次更改点：
- 在 Torch FSDP 参数管理中统一构造 optimizer sharding spec：原生 DTensor 复用原 global metadata；普通 TP/EP-local 参数利用
TPShardMetaInfo.mesh和 placements 恢复 logical global shape；普通 FSDP 参数沿用本地输入 shape。 - Muon 的 no-shard 路径复用统一的 local-tensor 访问逻辑，兼容普通 Parameter 和 DTensor。
- 多维 replicate relay broadcast 严格执行前一维 wait 后再启动下一维，最后一维保留异步以继续重叠。
- accuracy 用例增加 loss、grad norm、参数 gradient 日志以及 optimizer step 后 module 当前参数对拍。
迁移要求：master 分支的 fully_shard 实现也需要应用相同的 sharding spec 构造逻辑。迁移时应基于 master 自身的参数生命周期复用现有流程，不应直接假设传入参数 shape 是 global shape，也不应增加按模型或并行 case 分支的特例。
8月13日 关联了pull request：fix: restore optimizer layout for TP-sharded FSDP parameters
8月17日 关联了pull request：FSDP代码重构&支持参数非均匀切分
Checklist
问题描述
dualmode Trainer 生产路径中，TP/EP 处理后的参数在进入
fully_shard时是普通本地Parameter，不再携带DTensor的 global shape。FSDP 参数管理若直接把传入参数的 local shape 当作 optimizerLayout的 global shape，会造成元数据与 FSDP+TP/EP placements 不匹配。Muon 根据错误的 optimizer layout 恢复 Newton-Schulz 输入形状时，会在 TP+FSDP 场景触发
shape is invalid for input of size 16384。多维 replicate mesh 下，异步 relay broadcast 若未等待前一维通信完成，还可能让后一维从尚未收到有效数据的 rank 继续广播，污染 optimizer 参数。复现场景：8 卡 HSDP(2×2)+TP(2)、DP(2)+CP(2)+TP(2) 以及 HSDP+TP+EP MoE dualmode Trainer，用优化器执行参数更新并与 standalone 模型逐步对拍。
期望行为
FSDP 应根据
TPShardMetaInfo中已知的 mesh 和 placements，从 TP/EP-local shape 恢复 logical global shape，再构造统一的 FSDP+TP/EP optimizer layout。Muon 应同时支持普通 Parameter 与 DTensor，并保证多维 replicate relay 的维度间数据依赖。环境信息
HYPER_PARALLEL_PLATFORM=torchtrainer_dev