fix: RewardModelLinear bcast process_group (#249)

Co-authored-by: shidongxing <shidongxing@>
InternLM · Jun 12, 2024 · 2ac3dcf · 2ac3dcf
1 parent eaef99e
commit 2ac3dcf
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/internlm/model/modules/linear.py b/internlm/model/modules/linear.py
@@ -526,9 +526,10 @@ def __init__(
 
         # broadcast parameters for reward model head layer.
         parallel_mode = get_head_parallel_mode()
-        dist.broadcast(self.weight, gpc.get_ranks_in_group(parallel_mode)[0])
+        process_group = gpc.get_group(parallel_mode)
+        dist.broadcast(self.weight, gpc.get_ranks_in_group(parallel_mode)[0], process_group)
         if bias:
-            dist.broadcast(self.bias, gpc.get_ranks_in_group(parallel_mode)[0])
+            dist.broadcast(self.bias, gpc.get_ranks_in_group(parallel_mode)[0], process_group)
 
 
 def new_linear(