常见问题¶

TL;DR

常见采用问题的快速解答。每条条目均交叉链接到相关指南。如果你的问题不在此列，请开一个 GitHub Discussion 或查看故障排查。

DMuon 和 ZeRO-1 有什么区别？

ZeRO-1 将优化器状态分片到所有 rank，每个 rank 独立更新自己的 1/N 参数切片。对 Adam 来说这很高效：每个 rank 可独立更新自己的切片。

矩阵优化器的问题在于 Newton-Schulz 无法在参数切片上运行——它需要完整的 (m, n) 梯度矩阵才能计算有意义的正交更新。ZeRO-1 持有 1/N 行的 rank 在不先 all-gather 完整矩阵的情况下无法正确正交化。

专属所有权更进一步：单个 rank 拥有整个参数并在本地运行 NS，无需 any all-gather。代价是其他 rank 必须接收更新后参数的广播——但该广播可以隐藏在下一次前向传播中。

详见核心概念。

我需要 HSDP 吗？

单节点多 GPU — 1D shard-only mesh 就够用，更简单。将 mesh 传给 dedicate_params；省略 replicate_mesh。

多节点训练（replicate_size ≥ 2）— HSDP 的两阶段 reduce（shard → replicate）加上 DMuon 的异步后置广播能发挥作用：replicate 广播与下一次前向传播重叠，摊销了节点间 IB 的代价。在跨两个节点 16+ GPU 的场景下，异步隐藏通常是值得的。

将 replicate_mesh=hsdp["replicate"] 传给 dedicate_params，将 replicate_async=True（默认）传给 Muon 即可获得完整 HSDP 收益。

详见 HSDP 指南。

什么时候需要 hook_boundary_predicate？

默认启发式从参数的全限定名中查找 layers.N 或 blocks.N 来确定 hook 注册的层模块。对于标准的 Llama/Qwen model.layers.N.mlp.*_proj 结构和标准的 ViT visual.blocks.N.attn.*_proj 结构，这个方式有效。

在以下情况需要 hook_boundary_predicate：

VLA 模型：action head 位于主层堆栈之外
MoE 模型：expert 参数的 FQN 模式不同
嵌套多模态模型：视觉编码器和 LLM 有独立的层编号层级
自定义 adapter / LoRA 层：adapter 名称不匹配 layers.N

VLA action head 示例：

import dmuon

dmuon.dedicate_params(
    model,
    mesh,
    predicate=lambda n, p: "proj" in n,
    hook_boundary_predicate=lambda m: hasattr(m, "_is_action_layer"),
    hook_boundary_strict=True,
)

详见自定义 Hook 边界。

选 Z2 还是 Z3？

默认 Z3（reshard_after_forward=True）适用于 Muon 目标参数超出每 rank 可用显存预算的任何模型。每步通信代价为 3(N-1)/N · P_M 字节；前向后广播缓冲区被释放，每 rank 峰值显存较低。

Z2（reshard_after_forward=False）适用于可以在前向+反向过程中在每个 rank 保持 P_M 元素常驻的场景。通信代价降至 2(N-1)/N · P_M—— 理论上的 ring all-reduce 下界。最适合较小模型（≤ 3B 参数），此时 GPU 显存不是瓶颈。

粗略经验：当 Muon 目标参数的总字节数在考虑激活值和优化器状态后占每 GPU 显存的 20% 以下时使用 Z2。7B+ 参数几乎总是需要 Z3。

详见 Z2 与 Z3 模式。

DMuon 可以与 DeepSpeed 混合使用吗？

简短回答： ZeRO-3 目前不支持。DeepSpeed ZeRO-3 使用与 DMuon 的 dedicate_params + fully_shard 合约不兼容的参数存储机制（deepspeed.zero.Init + 自定义 hook）。

DeepSpeed ZeRO-0/1/2 集成在路线图中——专属所有权原语在原理上是兼容的，因为参数在存储层没有被碎片化。欢迎贡献；现有方法见集成方案。

当前推荐： 将 DMuon 与 PyTorch FSDP2 配合使用，这是主要的测试和支持配置。

有逐比特一致的收敛保证吗？

有。DMuon 在 4-GPU 测试环境（tests/distributed/test_hsdp_correctness.py）上从三个维度验证逐比特一致输出：

HSDP vs. shard-only： DMuon-HSDP（G=2, R=2）在 10 步训练中与 shard-only DMuon（G=4）产生相同的 loss 值。
非 TP HSDP 的异步 vs. 同步： replicate_async=True 与 replicate_async=False 产生相同输出；存在 TP-sharded dedicated 参数时当前会使用同步 publish。
检查点恢复： 从检查点恢复训练与不中断的连续训练在相同步数内产生相同 loss 值。

这些测试在每个 PR 上运行。如果观察到与单 GPU 基线的偏差，请查看故障排查。

DMuon 支持张量并行（TP）吗？

1D FSDP + TP 支持。先应用 TP，再应用 DMuon，最后 FSDP2：

from torch.distributed.tensor.parallel import parallelize_module
import dmuon
from torch.distributed.fsdp import fully_shard

for layer in model.layers:
    parallelize_module(layer.mlp, tp_mesh, {...})   # 先 TP

dmuon.dedicate_params(model, dp_mesh, ...)          # 再 DMuon

for layer in model.layers:
    fully_shard(layer, mesh=dp_mesh)                # 最后 FSDP2
fully_shard(model, mesh=dp_mesh)

在 DP 组内，所有 rank 共享相同的 TP 位置，因此广播 TP 分片是正确的。 DMuon 使用带 TP 感知的 Gram Newton-Schulz，通过 all-reduce Gram 矩阵实现 O(d²) TP 通信。

2D HSDP × TP（3D 并行）尚未验证。详见张量并行。

需要引用哪些相关工作？

Gram Newton-Schulz（使用默认 "gram" 后端时）：

@misc{GramNewtonSchulz,
  title  = {Gram Newton-Schulz},
  author = {Jack Zhang and Noah Amsel and Berlin Chen and Tri Dao},
  year   = {2026},
  url    = {https://dao-ailab.github.io/blog/2026/gram-newton-schulz/}
}

Muon 优化器（Jordan et al., 2024）：arXiv:2502.16982

Distributed Shampoo（Shi et al., 2023）和 ZeRO-1 （Rajbhandari et al., 2020）开创了 DMuon 所扩展的专属所有权原语。

常见问题¶

参见¶