Skip to content

Commit

Permalink
[pre-commit.ci] auto fixes from pre-commit.com hooks
Browse files Browse the repository at this point in the history
for more information, see https://pre-commit.ci
  • Loading branch information
pre-commit-ci[bot] committed Jan 6, 2024
1 parent 9b06a29 commit 376b49b
Showing 1 changed file with 4 additions and 3 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 376b49b

Please sign in to comment.