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

[HorizonChaser]

修复 _set_o_group_proj_grouped 在 PP>1 下导出时因 mg_attn=None 崩溃的问题.

问题分析

linear_o_group_proj (wo_a) 是 te.GroupedLinear, 存储多个 weight{i}, 无法通过 _set_state_dict 单一 key 路径处理.

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)

原实现直接访问 mg_attn.linear_o_group_proj 并手动拼接 tensor, 完全绕过了 bridge 的 PP 通信机制, 导致在 PP>1 时每个 stage 对于自己没持有的层, mg_attn = None, 从而出现如下错误:

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'

修复方案

重写 _set_o_group_proj_grouped, 复用 bridge 统一出口:

• 通过 _get_weight(params_list_or_None, mg_key) 实现导出, 自动处理 FP8 解析、scale 裁剪、PP 之间同步等等
• 通过 _set_weight(params_list, hf_weight, mg_key, hf_scale_inv=...) 实现导入
• 对于 peft_format 触发 return 跳过, 与非 FP8 路径 _set_state_dict 行为对齐, 防止导出冻结的基础权重
• 对于 merge_lora=True 时, 在合并后 unwarp 取出 group_proj.base_layer 获取合并后的真实权重

虽然 Deepseek V4 目前没有支持 TP, 但后续支持之后应该可以直接复用 _get_weight / _set_weight 中的 TP
gather/split 逻辑, 也不需要手动拼接.

局限性

对于 wo_a (也就是 megatron 的 linear_o_group_proj):

1. 没有实现 grouped LoRA adapter 导出
2. 没有实现 modules_to_save 的支持逻辑

考虑到 wo_a 本身是一个地址投影, 应该不会被注入 lora 训练, 目前会返回 NotImplementedError, 如果有需要后续可以再添加

验证

对于 DeepSeekV4 Flash FP8 权重下 lora 微调, PP=2, EP=CP=8, TP=ETP=1, 预先训练 3 个 step:

{'loss': 0.57271874, 'grad_norm': 0.1394829, 'learning_rate': 4.001e-05, 'seq_load_balancing_loss': 1.79930186, 'iteration': '1/3', 'elapsed_time': '4m 6s', 'remaining_time': '8m 13s', 'memory(GiB)': 65.39, 'train_speed(s/it)': 246.458667}                                   
{'loss': 0.65376574, 'grad_norm': 0.14328134, 'learning_rate': 1.444e-05, 'seq_load_balancing_loss': 1.79427707, 'iteration': '2/3', 'elapsed_time': '5m 43s', 'remaining_time': '2m 52s', 'memory(GiB)': 65.4, 'train_speed(s/it)': 171.505935}                                  
{'loss': 0.474004, 'grad_norm': 0.10214218, 'learning_rate': 1e-06, 'seq_load_balancing_loss': 1.73329628, 'iteration': '3/3', 'elapsed_time': '7m 17s', 'remaining_time': '0s', 'memory(GiB)': 66.74, 'train_speed(s/it)': 145.787802}                                           
{'eval_loss': 0.47187394, 'eval_seq_load_balancing_loss': 0.0, 'iteration': '3/3', 'elapsed_time': '9m 35s', 'remaining_time': '0s', 'memory(GiB)': 66.74, 'train_speed(s/it)': 191.692337}                                                                                       

后续 6 个 step 的续训, 从 adapter 加载之前的权重 (没有做到学习率和数据集的连续, 不过看 loss 对比最开始的 3 个 step 均有下降):

{'loss': 0.56651509, 'grad_norm': 0.12827173, 'learning_rate': 4.82e-05, 'seq_load_balancing_loss': 1.79181099, 'iteration': '1/6', 'elapsed_time': '4m 7s', 'remaining_time': '20m 36s', 'memory(GiB)': 65.66, 'train_speed(s/it)': 247.136432}{'loss': 0.64803648, 'grad_norm': 0.14245556, 'learning_rate': 4.001e-05, 'seq_load_balancing_loss': 1.78637731, 'iteration': '2/6', 'elapsed_time': '5m 44s', 'remaining_time': '11m 27s', 'memory(GiB)': 65.67, 'train_speed(s/it)': 171.774969}{'loss': 0.46325764, 'grad_norm': 0.08786696, 'learning_rate': 2.752e-05, 'seq_load_balancing_loss': 1.71304393, 'iteration': '3/6', 'elapsed_time': '7m 17s', 'remaining_time': '7m 17s', 'memory(GiB)': 67.12, 'train_speed(s/it)': 145.743129}{'loss': 0.48516139, 'grad_norm': 0.11561694, 'learning_rate': 1.444e-05, 'seq_load_balancing_loss': 1.71926117, 'iteration': '4/6', 'elapsed_time': '8m 49s', 'remaining_time': '4m 25s', 'memory(GiB)': 67.12, 'train_speed(s/it)': 132.334503}{'loss': 0.46950492, 'grad_norm': 0.09321651, 'learning_rate': 4.63e-06, 'seq_load_balancing_loss': 1.68137872, 'iteration': '5/6', 'elapsed_time': '10m 26s', 'remaining_time': '2m 5s', 'memory(GiB)': 67.12, 'train_speed(s/it)': 125.129293}
{'loss': 0.46142226, 'grad_norm': 0.13110261, 'learning_rate': 1e-06, 'seq_load_balancing_loss': 1.6710031, 'iteration': '6/6', 'elapsed_time': '12m 4s', 'remaining_time': '0s', 'memory(GiB)': 68.15, 'train_speed(s/it)': 120.736977}
{'eval_loss': 0.44165221, 'eval_seq_load_balancing_loss': 0.0, 'iteration': '6/6', 'elapsed_time': '14m 20s', 'remaining_time': '0s', 'memory(GiB)': 68.15, 'train_speed(s/it)': 143.276238}                                                                                     

同时经过检查, 导出后的 adapter 不包含 wo_a 相关权重, 结构完整.

[gemini-code-assist]

Code Review

This pull request updates the _set_o_group_proj_grouped method in deepseek_v4.py to be pipeline-parallel (PP) aware by utilizing _get_weight and _set_weight for distributed communication. It also adds checks and raises NotImplementedError for unsupported FP8 grouped linear projection configurations under PEFT. Feedback was provided regarding a potential runtime crash when calling dist.all_reduce on a boolean tensor without specifying the reduction operator under the NCCL backend; it is recommended to explicitly use dist.ReduceOp.MAX.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Using dist.all_reduce on a boolean tensor without specifying the reduction operator defaults to dist.ReduceOp.SUM. Under the NCCL backend, SUM is not supported on boolean tensors and will cause a runtime crash. Since we want to check if any rank has LoRA or modules to save, we should explicitly use dist.ReduceOp.MAX (which acts as a logical OR for boolean values).

Suggested change

	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 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, op=dist.ReduceOp.MAX, group=self.pp_group)
	is_lora, is_modules_to_save = state[0].item(), state[1].item()

[HorizonChaser]

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

Lines 472 to 478 in 38316cf

	if not to_mcore:
	state = torch.tensor([is_lora, is_modules_to_save], dtype=torch.bool, device='cuda')
	if is_expert and self.ep_pp_size > 1:
	dist.all_reduce(state, group=self.ep_pp_group)
	elif not is_expert and self.pp_size > 1:
	dist.all_reduce(state, group=self.pp_group)
	is_lora, is_modules_to_save = state

这里的已有实现与我写的模式完全一致, 没有显式提供 op 参数, 虽然确实提供比较好, 但是如果需要的话应该作为单独的 PR?

[Jintao-Huang]

#151