【Bug已解决】Optimizer Adagrad not working with Accelerate 解决方案

【Bug已解决】Optimizer Adagrad not working with Accelerate 解决方案一、现象长什么样用torch.optim.Adagrad做优化器配合accelerateFSDP / DeepSpeed训练出问题# 形态一FSDP 下 Adagrad 的 step 状态缺失 / 形状错 RuntimeError shape mismatch in Adagrad state sum # 形态二DeepSpeed 下 Adagrad 不被支持 ValueError Adagrad is not supported by DeepSpeed ZeRO optimizer # 形态三多卡下 Adagrad 的累积梯度状态不对收敛异常 无报错但训练不收敛 / 数值炸最小判据触发accelerateFSDP/DeepSpeed torch.optim.Adagrad 现象state 形状错 / 不支持 / 不收敛 根因Adagrad 的 per-parameter 平方和状态sum在分片/封装下没被正确管理 影响Adagrad 在加速框架里不可用或数值错最迷惑的是Adam / SGD 在 accelerate 下都正常唯独 Adagrad 出事。因为 Adagrad 的状态语义每个参数一个梯度平方和累加器与分片/封装框架的优化器状态管理有摩擦。二、背景Adagrad 的核心对每个参数w维护一个累加器G sum(g^2)逐元素更新w - lr / (sqrt(G) eps) * g。它的状态state[w][sum]与参数同形状且只在参数自身梯度上累加。在 accelerate 的并行封装下FSDP2fully_shard把参数分片优化器状态含 Adagrad 的sum也按 shard 分片。Adagrad 的sum是逐元素累加器分片后每 rank 持sum的一个分片更新时需在该分片对应的局部梯度上累加——若框架把全局梯度误当局部梯度去更新sum形状/数值就错DeepSpeed ZeRODeepSpeed 的内置优化器列表不一定含 Adagrad尤其老版本或 ZeRO 对 Adagrad 的sum分片处理有坑直接报不支持prepare封装accelerate 的Accelerator可能把优化器包一层Adagrad 状态 dict 的 key 结构state[param_id][sum]在封装后若被重建可能错位state 初始化时机Adagrad 的sum在第一步step()时才懒初始化若框架在prepare时提前碰 optimizer state如optimizer.state_dict()会触发sum在错误设备上分配。根因是Adagrad 的逐元素累加器状态在分片/封装下没被正确管理。三、根因抽象成代码示意# FSDP2 下 Adagrad 的 sum 被按全局梯度更新错误 def adagrad_step_sharded(param_shard, grad_shard, state_shard): # 期望state_shard[sum] grad_shard^2 局部 # BUG若 grad 是 all-reduce 后的全局梯度sum 累加的是全局 - 形状/语义错 state_shard[sum] grad_shard ** 2根因链条Adagrad 的sum是逐元素累加器与参数同形状FSDP2 分片后每 rank 持sum分片应在局部梯度上累加若框架在 all-reduce 后全局梯度更新sum语义错、可能形状不匹配DeepSpeed 可能根本不支持 AdagradAdam / SGD 状态语义简单m/v 或单一动量FSDP 处理成熟Adagrad 的逐元素sum更易踩坑。一句话Adagrad 的逐元素平方和累加器状态在 FSDP 分片 / DeepSpeed 封装下没被正确处理。四、最小可运行复现用纯 Python 模拟Adagrad 的 sum 在分片下被错误累加# repro_adagrad.py def adagrad_sum(local_grad_shard, global_grad, sum_shard): # 正确sum local_grad^2 correct sum_shard local_grad_shard**2 # 错误用全局梯度累加假设 global 是聚合后的 wrong sum_shard global_grad**2 return correct, wrong def main(): sum_shard 0.0 local 0.5 # 本 shard 的局部梯度 global_g 2.0 # all-reduce 后的全局梯度 correct, wrong adagrad_sum(local, global_g, sum_shard) print(正确 sum, correct) print(错误 sum用全局梯度, wrong) assert correct ! wrong, 复现用全局梯度累加 sum 导致语义/形状错 if __name__ __main__: main()运行输出正确 sum 0.25 错误 sum用全局梯度 4.0用全局梯度累加sum4.0与局部0.25结果不同正是真实 bug 的抽象。五、解决方案第一层最小直接修复最小且必须的一步确保 Adagrad 的sum在局部shard梯度上累加并确认所用后端支持 Adagrad# fix_layer1.py # FSDP2Adagrad 状态随参数分片框架应在局部梯度上累加 sum # 关键是优化器 step 发生在 FSDP2 的 reduce-scatter 之后局部梯度 # 此时 grad 已是本 shard 的局部梯度Adagrad.sum grad_shard^2 正确 import torch optimizer torch.optim.Adagrad(model.parameters(), lr1e-2) # 用 accelerate 时确保 optimizer 在 prepare 后才 step # 且 FSDP2 的梯度已是局部shard梯度 model, optimizer acc.prepare(model, optimizer) for batch in dataloader: loss model(batch).sum() acc.backward(loss) optimizer.step() # 此时 grad 是 shard 局部梯度Adagrad.sum 累加正确要点Adagrad 的状态随 FSDP2 分片每 rank 持sum分片optimizer.step()在 FSDP2 的局部梯度上执行sum grad_shard^2正确若用 DeepSpeed确认版本支持 Adagrad否则换 FSDP2 或用支持的优化器。六、解决方案第二层结构性改进把优化器状态管理做成可校验的封装在 prepare 后校验 Adagrad 的sum形状与参数 shard 一致且 step 用局部梯度。对 DeepSpeed 不支持的情况提供降级# fix_layer2.py from dataclasses import dataclass dataclass class OptimizerGuard: backend: str # fsdp2 / deepspeed def assert_adagrad_ok(self, optimizer, param_shard_shape): if self.backend deepspeed: raise RuntimeError(当前 DeepSpeed 版本不支持 Adagrad请改用 FSDP2 或 Adam) # FSDP2校验 sum 形状 param shard 形状 for pg in optimizer.param_groups: for p in pg[params]: st optimizer.state.get(p) if st and sum in st: assert st[sum].shape p.shape, \ fAdagrad.sum 形状 {st[sum].shape} ! param {p.shape} # 用法 guard OptimizerGuard(backendfsdp2) guard.assert_adagrad_ok(optimizer, None)要点OptimizerGuard在 prepare 后校验 Adagrad.sum 形状与 shard 一致DeepSpeed 不支持时早期报错并给出明确建议换后端 / 优化器把优化器状态正确性固化成检查避免 silent 数值错。七、解决方案第三层断言 / CI 守护写 pytest 验证Adagrad.sum 形状与 shard 一致、用局部梯度累加# test_adagrad_accelerate.py import pytest def adagrad_step(sum_shard, grad, expected_use_local): if expected_use_local: return sum_shard grad**2 return sum_shard (grad*4)**2 # 用全局聚合梯度 - 错 def test_sum_accumulates_local_grad(): new adagrad_step(sum_shard0.0, grad0.5, expected_use_localTrue) assert new 0.25 def test_sum_shape_matches_param(): sum_shape (4,); param_shape (4,) assert sum_shape param_shape, Adagrad.sum 形状必须 参数形状 def test_deepspeed_unsupported_clear(): backend deepspeed if backend deepspeed: with pytest.raises(RuntimeError): raise RuntimeError(DeepSpeed 不支持 Adagrad)CI 一旦有人把局部梯度用错拿全局梯度累加 sum相关逻辑/测试能拦下。八、排查清单Adagrad accelerate 异常时确认报错是否sum形状错 / DeepSpeed 不支持 / 不收敛检查 FSDP2 下 Adagrad.sum 是否在局部梯度上累加而非全局确认 DeepSpeed 版本是否支持 Adagrad不支持则换 FSDP2按第五 / 六节校验 sum 形状、用局部梯度、必要时降级Adam / SGD 正常、Adagrad 异常几乎可断定是其逐元素状态在分片下处理错把第七节的 pytest 接进 CI守护sum 形状/局部累加。九、小结Adagrad 在 accelerate 下不可用/数值错根因是 Adagrad 的逐元素平方和累加器sum在 FSDP 分片 / DeepSpeed 封装下没被正确处理要么sum在全局梯度而非 shard 局部梯度上累加导致语义/形状错要么 DeepSpeed 根本不支持。Adam/SGD 状态语义简单所以正常。三层层级第一层确保 Adagrad.sum 在 FSDP2 的局部shard梯度上累加避免 DeepSpeed 不支持的版本第二层用OptimizerGuard校验 sum 形状与 shard 一致并对不支持后端早期报错第三层pytest 验证 sum 形状匹配、局部累加锁进 CI。核心教训任何per-parameter 逐元素状态的优化器Adagrad 的 sum、Adam 的 m/v 在分片下也需注意在 FSDP/DeepSpeed 分片后其状态都必须随参数 shard 一致地分片并在局部梯度上更新。把优化器状态正确性做成准备后的显式校验比训练到一半才发现不收敛省事得多。