问题描述

昇腾910环境MindSpore训练大模型时出现内存分配失败，求解决方案

昇腾910平台上使用MindSpore进行大模型训练时遇到了一个棘手的内存问题，尝试了多种方法仍未解决，特来社区求助。

环境配置：

• 硬件：昇腾910 AI处理器

• 操作系统：Ubuntu 20.04 LTS

• MindSpore版本：2.7.1

• CANN版本：8.2.RC1

• Python版本：3.9

• 模型：基于Transformer架构的7B参数大模型

错误现象：

在训练启动后不久，程序报错退出，错误信息如下：

RuntimeError: Malloc device memory failed, drvRetCode=6
Memory resources are exhausted. Try to reduce model size or batch size.

已尝试的解决方案：

1. 检查了CANN与MindSpore版本配套性，确认版本匹配

2. 尝试减小batch_size从32降至8，问题依旧

3. 检查了设备内存使用情况，发现内存占用在训练开始后迅速达到上限

4. 尝试使用MindSpore的内存优化配置，包括设置context.set_context(memory_optimize_level="O1")

5. 验证了单卡推理可以正常运行，但多卡训练时出现此问题

问题分析：

从错误信息看，似乎是设备内存资源耗尽。但奇怪的是，同样的模型在GPU平台上训练时内存使用正常。我怀疑可能是昇腾平台上的内存管理机制或某些算子实现存在差异。

具体疑问：

1. 昇腾910上是否有特殊的内存优化配置或参数需要设置？

2. 是否有针对大模型训练的特定内存管理策略？

问题解答

昇腾910上使用MindSpore训练7B大模型时出现“Malloc device memory failed, drvRetCode=6”错误，这通常是由于设备内存（HBM）资源耗尽所致，尤其是在多卡分布式训练场景下。结合您已尝试的步骤，以下是针对性的解决方案：

核心解决方案

1. 检查并释放被占用的显存

首先确认是否有其他进程占用了昇腾910的显存。执行以下命令检查：

npu-smi info

查看HBM-Usage列，确认目标卡是否已被占用。如果被占用，可通过ps -ef | grep python找到进程ID并用kill -9 <PID>终止，或通过export DEVICE_ID=<空闲卡ID>切换到其他空闲卡。

2. 调整HCCL通信内存

多卡训练时，HCCL（昇腾集合通信库）需要额外的HBM进行通信。如果通信内存不足，会导致drvRetCode=6错误。

• 方法一：在训练脚本中降低max_device_memory的值（例如从默认的30GB降至25GB），为HCCL预留更多空间。

• 方法二：设置环境变量export HCCL_BUFFSIZE=<值>，降低通信域内存占用。建议从较小值（如1048576）开始尝试。

3. 启用更高级的内存优化配置

您已设置memory_optimize_level="O1"，但可以进一步组合以下配置：

import mindspore as ms
# 基础配置
ms.set_context(mode=ms.GRAPH_MODE, 
               device_target="Ascend",
               memory_optimize_level="O1",  # 启用SOMAS内存优化
               max_call_depth=2000,
               enable_reduce_precision=True)

# 内存池与复用配置（MindSpore 2.4+）
ms.set_context(mempool_block_size="1GB",
               enable_mem_reuse=True,
               enable_mem_scheduler=True)

# 如果内存仍然不足，启用offload功能（将部分数据临时卸载到Host内存）
offload_config = {
    "offload_path": "./offload",
    "offload_cpu_size": "512GB",
    "hbm_ratio": 0.9  # 设置HBM使用比例，保留10%给系统
}
ms.set_offload_context(**offload_config)  # 注意：此功能需在图编译等级O0下使用[15](@ref)

4. 启用混合精度训练（AMP）

混合精度训练可显著减少显存占用，同时利用昇腾910的FP16计算优势。

from mindspore import amp
# 在模型和优化器定义后
net = amp.auto_mixed_precision(net, 'O1')  # 'O1'模式自动将部分算子转为FP16
# 或使用更精细的控制
net = amp.build_train_network(net, optimizer, loss_fn, level="O2", loss_scale_manager=loss_scale_manager)