[bug] DeepSeekV4/V4-Flash 在 PP>1 + FP8 下训练时导出 lora adapter 时报错: _set_o_group_proj_grouped 未处理 mg_attn=None

[HorizonChaser]

Checklist / 检查清单

• I have searched existing issues, and this is a new bug report. / 我已经搜索过现有的 issues，确认这是一个新的 bug report。

Bug Description / Bug 描述

DeepSeek V4 Flash 使用 PP>1 + FP8 (fp8_param=True) 进行 lora 微调时，save_checkpoint 导出阶段崩溃:

File "mcore_bridge/model/gpts/deepseek_v4.py", line 514, in _set_o_group_proj_grouped
    param = getattr(mg_attn.linear_o_group_proj, f'weight{i}')
AttributeError: 'NoneType' object has no attribute 'linear_o_group_proj'

How to Reproduce / 如何复现

megatron sft \
    --model DeepSeek-V4-Flash \
    --fp8_recipe blockwise \
    --fp8_param_gather true \
    --pipeline_model_parallel_size 2 \
    --expert_model_parallel_size 8 \
    --tuner_type lora --save_safetensors true \
    --pipeline_model_parallel_layout 'Et*22|t*21mL' \
    --target_modules linear_qkv linear_proj linear_fc1 linear_fc2 \
    ...

训练时使用的并行是 EP=CP=8, PP=2, TP=ETP=1, 不论 merge_lora 值如何都会触发.

版本:

mcore-bridge @ git+https://github.com/modelscope/mcore-bridge.git@fdc13bbc9292862f14e58a83be0ebd374a663273
megatron-core @ git+https://github.com/NVIDIA/Megatron-LM.git@fd1121b8ff7e3a4f83a28d35aed172d7bc0260e1
ms-swift @ git+https://github.com/modelscope/ms-swift.git@6c7d381174609f432f32c59c3e4bffebb07574d7
torch==2.12.1+cu129
transformer-engine-torch==2.16.1

2x8xH20, CUDA 12.9.

Additional Information / 补充信息

我试着分析了一下相关逻辑:

mcore-bridge/src/mcore_bridge/bridge/gpt_bridge.py

Lines 1791 to 1813 in 38316cf

	while layer_idx < self.config.num_layers:
	lm_model = getattr(mg_model, 'language_model') if self.is_multimodal else mg_model
	if len(lm_model.decoder.layers) > 0:
	start_idx = lm_model.decoder.layers[0].layer_number - 1
	mg_layer_available = (start_idx <= layer_idx < lm_model.decoder.layers[-1].layer_number)
	else:
	mg_layer_available = False
	if mg_layer_available:
	mg_layer = lm_model.decoder.layers[layer_idx - start_idx]
	else:
	if to_mcore:
	layer_idx += 1
	prog_bar.update()
	continue
	else:
	mg_layer = None
	if not to_mcore and self.pp_size > 1:
	has_model = torch.tensor([mg_layer is not None], dtype=torch.bool, device='cuda')
	dist.all_reduce(has_model, group=self.pp_group)
	if not has_model:
	mg_model = next(mg_models) # compat vpp
	continue
	res = self._set_layer_state(mg_layer, hf_state_dict, f'{self.hf_layers_prefix}.', layer_idx, to_mcore)

PP>1 时, 不属于当前 PP stage 的 layer 的 mg_attn=None, _convert 在导出时将 mg_layer=None 传入 _set_layer_state, 这里应该是期望下游方法通过 _get_weight(None, ...) → _broadcast_ep_pp(None, ...) 正确参与集合通信.

mcore-bridge/src/mcore_bridge/model/gpts/deepseek_v4.py

Lines 500 to 519 in 38316cf

	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)
	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)

但 deepseekv4 相关的实现中 _set_o_group_proj_grouped 直接访问 mg_attn.linear_o_group_proj , 应该是没有考虑 PP>1 时每个 stage 会缺少一些层, 同时绕过了 bridge 的通用 PP-aware 权重通信机制:

• 无 _broadcast_ep_pp 导致不持有该层的 PP stage 无法获得权重
• 无 scale_inv padding 裁剪
• 无 _peft_format 检查 导致 lora adapter 导出时会尝试导出冻结的模型权重 (实际上这条路径会直接导出原始权重的 wo_a)

[HorizonChaser]

试着修复了这个问题 #149, 我这边测试正常

[Jintao-Huang]

#151

你试试这个是否正常