From 376b49ba59773b07c1012125156d40303e5c55b2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 6 Jan 2024 01:36:14 +0000 Subject: [PATCH] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../nlp/models/language_modeling/megatron_gpt_model.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py b/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py index be8b46c505fd..b8d5aa36df10 100644 --- a/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py +++ b/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py @@ -275,7 +275,7 @@ def __init__(self, cfg: DictConfig, trainer: Trainer): self.num_microbatches_in_previous_step = -1 self.microbatch_count = 0 self.is_prev_microbatch_training = True - + # configuration used for inference self._inference_config = None @@ -900,8 +900,9 @@ def fwd_output_and_loss_func(dataloader_iter, model, checkpoint_activations_all_ # compute fp8 weights first time a model chunk processes a microbatch pp_size = parallel_state.get_pipeline_model_parallel_world_size() forwards_per_step = vp_size * get_num_microbatches() - is_first_microbatch = (self.microbatch_count % forwards_per_step < vp_size*pp_size) and \ - (self.microbatch_count % pp_size == 0) + is_first_microbatch = (self.microbatch_count % forwards_per_step < vp_size * pp_size) and ( + self.microbatch_count % pp_size == 0 + ) else: is_first_microbatch = self.microbatch_count % get_num_microbatches() == 0 forward_args['is_first_microbatch'] = is_first_microbatch