Uh oh!
There was an error while loading. Please reload this page.
fix(dataloader): track distributed progress by global batch - #204
fix(dataloader): track distributed progress by global batch#204chen2021673 wants to merge 3 commits into
Conversation
| size_t batch_size_ = 0; | ||
| size_t batch_idx_ = 0; | ||
| size_t max_batch_idx_ = 0; | ||
| size_t global_batch_idx_ = 0; |
There was a problem hiding this comment.
这里变量不是直接表示batch维度的，为了防止理解混淆，建议改个名字，比如 dataloader_step_ ，另外其他变量也统一改下命名
There was a problem hiding this comment.
改动：
global_batch_idx_ -> dataloader_step_
num_global_batches_ -> num_dataloader_steps_
| 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_)) {} |
There was a problem hiding this comment.
建议先校验 dataset_ != nullptr 和 batch_size_ > 0，再计算 num_global_batches_。当前在初始化列表中直接访问 dataset_->Size()，空指针会导致非法解引用，batch_size_ == 0 也会引发除零。
There was a problem hiding this comment.
新增
CHECK(dataset_ != nullptr) 和 CHECK_GT(batch_size_, 0)
cf6fe96 to
ee324e4Compare- rename global batch tracking to dataloader steps - validate dataset and batch size before computing steps
| friend bool operator==(const DataLoaderIterator &lhs, const DataLoaderIterator &rhs); | ||
| size_t DataLoaderStep() const; | ||
| DataLoaderIterator &SeekDataLoaderStep(size_t dataloader_step); |
There was a problem hiding this comment.
“从第 xx 个 step 继续读取” 的状态控制应当是 Sampler 的职责：Sampler 根据恢复进度决定起始 sample/index，DataLoader 本身不需要额外维护一份可 seek 的状态。
Megatron 也是在 Sampler 层通过 consumed_samples 控制数据恢复位置：
https://github.com/NVIDIA/Megatron-LM/blob/f7f584d7a04c1891bc583a1213ebb1dbcac4e158/megatron/training/datasets/data_samplers.py#L172
建议暂时不要在 DataLoader / DataLoaderIterator 中引入额外的 dataloader_step 和 SeekDataLoaderStep 状态。现阶段继续在 main 侧通过迭代并跳过前 N 个 batch 的方式完成断点恢复，作为临时方案；同时保留 Sampler 的 TODO，后续引入 Sampler 抽象后，再将恢复位置的控制下沉到 Sampler。
概述
修复分布式训练中 DataLoader 使用局部 batch 进行数据划分和索引，可能导致 batch 访问越界的问题。
主要修改