Skip to content

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

Merged
kilinchange merged 2 commits into
masterfrom
feat/checkpoint-consumed-micro-batches
Aug 26, 2026
Merged

feat: track checkpoint progress in micro-batches#195
kilinchange merged 2 commits into
masterfrom
feat/checkpoint-consumed-micro-batches

Conversation

@JYMiracle305

@JYMiracle305JYMiracle305 commented Jul 31, 2026

Copy link
Copy Markdown
Contributor

主要修改

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

在不同流水线并行配置下,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 测试

@JYMiracle305
JYMiracle305 changed the base branch from master to feat/named-parametersAugust 3, 2026 08:04
@JYMiracle305
JYMiracle305 changed the base branch from feat/named-parameters to masterAugust 3, 2026 08:06
@JYMiracle305
JYMiracle305force-pushed the feat/checkpoint-consumed-micro-batches branch from 49d8d4b to 57ab888CompareAugust 3, 2026 08:27
@JYMiracle305
JYMiracle305 changed the base branch from master to feat/named-parametersAugust 3, 2026 09:52
@JYMiracle305
JYMiracle305force-pushed the feat/checkpoint-consumed-micro-batches branch from 57ab888 to 7689011CompareAugust 4, 2026 07:43
Comment threadinfini_train/include/checkpoint/checkpoint.h Outdated
@JYMiracle305
JYMiracle305force-pushed the feat/checkpoint-consumed-micro-batches branch 2 times, most recently from 59a9674 to 9d99035CompareAugust 6, 2026 08:44
@JYMiracle305
JYMiracle305force-pushed the feat/checkpoint-consumed-micro-batches branch from 9d99035 to 862a378CompareAugust 6, 2026 09:47
@kilinchange
kilinchangeforce-pushed the feat/checkpoint-consumed-micro-batches branch from 862a378 to f6f4715CompareAugust 7, 2026 02:18
@JYMiracle305
JYMiracle305force-pushed the feat/checkpoint-consumed-micro-batches branch 2 times, most recently from 9d30401 to 6799ce4CompareAugust 12, 2026 14:55

@Chamberlain0w0Chamberlain0w0 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

没啥问题了,还剩 stacked PR 的冲突解决

@JYMiracle305
JYMiracle305force-pushed the feat/checkpoint-consumed-micro-batches branch 2 times, most recently from 9d406c9 to 86aefcdCompareAugust 14, 2026 09:35
Base automatically changed from feat/named-parameters to masterAugust 20, 2026 08:49
@JYMiracle305
JYMiracle305force-pushed the feat/checkpoint-consumed-micro-batches branch from 86aefcd to 48ca7e4CompareAugust 20, 2026 08:56
@JYMiracle305

Copy link
Copy Markdown
ContributorAuthor

精度对比

image

性能对比

image

Comment threadexample/gpt2/main.cc
@JYMiracle305
JYMiracle305force-pushed the feat/checkpoint-consumed-micro-batches branch from 48ca7e4 to 3a1bd6dCompareAugust 26, 2026 08:09
@kilinchange
kilinchange merged commit e403ed4 into masterAug 26, 2026
2 checks passed
Sign up for freeto 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

@JYMiracle305@kilinchange@Chamberlain0w0