Add option to change batch size if needed (#11268)

* Add option to change batch size if needed * Apply isort and black reformatting Signed-off-by: BoxiangW <[email protected]> --------- Signed-off-by: BoxiangW <[email protected]> Co-authored-by: BoxiangW <[email protected]>
NVIDIA · Jan 8, 2025 · c0d7f2d · c0d7f2d
1 parent ee12d88
commit c0d7f2d
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 0 deletions.
diff --git a/nemo/collections/nlp/modules/common/megatron/megatron_init.py b/nemo/collections/nlp/modules/common/megatron/megatron_init.py
@@ -181,6 +181,7 @@ def initialize_model_parallel_for_nemo(
                     micro_batch_size=micro_batch_size,
                     data_parallel_size=app_state.data_parallel_size,
                     rampup_batch_size=rampup_batch_size,
+                    decrease_batch_size_if_needed=False,
                 )
             else:
                 if isinstance(_GLOBAL_NUM_MICROBATCHES_CALCULATOR, ConstantNumMicroBatchesCalculator):
@@ -201,6 +202,7 @@ def initialize_model_parallel_for_nemo(
                     micro_batch_size=micro_batch_size,
                     data_parallel_size=app_state.data_parallel_size,
                     rampup_batch_size=rampup_batch_size,
+                    decrease_batch_size_if_needed=False,
                 )
             else:
                 if isinstance(_GLOBAL_NUM_MICROBATCHES_CALCULATOR, ConstantNumMicroBatchesCalculator):

diff --git a/nemo/lightning/data.py b/nemo/lightning/data.py
@@ -103,6 +103,7 @@ def setup_microbatch_calculator(
                 micro_batch_size=micro_batch_size,
                 data_parallel_size=app_state.data_parallel_size,
                 rampup_batch_size=rampup_batch_size,
+                decrease_batch_size_if_needed=False,
             )
         else:
             if isinstance(_GLOBAL_NUM_MICROBATCHES_CALCULATOR, ConstantNumMicroBatchesCalculator):
@@ -121,6 +122,7 @@ def setup_microbatch_calculator(
                 micro_batch_size=micro_batch_size,
                 data_parallel_size=app_state.data_parallel_size,
                 rampup_batch_size=rampup_batch_size,
+                decrease_batch_size_if_needed=False,
             )
         else:
             if isinstance(_GLOBAL_NUM_MICROBATCHES_CALCULATOR, ConstantNumMicroBatchesCalculator):

diff --git a/nemo/lightning/megatron_init.py b/nemo/lightning/megatron_init.py
@@ -181,6 +181,7 @@ def initialize_model_parallel_for_nemo(
                     micro_batch_size=micro_batch_size,
                     data_parallel_size=app_state.data_parallel_size,
                     rampup_batch_size=rampup_batch_size,
+                    decrease_batch_size_if_needed=False,
                 )
             else:
                 if isinstance(_GLOBAL_NUM_MICROBATCHES_CALCULATOR, ConstantNumMicroBatchesCalculator):
@@ -201,6 +202,7 @@ def initialize_model_parallel_for_nemo(
                     micro_batch_size=micro_batch_size,
                     data_parallel_size=app_state.data_parallel_size,
                     rampup_batch_size=rampup_batch_size,
+                    decrease_batch_size_if_needed=False,
                 )
             else:
                 if isinstance(_GLOBAL_NUM_MICROBATCHES_CALCULATOR, ConstantNumMicroBatchesCalculator):