feat(parallel_context.py): add gqa process group to allreduce dkv (#54)

InternLM · Feb 23, 2024 · 8baa139 · 8baa139
1 parent 883bd5d
commit 8baa139
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/internlm/core/context/parallel_context.py b/internlm/core/context/parallel_context.py
@@ -564,7 +564,8 @@ def init_parallel_groups(self):
 
         # run initialization of different process groups
         initializers = []
-        initializers.append(pgroup_initializer.Initializer_GQA(*initializer_args))
+        if "gqa" in parallel_config and parallel_config["gqa"] is True:
+            initializers.append(pgroup_initializer.Initializer_GQA(*initializer_args))
         initializers.append(pgroup_initializer.Initializer_Weight(*initializer_args))
         initializers.append(pgroup_initializer.Initializer_Weight_Data(*initializer_args))
         initializers.append(pgroup_initializer.Initializer_Tensor(*initializer_args))