300i duo算力卡显存优化实战：基于npu-smi工具的高效监控与调优

1次阅读

共计 2486 个字符，预计需要花费 7 分钟才能阅读完成。

在深度学习训练任务中，300i duo 算力卡因其强大的 NPU（Neural Processing Unit）加速能力而备受青睐。然而，在实际应用中，显存（Memory）管理问题常常成为性能瓶颈。无论是 CV（Computer Vision）还是 NLP（Natural Language Processing）任务，开发者都可能遇到显存不足（OOM, Out Of Memory）报错、显存利用率波动大等问题。这些问题不仅影响训练效率，还可能导致任务中断，浪费宝贵的时间和资源。

与 NVIDIA 的 nvidia-smi 工具类似，华为的 npu-smi 工具是专为 NPU 设计的监控和管理工具。两者在功能上有一些差异：

nvidia-smi主要用于监控 GPU 的状态，包括显存使用情况、温度、功耗等。
npu-smi则专注于 NPU 的状态监控，提供了更多针对 NPU 优化的功能，如显存碎片整理、多进程显存共享等。

npu-smi monitor
功能：实时监控 NPU 的显存使用情况。
常用参数：
- -i：指定设备 ID。
- -d：监控间隔时间（秒）。
- -m：显示显存使用详情。
示例：npu-smi monitor -i 0 -d 1 -m
npu-smi config
功能：配置 NPU 的运行参数。
常用参数：
- --memory-optimize：开启显存优化模式。
- --fragmentation-threshold：设置显存碎片整理阈值。
示例：npu-smi config --memory-optimize on --fragmentation-threshold 0.8

以下是一个 Python 脚本示例，用于实时监控显存占用并在超过阈值时触发告警：

import subprocess
import time

def monitor_memory(device_id, threshold_gb, check_interval=1):
    while True:
        # 使用 npu-smi 获取显存使用情况
        cmd = f'npu-smi monitor -i {device_id} -m'
        result = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, text=True)
        output = result.stdout

        # 解析显存使用量
        memory_used = float(output.split('Memory Used:')[1].split('GB')[0].strip())

        # 检查是否超过阈值
        if memory_used > threshold_gb:
            print(f'警告：显存使用量超过阈值！当前使用量：{memory_used} GB')

        time.sleep(check_interval)

# 示例：监控设备 0，阈值设置为 8GB
monitor_memory(device_id=0, threshold_gb=8)

对于大模型训练任务，显存不足是一个常见问题。可以通过分段执行的方式解决：

import torch

# 假设模型太大，无法一次性加载到显存
model = LargeModel()
optimizer = torch.optim.Adam(model.parameters())

def train_in_segments(data_loader, segment_size):
    for batch_idx, (data, target) in enumerate(data_loader):
        # 分段处理数据
        for segment_start in range(0, len(data), segment_size):
            segment_end = segment_start + segment_size
            segment_data = data[segment_start:segment_end]
            segment_target = target[segment_start:segment_end]

            # 前向传播
            output = model(segment_data)
            loss = compute_loss(output, segment_target)

            # 反向传播和优化
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            # 清除中间变量，释放显存
            del segment_data, segment_target, output, loss
            torch.cuda.empty_cache()

显存碎片化（Memory Fragmentation）会导致显存利用率下降。以下是一些优化建议：