ai-infra-interview-305

第 60 题：如何实现自定义的分布式优化器？继承torch.optim.Optimizer的注意事项？

题目

如何实现自定义的分布式优化器？继承torch.optim.Optimizer的注意事项？

---

完整讲解

一、继承 Optimizer 要做什么？

torch.optim.Optimizer 要求子类：

• 在 __init__ 里调 super().__init__(params, defaults)，并保存 self.param_groups（及每个 group 的 param 列表和超参）。
• 实现 step(closure=None)：遍历 self.param_groups 里每个 param，用其 param.grad 按你的算法更新 param.data；若支持 closure（如 LBFGS），可执行 closure 再更新。
• 不在 step 里默认调 zero_grad()；由用户在 backward 后、step 前自己调，或你在 step 里可选地调。

分布式下：梯度通常已由 DDP/FSDP 等做 all-reduce，每个 rank 上的 param.grad 已是「全局梯度」；自定义优化器只需在 本 rank 上按梯度更新本 rank 的参数（FSDP 下参数是分片的，每卡只更新自己的分片）。所以「分布式」部分多数由 DDP/FSDP 解决，优化器侧主要是正确读 grad、写 data。

---

二、注意事项

• 状态与分片：若优化器有状态（如 Adam 的 m、v），在 FSDP 下状态也应按分片存：每卡只存本卡参数分片对应的状态；all-gather 参数时不同步状态（状态随分片走）。若用 DDP，每卡完整参数，状态也完整，与单机一致。
• 设备与 dtype：param.data 和 param.grad 的 device/dtype 要一致；混合精度时 grad 可能是 fp16，若优化器用 fp32 状态，要转成 fp32 再更、写回时再转（或用 master weight）。
• 梯度是否为 None：若某 param 未参与计算（如 find_unused_parameters），grad 可能为 None；step 里要 skip 这类 param，不要用 None 做运算。
• closure：若实现 step(closure)，应先执行 closure() 再更新，并返回 closure 的返回值（若需要）；多数优化器 closure 可为 None。
• 与 DDP/FSDP 顺序：先 backward()（DDP 会 all-reduce），再 optimizer.step()；若用 AMP，先 scaler.unscale_(optimizer)、检查 inf/nan，再 optimizer.step()、scaler.update()。

---

三、小结

• 子类实现 __init__(params, defaults) 和 step()；step 里遍历 param_groups、用 grad 更新 data。
• 分布式：梯度已由 DDP/FSDP 同步，优化器只做本地更新；FSDP 下状态要随分片、只更新本卡分片。
• 注意：grad 为 None 时 skip、设备与 dtype、与 AMP 的配合顺序。

---

面试要点

• 继承 Optimizer：init 调 super、保存 param_groups；step 里遍历 params、用 grad 更新 data。
• 分布式：梯度已 all-reduce，优化器只做本地更新；FSDP 下状态随分片。
• 注意：None grad 要 skip、设备/dtype、AMP 时 unscale 再 step。

---

记忆要点

1. 必须：super().init、实现 step；step 里用 param.grad 更新 param.data。
2. 分布式由 DDP/FSDP 管梯度；优化器管本地更新；FSDP 状态随分片。
3. 处理 None grad、设备/dtype、与 scaler 顺序。

返回模块

返回总览