Init for model_parallel == 1 (facebookresearch#577)

* gate by arch, not by mp size * add back mp > 1 conditional
sahajgg · Jan 1, 2023 · 511504b · 511504b
1 parent 59403be
commit 511504b
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/metaseq/distributed/utils.py b/metaseq/distributed/utils.py
@@ -165,7 +165,10 @@ def distributed_init(cfg: MetaseqConfig):
     if nodelist:
         logger.info(f"SLURM nodelist: {nodelist}")
 
-    if cfg.common.model_parallel_size > 1:
+    if (
+        getattr(cfg.model, "arch", None) == "transformer_lm_megatron"
+        or cfg.common.model_parallel_size > 1
+    ):
         try:
             from megatron.mpu import (
                 initialize_model_parallel,