Dobb: Shard-Local Preconditioning in Block-Diagonal Eigenbases
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.
Related
- A parameterized block LU preconditioning for double saddle point problems
- Robust augmented Block Preconditioning for two-phase Incompressible Inductionless MHD
- Block Preconditioning for All-at-Once Variable-Coefficient Fractional Evolution Equations via the GLT Analysis
- BLOKS: Scaling Kernel Solvers with Sketched Block Preconditioning
- Block diagonal preconditioners for arrow-shaped multiple saddle-point linear systems
BibTeX
@misc{fong2026dobb,
title = {{Dobb: Shard-Local Preconditioning in Block-Diagonal Eigenbases}},
author = {Shi-Yuk Fong and Shuhua Yu and Jiawei Zhang},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.39412},
eprint = {2610.39412},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.39412},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Dobb: Shard-Local Preconditioning in Block-Diagonal Eigenbases
AU - Shi-Yuk Fong
AU - Shuhua Yu
AU - Jiawei Zhang
PY - 2026
DA - 2026/10/29
PB - arXiv
DO - 10.48550/arxiv.2610.39412
UR - https://arxiv.org/abs/2610.39412
AB - 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.
ER -