作者​：昇腾实战派​知识地图​：https://blog.csdn.net/Lumos_Lovegood/article/details/161601003

背景概述

在 Ascend NPU 上使用 PyTorch 2.14.0 + TorchNPU 进行模型训练时，为了获得更好的性能，我们尝试启用 torch.compile 的 DVM（Dynamic Vectorization Module）后端。本文记录了在 Qwen3.5 模型上启用 DVM 编译过程中遇到的一系列问题及相应的解决方案，涉及 Triton 后端、TorchNPU Inductor 适配层以及模型代码的修改。

整个修改涉及三个主要目录：

• Triton 后端（triton/backends/ascend/）：2 个文件，5 处修改
• TorchNPU Inductor 适配（torch_npu/_inductor/）：3 个文件，3 处修改
• 模型代码（torchtitan/）：7 个文件，多处修改

---

一、Triton 后端修改

1.1 引入 Profiling 支持

文件：triton/backends/ascend/driver.py:389

在 NPU 头字符串生成函数 generate_npu_header_src() 中，新增了 Ascend Profiling 公共头文件的引入：

// 修改前
#ifndef TRITON_NPU_HEADERS
#define TRITON_NPU_HEADERS
#include 

// 修改后
#ifndef TRITON_NPU_HEADERS
#define TRITON_NPU_HEADERS
#include "aprof_pub.h"  // 新增
#include

原因：引入 aprof_pub.h 使得后续 NPU 扩展（PCH/so）可以调用 Ascend Profiling 相关 API。该头文件需要与 _get_ascend_path() 指向的 include/experiment/msprof 路径配合使用。这也是后续将编译标准从 C++17 升级到 C++20 的根本原因——aprof 头文件依赖 C++20 语法。

1.2 调整 ACL 头文件搜索路径

文件：triton/backends/ascend/utils.py

在预编译扩展（_precompile_npu_ext）和实际编译扩展（_build_npu_ext）两个函数中，均启用了之前被注释掉的 ACL 头文件路径：

# 修改前（被注释）
# cc_cmd += ["-I/usr/local/python3.12.13/lib/python3.12/site-packages/torch_npu/include/third_party/acl/inc"]

# 修改后（启用）
cc_cmd += ["-I/usr/local/python3.12.13/lib/python3.12/site-packages/torch_npu/include/third_party/acl/inc"]

原因：TorchNPU 自带的 ACL 头文件（包含 RI Cond API）必须优先于 CANN 自带的旧版 acl_rt.h，因为两者使用了相同的 include guard。原代码注释掉这行后，编译时命中的是 CANN 旧头，缺少 RI Cond API。启用后，将 TorchNPU 的 ACL include 路径排在最前，确保新头文件胜出。

1.3 升级编译标准到 C++20

文件：triton/backends/ascend/utils.py

在预编译和实际编译两个函数中，将编译标准从 C++17 升级到 C++20：

# 修改前
cc_cmd += ["-std=c++17", "-shared", "-fPIC", "-o", gch_path]

# 修改后
cc_cmd += ["-std=c++20", "-shared", "-fPIC", "-o", gch_path]

原因：配合引入的 aprof_pub.h 头文件，该头文件及其依赖的 msprof/experiment 头使用了 C++20 特性。PCH（预编译头）和后续 SO 必须使用相同标准，否则 -Winvalid-pch 会告警或失败。

---

二、TorchNPU Inductor 适配层修改

2.1 修复非张量参数处理

文件：torch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/inductor_patch/ir.py:543-544

在 _patch_externkernel_copy_input 函数中，新增了对非张量参数的守卫判断：

def _patch_externkernel_copy_input(x):
    if not hasattr(x, 'get_traced_graph'):  # 新增
        return x                             # 新增
    # 非张量参数（如 Subgraph）原样透传
    traced_graph = x.get_traced_graph()
    node_name = x.get_name()

原因：该函数原本无条件调用 x.get_traced_graph()，假设 x 一定是张量型 IR 对象。但 Inductor 在复制输入时也会传入非张量参数（如 Subgraph 这类高阶算子参数），这些对象没有 get_traced_graph 方法，会直接抛出 AttributeError。新增 hasattr 守卫后，对没有该方法的对象直接原样返回，跳过拷贝逻辑，避免崩溃。

2.2 关闭 Flex Attention 的自定义 Lowering

文件：torch_npu/_inductor/lowering_override_list.py:29-30

将 flex_attention 和 flex_attention_backward 从自定义 Lowering 列表中注释掉：

# 修改前
torch.ops.higher_order.flex_attention,
torch.ops.higher_order.flex_attention_backward,

# 修改后（注释掉）
# torch.ops.higher_order.flex_attention,
# torch.ops.higher_order.flex_attention_backward,

原因：LOWERIN_OVERRIDE_OP 列表中的算子会走 TorchNPU 的自定义 Lowering。注释掉这两个高阶算子后，它们将回落到上游 PyTorch 的默认实现。这是为了规避当前自定义 Lowering 在 flex_attention 上的兼容性适配差异，临时关闭以让模型能正常跑通。

2.3 补齐 Grid 维度

文件：torch_npu/_inductor/runtime/triton_heuristics.py:630-633

在 NPUCachingAutotuner._build_runtime_launch_args 方法中，新增了 grid 维度补齐逻辑：

def _build_runtime_launch_args(self, args, runtime_blocks: tuple[int, ...]):
    padded = (*runtime_blocks, 1, 1, 1)[:3]  # 新增
    return (*args, *runtime_blocks)

原因：Launcher 侧固定接收三维 grid 参数（_grid_0/_grid_1/_grid_2），而 runtime_blocks 可能只有 1 维或 2 维。需要像 Triton 运行时默认逻辑那样，把缺失的 grid 维度补为 1，避免因维度不匹配导致的运行时错误。

---

三、模型代码修改

3.1 放宽编译约束

文件：torchtitan/distributed/compile.py:56

将 fullgraph 参数从 True 改为 False：

# 修改前
transformer_block.compile(backend=backend, fullgraph=True)

# 修改后
transformer_block.compile(backend=backend, fullgraph=False)

原因：fullgraph=True 要求整张图无 Python 副作用、无动态控制流、不允许 graph break。在 NPU Inductor 环境下，图里容易出现 break 点导致编译直接失败。放宽约束后，允许图断裂，保证 TransformerBlock 能编译跑通。

3.2 添加本地数据集支持

文件：torchtitan/hf_datasets/multimodal/mm_datasets.py:288-294

新增本地数据集配置：

"cc12m-local": DatasetConfig(
    path="/path/to/cc12m",
    loader=lambda path: load_dataset(
        path, split="train", data_files={"train": "*.tar"}, streaming=True
    ),
    sample_processor=_process_cc12m_test_sample,
),

原因：原 "cc12m" 走默认/远端下载路径，本机已有一份本地 cc12m 数据（*.tar webdataset 格式）。新增 cc12m-local 指向本地路径，避免联网下载，加速数据加载。

3.3 支持 Scaled Dot-Product Attention

文件：torchtitan/models/common/decoder.py

在 Decoder.get_attention_mask 方法中，新增对 ScaledDotProductAttention 的支持：

if isinstance(inner_attn, FlexAttention.Config):
    return self._create_flex_attention_mask_for_document(positions, attn_config)
elif isinstance(inner_attn, VarlenAttention.Config):
    return create_varlen_metadata_for_document(positions)
elif isinstance(inner_attn, ScaledDotProductAttention.Config):  # 新增
    return None                                                  # 新增
else:
    raise TypeError(...)

原因：SDPA 原生通过 is_causal=True 处理 causal mask，不需要外部预编译的 flex block mask。原代码只识别 Flex/Varlen，遇到 SDPA 会抛出 TypeError。新增分支返回 None，让 SDPA 路径不再尝试构造 flex mask。

3.4 切换到 SDPA 注意力机制

文件：torchtitan/models/qwen3_5/__init__.py

将 Qwen3.5 的文本注意力和视觉注意力的 inner_attention 从 flex_attention 强制切换到 SDPA：

# 视觉 attention config - 新增 inner_attention 字段
proj=_linear(dim, dim),
inner_attention=ScaledDotProductAttention.Config(),

# 文本 attention - 直接使用 SDPA
# inner_attention = get_attention_config(attn_backend)
inner_attention = ScaledDotProductAttention.Config()

原因：这是将注意力机制从 flex_attention 切换到 SDPA 的关键修改点。SDPA 的切换发生在模型侧（这里），而不是 TorchNPU 的 LOWERIN_OVERRIDE_OP。

3.5 适配本地环境

文件：torchtitan/models/qwen3_5/config_registry.py:209,212

将权重路径和数据集名称改为本地配置：

# 修改前
hf_assets_path="./assets/hf/Qwen3.5-9B",
dataloader=_dataloader("cc12m"),

# 修改后
hf_assets_path="/path/to/hf/Qwen3.5-9B",
dataloader=_dataloader("cc12m-local"),

原因：本地环境适配，权重和 tokenizer 指向本机实际路径，dataloader 使用本地数据集。

3.6 视觉编码器 SDPA 适配

文件：torchtitan/models/qwen3_5/vision_encoder.py

在视觉编码器的 mask 构造逻辑中，增加 SDPA 探测：

use_sdpa = isinstance(
    next(iter(self.layers.values())).attn.flex_attention,
    ScaledDotProductAttention,
)
if use_sdpa:
    attention_mask = None
else:
    mask_mod = get_vision_block_mask_mod(num_patch)
    attention_mask = compiled_create_block_mask(
        mask_mod,
        num_vision,
        None,
        max_num_patch,
        max_num_patch,
        device=x.device,
    )

原因：视觉编码器之前无条件调用 compiled_create_block_mask（flex 专用）。切换到 SDPA 后不需要外部 mask。新增 use_sdpa 探测，SDPA 模式下跳过 mask 构造，避免白白编译/运行用不上的 flex mask，也避免 SDPA 下传入不相容的 mask 报错。

3.7 关闭确定性模式

文件：torchtitan/train.py:26-27

注释掉强制确定性模式的代码：

# 修改前
import torch._inductor.config
torch._inductor.config.deterministic = True

# 修改后（注释掉）
# import torch._inductor.config
# torch._inductor.config.deterministic = True

原因：之前强制 deterministic = True 是为了保证确定性，但在 NPU Inductor 下可能触发不支持的确定性算子，降低性能或直接报错。注释掉后恢复默认（非确定性），优先保证模型能跑通。

---

四、新增验证脚本

为了验证 Gated Delta Rule 算子在 NPU 上的正确性，新增了两个测试脚本：

4.1 DVM 编译路径验证

文件：demo_test/test_qwen35_gdr_dvm.py

对 _torch_recurrent_gated_delta_rule 进行 NPU 上的前向+反向验证，走 DVM/Inductor 编译路径：

fn = torch.compile(_torch_recurrent_gated_delta_rule, backend="inductor", fullgraph=True)
out = fn(q, k, v, g, beta)
loss.backward()

输入为 B=1, L=8, N=32, K=V=128 的 bf16 张量（g 为 fp32），打印输出形状、dtype、loss 及梯度的绝对值均值。

4.2 Eager 路径验证

文件：demo_test/test_qwen35_gdr_npu.py

同一算子的 NPU eager（不编译）版本，直接调用：

out = _torch_recurrent_gated_delta_rule(q, k, v, g, beta)
loss.backward()

目的：两个脚本构成对照基准，验证 Gated Delta Rule 算子在 NPU 上，torch.compile(inductor/fullgraph) 编译路径与 eager 路径的输出和梯度是否一致，用于定位 DVM 编译是否引入数值误差。

---

五、总结

本文记录了在 Qwen3.5 模型上启用 DVM 编译过程中遇到的一系列问题及解决方案。关键要点如下：

1. SDPA 切换发生在模型侧：真正将注意力机制从 flex_attention 切换到 SDPA 的修改在模型代码中（torchtitan/models/qwen3_5/__init__.py），而不是 TorchNPU 的 LOWERIN_OVERRIDE_OP 注释——后者只影响 device-dispatch 包装，不会自动让 attention 走 SDPA。
2. 编译标准升级：为了支持 Ascend Profiling 功能，需要将编译标准从 C++17 升级到 C++20。
3. 头文件路径优先级：TorchNPU 自带的 ACL 头文件必须优先于 CANN 旧版头文件，以确保 RI Cond API 可用。
4. 非张量参数处理：Inductor IR Patch 需要处理 FallbackKernel 等非张量参数，避免直接访问不存在的属性。
5. 编译约束放宽：在 NPU 环境下，fullgraph=True 过于严格，需要放宽到 fullgraph=False 以保证编译通过。

这些修改共同确保了 Qwen3.5 模型能够在 Ascend NPU 上成功启用 DVM 编译，并获得正确的训练结果。