实现5倍性能跃升!XTuner赋能昇腾超节点:大模型训练优化实践和创新
作者:GRJ, QJM, TZB from XTuner Group@Shanghai AI Lab;DYB, LQW from Huawei
随着 Qwen3.5 这类搭载 GDN 等创新算子的新一代大模型落地,基于昇腾 910C 开展训练面临新的工程挑战:新模型原生算子实现碎片化,直接使用 XTuner 启动训练时吞吐仅为单卡 4000 TGS,远达不到业务训练要求。
面对架构错配导致的性能断崖,我们摒弃传统的“抓单点慢算子”思路,基于 torch_npu profiler 构建了一套从时序全景透视到微秒级源码归因的性能诊断创新方法论。通过精准破除“算子微碎片化”、“通信伪重叠”与“Host 侧调度抖动”三大枷锁,在 XTuner 落地了 6 项软硬协同关键优化:零中心 RMSNorm 算子级融合、GDN 线性注意力 Ascend C 原生重构、SwiGLU 门控激活单 Pass 融合、Liger 分块交叉熵显存折叠、FSDP2 Unshard 通信归流与碎片治理、以及 Rank-NUMA 拓扑亲和性绑核调度。
最终,在昇腾 910C 跨超节点集群的大规模验证中,Qwen3.5-2B 单卡训练吞吐从基线 4000 TGS 跃升至 20000 TGS,实现 5 倍 的系统级性能突破。同时大幅削减了峰值激活显存,彻底消除了大 Batch 下的 OOM 隐患,在保障数值完全对齐与 Loss 稳定收敛的前提下,极致释放了 910C 超节点的原生算力潜能。
本文工作基于 XTuner 完成。XTuner 是上海人工智能实验室开源的新一代大模型训练引擎(github.com/InternLM/xtuner),面向预训练、指令微调与强化学习全链路,覆盖稠密/MoE 架构与语言/多模态模型,基于 FSDP 架构支撑大规模集群训练,并针对多元硬件芯片深度优化,是检验异构算力生态对新一代模型支持程度的典型载体。
一、面向新一代算力平台,大模型训练迎来性能新挑战
以 Qwen3.5 为代表的新一代大模型,正加速引入高密度门控与线性循环注意力(GDN)结构。这类原本依托 GPU 生态演进的动态计算范式,与昇腾 910C 的 Cube/Vector 异构计算单元、片上多级 Cache 及超节点拓扑存在显著错位:
-
算子碎片化开销高:门控与归一化算子缺乏融合,诱发海量 HBM 往返搬移;
-
通信流水难以掩盖:分布式反向通信切片细碎,无法被计算有效掩盖;
-
调度抖动引发木桶效应:CPU 端调度不均叠加设备端频繁起停,算力空转严重。
要破局,核心在于建立一套穿透框架调度与芯片底层的全链路性能透视标尺。
二、性能透视方法创新:构建从宏观全景到源码级归因的分析体系
我们基于 torch_npu profiler 建立了一套自顶向下、由宏观全景到微观切片、直达代码行的性能透视方法论:
[Step 1: 定步采样] ➔ 环境变量控制确定性采集,捕获 Kernel 明细与 Step 时序
│
[Step 2: 宏观分布] ➔ 耗时排序与分位数统计,识别“海量小 Kernel”形态
│
[Step 3: 时序全景] ➔ 甘特图逐流精确对齐,捕获 Host 调度时延与通信伪重叠
│
[Step 4: 指纹溯源] ➔ 40ms 识别执行周期 ➔ 4ms 剥离边界干扰 ➔ 1ms 锁定核心 Motif
│
[Step 5: 根因画像] ➔ 三维系统瓶颈对齐,精准驱动六大重构专项
2.1 第一步:定步采样——捕获窗口与分析对象
我们在训练脚本中通过环境变量将 torch_npu profiler 的采集窗口精准挂载在特定步(如 step-5 与 step-19)。其中 step-5 避开了初期冷启动抖动,已经进入稳定迭代状态,成为后续所有方案做单变量对照实验的黄金基准。
采集完成后,每个 Rank 均会输出一份标准的分析目录树:
outputs/training/<ts>/profiling_time/step-5/dev-<sn>_<pid>_<ts>_ascend_pt/
ASCEND_PROFILER_OUTPUT/
├── kernel_details.csv # 核心分析对象:包含每个算子的名称、执行时长、流 ID、起始时间戳
├── step_trace_time.csv # 宏观步级时间分解:划分计算、通信与空闲时间占比
└── op_statistic.csv # 算子调用频次与总时长汇总统计
在 16 卡超节点环境下,各 Rank 的行为特征高度一致。我们以 kernel_details.csv 为核心切入点,抽取其中单个典型 Rank 展开精细剖析。
2.2 第二步:宏观分布——打破“抓单点慢算子”的思维惯性
面对 Profile 数据,常规的调优习惯往往是按耗时降序排列,试图找到耗时占比最高的几个“慢算子”进行定向替换。然而在 Qwen3.5-2B 的基线剖析中,这种传统思路彻底失效了。
通过 Pandas 加载单 Rank 在 step-5 的明细数据,我们发现了一组反常的宏观指标:
-
单步 Kernel 实例激增至 19,252 个,端到端执行跨度 3,083ms,平均算子下发间隔仅~160μs;
-
耗时排名前五的通信算子累计仅占总时间的 9%,且步首 302.2ms 的集合通信算子本质是等待慢卡的屏障耗时。
这证实了核心瓶颈并非“单点算子过慢”,而是系统被近两万个极细粒度 Kernel 密集淹没。为此,我们创新性引入双对数耗时概率密度反演图(Log-Log Step Density Plot)展开全貌对比:
| 指标 / 耗时分位 | 全关基线(原始状态) | 全开优化(调优后) | 核心启示 |
|---|---|---|---|
| Kernel 实例数 | 19,252 | 9,608 | 算子实例缩减达 50.1%,碎片化现象腰斩 |
| p50 耗时 | 22.8 µs | 24.9 µs | 绝大多数计算由微秒级极小算子支撑 |
| p90 耗时 | 290 µs | 1180 µs | 优化后算子计算密度提升超 4 倍,聚合成块 |
| ≤ 50 µs 占比 | 62.4% | 56.1% | 显著减少硬件起停与切流边界开销 |
| ≤ 100 µs 占比 | 75.6% | 65.1% | 扭转 AI Core 频繁空转的非饱和状态 |

图 1:单 Rank 下 step-5 的 Kernel 耗时概率密度分布对比(橙线:优化前基线 run240;蓝线:六项优化全开 run200)
在双对数(Log-Log)坐标系下,横轴跨越了从 1 μs 到 300 ms 的 5 个数量级,纵轴反映了各耗时区间的出现概率密度。这张图将软硬件协同优化前后的“形态重塑”展现得淋漓尽致:
-
微秒区(1\~100µs)密度大幅沉降:基线中 ≤100µs 的极小算子占比高达 75.6%,高频硬件起停严重拖垮了流水线;优化后这部分被大幅压制;
-
毫秒区(10²\~10⁴µs)形成黄金融合带:优化后 p90 耗时拉升超 4 倍(从基线的 290 μs 至 1180 μs),原先散落的类型转换、标量偏移被合并为高计算密度的连续块;
-
右侧长尾收敛削峰:原本孤立突兀的 302ms 同步长尾被彻底拉平,最长通信耗时收敛至 150ms 以内。
2.3 第三步:时序全景——揭穿“伪重叠”与“调度抖动”
为了看清设备流水线究竟是如何被拖垮的,我们将单 Rank 在稳态步(step-5)的多流执行与空闲状态投影在同一物理时序轴上,绘制出宏观与局部的双尺度时序甘特图(Macro \& Zoom-in Timeline):

图 2:单 Rank 在 step-5 的宏观甘特图(上:整步全景;下:40ms 局部放大;右下:算子族耗时与频次图例)
我们将整个迭代步拆解为 Compute 流、Comm 流、Other 辅助流以及 Free(计算空闲)时间段,时序图呈现出三个致命问题:
-
步首 302ms 空转(红色整段):各 Rank 步启动极差高达 295ms,率先跑完的卡在 Barrier 处空等慢卡(Straggler),直指 Host CPU 跨 NUMA 调度竞争;
-
计算流呈现“发丝细纹”:缺乏宽平的大算子块,Cube 计算单元处于反复起伏、饥饿等待的亚健康状态;
-
算子结构严重倒挂:矩阵乘(MatMul/Attention)调用频次仅占 6.4%(耗时仅 677ms);而数据搬移(Glue)与逐元素操作(Elementwise)高达 13,238(5554 + 7684)次,吞噬近 70% 的调度开销;细碎的通信切片无法与计算交叠,形成大段裸露空洞。
在 40ms(t=1740∼1780 ms)局部视窗中,可以捕捉到一个极具价值的规律:这些碎小 Kernel 展现出高度同步的周期性脉冲特征(Motif)——亮蓝(matmul/attn)与金黄(GDN chunk)作为核心计算,却被大段密集的深灰(胶水搬移)与粉红(逐元素操作)反复穿插、分割,以几乎恒定的拓扑顺序呈现出成串的周期性脉冲式爆发,表明其对应特定神经网络层的重复调用,为微观溯源锁定了靶心。
2.4 第四步:指纹溯源——锁定微秒级算子指纹与源码破译
在宏观甘特图中捕捉到的周期性节拍指明了突破口。为了彻底摸清硬件流水线是如何被蚕食的,我们采用三级视窗逐层推进(40 ms → 4 ms → 1 ms),将观测尺度直接推到微秒级的物理算子层面:

图 3:逐级放大定位微观算子指纹(A:40ms 宏观切片;B:4ms 计算流截断;C:1ms 视窗逐算子名称与耗时指纹)
1. 显微镜推至 1ms 视窗:抓获微观“算子链条”
如上图 C 所示,在 t=1741.5∼1742.5 ms 这短短 1 ms(1000 μs)内,原本应由高吞吐矩阵乘主导的流水线,竟被整整 11 个碎片化微观 Kernel 强行塞满。其执行链条呈现出极度僵硬的串行时序:
Cast(32μs)→Pows(39μs)→Mean(41μs)→Adds(8μs)→Rsqrt(9μs)→Mul(62μs)→Mul(53μs)→Cast(96μs)→Silu(48μs)→Mul(138μs)→Cast(30μs)
2.对照数学公式:破译算子指纹(Gated RMSNorm 暴力拆解)
将这一长串碎片算子与模型算法逻辑进行逆向比对,其“算子指纹”瞬间与深度学习数学公式形成了绝对吻合——这正是 Gated RMSNorm 在 Eager 模式下的逐项暴力展开(Term-by-term Decomposition):
![]()
-
RMSNorm 主干计算链:上面的公式正好严格对应了图中的
Pows→Mean→Adds→Rsqrt→Mul;紧接着的第二个Mul(53μs)完成残差与主干的缩放。 -
门控与激活链:随后的 Silu(48μs) → Mul(138μs),构成了典型的门控激活相乘;
-
致命的精度强转“胶水”:由于缺乏 Kernel 融合,Eager 模式为了在 fp32 下保证方差数值稳定,在链条中被动插入了 3 次 Cast 算子(fp16 强转 fp32 算方差、转回、激活后再转),单是数据强转就白白浪费了 158 μs(占比接近 30%)。
归因结论:原本业界成熟的单个融合算子(Fused Gated RMSNorm)即可在几十微秒内解决的计算,在基准代码中被暴力拆解成 11 个细碎小算子,执行时间被硬生生拉长至 556 μs!
3. 周期复现:源码精准闭环归因
上述执行特征绝非局部偶发。沿着时序轴观察图 C 右侧:在完成一次常规的 MatMul(240 μs)与残差累加 Add(67 μs)后,仅推进至横坐标约 +948 μs 处,由 Cast(25μs) → Pows(54μs) → Mean(51μs)... 构成的算子链便再次严格按序出现。
-
微观开销积累:由 11 个细粒度 Kernel 构成的分解链在单步迭代内高频往复执行,直接诱发了成千上万次无谓的硬件起停开销、指令下发阻塞以及片外显存(HBM)的往返搬移;
-
源码精准定位:依据该算子调用拓扑所形成的特异性指纹,我们反向检索底层实现,直接定位到了缺乏 Kernel 融合的原始代码位置(
q35_npu_old.py:29-30)。
2.5 第五步:根因画像——确立六大重构专项
从宏观分布(2.2 节)、时序甘特图(2.3 节)到微观算子指纹(2.4 节),Profiler 数据逐层揭示了 Qwen3.5-2B 训练效率偏低的系统性原因:瓶颈并非来自某一个算子,而是 Host 侧调度延迟、Device 侧细粒度 Kernel 密集执行,以及通信与计算重叠不足等因素共同作用的结果。
针对上述问题,不能零敲碎打,必须推行“软硬协同、多管齐下”的系统性重构。我们量身定制了六大重构专项,形成了从代码级融合到系统级调度的完整应对策略:
| 诊断层级 | 捕获的 Profiling 关键指纹 | 根因机制判定 | 系统重构专项(优化项) |
|---|---|---|---|
| 算子拆解 | 周期性 11-Kernel 串联、fp32 Cast 频发 | RMSNorm(零中心与 Gated)逐项分解,4 次 fp32 显存往返 | 专项 1:零中心 RMSNorm 算子级融合 |
| 新型计算 | GDN Chunk 路径长、调度硬编码 | 线性注意力缺乏昇腾原生核函数,回退过渡期低效实现 | 专项 2:GDN 线性注意力 Ascend C 原生重构 |
| 门控激活 | 单步 5,554 次 elementwise,Silu+Mul 紧随 MatMul 之后 | Eager Autograd 导致激活与门控乘法分离,多余全量张量读写 | 专项 3:SwiGLU 门控激活单 Pass 融合 |
| 超宽张量 | Logits 产生超宽张量,反向衍生 2000 个小核 | 15.4 万超大词表导致全量张量显存往返,显存开销巨大 | 专项 4:Liger 分块交叉熵显存折叠 |
| 通信拓扑 | 535ms 纯通信暴露未被计算掩盖,碎片化通信切片 | FSDP/ZeRO 聚合通信切分粒度与流调度未对齐硬件流水 | 专项 5:FSDP2 Unshard 通信归流与碎片治理 |
| 主机调度 | 步首 302ms 巨量空转,各卡极差达 295ms,出现慢卡拖尾 | CPU 跨 NUMA 远端访存、线程争抢导致 Host 下发抖动严重 | 专项 6:Rank-NUMA 拓扑亲和性绑核调度 |
三、深入训练核心链路,多维优化释放昇腾超节点性能
围绕“算子级削减访存开销 → 通信级治理显存碎片 → 系统级消除长尾抖动”三维路径,我们在 XTuner 中落地了 6 项硬核专项重构。
3.1 算子优化:打破访存墙,释放新型架构硬件算力
3.1.1 零中心 RMSNorm 算子级融合
pr链接:https://github.com/InternLM/xtuner/pull/2119
在 §2.3 的时序图中,Compute 流上高频闪烁着一组橙色碎片(Norm 分解串),单步单 rank 累计高达 1,577 次调用。Qwen3.5 原生代码因缺乏算子融合,回退至 PyTorch 原生分解逻辑,使得全图的 158 条 Norm 链全部分解执行:

一条 Norm 链被暴力拆解为 10\~11 个小算子,经历 4 次 fp32 中间张量的全局显存进出,严重占用显存带宽。
【优化前:Eager 分解串 (10~11 个 Kernel)】
In (bf16) ─► Cast(fp32) ─► Pows ─► ReduceMean ─► Adds ─► Rsqrt ─► Mul ─► Cast(bf16) ─► Mul(1+w)
└────────────── 4 次 fp32 中间张量显存往返搬运 ──────────────┘
【优化后:PR #2119 算子级融合 (1 个 Kernel)】
In (bf16) ──────────────────► [ aclnnRmsNorm 融合核 ] ──────────────────► Out (bf16)
(将 1+w 缩放折叠进 gamma,全程单 pass)

图 4:零中心 RMSNorm 融合前后 6.5ms 微观对比(上:run191 分解基线;下:run192 融合后)
在 Compute 流上以“Norm 耗时最大”自动截取 6.5ms 满窗(图 4 上面板,run191):
-
固定节拍:橙色碎簇每隔 \~3.2ms 周期性重现,严密对应 61 个调用点 × 2 个 micro-batch;
-
特征核签名:括号高亮的
4 kernels / 77µs对应Pows→ReduceMean→Adds→Rsqrt序列,前后夹杂着大量Cast与Mul; -
显存墙本质:虽然 6.5ms 窗口内特征核仅耗时 0.24ms,但其强制触发的 4 次 fp32 显存往返才是吞吐锁死的主因。
PR #2119 新增 npu_zero_centered.py(+51 行):编写单 Pass npu_rms_norm,直接将 (1.0+w) 缩放等价折叠进 FP32 的 γ 权重张量中,从而直接调用昇腾底层的 aclnnRmsNorm 融合核,BF16 输出仅需一次截断,并设为默认开启。

图 5:零中心 RMSNorm 融合前后的算子耗时与端到端收益对比
对比融合后的图 4 下面板(run192)与图 5 聚合统计:
-
算子数量腰斩:10-Kernel 碎串缩减为单个
RmsNorm(前向 \~60µs / 反向RmsNormGrad196µs),窗口内算子数从 116 骤降至 75,全 Rank 算子调用净减近 2,800 次,特征核耗时(Pow(x^2)743ms 与Pows(g^2)127ms 累计高达 870ms )彻底归零; -
解除业务 OOM 死锁:端到端吞吐提升 +10.4%(149.4K → 165.0K TGS),Reserved 显存暴降 5.1GB(39.40GB → 34.29GB)。更关键的是,该优化彻底解决了在生产大 Batch 配置下启动即 OOM 的刚性阻塞。
3.1.2 GDN 线性注意力 Ascend C 原生重构
pr链接:https://github.com/InternLM/xtuner/pull/2120
在 §2.3 的全景时序图中,线性注意力核心 GDN(GatedDeltaNet)呈现出与硬件架构的严重割裂:
-
Causal Conv1d 碎片化:四重 for 循环切片强行拼接,单 Rank 制造了 104 对
StridedSlice / PadV3胶水算子; -
过渡版 Triton 核管线调度僵化:反向被粗暴压制为单个 3,642µs 的超长单体核,不仅引发片上计算流水频繁挂起,且过度依赖早期实验性编译后端,随底层 CANN 软件栈迭代极易出现性能劣化与兼容性问题;
-
输出侧 Gated Norm 呈碎片化分解执行:GDN 输出门控处的 RMSNorm 仍以 11 个小 kernel 的分解链形式运行。
针对上述问题,PR #2120 落地了包括原生三级派发、对接昇腾底层 fla_npu 的 Ascend C 专用算子栈、以及融合 Gated Norm 在内的完整重构。
【优化前:Gated Norm Eager 分解 (11 个 Kernel)】
Gate, X ─► [ Silu, Mul, Cast, Pow, ReduceMean, Rsqrt... ] ─► Out (多次 fp32 显存往返,反向多小核求导)
【优化后:PR #2120 门控归一化融合 (1 个 Kernel)】
Gate, X ────────────────► [ Fused Gated RmsNorm / RmsNormGrad ] ────────────────► Out
(门控与 Norm 深度折叠,单 Pass 片上完成,反向单核承接)

图 6:Gated RMSNorm 融合前后 9ms 窗口对比(上:未融合;下:融合后)
图 6 的 9ms 窗口真实捕获了两种 Norm 的异构同框对比:
-
上部未融合态:左侧零中心 Norm 已收敛为规整的单块(
RmsNorm\~60µs),而右侧 Gated Norm 仍处于由3 kernels / 106µs特征核及大量Cast/Mul构成的 11-Kernel 分解泥潭中; -
下部融合后:开启
FUSED_GATED_NORM=1后,9ms 窗口内核数量由 101 个腰斩至 50 个,反向统一由单个RmsNormGrad(197µs) 承接,两种 Norm 在图谱形态上完全归一。

图 7:GDN Chunk 路径微观对比(上:基线 Triton 实现;下:PR #2120 Ascend C 原生实现)
如图 7 所示,底层计算核心完成了向 Ascend C 专用算子栈的形态重塑:
-
解构粗粒度单体长核:优化前 3642µs 的超长单体核被彻底解构,重塑为形态规整、与昇腾 Cube/Vector 计算管线严密契合的 Ascend C 原生核家族:
-
ChunkBwdDqkwg(2106µs):专门负责梯度投影与权重匹配; -
ChunkGatedDeltaRuleBwdDhu(974µs):负责状态转移链的反向传播; -
ChunkGatedDeltaRuleFwdH(835µs):交织执行的前向状态更新; -
ChunkBwdDvLocal(401µs):本地局部变量更新。
-
-
规整硬件亲和流水:每个原生核与各自算力负载严密契合,消除了超长单体核内部的流水挂起,彻底摆脱了过渡期 Triton 编译器在固件迭代下的运行脆弱性。

图 8:PR #2120 优化后 GDN 算子栈单步耗时构成解剖(单 rank Σ480ms / 2.99s 步长)
图 8 展示了在开启优化情况下,单 rank 在 Step-5 中 GDN 各组成算子的实际耗时分布。优化后整个 GDN 核心栈累计耗时约 480ms(仅占整步耗时的 \~16%):
-
数百对胶水算子彻底消失,全由原生封装的
causal conv1d bwd(88.0ms) 与fwd(28.6ms) 承接; -
致密的 Ascend C 算子家族(
dqkwg75.6ms、delta rule (H)47.6ms、dHu34.9ms 等)构成计算绝对主力,中间仅保留经过深度验证的高效辅助算子(如solve_tril71.8ms)。

图 9:Gated RMSNorm 融合前后的算子耗时与端到端指标对比
如图 9 所示,全 rank 维度的特征算子 Pow(x^2) 耗时从 233ms 彻底归零,Div 从 153ms 降至 97ms;单 rank 内核数从 10,507 降至 9,535,端到端吞吐提升 +4.0%(158.7K → 165.0K TGS),Reserved 显存净省 4.5GB。
3.1.3 SwiGLU 门控激活单 Pass 融合
pr链接:https://github.com/InternLM/xtuner/pull/2128
在 §2.3 的时序图中,Elementwise 家族以 Σ554ms × 5554 的高频发射位列全图第二,其中单算子 aclnnMul 累计耗时超 300ms。排查发现,Dense MLP 层的激活计算采用了 PyTorch 原生 Eager 实现:前向执行 silu(gate) * up,反向对应两遍梯度求导。这一过程被拆解为 SiLU + Mul + SiluGrad + MulGrad 等 4\~5 个离散小算子,伴随 3 遍全尺寸张量的全局显存搬运,单步全 rank 累计在逐元素计算上空耗近 2.9 秒。
【优化前:Eager 分解 (4~5 个 Kernel)】
Gate, Up ─► [ SiLU ] ─► Temp ─► [ Mul ] ─► Out (反向对应多次 Grad 算子,3 遍显存往返)
【优化后:PR #2128 Triton 单 Pass 融合 (2 个 Kernel)】
Gate, Up ──────────────────► [ Fused SwiGLU Fwd/Bwd ] ──────────────────► Out (单 Pass 显存进出)

图 10:SwiGLU 融合前后 4ms 微观对比(上:run221 Eager 未融合;下:run200 融合后)
图 10 以 4ms 窗口捕获了一个典型的 MLP 位点(全网 24 层 × 2 micro-batch = 48 个位点):
-
优化前(上面板,run221):Down 投影梯度计算(
MatMulV3 1320µs)完成后,进入离散的激活反向链,橙色高亮的Swish (404µs)与紧邻的多个 Mul 碎屑交织,形成典型的“大 MatMul — 逐元素碎片 — 大 MatMul”断层;单点耗时虽小,但在整步 48 个位点的高频累积下,极度挤占显存带宽; -
优化后(下面板,run200):4ms 窗口内算子数从 12 个骤减至 3 个,原先分散的逐元素碎片被单内核
_swiglu_bwd(1281µs)一体化接管,计算图紧凑收敛。
PR #2128 新增 ops/fused_swiglu.py,采用纯 Triton 实现单 Pass 前反向算子,不依赖特定 NPU 底层私有库(兼具跨硬件平台可移植性),并在 dense_decoder_layer.py 中挂载使能。

图 11:SwiGLU 融合前后全 rank 算子耗时与端到端收益对比
对比基线(run221)与开启 XTUNER_NPU_FUSED_SWIGLU=1(run200):
-
碎小算子大范围出清:全 rank 口径下,
Mul(gate*up)耗时从 1734ms 骤降至 415ms(压减 76%),SiLU前向与反向耗时分别压减 50%(532ms → 267ms)与 80%(641ms → 134ms);激活算子开销整体折叠为fused swiglu fwd/bwd(498ms / 868ms); -
端到端加速与显存削减:单 rank 单步内核数从 11,050 降至 9,608;训练吞吐提升 +3.3%(173.2K → 178.9K TGS),Reserved 显存净省 1.0GB(55.78GB → 54.78GB)。
3.1.4 Liger 分块交叉熵显存折叠
pr链接:https://github.com/InternLM/xtuner/pull/2129
在 §2.3 的全景时序图中,Loss Chain 表面上仅占 Σ91ms × 80,极易被误判为微不足道的开销。然而其背后操作的是全模型物理尺寸最宽的张量——高达 154,880 的词表维度(mbs×seq×vocab)。
在原生 Chunk CE 模式下,尽管前向激活做了切片,但反向链依然依赖 PyTorch 原生算子逐个求导,导致全尺寸 Logits 梯度(d_logits)及其 fp32 中间体在全局显存中频繁搬运物化,造成极大的显存带宽浪费与 Kernel Launch 开销。
【优化前:Eager Chunk CE 反向 (多 Kernel 离散往返)】
Loss ─► [ NLLLossGrad ] ─► [ LogSoftmaxGrad ] ─► [ Cast(fp32) ] ─► d_logits (中间张量全局显存物化)
【优化后:PR #2129 Liger 融合反向 (单 Kernel 片上折叠)】
Loss ────────────────► [ liger_cross_entropy_backward (Triton) ] ────► d_weight / d_input
(片上分块即时累加,全尺寸 d_logits 零物化)

图 12:Loss 路径单个 Chunk 反向微观对比(上:run223 原生 Chunk CE;下:run200 Liger 融合后)
图 12 截取了处理单个 Chunk 反向计算的 8ms 典型窗口:
-
优化前(上面板,run223):窗口内挤满了19个离散算子,括号高亮区域的单 Chunk 反向链达 5796µs,包含
LogSoftmaxV2(2455µs) 与NLLLossGrad(650µs),紧随其后的是耗时达 1253µs 的Cast算子(用于将梯度回写至 fp32 缓冲); -
优化后(下面板,run200):全窗口紧凑收敛为单个
liger_cross_entropy_backward(3795µs)内核。Liger 将梯度推导与分块累加直接在片上完成,消除了全尺寸中间体落地 HBM 的沉重负担。
针对原版 Triton 算子在 NPU 多 Rank 下计算 loss_1d 偶发数值异常的缺陷,PR #2129 在 loss/liger_npu.py 中重构了专属于 NPU 的 NPUFusedLinearCrossEntropyLoss,驱动修复后的 Liger-Ascend Triton 核函数并默认使能。

图 13:Loss 路径全 rank 耗时构成与吞吐收益对比
图 13 汇总了全 rank(16 卡求和)下单步 Loss 链路的耗时与端到端对比:
-
反向耗时断崖式下跌:原生模式下耗时巨大的
LogSoftmaxGrad(1,371ms)与NLLLoss + grad(349ms)彻底出清,由单一融合内核liger fused CE bwd(482ms)取代,反向链路耗时直降 72%(1,720ms → 482ms),整条 Loss 链总耗时由 2,927ms 压缩至 1,642ms; -
小算子发射骤降:单 Rank 单步 Kernel 数由 11,633 降至 9,608,一次性净减 2,025 次微算子发射;
-
单 PR 最大吞吐跃升:在基线标准的生产配置下,全局训练吞吐由 155.7K 飙升至 178.9K tok/s(净增 +14.9%),成为本阶段性能增益最显著的单一改动。
3.2 通信优化:突破超节点规模扩展瓶颈
3.2.1 FSDP2 Unshard 通信归流与碎片治理
pr链接:https://github.com/InternLM/xtuner/pull/2130
在 §2.3 的全景时序图中,计算流(Compute Stream)存在长达 593ms 的被动等待空档,该区间通信流上的 All-Gather 虽在全速传输,但下游计算强依赖权重,导致通信与计算本质上是一种“伪重叠”(Pseudo Overlap)。

图 14:基线 run240 下 24ms 窗口内的 Unshard 伪重叠现象(Compute 挂起 18ms,Comm 独占 24ms)
图 14 在 24ms 的高精度窗口下揭示了伪重叠的本质形态:
-
“伪重叠”的物理本质:计算流(Compute 行)在阴影带内整段空白(Free 行为等宽的 18ms 红条),而通信流(Comm 行)满载运行 24ms。由于下游算子强依赖该权重,计算流只能被动挂起。独立通信流提前发射并没有压缩端到端时间,只是掩盖了计算流的等待事实;
-
Per-Stream 显存滞留:
torch_npu的显存分配器按流分桶管理。FSDP2 默认将参数 All-Gather 派发在独立的all_gather_stream上,导致大块参数常驻于通信流内存池,主计算流无法跨流复用,即使用户执行empty_cache()也无法越流回收,固化为约 1.0GB 的显存碎片。
PR #2130 引入 XTUNER_FSDP_UNSHARD_ASYNC_OP=1,重构了 FSDP2 的 Unshard 异步预取与等待流水,规范了通信流与计算流的时序协同,并彻底消除了跨流内存池隔离:
【优化前:双流异步伪重叠 (Per-Stream 显存孤岛)】
Compute Stream ──► [ 空闲挂起等待 (18ms) ] ──────────────────► [ 下一算子计算 ]
▲
Comm Stream ──► [ Unshard All-Gather (24ms) ] ─┘
(权重常驻 all_gather_stream 独立内存池,产生 1GB 无法回收的跨流碎片)
【优化后:PR #2130 同步归流调度 (显存池单流共享)】
Compute Stream ──► [ Unshard All-Gather ] ──► [ 前向预取计算 ] ──► Out
(同流复用计算内存池,消除跨流内存死锁,零时间惩罚释放 1GB 显存)

图 15:Unshard 异步优化前后的 Reserved 显存曲线与 Step-5 耗时分解
图 15 将优化前后的长期显存走势与单步耗时结构进行了直观映射:
-
典型碎片消除:左图显示,基线(run214)Reserved 显存从 Step-2 起便恒定高出约 1.0GB,且全程无斜率变化,属于经典的跨流分配器碎片指纹;
-
零代价释放显存:右图表明,优化后(run213)单步计算(2.76s vs 2.77s)与总吞吐(176.4K vs 177.1K tok/s)完全持平,在零性能损失下稳态释放 1.0GB 显存(56.74GB → 55.70GB),有效拓宽了更大 Batch Size 的安全边界。
3.3 系统优化:提升训练过程整体效率
3.3.1 Rank-NUMA 拓扑亲和性绑核调度
pr链接:https://github.com/InternLM/xtuner/pull/2131
在同步分布式训练中,整个集群的步调节奏完全由最慢的节点决定。在 §2.3 的全景时序图中,基线在每一步起始均出现长达 302ms 的 hcom_allReduce 巨型通信块,伴随完全等宽的 302.6ms Free 空档。

图 16:基线 run240 步首 700ms 窗口微观图谱(302ms 巨型 AllReduce 与等宽 Free 空档同框)
图 16 还原了步首的调度阻塞真相:
-
假性通信阻塞:步首执行极少量的 NaN-Check 后,通信流挂起长达 302ms 的 AllReduce,计算流完全陷入停滞(Free 表现为等宽的 302.6ms 红条);
-
根因定位:跨卡 AllReduce 物理传输仅需毫秒级,300ms 纯属“先到者等最后一名”。在未绑核的宿主环境中,16 个 Rank 的 Python 主进程、ACL 派发及 HCCL 伴生线程在跨 NUMA 节点间随机漂移,CPU 跨片访存与核心争抢导致个别卡启动严重滞后(基线各 Rank 首核发射最大离散度达 295ms)。
PR #2131 新增 utils/npu_cpu_binder.py,完全基于系统原生接口,将每个 Rank 主进程及其 ACL/Release/HCCL 线程刚性绑定至对应 NPU 直连的 NUMA 本地 8 核心片上,在 trainer.py 中全流程使能。
【优化前:无序调度与 NUMA 漂移 (Straggler 拖累全局)】
Rank 0~15 ──► [ OS 随机分派 / 跨 NUMA 访存争抢 ] ──► 启动时刻散布 104.5ms
│
全卡时序 ──► [ 最慢 Rank 拖累集合通信 (步首等宽 Free 挂起 302ms) ] ◄─┘
【优化后:PR #2131 纯原生 NUMA 拓扑绑核对齐】
Rank i ──► [ 锁定 NPU 直连 NUMA 节点的 8 核切片 ] ──► 启动时刻散布压缩至 70.2ms
│
全卡时序 ──► [ 步首 AllReduce 毫秒级即时汇聚,消除调度抖动 ] ◄───────┘

图 17:Step-5 跨 16 Rank 首个 Kernel 启动偏移散点对比(run224 vs run225)
图 17 的散点分布量化呈现了各卡在步首启动时差上的收敛效果,直观反映了硬件亲和性绑定对系统抖动的平抑作用:
-
调度离散度收敛:Step-5 中各卡相对最早 Rank 的首核启动延迟,由未绑核时的 104.5ms 显著收窄至 70.2ms(收敛 32.8%),消除了 CPU 级无序竞争;
-
超节点规模收益:在单机 16 卡环境下优化体现为抖动平抑与稳态对齐;在 256 Rank(16 节点 × 16 卡超节点)、100 步配对规模化验证中,消除调度抖动直接换取了 +1.7% 的全局端到端吞吐净提升。
四、性能跃升:从单点优化到系统级加速
在完成算子融合、通信归流与拓扑亲和绑核后,整个训练流水线完成了从“细碎停滞”到“致密连续”的蜕变。第三章落地的各项技术重构,最终在全局图谱和系统级性能上形成了合力汇聚。
4.1 全景时序重塑:消除碎片与空泡
将优化后的全景时序(run200,图 18)与 §2.3 的基线时序(run240,图 2)进行对比,微观图谱呈现出四大本质重塑:

图 18:全 PR 使能后(run200)单 rank 单步全景执行时序
1. 从“碎细斑驳”到“宽幅主导”(Compute 行):
- 基线(图 2):基线被 5,554 次逐元素碎片、7,684 次胶水切片及 1,577 次分解 Norm 稀释,核心 MatMul 算子严重碎片化;
- 优化后(图 18):胶水算子与小核被深度融合出清,计算流蜕变为以大块连续的蓝色 MatMul(1,534ms,时长提升 2.3 倍)和黄色 Ascend C 专用 GDN 核(411ms)为主导的高算力充盈态;
2. 从“零敲碎打”到“整段掩盖”(Comm 行):
- 通信流由基线 418 段离散竖条汇聚为规则大段的聚合 All-Gather 通信块(如 600\~1,700ms、2,050\~2,750ms 区间),实现了计算与通信在物理时空上的成块深度重叠;
3. 从“步首空转+全流程停顿”到“极致收敛”(Free 行):
- 步首 300ms 级的跨卡同步空泡彻底绝迹,整步仅剩零星离散的缝隙;
4. 算子总数腰斩出清:
- 单 Rank 单步发射的底层 Kernel 总量由基线的 19,252 个骤降至 9,608 个(压减 50.1%),根治了 CPU 侧调度排队与下发拥堵瓶颈。
4.2 端到端性能跃升:解锁单卡5倍算力释放
如图 19 所示,在 Step-5 耗时构成中,非重叠通信耗时由 535ms 压缩至 178ms,非重叠 Free 空泡由 123ms 降至 33.5ms(降幅达 72.8%):

图 19:Step-5 耗时结构对比(run240 基线 vs run200 优化全开)
这一系列系统级优化不仅压缩了绝对耗时,更彻底解锁了原本因显存受限而无法落地的生产级并行配置:
-
同配置等价加速(GA=2 / mbs=1):在基线相同硬件与并行配置下,全局训练吞吐从基线的 43.0K 跃升至 165.0K tok/s,纯算法与内核级重构直接带来 3.8 倍(+283.7%) 的算力释放;
-
生产配置解锁增益(GA=4 / mbs=2):基线版本在该高效配置下因显存碎片与张量未融合直接 OOM。优化全开后成功释放超 5GB 显存,顺利解锁生产配置,全局吞吐进一步攀升至 178.9K tok/s,达成累计 4.2 倍 的端到端系统级加速。
我们进一步在多 MBS、多节点集群上完成扩缩验证,结果如下表所示。对比早期基线单卡 4000 TGS 的水平,优化后单节点稳定跑到 20000+ TGS,单卡算力密度提升 5 倍;即使扩容至 96 节点,单卡 TGS 依旧保持 17000 的高位,具备优秀的分布式扩展性。

五、面向未来:打造开放高效的大模型训练生态
本次实践是 XTuner 高性能训练框架 与 昇腾 910C 超节点算力底座 深度耦合、系统级创新的落地范式。它有力验证了:国产软硬全栈基础设施不仅能高效承接传统主力模型,更具备支撑乃至引领混合注意力(Hybrid Attention)、线性注意力(GDN)等下一代非传统模型架构演进的工程硬实力。
立足当前突破,XTuner 与昇腾 910C 将持续深化技术协同,重点发力以下方向:
-
破局架构演进,攻坚原生非标准算子:大模型正快速向长序列状态网络、万亿 MoE 及混合注意力架构演进。XTuner 将依托昇腾 910C 的高带宽互联与片上计算特性,持续攻关下一代复杂拓扑的原生 Ascend C 算子与定制并行范式,为前沿架构探索筑牢算力底座。
-
深化软硬协同,释放超节点集群极致效能:深入 910C 硬件微架构与超节点拓扑,XTuner 将持续演进原生 FSDP2、跨流无损掩盖及全局显存分桶复用机制,消除跨节点通信与调度抖动瓶颈,保障超大规模集群下的线性扩展能力与长训稳定性。
-
践行开源共赢,繁荣自主智算软件生态:本次沉淀的原生算子重构、通信归流与拓扑绑核方案均全面反哺开源社区(合并至 XTuner 上游)。团队将持续把硬件调优沉淀转化为模块化、开箱即用的标准能力,大幅降低国产先进算力的使用门槛,赋能更多开发者高效创新。
更多推荐


所有评论(0)