⚠ This page is served via a proxy. Original site: https://github.com
This service does not collect credentials or authentication data.
Skip to content

Conversation

@Chamberlain0w0
Copy link
Contributor

@Chamberlain0w0 Chamberlain0w0 commented Dec 24, 2025

通过 --use_distributed_optimizer 来开启基于分布式优化器的 ZeRO-1 优化。Stage 1 的优化对 SGD 没有效果,但是依然在 GPT-2 训练脚本中加了相关 option,为的是之后 stage 2/3 的扩展。

实现细节:

  1. 添加 ParamAndGradBucket/Group/Buffer 的基建,采用了与 Megatron-LM 类似的实现思路:所有 param/grad 都连续放在一个一维 buffer 中,并分 bucket 按组进行通信;
  2. 修改 DistributedDataParallel 的逻辑,需要在构造函数中创建上述 buffer、划分 bucket group,并完成 hook 注册;同时保留了原 reducer 逻辑的实现分支,不开启 --use_distributed_optimizer 时则走原先的逻辑;
  3. 添加 DistributedOptimizer 类,继承自 Optimizer;同时为每种优化器包了一个仅接受 models->Parameters() 的构造 function 方法,DistributedOptimizer 在构造时会接受这个 function,并在其构造函数内部创建 base_optimizer;
  4. 修改了 PP 的构造函数参数,把 optimizer 的传入延后到了 Module::TrainStep,这样 PP 对象可以完全不持有 optimizer 对象;
  5. 在训练循环添加了本轮峰值占用/预留显存信息的输出。

详细实现流程可以参考叙述:https://gxtctab8no8.feishu.cn/wiki/XQbGwXSsZi3MutkZuhXcKWsinnY#share-QU0fdM1cYoT06vxmwMbcnLJUn1b

两个比较重要的细节:

  1. DDP 对象在构造时,会一同构造模型对应的 buffer/bucket_group 等;同时 DistOpt 在构造时会接受 DDP 对象的 buffer/group;
  2. 由于现在 PP + DDP 的实现,会单独将每个 chunk 构造为 DDP 对象,所以传给 DistOpt 的 buffer/bucket_group 等,需要额外在 chunk 之间做一道汇总,得到一个总 list 再传进去。

@Chamberlain0w0 Chamberlain0w0 force-pushed the feat/distributed_optimizer branch from be321ca to 4afb235 Compare January 12, 2026 06:27
@Chamberlain0w0 Chamberlain0w0 changed the title [WIP] feat: Add DistributedOptimizer and support ZeRO-1 feat: Add DistributedOptimizer and support ZeRO-1 Jan 12, 2026
@kilinchange
Copy link
Collaborator

把 ddp 相关的文件在 parallel 里单独放一个 ddp 文件夹吧。

@Chamberlain0w0
Copy link
Contributor Author

记录下写入飞书表格脚本的更新:

  1. 表格中新增两列显存占用的信息:
image 2. 之后扩充测例的时候支持自动创建 sheet(sheet 名称与 test_config 中指定的测例 id 一致),无需手动创建,执行逻辑大致类似于: image

@Chamberlain0w0
Copy link
Contributor Author

image image

@Chamberlain0w0
Copy link
Contributor Author

image image

@kilinchange kilinchange merged commit c4703e1 into master Jan 30, 2026
2 checks passed
@kilinchange kilinchange deleted the feat/distributed_optimizer branch January 30, 2026 01:46
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants