Runtime 在整个流程里扮演什么角色

在说 Runtime 之前，先回顾一下你的 PyTorch 模型在昇腾 NPU 上跑的时候，数据是怎么流的：

你的 PyTorch 代码
        ↓ torch_npu 插件（Framework Adaptor）
    GE 图引擎（编译优化）
        ↓ 优化后的计算图
    Runtime 运行时 ← 这里
        ↓ 具体算子执行指令
    昇腾 NPU 硬件（AI Core + HBM）

GE 把计算图优化完了，Runtime 接过来，负责真正把这些算子扔到硬件上去执行。你可以理解为：GE 是图纸设计师，Runtime 是施工队长——图纸画得再好，施工队长安排不对，工地还是乱成一锅粥。

Runtime 的上游是 GE 的输出（优化后的静态图），下游是昇腾 NPU 的硬件接口。它的核心职责是：把 GE 给的计算图，转换成可执行的任务序列，然后调度到硬件上去跑，同时管好硬件资源（显存、计算单元、多进程/多任务并发）。

Runtime 管的三件事

Runtime 的工作很多，但核心可以分成三块：资源管理、任务调度、执行协调。一件一件说。

第一件事：资源管理——显存怎么分配

你把模型扔到 NPU 上之前，Runtime 要先给模型分配一块显存来存放权重参数、中间激活值、算子临时 buffer。这个过程叫 HBM 显存分配。

这块工作听起来简单，但坑很多。

第一个坑：分配粒度

GPU 的显存分配相对灵活，你申请多少字节，系统给你分配多少。昇腾 NPU 的 Runtime 对显存分配有一些对齐要求——不是说你要 100MB 就给你 100MB，可能是 128MB 一块，不够的部分也要算。这意味着你用 nvidia-smi 看显存的时候，看到的数字总是 2 的幂次，不是因为硬件要求，是因为 Runtime 的分配策略。

第二个坑：多任务显存抢占

如果你同时跑了两个推理任务在同一张卡上，Runtime 有一个显存分配策略的问题。两个任务都申请显存，Runtime 要决定怎么分：是静态切分（一人一半）、还是动态抢占（谁先申请谁先用）、还是配额制（每个任务最多用多少）。

默认的动态抢占策略在大多数场景下没问题，但如果你的两个任务一个是短期小请求、一个是长期大模型推理，抢占策略可能会导致短期请求因为显存不够而排队。

解法是用 显存配额（Memory Quota） 限制每个任务的最大显存占用：

import torch_npu

# 为当前进程设置最大显存占用为卡总显存的 80%
torch.npu.set_per_process_memory_fraction(0.8)

# 为特定模型单独分配一个显存池
with torch.npu.npu.stream(priority=0):
    output = model(input)   # 这个流用 80% 显存

第三个坑：显存泄漏

深度学习模型在推理过程中，如果某些中间张量没有被正确释放（Python 层面没有删除引用、Runtime 层面的资源计数没有归零），显存会越用越多。Runtime 本身有 GC 机制，但如果你在循环里不断创建新张量而不手动 del 掉旧的，GC 来不及回收，显存会爆。

排查显存泄漏有一个技巧：跑推理的时候，用 profiling 看显存使用曲线：

import torch
from torch_npu.contrib import profile

with profile(profile_dir="./runtime_profile", with_stack=True):
    for i in range(100):
        output = model(inputs[i])
        if i % 10 == 0:
            # 每 10 步打印一次显存占用
            memory_allocated = torch.npu.memory_allocated() / 1024**3
            memory_reserved = torch.npu.memory_reserved() / 1024**3
            print(f"Step {i}: allocated={memory_allocated:.2f}GB, reserved={memory_reserved:.2f}GB")

如果 reserved 一直涨，说明显存泄漏了。如果 allocated 涨但 reserved 不涨，说明是 Runtime 在复用 buffer，属于正常现象。

第二件事：任务调度——算子怎么排队执行

显存分配完了，下一步是执行。Runtime 有一个任务调度器，决定每个算子什么时候执行、按什么顺序执行。

这里有个关键概念：Stream（流）。

你可以把 NPU 想象成一个工厂，工厂里有很多工作台（AI Core），每个工作台可以同时处理一个任务。但工作台不是随时都有空，你需要有人安排任务的排队和调度。Stream 就是干这个的——它是一条任务队列，告诉你某个算子该排到哪个位置、什么时候轮到它执行。

在 GPU 上，默认情况下所有算子都在同一个 Stream 里按顺序执行。你也可以创建多个 Stream，让某些算子并行排队：

import torch

# 创建两个独立流
stream1 = torch.npu.Stream(0)   # 流 0
stream2 = torch.npu.Stream(1)   # 流 1

# 在流 1 上执行计算 A
with torch.npu.stream(stream1):
    result_a = model_a(input_a)

# 在流 2 上执行计算 B（跟 A 并行，不等待 A 完成）
with torch.npu.stream(stream2):
    result_b = model_b(input_b)

# 等待两个流都完成
torch.npu.synchronize()

为什么要用多 Stream？通常有几种场景：

• Preload 模式：流 1 跑当前帧推理的同时，流 2 已经在加载下一帧的数据，算和搬重叠，减少等待
• 异构计算：某些算子（比如数据预处理）适合在 CPU 上做，跟 NPU 计算并行，减少总延迟

昇腾 NPU 的 Runtime 支持多 Stream 调度，但需要注意一点：不同 Stream 之间如果涉及到同一个张量的读写，需要显式加同步（synchronize），否则可能出现 data race。这个跟 CUDA 的多 Stream 逻辑是一样的。

第三件事：执行协调——Host 和 Device 怎么配合

昇腾 NPU 有两个部分：Host（CPU 端）和 Device（NPU 端）。你的 Python 代码跑在 Host 上，真正的计算跑在 Device 上。Runtime 负责这两端之间的协调：数据什么时候搬过去、结果什么时候搬回来、搬的时候走哪条总线、带宽够不够。

这个过程叫 数据搬运（Data Transfer），是影响延迟的重要因素。

一个典型的坑是：频繁的小数据搬运。假设你每次推理只需要一个 1MB 的输入，但你每帧都从 CPU 搬到 NPU，算完再搬回来，1000 帧就是 2GB 的搬运量。虽然总线带宽不小，但搬运本身有启动开销（setup time），如果每次只搬 1MB，总线的启动开销会抵消掉带宽的优势。

解法是批量搬运：把多帧数据攒在一起，一次性从 CPU 搬到 NPU，NPU 上批量算完，再一次性搬回来。Runtime 提供了 torch.npu.H2D 和 torch.npu.D2H 的异步拷贝接口：

import torch

# 同步方式（每帧单独搬）
for frame in frames:
    frame_npu = frame.npu()   # CPU → NPU
    result = model(frame_npu)
    result_cpu = result.cpu() # NPU → CPU

# 异步批量方式（攒在一起搬）
batch_size = 8
for i in range(0, len(frames), batch_size):
    batch = torch.stack(frames[i:i+batch_size]).npu()  # 批量 CPU → NPU
    results = model(batch)                               # 批量推理
    results_cpu = results.cpu()                          # 批量 NPU → CPU

Runtime 在做数据搬运的时候会尽量用 DMA（Direct Memory Access，直接内存访问）来减少 CPU 的参与。但有些场景下（尤其是小数据量），同步拷贝反而比异步 DMA 更快，因为异步 DMA 有调度开销。实际使用时可以 profiling 看哪种方式更快。

一个真实场景：多卡训练时 Runtime 在干什么

说完三件事，来一个具体场景，看看 Runtime 在多卡训练时是怎么工作的。

import torch
import torch.distributed as dist
import torch_npu

# 初始化分布式训练
dist.init_process_group(backend='hccl', init_method='env://')
local_rank = int(os.environ['RANK'])
torch.npu.set_device(local_rank)

model = MyModel().npu(local_rank)
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])

# 训练循环
for data, target in dataloader:
    data = data.npu(local_rank)
    target = target.npu(local_rank)

    optimizer.zero_grad()
    output = model(data)
    loss = criterion(output, target)
    loss.backward()          # 反向传播
    optimizer.step()         # 参数更新

在这个多卡训练场景下，Runtime 要处理以下事情：

第一步：进程资源分配

每个进程（对应一张卡）调用 torch.npu.set_device 时，Runtime 会在那张卡上建立一个独立的资源上下文。这个上下文里有独立的显存池、独立的任务队列。不同进程之间不会互相干扰。

第二步：算子执行

每个进程内部的 model(data) 调用，在本地的 Runtime 调度下执行。Runtime 按照 GE 优化后的图，把每个算子排到 Stream 里执行。

第三步：梯度同步（多卡通信）

loss.backward() 完成后，每个进程的梯度在各自卡上算出来了。但模型参数需要同步，不然每张卡的权重会越来越不一致。

这一步调用的是 HCCL（集合通信库）。Runtime 负责把梯度同步的请求发到 HCCL，HCCL 在多张卡之间做 AllReduce，把每张卡的梯度累加、再平均、分发回去。同步完成后，每张卡的模型参数保持一致。

坑点：梯度同步卡死

多卡训练最常踩的一个坑是梯度同步卡死。常见原因有两个：

1. 某张卡的某个算子执行超时：比如某个自定义算子在某张卡上跑得特别慢，其他卡已经到同步点了，这张卡还没算完，所有卡都在等它。解法是用 torch.distributed.barrier() 加超时检测。
2. 通信带宽被打满：如果模型的梯度 tensor 很大（模型参数多、batch 大），AllReduce 的通信时间会很长。如果通信时间和计算时间没做好 pipeline，部分卡会在通信阶段阻塞。解法是调整通信的时机，或者用梯度压缩减少通信量。

怎么调试 Runtime 相关的问题

如果你的模型跑起来有问题（OOM、卡死、延迟高），有几个实用的调试方法：

方法一：打开 Runtime 日志

CANN 的 Runtime 有详细的执行日志，打开方式：

# 在运行前设置环境变量
export ASCEND_GLOBAL_LOG_LEVEL=1
export HCCL_DEBUG_MODE=1

python your_script.py

日志里会记录每个算子的执行时间、显存分配释放、Stream 调度情况。

方法二：用 profiling 工具看算子耗时分布

from torch_npu.contrib import profiler

with profiler.profile(
    activities=[
        profiler.ProfilerActivity.CPU,
        profiler.ProfilerActivity.NPU,
    ],
    record_shapes=True,
    profile_memory=True,
    with_stack=True,
) as prof:
    model(input)

# 打印耗时 Top 10 的算子
print(prof.key_averages().table(sort_by="npu_time_total", row_limit=10))

方法三：检查显存分配详情

print(torch.npu.memory_summary(device=0))

这条命令会打印当前卡的显存使用详情：总共分配了多少、当前占用多少、峰值多少、各个算子的显存分配记录。

总结

Runtime 运行时是 CANN 里直接跟硬件打交道的那一层。它的核心工作有三件：资源管理（显存怎么分配）、任务调度（算子怎么排队）、执行协调（Host 和 Device 怎么配合）。

理解 Runtime 的工作方式，对排查以下几类问题特别有用：

• OOM 问题：看显存分配策略是不是合理，多任务是否有抢占
• 延迟问题：看数据搬运是否成为瓶颈，Stream 调度是否充分并行
• 多卡训练卡死：看梯度同步是不是在正确时机发生，通信带宽是否被打满

GE 负责把图优化好，Runtime 负责把优化后的图执行好。两者配合好，昇腾 NPU 的性能才能真正跑出来。去看 cann-recipes-infer 和 runtime 源码，能帮助你更深入地理解这两层是怎么协作的。

本篇文章涉及的相关仓库：

• Runtime 源码：https://atomgit.com/cann/runtime
• 推理最佳实践：https://atomgit.com/cann/cann-recipes-infer
• 集合通信库源码：https://atomgit.com/cann/hccl