Add pg_collection when init PackedPerExpertLinear

Signed-off-by: Hollow Man <hollowman@opensuse.org>

Author: HollowMan6

src/megatron/bridge/peft/utils.py

@@ -2634,13 +2634,15 @@ def __init__(
         init_method: Optional[Callable] = None,
         dtype: Optional[torch.dtype] = None,
         device: Optional[torch.device] = None,
+        pg_collection: ProcessGroupCollection | None = None,
     ):
         super().__init__()
         if not hasattr(torch, "_grouped_mm"):
             raise RuntimeError("PackedPerExpertLinear requires torch._grouped_mm (torch >= 2.9).")
         self.num_local_experts = num_local_experts
         self.in_features = in_features
         self.out_features = out_features
+        self.pg_collection = pg_collection
         weight = torch.empty(num_local_experts, out_features, in_features, dtype=dtype, device=device)
         if init_method is not None:
             for e in range(num_local_experts):
@@ -2670,7 +2672,11 @@ def sharded_state_dict(
         key = f"{prefix}weight"
         return {
             key: _make_grouped_expert_sharded_tensor(
-                self.weight.data, key, tp_axis=None, sharded_offsets=sharded_offsets
+                self.weight.data,
+                key,
+                tp_axis=None,
+                sharded_offsets=sharded_offsets,
+                pg_collection=self.pg_collection,
             )
         }
 
@@ -2770,6 +2776,7 @@ def __init__(
                 init_method=row_init,
                 device=params_device,
                 dtype=params_dtype,
+                pg_collection=self.pg_collection,
             )
         else:
             # Per-expert A (intermediate → rank); shared B (rank → hidden).
@@ -2780,6 +2787,7 @@ def __init__(
                 init_method=column_init,
                 device=params_device,
                 dtype=params_dtype,
+                pg_collection=self.pg_collection,
             )
             self.linear_out = RowParallelLinear(
                 dim,