Fix checkpoint converter missing parallel group initialization (#3217)

Author: yashaswikarnati

pyproject.toml

@@ -64,7 +64,7 @@ Download = "https://github.com/NVIDIA/Megatron-LM/releases"
 Homepage = "https://github.com/NVIDIA/Megatron-LM/megatron/core"
 
 [project.optional-dependencies]
-mlm = ["flask-restful", "sentencepiece", "tiktoken", "wandb", "transformers"]
+mlm = ["flask-restful", "sentencepiece", "tiktoken", "wandb", "transformers", "accelerate"]
 
 dev = [
     "nvidia-modelopt[torch]; sys_platform != 'darwin'",

tools/checkpoint/saver_base.py

@@ -170,9 +170,15 @@ def initialize_megatron_env(self):
 
         # For backward compatibility during local parallel states refactoring
         fake_tp_group = _ConverterFakeProcessGroup(size=self.args.target_tensor_parallel_size)
+        fake_pp_group = _ConverterFakeProcessGroup(size=self.args.target_pipeline_parallel_size)
         fake_ep_group = _ConverterFakeProcessGroup(size=self.args.target_expert_parallel_size)
+        fake_dp_group = _ConverterFakeProcessGroup(size=1)
         mpu._TENSOR_MODEL_PARALLEL_GROUP = fake_tp_group
+        mpu._PIPELINE_MODEL_PARALLEL_GROUP = fake_pp_group
         mpu._EXPERT_MODEL_PARALLEL_GROUP = fake_ep_group
+        mpu._DATA_PARALLEL_GROUP = fake_dp_group
+        mpu._DATA_PARALLEL_GROUP_WITH_CP = fake_dp_group
+        mpu._INTRA_PARTIAL_DATA_PARALLEL_GROUP_WITH_CP = fake_dp_group
         fused_kernels.load(self.margs)
 
         try: