Level 27 | ZeRO Optimizer Sim

Mission 1 TODO 1 / TODO 2

ZeRO：把优化器状态切给不同 GPU

先补的知识

Adam 等优化器会为参数保存额外状态。
每张卡都保存完整状态会浪费显存。
ZeRO 把参数或优化器状态按 shard 分给多张卡。

图解原理

多人共同保管一本大账本，每个人只负责自己的几页，需要时再同步完整结果。

GPU0: 参数 0-3GPU1: 参数 4-7GPU2: 参数 8-11

语法热身：用参数列表划分负责范围

params = list(model.parameters())
half_idx = len(params) // 2
gpu_partitions = {
    0: params[:half_idx],
    1: params[half_idx:],
}

闯关题

ZeRO 分片的直接目的是什么？

减少每张 GPU 的显存占用让 loss 不需要反传把模型改成 CNN

回到 notebook 的作业

这里不直接写答案。你已经拿到足够输入，最后用 notebook 的 TODO 做举一反三。

把 model_params 转成参数列表，并按列表切成两份。为每张 GPU 初始化局部 optimizer state。确认局部状态长度小于完整参数长度。

Mission 2 TODO 3 / TODO 4

局部更新：单机模拟里原位修改就是同步

先补的知识

每张卡只遍历自己负责的参数列表。
TODO 3 要更新局部动量并用动量更新参数。
TODO 4 在本 notebook 中不需要写代码，只要理解真实多卡需要把更新后的权重同步出去。

图解原理

本练习里每个 p 都是同一个 Parameter 的引用；你原位改 p.data，模型里的参数也同步变了。

GPU0 负责 fc1GPU1 负责 fc2p.data 原位更新真实多卡才需要广播

语法热身：遍历本 GPU 参数和梯度

for p, g in zip(params, grads):
    momentum = states[id(p)]
    momentum = momentum + g
    states[id(p)] = momentum
    p.data = p.data - lr * momentum

闯关题

为什么 notebook 的 TODO 4 不需要额外写 All-Gather 代码？

因为梯度会自动消失 p.data 是原位引用，单机模拟里已经同步因为 ZeRO 不需要参数

回到 notebook 的作业

这里不直接写答案。你已经拿到足够输入，最后用 notebook 的 TODO 做举一反三。

只遍历当前 GPU 负责的 params 和对应 grads。更新 states[id(p)] 中保存的动量。用 p.data 原位更新参数，并理解 TODO 4 在真实多卡中的同步意义。