support option to disable creating gloo pg

Signed-off-by: Ananth Subramaniam <[email protected]>

Author: ananthsub

nemo/tron/config.py

@@ -212,6 +212,9 @@ class DistributedInitConfig:
     use_tp_pp_dp_mapping: bool = False
     """If set, distributed ranks initialize order is changed from tp-dp-pp to tp-pp-dp. Make sure EP and CP aren't used with this option enabled"""
 
+    use_gloo_process_groups: bool = True
+    """If set, create Gloo process groups for communications."""
+
 
 @dataclass
 class ProfilingConfig:
@@ -543,6 +546,14 @@ def __post_init__(self):
             self.model_config.use_cpu_initialization or self.dist_config.lazy_mpu_init
         )
 
+        # Make sure all functionality that requires Gloo process groups is disabled.
+        if not self.dist_config.enable_gloo_process_groups:
+            if self.optimizer_config.use_distributed_optimizer:
+                # If using distributed optimizer, must use distributed checkpointing.
+                # Legacy checkpointing uses Gloo process groups to collect full distributed
+                # optimizer state in the CPU memory of DP rank 0.
+                assert self.checkpoint_config.ckpt_format == "torch_dist"
+
         # Scheduler
         if self.scheduler_config.lr_decay_iters is None:
             self.scheduler_config.lr_decay_iters = self.train_config.train_iters

nemo/tron/init.py

@@ -229,6 +229,7 @@ def _initialize_distributed(
                 ),
                 get_embedding_ranks=get_embedding_ranks,
                 get_position_embedding_ranks=get_position_embedding_ranks,
+                create_gloo_process_groups=cfg.dist_config.use_gloo_process_groups,
             )
             if get_rank_safe() == 0:
                 print(

nemo/tron/optim.py

@@ -19,7 +19,7 @@
 
 
 def setup_optimizer(cfg: ConfigContainer, model, no_weight_decay_cond=None, scale_lr_cond=None, lr_mult=1.0):
-    optimizer = get_megatron_optimizer(cfg.optimizer_config, model, no_weight_decay_cond, scale_lr_cond, lr_mult)
+    optimizer = get_megatron_optimizer(cfg.optimizer_config, model, no_weight_decay_cond, scale_lr_cond, lr_mult, use_gloo_process_groups=cfg.dist_config.use_gloo_process_groups)
     scheduler = _get_scheduler(cfg, optimizer)
 
     return optimizer, scheduler