Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Fix progress bar to update the total steps with trainer.max_steps #8499

Merged
merged 2 commits into from
Feb 26, 2024
Merged

Fix progress bar to update the total steps with trainer.max_steps #8499

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d910964
Fix progress bar to update the total steps with trainer.max_steps
athitten Feb 23, 2024
3559db9
Merge branch 'main' into athitten/fix_progress_bar
athitten Feb 23, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
9 changes: 3 additions & 6 deletions nemo/collections/nlp/parts/nlp_overrides.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1462,12 +1462,9 @@ def init_train_tqdm(self):
return self.bar

def on_train_epoch_start(self, trainer, *_):
if trainer.max_steps > 0 and (trainer.ckpt_path is not None):
# while resuming from a ckpt use trainer.max_steps as the total for progress bar as trainer.num_training_batches
# is truncated to max_steps - step being resumed at
num_training_batches = trainer.max_steps
else:
num_training_batches = trainer.num_training_batches
# Use trainer.max_steps as the num_training_batches since len(dataloader) aka num_training_batches is returned as the total num of micro batches
# instead of total num of global batches with this PR: https://github.com/NVIDIA/NeMo/pull/8426
num_training_batches = trainer.max_steps
self.train_progress_bar.reset(num_training_batches)
self.train_progress_bar.initial = 0
self.train_progress_bar.set_description(f"Epoch {trainer.current_epoch}")
Expand Down