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

1. 背景概述

随着大模型训练对计算效率和显存利用率的要求不断提高，PyTorch 的 torch.compile 机制成为优化模型执行的关键技术。DeepSeek V4 作为新一代大规模语言模型，在 Ascend NPU 上进行训练时，需要适配 torch_npu 的 DVM（Dynamic Vector Machine）编译后端，以充分发挥硬件性能。本文记录了在适配过程中遇到的一系列问题、分析思路及解决方案，旨在为类似场景下的 compile 模式适配提供参考。

2. 环境与验收配置

本文结论基于以下环境和最终验收配置：

项目	版本或配置
架构	Linux aarch64
NPU	8 张 Atlas 800I A2
CANN	9.1.0
Python	3.12.13
PyTorch	2.12.0+cpu
torch_npu	2.12.0+gitf380339
ops-transformer	cd1734e84b19c130224db087e85bead27a3d31e8
模型	deepseek_v4 / v4_13b
并行	EP8 + FSDP8
训练形状	sequence length 2048，global batch size 16
dtype	BF16
attention	auto，实际命中 Vendor SparseFlashMla
compile backend	DVM

3. 最终结论

在固定 step-0 checkpoint、完整 C4 数据集、debug.seed=42 的 8 卡 10-step 验收中，DVM compile 已同时通过功能、数值、性能和容量门禁：

指标	结论
step 1 / step 10 loss	与 eager 接近
10-step 最大 loss / grad norm 相对差	低于既有 Vendor backward 波动参考
中位 step time 与中位吞吐	DVM 明显优于 eager
rank0 active peak 与 8-rank 最坏 reserved peak	DVM 更低
allocator retry / OOM	0 / 0，通过

编译产物提供了后端证据：80/80 个 output_code.py 包含 @dvm.kernel，没有 @triton.jit；其中 64 个仍保留 cann_ops_transformer 或 sparse_flash_mla custom-op 调用。这比仅看到 torch.compile 启用更可靠，因为“启用了 compile”不等于“实际走了预期 backend”。

最终落地依赖四个关键条件：

1. Vendor custom op 必须具备可捕获的 autograd、Fake/Meta 和副作用边界。
2. DVM 必须来自实际包含 DVM 实现的 torch_npu build。
3. DVM reduction 与 input/output alias 的错误融合必须修复。
4. v4_13b 必须显式设置 training.dtype=bfloat16。

4. 推荐的排查分层

本次适配最有效的做法不是反复启动完整 Trainer，而是按编译流水线逐层缩小问题范围：

1. Dynamo/AOTAutograd 层：先排查 Python import、日志、集合写入、autograd 多输出合同等 fullgraph 阻断点。
2. FakeTensor/Meta 层：若真实 kernel 尚未执行就失败，直接检查 Meta implementation 的 shape、dtype 和 schema，不先怀疑设备算子。
3. Lowering/codegen 层：检查堆栈实际落在哪个 backend，并查看生成代码，不能根据启动参数推断后端。
4. 单算子 A/B 层：固定输入、cotangent 和 seed，对 eager/compile 做前向、反向、重复哈希及冷/热缓存对比。
5. 融合边界层：原子表达式都正确但组合错误时，逐步增加融合前缀，或用精确 fallback 切开节点，定位首个错误融合。
6. 完整 Trainer 层：最后才做至少 10 step 的数值、显存和稳态性能验收。

这套顺序能区分四类经常被混为一谈的问题：图捕获失败、编译失败、编译成功但静默算错，以及算法正确但运行配置导致 OOM。

5. 闭环一：让 Vendor attention 成为 compile-safe 的 PyTorch 边界

5.1 现象与根因

Vendor SparseFlashMlaGrad 只支持 attention output 的梯度，不支持 LSE 梯度；原 adapter 却将 output 和 LSE 都声明为可微输出。AOTAutograd 即使看到调用方对 LSE 做了 detach，仍可能向 raw autograd.Function.backward 传入 synthetic grad_lse。只调用 mark_non_differentiable() 不够，backward 签名仍必须接收该占位参数。

同一条路径还有两个 fullgraph 副作用：

• cann_ops_transformer 在第一次 forward 内延迟 import，模块加载和算子注册无法作为张量图捕获；
• Vendor 首次调用会执行 logger.info() 并写 _LOGGED_OPS 集合。

这三个问题都发生在 Python/PyTorch 集成边界，不是 Vendor kernel 的数值问题。

5.2 修复

LSE 在 NPU Vendor 路径中明确设为不可微，backward 接收但忽略 synthetic LSE 梯度。模型中唯一使用 LSE 的 indexer auxiliary loss 原本就显式 detach，因此没有改变训练梯度语义。

首次 import 被移动到并行化完成后、activation checkpoint 和 compile 包装之前；同时只在真实会命中 Vendor contract 的 NPU 配置下预加载，避免污染非目标后端（如 CPU、Fusion）或仅编译 loss 的路径。三个修复点集中如下：

--- a/torchtitan/models/deepseek_v4/ops_transformer_attention.py
+++ b/torchtitan/models/deepseek_v4/ops_transformer_attention.py
@@ _forward_impl(...)
-    return out_BSND, native_lse_B1SN.squeeze(1)
+    lse_BSN = native_lse_B1SN.squeeze(1)
+    ctx.mark_non_differentiable(lse_BSN)
+    return out_BSND, lse_BSN

@@ class _SparseFlashMlaFunction(torch.autograd.Function):
-    def backward(ctx, grad_out_BSND, grad_lse_BSN):
-        if grad_lse_BSN is not None:
-            raise RuntimeError("SparseFlashMlaGrad does not support grad_lse")
+    def backward(ctx, grad_out_BSND, _grad_lse_BSN):
         return _backward_impl(ctx, grad_out_BSND)

@@ provider loader
-def _ops_module():
+def load_ops_transformer_attention():
     ...

@@ def _run_vendor_op(name, *args, **kwargs):
-    _ops_module()
+    load_ops_transformer_attention()
-    if name not in _LOGGED_OPS:
+    if not torch.compiler.is_compiling() and name not in _LOGGED_OPS:
         logger.info("DeepSeek V4 called vendor op cann_ops_transformer::%s", name)
         _LOGGED_OPS.add(name)

--- a/torchtitan/models/deepseek_v4/parallelize.py
+++ b/torchtitan/models/deepseek_v4/parallelize.py
@@ after TP/EP/CP setup, before AC and compile wrapping
+    if model_compile_enabled and parallelism.spmd_backend == "default":
+        _preload_ops_transformer_attention_for_compile(model)

5.3 验证

• eager/fullgraph 前反向回归覆盖 output、LSE 和四类输入梯度合同；
• import、activation checkpoint、compile 顺序及非适用配置隔离测试通过；
• 最终 v4_13b 8 卡 10-step 中，64 个生成 wrapper 保留 Vendor custom op，前向、反向和正常 teardown 均完成。

经验是：custom op 的 compile 适配不只需要 schema 和 kernel；首次 import、日志、Python 容器写入和多输出 autograd 合同都属于可编译边界的一部分。

6. 闭环二：FakeTensor 报 float shape，实际是运行时加载了旧 Meta 实现

6.1 保存的错误日志

失败日志：outputs/logs/dsv4_v4_13b_dvm_10step/

核心堆栈如下：

failed while attempting to run meta for
cann_ops_transformer.sparse_flash_mla.default

File "/usr/local/Ascend/cann-9.1.0/python/site-packages/
cann_ops_transformer/ops/sparse_flash_mla.py", line 101
  softmax_lse = torch.empty([
      q.shape[0],
      ori_kv.shape[2],
      q.shape[1],
      q.shape[2] / ori_kv.shape[2],
  ], ...)
TypeError: empty(): argument 'size' ...
"type must be tuple of ints,but got float"

6.2 错误分析过程

1. 错误来自 FakeTensor 执行 Meta kernel，真实 NPU kernel 尚未运行，因此先排除 DVM kernel、HCCL 和显存。
2. traceback 显示导入路径是 CANN 系统目录，而不是固定 commit 的 overlay。
3. 系统 Meta 用 Python / 计算 q_heads / kv_heads，结果是 float；C++ schema 和张量 shape 要求整数。
4. 固定 ops-transformer commit 已使用 //，所以不是模型传入了非法 shape，而是 Python import precedence 选择了旧包。

6.3 修复

没有直接修改系统 site-packages，而是从固定 commit 创建独立 Python overlay，并给 JIT extension 单独的缓存目录。准备脚本同时校验源码和 Meta 哈希，避免“目录存在但版本不对”。

Meta 层的关键差异是：

 softmax_lse = torch.empty(
     [
         q.shape[0],
         ori_kv.shape[2],
         q.shape[1],
-        q.shape[2] / ori_kv.shape[2],
+        q.shape[2] // ori_kv.shape[2],
     ],
     dtype=torch.float32,
     device=q.device,
 )

运行时环境必须让 overlay 排在系统包之前：

PYTHONPATH=/tmp/ops_transformer_overlay/torch_extension:/path/to/torchtitan \
TORCH_EXTENSIONS_DIR=/tmp/ops_transformer_overlay/torch_extensions \
python3 ...

6.4 验证

• Meta/FakeTensor 和 raw opcheck 通过；
• ratio 1/4/128 的 production Vendor fullgraph 前反向通过；
• 相同 v4_13b 命令不挂 overlay 时稳定复现 float shape，挂 overlay 后完成 10 step。

经验是：FakeTensor 的错误信息看似是模型 shape 问题时，先看 traceback 中真正加载的 Meta 文件。Python Meta、C++ extension 和系统 CANN 包可能来自三个不同版本。

7. 闭环三：确认 DVM backend 真实存在并被实际选择

7.1 现象与排查

旧系统 wheel 接收 options={"npu_backend": "dvm"} 后，仍进入 NPU Triton range-tree 初始化并报 KeyError: 'r'。仅凭参数名容易误判为“DVM 初始化失败”。

源码核验发现该 wheel 对应 commit fa0f83fe，不包含以下任一实现：

• torch_npu/_inductor/dvm；
• _load_dvm_backend；
• _C.dvm binding。

因此真实行为是未知 backend 回落到 Triton loader，错误来自 Triton，不是 DVM。

7.2 修复

从官方 v2.12 分支 commit f38033913dbefc3a74dbdc4d4838794790d834f8 构建隔离 wheel，先在 PYTHONPATH overlay 中验证，再将已验证 wheel 安装到系统环境。安装前后对 Fusion 和 Vendor eager 做同输入回归，排除了 torch_npu commit 升级带来的前向语义回归。

最终环境为：

torch       2.12.0+cpu
torch_npu   2.12.0+gitf380339
CANN        9.1.0

DVM 的启动条件是同时设置 backend 和启用 TorchTitan compile：

TORCHINDUCTOR_NPU_BACKEND=dvm \
./run_train.sh \
  --compile.enable \
  ...

7.3 验证

• 20 轮 Sinkhorn 生成代码出现 @dvm.kernel 和 dvm_fused_*.run；
• 没有 @triton.jit；
• 前向与梯度最大误差均处于极小量级（远低于容差阈值）；
• 最终全模型 80/80 个 wrapper 均为 DVM。

经验是：backend 选择必须用“实现存在 + 生成代码证据”验证。环境变量和 options 只能说明意图，不能说明实际执行路径。

8. 闭环四：Triton 路线为何被放弃，DVM 路线为何成立

本节记录一条已经闭环但不属于当前生产依赖的诊断路线。它的价值是说明如何从“编译失败”继续追到“编译成功但静默算错”，以及为什么最终选择 DVM，而不是继续堆叠 Triton workaround。

8.1 常量索引被错误生成为 rank-1 load

保存日志：outputs/logs/dsv4_debug_layer6_compile_s128/

tmp1 = tl.load(in_ptr1 + tl.arange(0,1) + (2))
tmp2 = tl.broadcast_to(tmp1, [1, 1])
       ^
ValueError: Cannot broadcast, rank mismatch: [1], [1, 1]

hc_scale[2] 语义上是标量，NPU codegen 却用 tl.arange(0, 1) 把它变成长度为 1 的向量。修复是让整数索引直接执行标量 load：

@@ integer index lowering
-    line = f"tl.load({var} + tl.arange(0,1) + ({original_index}))"
+    line = f"tl.load({var} + ({original_index}))"

8.2 HFusion 不支持 bool/bool unsigned-greater

标量 rank 修复后，backward 在 aten.max.dim 重建 argmax 时生成 bool/bool a > b，最终映射到 HFusion 尚不支持的 CompareFn::vugt。布尔比较可等价写为“a 为真且 b 为假”，所以只对 bool/bool 做规范化，其他 dtype 保持上游实现：

@@ class NPUTritonKernelOverrides(TritonKernelOverrides):
+    @staticmethod
+    def gt(a, b):
+        if triton_arg_dtype(a) == torch.bool and triton_arg_dtype(b) == torch.bool:
+            return f"({a} & ({b} == 0))"
+        return TritonKernelOverrides.gt(a, b)

8.3 编译成功后 Sinkhorn 仍静默算错

前两个 codegen 问题修复后，MHC Sinkhorn 可以编译，但固定输入上的融合前缀出现约 0.08 最大误差，完整 20 轮还有冷/热漂移和 backward NaN。排查顺序为：

1. 分别执行 affine、row max、exp、sum 和 div，原子结果正确；
2. 逐步增加融合前缀，错误只在归约进入融合 kernel 后出现；
3. FP32/BF16、bool 比较形式、autotune winner 和 partial tile 都不是必要条件；
4. 精确 fallback aten.max.dim 只修最小前缀，完整 Sinkhorn 仍错误；
5. 同时 fallback aten.max.dim,aten.sum.dim_IntList 后，冷/热输出逐位一致，梯度最大绝对误差处于极小量级。

这个 A/B 证明错误边界位于 NPU Triton 融合归约，而不是 Sinkhorn 数学公式。

排查还发现用户 fallback 注册过早，会被后续 NPU custom lowering 覆盖。通用修复是先保存用户选择，所有 custom lowering 完成后再 make_fallback()：

@@ def _register_npu_inductor_fallbacks():
+    user_fallback_ops = []
     ...
-            FALLBACK_LIST.append(op)
+            user_fallback_ops.extend(
+                getattr(op, overload) for overload in op.overloads()
+            )
     ...  # built-in and NPU-specific lowerings are registered here
+    for op in user_fallback_ops:
+        if op not in decompositions:
+            make_fallback(op)
+            FALLBACK_LIST.append(op)

8.4 路线决策

该 fallback 每个 block 每步会产生 2 次 max 和 123 次 sum extern call，破坏 Sinkhorn 融合，不能作为高性能最终方案。DVM 在相同算子上通过数值门禁；随后在同一 torch_npu build、NPU0、固定 shape、FP32、seed 42 和 strict deterministic 下按 warmup 10、采样 30、独立重复 3 轮取得：

范围	结论
forward 设备耗时中位数	DVM 明显低于 eager
forward+backward 设备耗时中位数	DVM 明显低于 eager

所以最终撤销 Triton 实验补丁和 fallback 环境变量，选择 DVM。闭环不是“把所有 Triton bug 都变成生产补丁”，而是通过最小修复和精确 fallback 完成根因证明，再用数值与性能证据选择更合适的 backend。

9. 闭环五：DVM spec reduction 与 input/output alias 导致静默错误

9.1 现象

Fusion eager 与全模型 DVM 对比时，dsa_indexer_loss/mean 首步出现明显跳变。真实 layer 2、ratio=4 输入中，eager 与 DVM 的 _acc 相差约三个数量级。

这不是异常退出，而是编译成功后的错误结果，因此普通错误日志没有堆栈。保存的 DVM 生成代码位于：

/tmp/dvm_alias_debug/.../model__2_forward_6.3/output_code.py

9.2 错误分析过程

FX 图的语义只是：

raw_sum = (...).mean()
new_acc = old_acc + raw_sum.detach()
old_acc.copy_(new_acc)

错误生成代码却把 reduction、add 和 copy 融成一个 spec kernel，并在调用端让同一 primals_6 同时作为输入和输出：

@dvm.kernel(ktype="spec", dyn_shape=False)
def dvm_fused_add_copy__div_mean_mul_squeeze_5_build(k):
    arg2_1 = k.load([], dvm.float32)
    ...
    sum_1 = k.sum(mul, [0, 1], False)
    k.spec_next()
    add = k.add(arg2_1, div_2)
    copy = add
    k.store(copy, dvm.float32)

dvm_fused_add_copy__div_mean_mul_squeeze_5.run(
    buf29, buf33, primals_6, primals_6, stream=stream0
)

直接调用生成 kernel 做 alias A/B：旧 accumulator 与归约项取典型示例值时，独立 output 得到正确结果；input/output alias 时结果异常偏离。其中整数部分正是 spec_next() 前第一阶段归约的和。

由此确认：spec kernel 会先借输出 buffer 保存阶段性 reduction，覆盖与其同地址的旧 accumulator；第二阶段再读取的已经不是旧值。

9.3 修复

最初的诊断补丁禁止所有 reduction 与 mutation 纵向融合，证明拆分后结果正确；最终补丁收窄为真实 alias 判定：只有 mutation target 同时以 MemoryDep 被融合组读取时，才拒绝融合。仅用于更新顺序的 StarDep 不算读取旧值。

+from torch._inductor.dependencies import MemoryDep

+def _has_mutated_input_alias(*nodes):
+    mutations = {
+        name for node in nodes for output in node.get_outputs()
+        for name in output.get_mutations()
+    }
+    inputs = {
+        dep.name for node in nodes for dep in node.read_writes.reads
+        if isinstance(dep, MemoryDep)
+    }
+    return bool(mutations & inputs)

@@ def can_fuse_vertical(self, node1, node2):
+    if (node1.is_reduction() or node2.is_reduction()) and _has_mutated_input_alias(
+        node1, node2
+    ):
+        return False

修复后生成代码明确分成两个 kernel：

-dvm_fused_add_copy__div_mean_mul_squeeze_5.run(
-    buf29, buf33, primals_6, primals_6, stream=stream0
-)
+buf34 = empty_strided((), (), device="npu", dtype=torch.float32)
+dvm_fused_mean_mul_squeeze_5.run(buf29, buf33, buf34, stream=stream0)
+dvm_fused_add_copy__div_mean_mul_squeeze_7.run(
+    primals_6, buf34, primals_6, stream=stream0
+)

9.4 验证

• 真实 ratio=4 输入的 DVM _acc 恢复至与 eager 一致，相对绝对差处于极小量级；
• carrier output 和 carrier gradient 逐位一致；
• index-score gradient 最大绝对差处于极小量级；
• 安全反例 dst.copy_(x.sum()) 仍保留一个 spec kernel，结果与 eager 都为 63.5，证明没有全局关闭 reduction+copy 融合；
• 最终 8 卡 10-step 无该 metric 数量级异常。

经验是：带 spec_next() 的多阶段 kernel 不能只检查拓扑依赖，还必须检查输出 buffer 是否与后续阶段需要读取的旧输入别名。

10. 闭环六：明确 reduction 的 deterministic 生效时机

10.1 现象与根因

DVM Sinkhorn 默认模式满足数值容差，但相同输入和 cotangent 连续执行时，hc_scale、hc_base 梯度哈希会变化。最小 broadcast-gradient 复现中，默认模式的 base gradient 运行 10 次得到 10 个哈希；主输入和非并行归约梯度稳定。

DVM 默认 sum 允许非固定的并行归约顺序。TorchKernelPy 在构造 kernel 时读取 PyTorch 全局 deterministic 状态；严格模式下，DVM ReduceOp 才会选择额外 workspace 和固定访问路径。因此开关必须在 DVM kernel 构造，也就是 torch.compile() 之前生效。

10.2 修复与验证

torch.use_deterministic_algorithms(True)
compiled = torch.compile(fn, options={"npu_backend": "dvm"})

TorchTitan 的 --debug.deterministic 已在模型构建和 parallelize_fn 之前设置该状态，不需要增加 DeepSeek V4 专属配置，也不能使用 --debug.deterministic_warn_only。

严格模式下：

• 最小归约的三类梯度各只有一个哈希；
• 完整 20 轮 Sinkhorn 在两个独立冷缓存进程、两张 NPU 上得到完全相同的输出和梯度哈希；
• 输出与梯度最大误差均处于极小量级（远低于容差阈值）。

这里的严格确定性结论只适用于支持 deterministic 的 DVM 路径；Vendor backward 的限制和后续验收条件统一列在文档末尾的开放问题中。

11. 闭环七：保住显式的 BF16 舍入和 RMSNorm 语义边界

11.1 Query RMS 的首个差异

原 Query RMS 使用手写表达式，DVM 把它与前序 matmul 融合并重排中间计算。真实 layer-0 BF16 query 上，相对 eager 的 output 与 input gradient 最大绝对误差均处于极小量级，relative L2 远低于门禁阈值。分阶段捕获证明首个差异发生在 attention_q_rms，早于 Vendor attention。

F.rms_norm 明确表达 PyTorch RMSNorm 合同，在 BF16 输入上使用 FP32 reduction；它的 eager/DVM 前向和输入梯度逐位一致，也与 npu_rms_norm 前向逐位一致。因此选择最小的 PyTorch-native 修改：

-q = q * torch.rsqrt(q.square().mean(-1, keepdim=True) + self.norm_eps)
+q = F.rms_norm(q, (self.head_dim,), eps=self.norm_eps)

11.2 Indexer aux 的隐式 BF16 store 被融合穿透

aux 路径的 gather、einsum 和独立保存的 BF16 selected_score 都逐位一致；首个差异出现在 selected_score.float() - attn_lse。原表达式先在 BF16 中乘 scale，eager 会在 .float() 前经历 BF16 舍入；DVM 融合后把 BF16 scalar-mul 的内部值直接送到 FP32 cast/sub，跳过了该隐式 store/round 边界。

修复不是强行插入 materialization，而是明确目标分布用 FP32 计算：

 selected_score = (
-    torch.einsum("bshd,bskd->bhsk", q, selected_kv) * self.softmax_scale
+    torch.einsum("bshd,bskd->bhsk", q, selected_kv).float()
+    * self.softmax_scale
 )
 selected_prob = torch.exp(
-    selected_score.float()
-    - attn_lse.transpose(1, 2).unsqueeze(-1).float()
+    selected_score
+    - attn_lse.transpose(1, 2).unsqueeze(-1).float()
 )

11.3 验证

• attention_q_linear、attention_q_rms、Vendor 输入输出和 attention 最终输出均逐位一致；
• aux raw loss 绝对差显著下降至极小量级；
• index-score gradient 最大绝对差显著下降至极小量级；
• focused CPU tests 及 8 卡 Trainer smoke 通过；
• F.rms_norm 与 custom op 的 DVM 性能接近，前反向中 F.rms_norm 略快，且没有新增不透明 dispatcher 边界。

经验是：compile 融合可以合法改变归约树，但不能穿过模型明确依赖的 dtype 舍入合同。需要先用阶段捕获找到首个差异，再决定是表达式缺少稳定算子语义，还是编译器错误跨越了显式 cast/store 边界。

12. 闭环八：EP 动态 token shape 不应强建 singleton guard

12.1 保存的错误日志

失败日志：outputs/logs/dsv4_dvm_manual/

GuardOnDataDependentSymNode:
Could not guard on data-dependent expression
Eq(u10 + u11 + u12 + u13 + u14 + u15 + u8 + u9, 1)

Caused by:
torch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/
inductor_patch/lowering.py:569 in broadcast_symbolic_shapes

target: aten.add.Tensor
user stack:
torchtitan/models/common/token_dispatcher.py:550 in _permute

12.2 错误分析过程

EP8 中每个 expert/rank 的 token 数来自 runtime collective，shape 表达式是无 hint 的 u8 + ... + u15。token_dispatcher 构造的两个张量 shape 完全相同，本应直接进入 elementwise add；旧 torch_npu lowering 却先调用 evaluate_expr(Eq(size, 1)) 判断是否为 singleton。没有 hint 的数据依赖 symbol 无法静态证明 size == 1，于是 guard 构造失败。

对照配套 PyTorch 2.12 后确认 torch_npu 保存的是旧版 broadcast helper。当前 PyTorch 的处理有两个关键点：

1. 两个 shape 完全相等时直接返回，不做 singleton 判断；
2. 必须判断时使用 is_size_one_or_false()，对未知动态尺寸保守返回 false。

因此根因在共享 torch_npu lowering，不在 DSV4 MoE 数学或 token dispatcher。

12.3 修复

两份 torch_npu broadcast helper 同步到配套 PyTorch 2.12 语义：

@@ def broadcast_symbolic_shapes(a, b):
+    b = tuple(b)
+    if not a or a == b:
+        return b
     ...
-        if V.graph.sizevars.shape_env.evaluate_expr(
-            sympy.Eq(y, 1), size_oblivious=True
-        ):
+        if V.graph.sizevars.is_size_one_or_false(y):
             output.append(x)
-        elif V.graph.sizevars.shape_env.evaluate_expr(
-            sympy.Eq(x, 1), size_oblivious=True
-        ):
+        elif V.graph.sizevars.is_size_one_or_false(x):
             output.append(y)

@@ def broadcast_tensors(*inputs):
-    if len(inputs) == 1 and isinstance(inputs[0], (list, tuple)):
-        return broadcast_tensors(*inputs[0])
+    if len(inputs) == 1:
+        if isinstance(inputs[0], (list, tuple)):
+            return broadcast_tensors(*inputs[0])
+        return inputs
     ...
-        sizes = x.get_size()
-        if len(sizes) != len(target) or any(...):
+        if (sizes := tuple(x.get_size())) == target:
+            pass
+        elif len(sizes) != len(target) or any(
+            V.graph.sizevars.is_size_one_or_false(a)
+            != V.graph.sizevars.is_size_one_or_false(b)
+            for a, b in zip(sizes, target)
+        ):
             x = expand(x, target)

相同逻辑同时修复：

• torch_npu/_inductor/lowering_fx.py；
• torch_npu/_inductor/ascend_npu_ir/.../inductor_patch/lowering.py。

12.4 验证

• 最小 A/B 用八个 Tensor.item() 构造无 hint 的 u0 + ... + u7；旧环境稳定失败，修复后连续两次精确得到 [1, 2, ..., 8]；
• 普通等 shape 与 singleton broadcast 回归通过；
• 8-rank v4_13b、EP8、FSDP8、S2048 一步功能验证通过；
• 同一路径最终完成 10 step，证明动态 EP shape 不只在首图通过。

经验是：动态 shape 的相等性和“是否等于 1”是两个不同问题。对相同 shape 先短路，对未知 symbol 使用保守 predicate，不能通过 evaluate_expr() 强建数据依赖 guard。

13. 闭环九：第二步 OOM 是 dtype 配置错误，不是单卡占用或 DVM 回归

13.1 保存的错误日志

失败日志：outputs/logs/dsv4_v4_13b_eager_10step/

[rank7] NPU out of memory. Tried to allocate 7.00 GiB
NPU 7; 60.96 GiB total capacity
53.29 GiB already allocated
6.66 GiB free
53.44 GiB reserved in total by PyTorch

File "torchtitan/models/deepseek_v4/attention.py", line 132,
in _selected_main_attn_dist
    selected_kv = torch.gather(...)

13.2 错误分析过程

最初怀疑卡 2 被其他进程占用，但证据不支持：

1. rank 2、3、6、7 都出现过同类失败，示例堆栈来自 rank 7；
2. reserved 与 allocated 很接近，不符合严重 allocator 碎片；
3. 第一步前反向成功，失败发生在 AdamW state 建立后的第二步；
4. 调用点是 indexer aux gather，确实需要一块较大的临时张量；
5. 启动命令遗漏 training.dtype=bfloat16，实际使用 v4_13b 默认 FP32。

FP32 下模型常驻显存约为 BF16 的两倍。optimizer state 建立后，额外的模型/梯度容量正好使第二步那块较大的 gather 越过上限。因此根因是验收命令没有满足原定 full-BF16 配置，不是 NPU2 外部占用，也不是 DVM 生成代码泄漏。

13.3 修复与验证

修复只有一个参数：

--training.dtype=bfloat16

补齐后 eager 和 DVM 都完成 10 step，DVM 的最坏 reserved peak 低于 eager，allocator retry 和 OOM 均为 0。

经验是：第二步 OOM 要先检查 optimizer state 生命周期和实际 dtype。第一步能跑通只证明模型、activation 和 gradient 能容纳，不代表 AdamW 一、二阶矩建立后仍有足够空间。

14. 跨 backend 数值门禁应按计算合同设置

14.1 为什么不能要求 Vendor eager/DVM 逐位一致

同一 backend、同一实现的非计算改动仍应使用 debug.seed=42 和 strict deterministic 做逐位门禁。但 eager 与 compile 可能改变融合、归约树和中间精度；跨 backend 的低位差异本身不能证明错误。

更重要的是，当前 Vendor backward 明确拒绝 strict deterministic。保存日志：outputs/logs/dsv4_v4_13b_eager_10step/

File "ops_transformer_attention.py", line 345, in _backward_impl
  raise RuntimeError(
      "SparseFlashMlaGrad does not support deterministic mode"
  )
RuntimeError: SparseFlash

14.2 分层门禁

跨 backend 计算变更使用以下顺序验收：

1. shape、dtype、layout、mutation/alias 和显式舍入边界一致；
2. output、loss、gradient 和 metric 全部 finite；
3. 同一 backend 在其支持的 deterministic 合同内可复现；
4. 单算子使用 max absolute error、relative L2、cosine 和 reference；
5. 离散路由检查 top-k、expert load 和边界翻转；
6. 代表性数据上比较平滑 loss、收敛趋势、grad norm 和非有限值；
7. 性能至少运行 10 step，并排除冷编译和 optimizer 初始化阶段。

在模型适配过程中，以下现象必须作为硬性指标进行排查，不能简单归因于“浮点非结合性”：

• NaN/Inf：任何非有限值的出现都需立即定位根因
• 同输入重复执行漂移：相同输入下多次执行结果不一致
• 数量级误差：输出结果与参考值存在数量级差异
• 显式 BF16 cast/store 边界被融合穿透：编译优化破坏了显式精度控制
• 异常 gradient spike：梯度出现异常尖峰
• 路由或 expert load 异常：MoE 路由分配出现异常模式
• 超出 eager 自身波动的持续收敛退化：训练损失持续偏离 eager 模式的正常波动范围

可复现启动命令

以下为实际完成 v4_13b 模型、EP8、FSDP8、S2048、BF16、Vendor attention、DVM compile 10-step 的命令结构。复用时需更换 dump、log 和 cache 目录，避免覆盖已有产物。

cd /path/to/torchtitan

NGPU=8 \
LOG_RANK=0 \
MODULE=deepseek_v4 \
CONFIG=v4_13b \
LOG_DIR=/path/to/torchtitan/outputs/logs/dsv4_v4_13b_dvm \
PYTHONPATH=/tmp/ops_transformer_overlay/torch_extension:/path/to/torchtitan \
TORCH_EXTENSIONS_DIR=/tmp/ops_transformer_overlay/torch_extensions \
TORCHINDUCTOR_NPU_BACKEND=dvm \
TORCHINDUCTOR_CACHE_DIR=/tmp/dvm_cache \
./run_train.sh \
  --dump_folder=/path/to/torchtitan/outputs/dsv4_v4_13b_dvm \
  --model.attention_backend=auto \
  --parallelism.data_parallel_shard_degree=8 \
  --parallelism.expert_parallel_degree=8 \
  --training.dtype=bfloat16 \
  --training.global_batch_size=16 \
  --training.seq_len=2048 \
  --training.steps=10 \
  --dataloader.dataset=c4 \
  --dataloader.dataset_path=/path/to/datasets/c4 \
  --checkpoint.enable \
  --checkpoint.load_only \
  --checkpoint.initial_load_path=/path/to/torchtitan/outputs/flash_seed/seed_checkpoint/13b/step-0 \
  --checkpoint.initial_load_model_only \
  --debug.seed=42 \
  --metrics.enable_tensorboard \
  --metrics.log_freq=1 \
  --compile.enable

注意事项：

• 当前 Vendor backward 不支持 --debug.deterministic，启用会按合同显式失败
• 若切换到支持 strict deterministic 的 Fusion/DVM 路径，应在 compile 前启用严格模式，且绝不使用 --debug.deterministic_warn_only

复用检查清单

编译前检查

1. 环境固定：确认 PyTorch、torch_npu、CANN、SoC 和 ops-transformer commit 版本
2. 导入路径：确认实际 Python import path，尤其是 Meta implementation
3. 编译配置：确认 TORCHINDUCTOR_NPU_BACKEND=dvm 与 --compile.enable 同时存在
4. 模型精度：确认模型实际 dtype，不依赖 config 名称推断
5. 确定性设置：在 compile 前设置 deterministic 状态
6. 依赖管理：将 optional provider 的 import/registration 移出 fullgraph

编译失败排查

• Dynamo/AOT 阶段 traceback：检查 Python 副作用和 autograd 合同
• FakeTensor/Meta 阶段 traceback：检查 schema、shape 类型和加载版本
• lowering 阶段 traceback：对照配套 PyTorch 实现，避免旧 fork 对动态 shape 强建 guard
• codegen 阶段 traceback：检查标量/向量 rank、dtype lowering 和实际 backend
• 原则：不要用全局关闭融合代替根因定位，先做最窄 A/B 测试

编译成功验证

1. 代码检查：检查 output_code.py，确认 DVM kernel 数量和 custom-op 保留情况
2. 固定输入测试：做 eager/compile 前向、反向和重复哈希对比
3. 特殊场景测试：对 reduction、mutation 和 alias 组合单独做反例测试
4. 差异定位：用阶段捕获找第一个差异，不从最终 loss 反推算子
5. 完整训练验证：至少运行 10 step，性能统计排除冷编译和 optimizer 初始化
6. 指标记录：同时记录 loss、grad norm、aux、step time、throughput、active/reserved peak、allocator retry 和 OOM

关联实现与证据

主要模型实现文件

• torchtitan/models/deepseek_v4/ops_transformer_attention.py
• torchtitan/models/deepseek_v4/parallelize.py
• torchtitan/models/deepseek_v4/attention.py

可复用补丁和复现脚本

• scripts/prepare_ops_transformer_overlay.sh
• scripts/torch_npu_triton_integer_index_scalar_load.patch
• scripts/torch_npu_triton_bool_gt.patch
• scripts/torch_npu_user_fallback_registration_order.patch
• scripts/torch_npu_dvm_reject_spec_mutation_fusion.patch
• scripts/torch_npu_sync_symbolic_broadcast.patch
• scripts/repro_dsv4_mhc_sinkhorn_compile.py
• scripts/repro_dsv4_full_model_block_compare.py
• scripts/repro_torch_npu_dynamic_broadcast.py

最终训练证据目录

• Eager 模式：outputs/dsv4_v4_13b_eager_10step/
• DVM 模式：outputs/dsv4_v4_13b_dvm_overlay_10step/
• DVM 日志：outputs/logs/dsv4_v4_13b_dvm_overlay_10step/
• DVM debug 图：/tmp/dsv4_dvm_debug_overlay/

开放问题

以下问题尚未形成“根因、生产修复、验证”闭环。除特别说明外，它们不影响本文 v4_13b 默认配置的 8 卡 10-step 验收结论。

当前默认路径仍需补充的门禁

1. Vendor backward 不支持 strict deterministic

SparseFlashMlaGrad 会在严格确定性模式下显式失败，因此当前 Vendor eager/DVM 门禁使用的是同类 Vendor eager 重复运行的经验波动范围，而非逐位一致性。该波动范围来自 debug_layer6，尚未在相同 v4_13b 配置下完成重复统计。

闭环条件：Vendor backward 提供严格确定性实现，并用相同 v4_13b 配置完成独立重复的 10-step output、loss、grad norm 和梯度哈希验证。

2. 长程、多 seed 收敛尚未验证

当前最终证据只覆盖完整 C4 上的单个 seed、10 step，可以验证功能、短程数值、容量和稳态性能，但不能替代长程收敛结论。

闭环条件：在代表性 C4 训练长度上运行多个 seed，对比 eager/DVM 的平滑 loss、grad norm、路由负载、expert load 和非有限值，并以 eager 自身的多 seed 波动作为参考。

3. EP8 持久 FX graph cache 的 key 跨进程不稳定

DeviceMesh 循环引用导致的 pickle 失败已经修复，同一进程也能 cache hit；但两个独立进程的首图 key 仍从 fdpe... 变为 f5df...，热进程继续 cache miss，说明完整缓存键中仍有跨进程不稳定字段。证据位于 outputs/logs/dsv4_ep8_cold/ 和 outputs/logs/dsv4_ep8_warm/。

闭环条件：独立冷、热进程生成相同首图 key，第二个进程出现持久 cache hit，并确认生成结果和数值不变。

4. Full activation checkpointing 的 DVM 冷编译成本未验收

FullAC 任务在生成 model__2_backward_13 后持续冷编译超过 24 分钟，未进入 step 1，也没有错误堆栈。该问题仅说明该组合的冷编译成本过高，并非正确性失败；FullAC 已移出默认配置门禁。

闭环条件：分别记录冷、热缓存编译时间，并完成至少 10-step FullAC 训练的数值、显存和稳态性能验证。

5. Fusion strict deterministic 下的精确差异尚未定位到具体 lowering

已确认差异至少来自两组独立路径：MHC 主要集中在 layer 0–2；attention 的 loss 差异集中在 ratio=4 的 layer 2/4，ratio=128 的 layer 3/5 还存在 grad norm 差异。现有证据只定位到子模块和层分支，没有继续缩小到具体算子，也没有生产修复。

闭环条件：对命中层做逐阶段前向、反向捕获，定位第一个错误 operator/lowering，再以高精度 reference、同后端重复和 10-step 训练证明修复有效。

非默认或已放弃路线的开放项

6. Triton 外层图加 DVM Sinkhorn 尚无生产 custom-op ABI

/tmp/dvm_custom_op_probe.py 已证明预生成 DVM runtime 可以通过 opaque custom op 保留在外层 Triton fullgraph 中，但原型依赖缓存文件路径和固定 shape，不能作为生产接口。

闭环条件：提供稳定 schema、Fake/Meta、autograd、动态 shape 和打包加载合同，再完成全模型数值与缓存复用验证。

7. strict deterministic 下的 Triton pointwise autotune 没有生产修复

当前只有 /tmp 原型通过 monkeypatch 选择第一个成功编译的配置；它足以验证 custom-op 路线可行，但未建立可长期维护的 benchmark 合同。

闭环条件：torch_npu 为安全 pointwise benchmark 提供 vetted deterministic 合同，或正式提供固定配置路径，并完成冷、热缓存重复验证。

8. Triton 仿射索引 KeyError: y0 + 1 未修复

根因只缩小到 detect_flattened_axis 用完整仿射表达式查询轴变量映射；最小触发算子和通用修复尚未确认。证据位于 outputs/logs/dsv4_debug_layer6_compile_s128_14a/ 和 /tmp/dvm_debug/。

闭环条件：提取最小复现，修复通用仿射索引解析或确认精确 fallback 边界，并验证相关 sibling callsites。该问题属于已放弃的 Triton+fallback 路线，不影响当前全 DVM 默认路径。