modelscope / modelscope/ms-swift

优化 npu_swiglu_forward:拼接投影权重而不是拼接激活以降低内存并提升 NPU 性能

Open
#10,108 2 comments 0 reactions 1 assignee View on GitHub

@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 / 问题描述

位置

问题综述
当前实现先对 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)是否存在。实现时应加健壮性检查并在不可行时回退到现有实现。
  • 在多设备或分布式场景要注意权重合并时的设备/数据布局。

建议的验收条件

  1. 提交实现后在至少一个 NPU 机器上跑对比基准,验证峰值显存降低和/或前向延迟下降。
  2. 在 CPU/GPU 环境提供功能等价的 fallback(例如未能读取权重时使用原实现)。
  3. 添加单测验证合并权重实现与原实现的输出在数值上接近(允许小数值误差)。

可选扩展

  • 将权重合并逻辑缓存为模块属性(例如 self._cached_cat_weight)避免每次前向重新 concat 权重。注意训练时要在权重更新后刷新缓存。
  • 在模块初始化或迁移脚本中一次性合并权重,并在训练/微调时按需恢复为可训练的拆分权重(如果需要单独优化)。

我建议由实现者在 npu_patch 中加入此优化分支,并在 PR 描述中附上基准对比图表。

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.