Skip to content

fix(dataloader): track distributed progress by global batch - #204

Merged
kilinchange merged 4 commits into
masterfrom
fix/dataloader-global-batches
Sep 11, 2026
Merged

fix(dataloader): track distributed progress by global batch#204
kilinchange merged 4 commits into
masterfrom
fix/dataloader-global-batches

Conversation

@chen2021673

Copy link
Copy Markdown
Contributor

概述

修复分布式训练中 DataLoader 使用局部 batch 进行数据划分和索引,可能导致 batch 访问越界的问题。

主要修改

  • 统一使用 global batch 作为分布式数据划分和读取单位。
  • 按 data parallel rank 计算各进程应读取的数据范围。
  • 增加 DataLoader 单元测试,覆盖多 rank 下的 batch 划分和边界行为。

Comment thread infini_train/include/dataloader.h Outdated
size_t batch_size_ = 0;
size_t batch_idx_ = 0;
size_t max_batch_idx_ = 0;
size_t global_batch_idx_ = 0;

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.

这里变量不是直接表示batch维度的,为了防止理解混淆,建议改个名字,比如 dataloader_step_ ,另外其他变量也统一改下命名

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

改动:
global_batch_idx_ -> dataloader_step_
num_global_batches_ -> num_dataloader_steps_

Comment thread infini_train/src/dataloader.cc Outdated

DataLoader::DataLoader(const std::shared_ptr<Dataset> &dataset, size_t batch_size)
: dataset_(dataset), batch_size_(batch_size), max_batch_idx_((dataset_->Size() + batch_size_ - 1) / batch_size_) {}
: dataset_(dataset), batch_size_(batch_size), num_global_batches_(CheckedCeilDiv(dataset_->Size(), batch_size_)) {}

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.

建议先校验 dataset_ != nullptr 和 batch_size_ > 0,再计算 num_global_batches_。当前在初始化列表中直接访问 dataset_->Size(),空指针会导致非法解引用,batch_size_ == 0 也会引发除零。

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

新增
CHECK(dataset_ != nullptr) 和 CHECK_GT(batch_size_, 0)

@chen2021673
chen2021673 force-pushed the fix/dataloader-global-batches branch from cf6fe96 to ee324e4 Compare August 27, 2026 02:36
- rename global batch tracking to dataloader steps
- validate dataset and batch size before computing steps
Comment thread infini_train/include/dataloader.h Outdated
- resume training by skipping consumed batches in main
- keep resume positioning as a future sampler responsibility
- preserve distributed global-batch partitioning
@kilinchange

Copy link
Copy Markdown
Collaborator

麻烦贴一下测试通过截图,并备份 log 和写入飞书性能表格。

@chen2021673

chen2021673 commented Sep 11, 2026

Copy link
Copy Markdown
Contributor Author

麻烦贴一下测试通过截图,并备份 log 和写入飞书性能表格。

单元测试:
image

精度测试:
image

性能测试:
image

@kilinchange
kilinchange merged commit 39889aa into master Sep 11, 2026
2 checks passed
@kilinchange
kilinchange deleted the fix/dataloader-global-batches branch September 11, 2026 02:07
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