提出分片局部块对角特征基预条件优化方法Dobb
原文：Dobb: Shard-Local Preconditioning in Block-Diagonal Eigenbases
训练参数量和优化器状态超出单节点内存的语言模型，需要将这些张量分片到多个设备上。在张量并行（TP）或全分片数据并行（FSDP）下，诸如Muon之类的矩阵感知优化器会引入通信，以跨分片执行全局矩阵运算。我们提出了Dobb（Distributed Optimizer with Block-diagonal Eigenbases），一种矩阵感知优化器，设计用于在参数分片内部本地计算更新。Dobb用对角块近似输入激活二阶矩，并在每个块内跟踪一个完整基。它将动量投影到这些基上，对系数应用符号归一化，并重建更新。当在TP或FSDP下每个块位于一个分片内时，Dobb可在本地应用其矩阵更新，除激活二阶矩的数据并行同步外无需任何通信。我们将Dobb与Dion进行基准对比，后者是一种分布式矩阵感知优化器，旨在保留正交归一化更新的优势同时减少通信。在GPT-2 XL（15.6亿参数）上，使用16个块的Dobb达到2.476的验证损失，而Dion在秩比例为1/16时为2.511。在32块H100 GPU上运行FSDP时，Dobb将摊销优化器调用时间从Dion的124.1毫秒降至9.2毫秒。作者：Shi-Yuk Fong, Shuhua Yu, Jiawei Zhang 类别：cs.LG,cs.DC,math.OC