ZeRO-1 风格的 Distributed Optimizer

nanotron 没用 DeepSpeed / Megatron-Core 的 distributed optimizer，自己实现了一个 ZeRO-1（仅切 optimizer state）。本章看它的设计。

6.1问题：Adam state 占大头

fp32 mixed precision 训练时，每个参数的内存账：

其中 master + m + v 共 12 bytes（75%）都是 fp32 optimizer state。如果把它沿 DP 切，单卡显存减 ~75%。

6.2nanotron 的 ZeRO-1

看 src/nanotron/optim/zero.py（约 300 行）：

class ZeroDistributedOptimizer(Optimizer):
    def __init__(self, named_params, optimizer_builder, dp_pg):
        self.dp_pg = dp_pg
        dp_size = dp_pg.size()
        dp_rank = dp_pg.rank()

        # 1) 把所有 parameter 平分到 DP rank
        self.local_params = self._shard_params(named_params, dp_rank, dp_size)

        # 2) 只为本地那份建 optimizer state
        self.optimizer = optimizer_builder(self.local_params)

    @torch.no_grad()
    def step(self):
        # 3) 每张卡只更新自己那份
        self.optimizer.step()

        # 4) all-gather 更新后的 weight
        for p_full, p_shards in self.param_groups:
            dist.all_gather(p_shards, p_local, group=self.dp_pg)

6.3shard 划分策略

nanotron 把每个 parameter 单独切（不像 DeepSpeed flatten 一坨）。具体：

• 对一个 weight tensor，按总 numel 平分给 dp_size 个 rank；
• 每 rank 保有自己那 1/N（fp32 master + m + v）；
• 训练前 forward / backward 用完整 bf16 weight；
• 反向后 reduce-scatter 梯度到自己负责的 shard；
• step 后 all-gather 更新好的 bf16 weight。

6.4显存对比（7B 模型）

6.5跟 ZeRO-2 / ZeRO-3 区别

nanotron 设计上是TP + PP + ZeRO-1，跟 Megatron 一致。要 ZeRO-3 的人通常该选 ColossalAI / DeepSpeed。

6.6启用

YAML 里：

optimizer:
  zero_stage: 1                # 1 启用 ZeRO-1；0 关闭
  weight_decay: 0.1
  clip_grad: 1.0
  optimizer_factory:
    name: adam_w
    adam_beta1: 0.9
    adam_beta2: 0.95
    adam_eps: 1.0e-8
    torch_adam_is_fused: true

6.7梯度累积时怎么办

nanotron 的 grad accumulation 在 DP all-reduce 之前：

1. 多次 micro_batch 前向 + 反向，梯度累加到 .grad；
2. 所有 micro_batch 跑完，再 DP all-reduce；
3. ZeRO optimizer step。

6.8optim/ 模块源码定位

6.9这章你需要带走的

• nanotron 自带 ZeRO-1 distributed optimizer，~300 行；
• 核心：每 parameter 切 + 本地 step + all-gather 更新后 weight；
• 7B 模型显存 84GB → 10GB（DP=8）；
• 没实现 ZeRO-2/3，因为 TP+PP 是另一条路径；
• YAML 用 optimizer.zero_stage: 1 启用。