Quick fix for pipeline tasks to get learning rate correct

6e83649f · Jared Casper · Deepak Narayanan · 25c07e14 · 6e83649f
Commit 6e83649f authored Dec 18, 2020 by Jared Casper Committed by Deepak Narayanan Dec 19, 2020
Hide whitespace changes
Inline Side-by-side

Showing with 6 additions and 2 deletions

tasks/finetune_utils.py tasks/finetune_utils.py +6 -2

No files found.
--- a/tasks/finetune_utils.py
+++ b/tasks/finetune_utils.py
@@ -129,10 +129,14 @@ def _build_train_valid_dataloaders(train_dataset, valid_dataset):
                                          args.num_workers, not args.keep_last)
    valid_dataloader = _build_infinite_size_dataloader(valid_dataloader_)
-    # Now that we've built the data loaders, set args.micro_batch_size to
+    # Now that we've built the data loaders, set batch_size arguments
-    # the actual batch size the model will see for this dataset
+    # to the actual batch size the model will see for this dataset.
+    # This is necessary so pipeline transfers know what size they are
+    # and the LR schedule, which is based on samples seen, gets set
+    # correctly.
    if hasattr(train_dataset, 'sample_multiplier'):
        args.micro_batch_size *= train_dataset.sample_multiplier
+        args.global_batch_size *= train_dataset.sample_multiplier
    return train_dataloader, valid_dataloader