Dobb optimizer speeds up local updates for large language models
Dobb: Shard-Local Preconditioning in Block-Diagonal Eigenbases
Machine LearningDistributed, Parallel, and Cluster Computing
Summary
Training very large language models requires splitting their data across many computers, which can slow down training because of the need to share information frequently. The authors propose Dobb, a new way to update model parameters locally within each computer shard without extra communication, making the process faster. Dobb divides the model’s data into blocks and works within each block, keeping updates efficient and reducing synchronization time. Tests show Dobb achieves better accuracy and much faster updates compared to previous methods on a large GPT-2 model.
What this means in practice
- •For ai infrastructure teams: Speed up training large language models by reducing communication overhead in optimizer updates across GPUs.
- •For high performance computing teams: Improve distributed training efficiency by applying local optimizer updates within device shards for large-scale matrix computations.
Authors
Shi-Yuk Fong, Shuhua Yu, Jiawei Zhang
Abstract
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.