[bugfix] fix DeepSeek V4 FP8 _set_o_group_proj_grouped crash when exporting LoRA adapter under PP>1

src/mcore_bridge/model/gpts/deepseek_v4.py

@@ -102,7 +102,7 @@ def get_query_key_value_tensors(
         """
         # s = sequence length, b = batch size, h = hidden size, n = num attention heads
         # Attention heads [s, b, n*h]
-        assert (hidden_states.ndim == 3), f"hidden_states should be 3D, [s, b, n*h], got {hidden_states.ndim}D"
+        assert (hidden_states.ndim == 3), f'hidden_states should be 3D, [s, b, n*h], got {hidden_states.ndim}D'
         if packed_seq_params is not None:
             assert (packed_seq_params.local_cp_size
                     is None), 'dynamic_context_parallel is not supported with MLA yet and is planned for future. \
@@ -495,31 +495,60 @@ def _set_o_group_proj_grouped(self, mg_attn, hf_state_dict, to_mcore):
 
         HF stores a single wo_a.weight of shape [G*R, D].
         GroupedLinear stores per-gemm weight{i} each of shape [R, D].
+
+        PP-aware: uses _get_weight/_set_weight for correct distributed communication.
         """
+        import torch.distributed as dist
+        from peft.utils import ModulesToSaveWrapper
+
+        from mcore_bridge.tuners import LoraParallelLinear
+
         o_groups = self.config.o_groups
+        group_proj = None if mg_attn is None else mg_attn.linear_o_group_proj
+
+        is_lora = isinstance(group_proj, LoraParallelLinear)
+        is_modules_to_save = isinstance(group_proj, ModulesToSaveWrapper)
+
+        if not to_mcore and self.pp_size > 1:
+            state = torch.tensor([is_lora, is_modules_to_save], dtype=torch.bool, device='cuda')
+            dist.all_reduce(state, group=self.pp_group)
+            is_lora, is_modules_to_save = state[0].item(), state[1].item()
+
+        if is_modules_to_save and self._peft_format:
+            raise NotImplementedError('FP8 grouped linear_o_group_proj as modules_to_save is not yet supported.')
+
+        if is_lora and self._peft_format:
+            raise NotImplementedError('FP8 grouped linear_o_group_proj LoRA adapter export is not yet supported. '
+                                      'Use merge_lora=True or exclude linear_o_group_proj from target_modules.')
+
+        if self._peft_format:
+            return
+
         if to_mcore:
             hf_weight = hf_state_dict['wo_a.weight'].load()
             hf_scale_inv = None
             if 'wo_a.weight_scale_inv' in hf_state_dict:
                 hf_scale_inv = hf_state_dict['wo_a.weight_scale_inv'].load()
-            weights = hf_weight.chunk(o_groups, dim=0)
-            scale_invs = hf_scale_inv.chunk(o_groups, dim=0) if hf_scale_inv is not None else [None] * o_groups
-            for i, (w, s) in enumerate(zip(weights, scale_invs)):
-                param = getattr(mg_attn.linear_o_group_proj, f'weight{i}')
-                self._set_param(param, w, s)
+
+            base_proj = (group_proj.base_layer if isinstance(group_proj, LoraParallelLinear) else group_proj)
+            params = [getattr(base_proj, f'weight{i}') for i in range(o_groups)]
+
+            self._set_weight(
+                params, hf_weight, 'linear_o_group_proj.weight', is_expert=False, hf_scale_inv=hf_scale_inv)
         else:
-            weights = []
-            scale_invs = []
-            for i in range(o_groups):
-                param = getattr(mg_attn.linear_o_group_proj, f'weight{i}')
-                if self._is_fp8_param(param):
-                    weights.append(param._rowwise_data)
-                    scale_invs.append(param._rowwise_scale_inv)
-                else:
-                    weights.append(param.data)
-            hf_state_dict['wo_a.weight'] = torch.cat(weights, dim=0).view(torch.float8_e4m3fn)
-            if scale_invs:
-                hf_state_dict['wo_a.weight_scale_inv'] = torch.cat(scale_invs, dim=0)
+            base_proj = None
+            if group_proj is not None:
+                base_proj = (group_proj.base_layer if isinstance(group_proj, LoraParallelLinear) else group_proj)
+
+            params = None
+            if base_proj is not None:
+                params = [getattr(base_proj, f'weight{i}') for i in range(o_groups)]
+
+            weight, scale_inv = self._get_weight(params, 'linear_o_group_proj.weight', is_expert=False)
+            if weight is not None:
+                hf_state_dict['wo_a.weight'] = weight
+            if scale_inv is not None:
+                hf_state_dict['wo_a.weight_scale_inv'] = scale_inv
 
     def _convert_hf_state_dict(self, hf_state_dict, to_mcore):
         res = super()._convert_hf_state_dict(hf_state_dict, to_mcore)