Skip to content

[WIP] refactor: unify pretrain entry for GPT model - #243

Open
Chamberlain0w0 wants to merge 2 commits into
masterfrom
refactor/pretrain-gpt
Open

Chamberlain0w0 wants to merge 2 commits into
masterfrom
refactor/pretrain-gpt

Conversation

@Chamberlain0w0

@Chamberlain0w0 Chamberlain0w0 commented Oct 8, 2026 •

Copy link
Copy Markdown
Collaborator

GPT-2、LLaMA3、Qwen3 和 Mixtral 的入口包含大量重复的训练逻辑。本次将公共训练生命周期收敛到 infini_train::training::Pretrain(),由 example/pretrain_gpt.cc 统一组装模型、数据和 forward/loss 回调。

  • 框架训练层:统一管理并行初始化、模型包装、梯度累积、优化器、学习率调度和 checkpoint 保存恢复。
  • 模型接入:各模型通过 provider 注册配置与权重加载逻辑,通过独立的 train.flags 提供运行预设。保留四个命名 executable,并支持命令行覆盖参数。
  • 并行执行:普通训练与 PP 共用 forward/loss 协议,通过 FinalizeModelGrads() 完成梯度同步;MoE 根据模型配置处理未命中专家及 EP/EDP 通信。
  • 配置与构建:接入 activation recompute,兼容 Mixtral 原有 batch 参数;模型目录通过统一 CMake helper 声明源码和预设,由示例构建层自动发现。
  • 开发文档:补充职责划分、模型扩展方式、配置部署和验证结果。

验证覆盖四模型 CPU/CUDA、命名入口与统一入口、recompute、checkpoint 恢复,以及 dense 模型的 DP/TP/SP/PP/VPP。Mixtral 的 DP、EP、EP+EDP 数值对照通过,PP/VPP/EP+PP 完成训练,EP+ZeRO 0/1/2 与 master 原训练逻辑结果一致。

@Chamberlain0w0 Chamberlain0w0 self-assigned this Oct 8, 2026
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.

1 participant