remove assertion (NVIDIA#8302) (NVIDIA#8321)

Signed-off-by: dimapihtar <[email protected]> Co-authored-by: Dmytro Pykhtar <[email protected]> Signed-off-by: Zeeshan Patel <[email protected]>
zpx01 · Mar 8, 2024 · 63b5d5a · 63b5d5a
1 parent 62452bd
commit 63b5d5a
Showing 1 changed file with 0 additions and 5 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
@@ -1303,11 +1303,6 @@ def setup(self, stage=None):
         self.init_global_step = self.trainer.global_step
 
         if self.rampup_batch_size:
-            optimizer = self.cfg.optim.get('name', None)
-            assert (
-                optimizer == 'fused_adam'
-            ), f'{optimizer} optimizer is not supported yet with rampup batch size. Please, use fused_adam optimizer instead.'
-
             num_microbatch_calculator = apex.transformer.pipeline_parallel.utils._GLOBAL_NUM_MICROBATCHES_CALCULATOR
             num_microbatch_calculator.update(self.init_consumed_samples, consistency_check=False)
             self.prev_consumed_samples = self.init_consumed_samples