slym/sharp (#7391)

Author: erhoo82

nemo/collections/nlp/models/language_modeling/megatron_base_model.py

@@ -132,6 +132,7 @@ def __init__(self, cfg: DictConfig, trainer: Trainer, no_lm_init=True):
             global_batch_size=cfg.get('global_batch_size'),
             rampup_batch_size=cfg.get('rampup_batch_size'),
             use_fp8=cfg.get('fp8', False),
+            use_sharp=cfg.get('sharp', False),
             init_mpi_proc_group=cfg.get('ub_tp_comm_overlap', False),
             seed=self.cfg.get('seed', 1234),
             apex_transformer_log_level=self.cfg.get('apex_transformer_log_level', 30),

nemo/collections/nlp/modules/common/megatron/megatron_init.py

@@ -67,6 +67,7 @@ def initialize_model_parallel_for_nemo(
     global_batch_size=None,
     rampup_batch_size=None,
     use_fp8=False,
+    use_sharp=False,
     init_mpi_proc_group=False,
     seed=1234,
     apex_transformer_log_level=30,
@@ -84,6 +85,7 @@ def initialize_model_parallel_for_nemo(
     app_state.pipeline_model_parallel_size = pipeline_model_parallel_size
     app_state.virtual_pipeline_model_parallel_size = virtual_pipeline_model_parallel_size
     app_state.use_fp8 = use_fp8
+    app_state.use_sharp = use_sharp
     app_state.init_mpi_proc_group = init_mpi_proc_group
     (
         app_state.tensor_model_parallel_rank,

nemo/collections/nlp/parts/nlp_overrides.py

@@ -169,6 +169,7 @@ def init_model_parallel(self, global_rank: int, world_size: int) -> None:
                     virtual_pipeline_model_parallel_size=app_state.virtual_pipeline_model_parallel_size,
                     pipeline_model_parallel_split_rank=app_state.pipeline_model_parallel_split_rank,
                     use_fp8=app_state.use_fp8,
+                    use_sharp=app_state.use_sharp,
                 )
 
                 # assert that fake tp and pp rank match after model parallel init

nemo/utils/app_state.py

@@ -55,6 +55,7 @@ def __init__(self):
         self._data_parallel_group = None
         self._megatron_checkpoint_version = None
         self._use_fp8 = False
+        self._use_sharp = False
         self._init_mpi_proc_gruop = False
 
         self._random_seed = None
@@ -364,6 +365,22 @@ def use_fp8(self, use_fp8):
         """
         self._use_fp8 = use_fp8
 
+    @property
+    def use_sharp(self):
+        """ Property returns the use of sharp.
+            Returns:
+                Use of sharp.
+        """
+        return self._use_sharp
+
+    @use_sharp.setter
+    def use_sharp(self, use_sharp):
+        """ Property sets the use of sharp.
+            Args:
+                use_sharp:  Use of sharp.
+        """
+        self._use_sharp = use_sharp
+
     @property
     def init_mpi_proc_group(self):
         """ Property sets the initialization of mpi process group.