Dobb 优化器让分片训练中的矩阵更新无需跨设备通信
Dobb: Shard-Local Preconditioning in Block-Diagonal Eigenbases
训练超过单机显存的模型时，Muon 这类优化器会被跨卡通信拖慢，这篇论文的 Dobb 把矩阵更新改成块内本地计算，H100 上优化器调用从 124ms 降到 9.2ms，做分布式训练的可以看看。
论文提出 Dobb（Distributed Optimizer with Block-diagonal Eigenbases），一种面向参数分片场景的矩阵感知优化器。它用对角块近似输入激活的二阶矩，在每个块内跟踪完整基，将动量投影到这些基上并做符号归一化，使得每个块落在同一分片内时更新可完全本地计算。在 GPT-2 XL（1.56B 参数）上，Dobb 用 16 个块达到验证损失 2.476，优于 Dion 在 rank fraction 1/16 下的 2.511。在 32 张 H100 GPU 的 FSDP 环境下，Dobb 将摊销的优化器调用时间从 Dion 的 124.1 ms 降到 9.2 ms。
Dobb: Shard-Local Preconditioning in Block-Diagonal Eigenbases
Training language models whose parameters and optimizer state exceed a single node's memory requires sharding these tensors across devices. Under tensor parallelism (TP) or fully sharded data parallelism (FSDP), matrix-aware optimizers such as Muon introduce communication to perform global matrix operations across shards. We propose Dobb (Distributed Optimizer with Block-diagonal Eigenbases), a matrix-aware optimizer designed to compute updates locally within parameter shards. Dobb approximates the input-activation second moment with diagonal blocks and tracks a complete basis within each block. It projects momentum into these bases, applies sign normalization to the coefficients, and reconstructs the update. When each block lies within one shard under TP or FSDP, Dobb applies its matrix update locally, requiring no communication beyond data-parallel synchronization of the activation second moments. We benchmark Dobb against Dion, a distributed matrix-aware optimizer designed to retain the benefits of orthonormalized updates while reducing communication. On GPT-2 XL (1.56B parameters), Dobb with 16 blocks reaches a validation loss of 2.476, compared with 2.511 for Dion at rank fraction 1/16. Under FSDP on 32 H100 GPUs, Dobb reduces the amortized optimizer-call time from 124.1 ms for Dion to 9.2 ms.