Pull Request已成功合入, 合并人@ascend-robot
（感谢 guihaowen666 的贡献）
变更摘要
此 PR 引入了一个新的 --fp8-quant-fb-split 参数，用于将 FP8 量化操作在前向传播和反向传播中进行拆分。当启用该模式时，前向传播不再将量化后的权重/激活张量保存到 ctx 上下文中供反向传播复用，而是在反向传播时重新计算量化值。这主要影响 MoE 分组矩阵乘法（MXFP8GMMFunction）和通用矩阵乘法（MXFP8MatMul）两条路径。
主要改动
-
新增
--fp8-quant-fb-split参数：在transformer_engine_basic.py中注册该命令行参数（默认False），并添加与 FP8 训练的依赖校验，仅在 FP8 启用时有效。 -
MXFP8GMMFunction.op_forward单轴量化分支：当fp8_quant_fb_split为True时，权重量化从双轴（npu_dynamic_mx_quant_with_dual_axis）切换为仅行方向单轴（npu_dynamic_mx_quant，axis=-2），不再生成列量化权重，也不再向ctx保存w_quant。 -
MXFP8GMMFunction.op_dx反向重算权重量化：反向传播时通过hasattr(ctx, 'w_quant')判断前向是否已保存量化权重；若未保存（即拆分模式），则调用npu_dynamic_mx_quant（axis=-1）重新计算权重量化值。 -
MXFP8MatMul跳过前向量化保存：在mxfp8_scaling_recipe.py中，当fp8_quant_fb_split为True时，前向不再将ctx.x和ctx.x_scale保存到上下文中（跳过npu_dynamic_mx_quant_with_dual_axis的额外输出），使反向传播自行处理量化。
代码审查
审查结论
经过对全部 3 个文件的详细审查：
审查总结
| 优先级 | 数量 | 说明 |
|---|---|---|
| P0 | 1 | grouped_matmul_util.py:211 — fp8_quant_fb_split=True 时 weight_col_mxfp8 未定义导致 NameError，训练直接崩溃 |
| P2 | 1 | transformer_engine_basic.py:116-117 — 错误消息引用了错误的参数名（fp8_reuse_quantized_weight 应为 fp8_quant_fb_split） |
| P3 | 1 | grouped_matmul_util.py:194-200 — fp8_quant_fb_split 分支缺少 reuse_identity 参数（可能有意，但无注释说明） |
逐文件确认：
mindspeed/core/transformer/moe/grouped_matmul_util.py：发现 2 个问题（P0 + P3），属于正确性回归和潜在功能遗漏。mindspeed/features_manager/megatron_basic/transformer_engine_basic.py：发现 1 个问题（P2），错误消息复制粘贴错误会误导用户排查。mindspeed/te/pytorch/fp8/recipes/mxfp8_scaling_recipe.py：无问题。正向/反向量化轴选择与hasattr(ctx, 'x')分发逻辑一致，ctx.save_for_backward/ctx.output_dtype在各分支均正确设置。
整体风险判断： P0 问题导致该特性完全不可用——一旦启用 --fp8-quant-fb-split，训练会立即因 NameError 崩溃。建议在合入前必须修复 P0 和 P2 两个问题；P3 建议补充注释或补传参数。
| 类型 | 数量 |
|---|---|
| 🔴 阻塞 | 1 |
| 🟡 建议 | 1 |
⛔ 需要修改
The MR can not be merged, because of CodeReview discussion not resolved
If you want to solve this problem, you can click here to do it in the FAQs.
Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.
What this PR does / why we need it?
When doing low-precision matmul and low-precision GMM calculations, using the double-axis quantization method saves the quantized weights and scales for the model's backward pass, which takes up a lot of GPU memory. This PR changes the double-axis quantization to single-axis quantization, so the model doesn’t save the quantized weights and scales during the forward pass but recalculates them during the backward pass, saving GPU memory.
Does this PR introduce any user-facing change?
This is a newly developed feature. Use the --mxfp8-defer-backward-quant parameter to control whether to enable this feature. It doesn't affect other features and involves no user-facing changes.
How was this patch tested?
Local checks and CI checks passed