如何解决昇腾 NPU 上频繁创建张量导致的内存碎片问题

在将模型从 CUDA 迁移到昇腾 NPU（Ascend）时，很多开发者会遇到一个诡异现象：通过 nvidia-smi 类似的工具观察，显存（HBM）占用并没满，但程序却频繁报出 Out of Memory (OOM)。

这通常不是因为物理内存不足，而是由 内存碎片（Memory Fragmentation） 导致的。本文将揭开昇腾显存管理的黑盒，并给出实操优化方案。

1. 昇腾 NPU 的显存分配机制

昇腾 NPU 的底层架构 CANN（Compute Architecture for Neural Networks）为了追求极致的搬运效率，使用了自研的显存池管理机制。与通用 CPU 内存管理不同，NPU 的显存分配有以下特点：

1. 块对齐（Block Alignment）：NPU 要求张量的起始地址和大小通常需要按 512 字节或更高字节对齐。
2. 显存池复用：PyTorch NPU 插件（torch_npu）会预先申请大块显存作为池，再切分给用户。
3. 非紧凑分配：如果代码中存在大量生命周期不一致、形状大小不一的临时张量，显存池会产生大量无法合并的“空隙”。

2. 典型错误示例：频繁创建临时张量

以下是一个导致内存碎片的典型负面示例。在训练循环中，如果根据输入动态地创建 mask 或进行不规则的切片操作，会迅速耗尽连续内存空间。

import torch
import torch_npu

def bad_loop(model, data_loader):
    for i, (input_data, target) in enumerate(data_loader):
        # 错误做法：在循环中频繁创建不同形状的临时 Tensor
        # 每次循环都会在显存池中挖一个新坑，且大小可能不同
        mask = torch.randn(input_data.shape, device='npu:0') > 0.5

        # 复杂的不规则切片也会产生临时碎块
        temp_result = input_data[mask]

        output = model(temp_result)
        loss = criterion(output, target)
        loss.backward()

3. 解决策略与实操方案

策略 A：张量复用（Static Buffer）

最有效的方法是提前分配好 Buffer，在循环中只进行数据填充。

# 优化方案：预分配固定大小的 Buffer
buffer_mask = torch.empty((BATCH_SIZE, CHANNELS, H, W), device='npu:0')

def optimized_loop(model, data_loader):
    for input_data in data_loader:
        # 使用 inplace 操作或 copy_ 来复用内存
        # .bernoulli_ 是原位操作，不会触发新显存申请
        buffer_mask.resize_(input_data.shape).bernoulli_(0.5)

        # 尽量保持 Shape 为 16 或 32 的倍数，利于对齐
        output = model(input_data * buffer_mask)

策略 B：手动触发 Cache 清理

如果业务逻辑必须产生大量临时变量，可以在迭代周期结束时手动清理显存池碎片。

# 在每个 Epoch 结束或每 100 个 Step 清理一次
if step % 100 == 0:
    torch_npu.npu.empty_cache()

策略 C：环境变量调优

可以通过 CANN 的环境变量优化显存块的切分粒度。例如，通过设置 PYTORCH_NPU_ALLOC_CONF 来限制小块内存的拆分。

# 在启动脚本中设置，减少小块碎片的产生
export PYTORCH_NPU_ALLOC_CONF=\"max_split_size_mb:128\"

4. 总结

昇腾 NPU 的显存管理对地址对齐和连续性极其敏感。解决碎片的关键在于：
1. 能复用则复用：尽量避免在 for 循环内部 new 张量。
2. 形状对齐：输入形状尽量保持稳定，避免 Padding 频繁变动。
3. 监控现状：使用 torch_npu.npu.memory_summary() 打印显存状态，观察 segment 的数量，如果 segment 数量持续攀升而 active 内存没变，即说明碎片化严重。”, “tags”: [“npu”, “昇腾”, “pytorch”, “内存优化”, “CANN”], “summary”: “本文深入剖析了昇腾 NPU 显存池的工作机制，揭示了频繁创建变长张量引发内存碎片导致 OOM 的原因，并提供了预分配与显存清理等实操解决方案。”}