Make sure dataloader state is the same after checkpoint is loaded

cd4822f1 · Deepak Narayanan · c671de3e · cd4822f1
Commit cd4822f1 authored Nov 12, 2020 by Deepak Narayanan
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

megatron/training.py megatron/training.py +2 -2

No files found.
--- a/megatron/training.py
+++ b/megatron/training.py
@@ -870,8 +870,8 @@ def build_train_valid_test_data_iterators(

    # Shift the start iterations.
    if train_dataloader is not None:
-        train_dataloader.batch_sampler.start_iter = args.iteration % \
-            len(train_dataloader)
+        train_dataloader.batch_sampler.start_iter = \
+            (args.iteration * args.num_microbatches_in_minibatch) % len(train_dataloader)
        print_rank_0('setting training data start iteration to {}'.
                     format(train_dataloader.batch_sampler.start_iter))
    if valid_dataloader is not None: