modelscope / modelscope/ms-swift
优化 npu_swiglu_forward:拼接投影权重而不是拼接激活以降低内存并提升 NPU 性能
Open
@hazelduan is already working on this.
Since Sep 15, 2026.
question
- Dominant language
- Python
- Stars
- 15.7k
- Forks
- 1.7k
- Avg merge
- 1d 16h
- Merged PRs (30d)
- 136
Description
Checklist / 检查清单
- I have searched existing issues, and this is a new question or discussion topic. / 我已经搜索过现有的 issues,确认这是一个新的问题与讨论。
Question Description / 问题描述
位置
- 文件: swift/model/npu_patch/model.py
- 函数: npu_swiglu_forward
- 链接: https://github.com/modelscope/ms-swift/blob/main/swift/model/npu_patch/model.py#L193
问题综述
当前实现先对 hidden_states 分别进行 gate_proj 和 up_proj,得到两个激活张量后再在最后一维上做 torch.cat,然后调用 torch_npu.npu_swiglu 进行融合激活,最后再做 down_proj:
gate = self.gate_proj(hidden_states)
up = self.up_proj(hidden_states)
cat = torch.cat((gate, up), dim=-1)
activated = torch_npu.npu_swiglu(cat, dim=-1)
return self.down_proj(activated)
建议与理由
建议改为“拼接投影层的权重(和 bias)而不是拼接激活”。也就是说,预先把 gate_proj 和 up_proj 的权重在输出维度上合并为一个权重(以及对应 bias),然后用一次线性变换得到等价的 concatenated activation,再调用 npu_swiglu,如:
# 合并投影层的权重/偏置
W_cat = torch.cat([self.gate_proj.weight, self.up_proj.weight], dim=0) # output dim concat
b_cat = None
if self.gate_proj.bias is not None or self.up_proj.bias is not None:
b1 = self.gate_proj.bias if self.gate_proj.bias is not None else torch.zeros(self.gate_proj.out_features, device=W_cat.device, dtype=W_cat.dtype)
b2 = self.up_proj.bias if self.up_proj.bias is not None else torch.zeros(self.up_proj.out_features, device=W_cat.device, dtype=W_cat.dtype)
b_cat = torch.cat([b1, b2], dim=0)
# 直接一次线性变换得到 cat 激活,避免分配两个中间激活
cat = F.linear(hidden_states, W_cat, b_cat)
activated = torch_npu.npu_swiglu(cat, dim=-1)
return self.down_proj(activated)
主要好处
- 降低内存峰值:避免同时存在 gate 和 up 两个中间激活,尤其在 batch 大或序列长时效果明显。
- 减少内存拷贝与张量分配:在 NPU 上能够减少不必要的数据移动,提升性能。
- 更好地控制 dtype/device:在构建合并权重时可以一次性对齐 dtype/device,减少隐式转换。
- 保持与现有数学等价:合并权重后的一次线性计算与分别计算再 concat 在数值上等价(在无非线性/状态依赖时)。
注意事项与实现细节
- 合并权重需保证两层的输入维度一致(通常是相同的),并按输出维度拼接(output dim)。
- 如果 gate_proj/up_proj 是带有特殊 forward 的模块(比如有参数化 dropout、噪声等),需要确认直接拼权重是否仍然可行。
- 如果模块被包装(例如 nn.Linear 被替换或按列分块),需要检查属性名(weight, bias)是否存在。实现时应加健壮性检查并在不可行时回退到现有实现。
- 在多设备或分布式场景要注意权重合并时的设备/数据布局。
建议的验收条件
- 提交实现后在至少一个 NPU 机器上跑对比基准,验证峰值显存降低和/或前向延迟下降。
- 在 CPU/GPU 环境提供功能等价的 fallback(例如未能读取权重时使用原实现)。
- 添加单测验证合并权重实现与原实现的输出在数值上接近(允许小数值误差)。
可选扩展
- 将权重合并逻辑缓存为模块属性(例如 self._cached_cat_weight)避免每次前向重新 concat 权重。注意训练时要在权重更新后刷新缓存。
- 在模块初始化或迁移脚本中一次性合并权重,并在训练/微调时按需恢复为可训练的拆分权重(如果需要单独优化)。
我建议由实现者在 npu_patch 中加入此优化分支,并在 PR 描述中附上基准对比图表。
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.