2026/1/15 13:36:07
网站建设
项目流程
青岛个人接网站建设,淘宝买cdk自己做网站,网店美工分为几个级别,网站分析与优化GMM NZ 全流程详解实战#xff1a;FSDP MOE 训练加速一、从推理到训练#xff1a;性能优化的新起点
在昇腾生态的推理实践中#xff0c;Grouped MatMul#xff08;简称 GMM#xff09;结合 NZ#xff08;Fractal_NZ#xff09;格式已被证明能带来显著性能收益。 但当我们…GMM NZ 全流程详解实战FSDP MOE 训练加速一、从推理到训练性能优化的新起点在昇腾生态的推理实践中Grouped MatMul简称 GMM结合 NZFractal_NZ格式已被证明能带来显著性能收益。 但当我们试图把这项优化迁移到训练端各种各样的问题接踵而至——框架支持差异、反向计算适配、内存搬运开销等使得原本的“推理提速方案”在训练阶段几乎无法使用。本次实战的目标就是在FSDP2 框架下训练 Qwen3 235B MOE 模型时使能 GMM NZ 格式并尽可能获得真实的性能提升。小示例如何构造一个最小 GMM 输入import torch import torch_npu # 构造一个最小化 GMM 输入 x torch.randn(4, 16).npu() # 模拟 4 tokensdim 16 weight torch.randn(2, 16, 32).npu() # 模拟 2 个专家每个 Expert 的 W: 16→32 tokens_per_expert torch.tensor([2, 2], dtypetorch.int32).npu() outs torch_npu.npu_grouped_matmul( [x], [weight], group_listtokens_per_expert, group_type0, group_list_type1, split_item2 )[0] print(outs shape:, outs.shape) print(outs dtype:, outs.dtype)二、GMM 算子适配从前向到反向的全链路打通1. PTA 框架下的 Grouped MatMul昇腾提供的torch_npu.npu_grouped_matmul能够批量执行矩阵乘法将相同形状的计算统一处理以减少内存访问和调度开销。 对于 MOE 结构它能自然支持按专家维度Expert分组计算。比较经典的调用如下outs torch_npu.npu_grouped_matmul( [x], [weight], group_listtokens_per_expert, group_type0, group_list_type1, split_item2 )[0]这段代码把输入x按 token 数切分为多个专家子块与对应的权重分片分别执行 MatMul然后合并输出。2. 手动适配反向传播然而npu_grouped_matmul并未原生支持反向计算。 为此我们基于 PyTorch 的Function机制手动实现了前反向逻辑class NpuGMMOp(Function): staticmethod def forward(ctx, weight, x, tokens_per_expert): ctx.save_for_backward(weight, x, tokens_per_expert) outs torch_npu.npu_grouped_matmul( [x], [weight], group_listtokens_per_expert, group_type0, group_list_type1, split_item2 )[0] return outs staticmethod def backward(ctx, grad_output): weight, input_tensor, tokens_per_expert ctx.saved_tensors grad_input torch_npu.npu_grouped_matmul( [grad_output], [weight.transpose(1,2)], group_listtokens_per_expert, group_type0, group_list_type1, split_item2 )[0] grad_weight torch_npu.npu_grouped_matmul( [input_tensor.T], [grad_output], group_listtokens_per_expert, split_item3, group_type2, group_list_type1 )[0] return grad_weight, grad_input, None到此GMM 的正反向计算路径已在 PTA 框架中打通了。验证# 验证自定义 GMM 是否能够正常完成前反向 x torch.randn(4, 16, dtypetorch.float16).npu().requires_grad_(True) weight torch.randn(2, 16, 32, dtypetorch.float16).npu().requires_grad_(True) tokens torch.tensor([2, 2], dtypetorch.int32).npu() out NpuGMMOp.apply(weight, x, tokens) loss out.sum() loss.backward() print(out shape:, out.shape) print(x grad:, x.grad.shape) print(weight grad:, weight.grad.shape)那就是打通啦嘻嘻。三、ND → NZ 转换带来的新开销虽然理论上权重转为 NZ 格式应能减少访存并提升并行效率。 但在训练中如果我们在前向前执行如下操作weight torch_npu.npu_format_cast(weight, 29) # ND → NZ那么每次调用都引入额外的transdata开销。 Profiling 结果显示这部分耗时足以抵消 GMM NZ 本身带来的性能收益。 于是如何在不显著增加数据搬运的情况下获得 NZ 格式是新的方向。四、sliceNZ一次搬运双倍收益1. 设计动机在 FSDP2 框架下MOE 权重会经历如下流程初始化 → AllGather → Slice → 进入 GMM 前向slice 操作与 ND→NZ 转换都涉及显存HBM搬运。 因此我们尝试将这两个步骤融合为一个算子即sliceNZ。 这样一来只需一次内存访问就能完成切分与格式转换。2. 算子原型sliceNZ 算子的核心定义如下aclnnSliceNzGetWorkspaceSize( const aclTensor *in, uint64_t dim, uint64_t start, uint64_t end, aclTensor *output )output通常是三维张量[num_expert, n, k]当专家数为 1 时为二维。 输出结果将直接进入 GMM 的前反向运算。3. PTA 注册与调用在op_plugin_functions.yaml中新增接口- func: npu_special_slice(Tensor self, int dim, int start, int end, Tensor(a!) output) - () op_api: [v2.1, newest] gen_opapi: exec: aclnnSliceNz, self, dim, start, end, output随后在 PTA 源码中替换原有 slice 实现逻辑torch.library.impl(lib, split_with_sizes_copy, PrivateUse1) def split_with_sizes_copy(all_gather_output, all_gather_input_split_sizes, dim, out): if len(all_gather_input_split_sizes) 1 and out[-1].shape[0] * out[-1].shape[1] 128 * 4096 * 1536: from special_op import npu_special_slice num_exp 128 in_dim1, out_dim1 1536, 4096 in_dim2, out_dim2 4096, 3072 out[-1].resize_(num_exp, out_dim1, in_dim1) out[-2].resize_(num_exp, out_dim2, in_dim2) npu_special_slice(all_gather_output, dim, weight_1_start, total_size, out[-1]) npu_special_slice(all_gather_output, dim, weight_2_start, weight_1_start, out[-2])经过编包、安装并重新加载 PTA即可在 FSDP 训练过程中调用 sliceNZ 实现权重的自动转 NZ。五、Profiling验证NZ 启用成功但性能仍未提升当我们初步完成 sliceNZ 接入并在 FSDP2 框架中跑通训练流程后第一时间使用msprof进行了全流程性能采样。msprof --applicationpython train_qwen_moe.py --output./profile_nz采样结果中可以明显看到以下现象sliceNZ 输出张量的格式标识为 FRACTAL_NZ说明格式转换成功GroupedMatmul 的入参依旧被识别为 ND 格式没有走到高性能 NZ 分支整体 step 耗时与 ND baseline 基本持平未见预期加速。这显然是一个“成功了一半”的结果。进一步对比msprof的 operator-level 统计发现sliceNZ 自身执行耗时约 1.7 ms比原 slice transdata 的组合方案略优节省约 0.4 ms但在 GMM 前反向中出现了多次冗余的 Transpose操作累计耗时 3~4 ms每个 step 的计算图中GMM 权重在正反向之间频繁出现 ND - NZ 转换迹象。关键问题 由于 PTA 框架对 tensor 的格式识别仍停留在 ND 层面即便底层张量已经是 FRACTAL_NZ框架仍会强制执行一次 transpose以符合 ND 算法接口要求。这直接抵消了 NZ 格式的加速收益。Profiling 数据片段示例算子名调用次数平均耗时 (ms)格式备注npu_special_slice11.72NZsliceNZ 成功执行npu_grouped_matmul43.15ND未识别 NZnpu_transpose83.92ND → NZ冗余操作total_step_time-22.1-与 baseline 基本持平由此可见瓶颈不在 sliceNZ 本身而在格式识别与矩阵布局转换。Profiling 采集与关键指标分析在验证 sliceNZ 的实际效果时可以使用msprof或torch_npu.profiler双路径采样。 下面是完整代码模板兼容昇腾 NPU 环境import torch import torch_npu from torch_npu.profiler.analysis.profiler_config import ProfilerConfig from torch_npu.profiler import profiler # GMM 训练循环 def train_step(weight, x, tokens_per_expert): out NpuGMMOp.apply(weight, x, tokens_per_expert) loss out.sum() loss.backward() return loss.item() # 配置 NPU profiler config ProfilerConfig( output_path./profile_data, record_shapesTrue, with_stackTrue, sample_rate1, ) with profiler(config): for i in range(10): loss train_step(weight, x, tokens_per_expert) if i % 2 0: print(fStep {i} loss: {loss:.4f}) # 结果分析命令 # msprof --view ./profile_data --report report.html检查 tensor 格式的辅助代码Profiling 结果只是宏观确认我们也可以在训练中动态打印 tensor 格式验证def debug_format_check(tensor, nametensor): fmt torch_npu.get_npu_format(tensor) print(f[DEBUG] {name} format: {fmt}) # 调用 debug_format_check(weight, sliceNZ weight) debug_format_check(out, GMM output)输出[DEBUG] sliceNZ weight format: FRACTAL_NZ [DEBUG] GMM output format: ND六、瓶颈优化提前转置消除额外开销在分析性能热点后我们锁定问题根源——GroupedMatmul 的输入维度约定。1. 问题剖析在常规的线性层中我们的计算逻辑是其中的 shape 通常为(out_feature, in_feature)。而在torch_npu.npu_grouped_matmul实现中权重输入的 shape 被定义为这意味着当用户直接传入 PyTorch 风格的权重时系统会自动执行一次transpose(in_feature, out_feature)来匹配算子接口。问题在于在 FSDP 训练中这种自动 transpose 每个 step 都会触发 结合反向传播阶段的多次权重复用每个 GMM 模块的 transpose 累积耗时甚至超过了算子本身的计算时间。2. 解决思路将问题前移既然 transpose 是为满足算子入参约定那就不妨提前在权重加载阶段完成一次性转置让后续计算直接以目标布局运行。在 Qwen3 235B MOE 中涉及 GMM 的主要模块包括gate_projup_projdown_proj其中gate_proj与up_proj共享输入维度并在输出轴拼接因此我们可以在模型初始化时对其权重执行def prepare_weight(weight): # 仅在模型加载阶段执行一次 return weight.permute(1, 0).contiguous() model.gate_proj.weight.data prepare_weight(model.gate_proj.weight.data) model.up_proj.weight.data prepare_weight(model.up_proj.weight.data) model.down_proj.weight.data prepare_weight(model.down_proj.weight.data)这样在训练阶段 GMM 调用时权重与算子接口天然对齐无需额外 transpose。3. 效果验证优化后再次采集 Profiling性能表现如下指标项优化前 (ND)启用 sliceNZ启用 sliceNZ 提前转置每 step 耗时22.1 ms21.7 ms17.4 msGMM 模块计算耗时8.6 ms8.4 ms5.2 mstranspose 调用数8 次6 次2 次内存带宽占用率88%86%79%最终结果表明sliceNZ 仅能带来小幅收益结合提前转置后GMM NZ 的性能提升约21%整体 FSDP2 训练吞吐tokens/s提升约18.6%。从单步优化到全流程增益关键就在于减少不必要的数据重排让计算更贴近硬件本地格式。4. 小总结我们发现算子级性能提升的瓶颈往往不在算子本身而在框架与算子之间的“接口层”。 NZ 格式在推理阶段表现优异是因为数据格式在编译前已固化 而训练中动态反向路径和多进程通信导致格式频繁切换。提前转置、格式冻结、算子融合正是我们将推理优化思路迁移到训练端的重点。5. 代码模型加载阶段提前转置可直接嵌入Qwen或MOE模块def prepare_weight(weight): # 仅在模型初始化或权重加载阶段执行 return weight.permute(1, 0).contiguous() def apply_pre_transpose(model): # 识别并处理所有包含 GroupedMatmul 的模块 target_layers [gate_proj, up_proj, down_proj] for name, module in model.named_modules(): for t in target_layers: if t in name and hasattr(module, weight): module.weight.data prepare_weight(module.weight.data) print(f[INFO] Transposed weight in {name})在训练启动前调用apply_pre_transpose(model)这样后续执行 GMM 算子时权重布局已经与算子定义完全一致。6. 代码对比测试小脚本观察 transpose 调用差异你可以通过计时测试不同方案的耗时差import time def test_transpose_cost(): x torch.randn(4096, 1536, dtypetorch.bfloat16).npu() w torch.randn(1536, 4096, dtypetorch.bfloat16).npu() # baseline: 每次临时转置 t1 time.time() for _ in range(100): y torch.matmul(x, w.T) torch.npu.synchronize() t2 time.time() print(临时转置耗时:, t2 - t1, s) # 优化后预转置一次 w_opt w.T.contiguous() t3 time.time() for _ in range(100): y torch.matmul(x, w_opt) torch.npu.synchronize() t4 time.time() print(提前转置耗时:, t4 - t3, s) test_transpose_cost()实测结果直观展示提前转置能减少 15%~25% 的整体矩阵乘耗时。7. 代码训练性能统计脚本最后可以附带一段性能采集逻辑用于评估整体吞吐import time def benchmark_training(model, dataloader, steps100): start time.time() for i, batch in enumerate(dataloader): if i steps: break loss model(batch) loss.backward() torch.npu.synchronize() end time.time() print(fAvg step time: {(end - start) / steps * 1000:.2f} ms) # 比较两种方案 print( Baseline (ND) ) benchmark_training(model_nd, dataloader) print( With sliceNZ Pre-Transpose ) benchmark_training(model_nz, dataloader)配合tokens/s统计就能量化最终性能收益。七、总结与启示整个 GMM NZ 训练使能过程从框架底层到算子适配经历了多次性能权衡与调试。 最终的经验可总结为一次搬运双重优化sliceNZ 的核心价值在于合并 slice 与格式转换格式识别是关键PTA 层需正确识别 NZ 格式才能真正走高性能路径冗余操作要前移处理提前转置可显著减少训练时的算子级开销验证手段要量化profiling 数据是判断优化是否有效的唯一依据。在推理端验证的性能思路迁移到训练端有时候并不是简单的复用。 但是通过算子融合、前后向一体化设计与格式识别优化 GMM NZ 的潜力在训练端才被真正激活。未来随着 PTA 主线对 NZ 支持的持续完善 这套方案也将更容易被开发者直接复用在更大规模的模型训练中。注明昇腾PAE案例库对本文写作亦有帮助。