Skip to content

feat: track checkpoint progress in micro-batches - #195

Open
JYMiracle305 wants to merge 1 commit into
masterfrom
feat/checkpoint-consumed-micro-batches
Open

feat: track checkpoint progress in micro-batches#195
JYMiracle305 wants to merge 1 commit into
masterfrom
feat/checkpoint-consumed-micro-batches

Conversation

@JYMiracle305

@JYMiracle305 JYMiracle305 commented Jul 31, 2026

Copy link
Copy Markdown
Contributor

主要修改

该 PR 统一了 checkpoint 中训练进度的计量语义,将原先含义模糊的 consumed_batches 重命名为 consumed_micro_batches。

在不同流水线并行配置下,DataLoader 单次迭代包含的 micro-batch 数量不同。该 PR 将 checkpoint 中保存的进度统一为 micro-batch 数量,并在恢复训练时根据当前 PP 配置换算为对应的 DataLoader 位置,同时正确处理 DDP rank 和步长。

  • 将 checkpoint 数据结构、保存/恢复接口及 JSON 字段统一重命名为 consumed_micro_batches
  • PP 场景保存 checkpoint 时,将 DataLoader batch 进度转换为 micro-batch 进度
  • 恢复训练时,根据当前 PP 配置换算回 DataLoader batch 位置
  • 增加 PP batch 边界及 DDP stride 合法性检查
  • 同步修改 GPT-2、LLaMA3 训练入口和 checkpoint 测试

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