Add option to change batch size if needed

NVIDIA · Nov 13, 2024 · e5e2dc2 · e5e2dc2
1 parent 085e957
commit e5e2dc2
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):