MindSpore训练模型时报错”RuntimeError: Memory allocation failed”分析和解决

经验分享 Tech Blogs 模型训练-Model Training
1

1 系统环境

  • 硬件环境(Ascend/GPU/CPU): Ascend910
  • MindSpore版本: mindspore=1.8.1
  • 执行模式（PyNative/ Graph）: 不限
  • Python版本: Python=3.9
  • 操作系统平台: Linux

2 报错信息

2.1 问题描述

自定义的Transformer模型（基于MindSpore 1.8.1版本），模型结构相对复杂，包含12层编码器和6层解码器，隐藏维度为1024。在训练到第3个epoch时，程序突然崩溃，报错信息显示内存分配失败。奇怪的是，前两个epoch训练完全正常，loss也在稳步下降。

2.2 脚本信息

import mindspore as ms
from mindspore import nn, ops
from my_transformer import CustomTransformer

# 模型初始化
model = CustomTransformer(
    src_vocab_size=50000,
    tgt_vocab_size=50000,
    hidden_size=1024,
    num_encoder_layers=12,
    num_decoder_layers=6,
    num_heads=16
)

# 优化器配置
optimizer = nn.Adam(model.trainable_params(), learning_rate=0.0001)

# 训练步骤
def train_step(data, label):
    def forward_fn(data, label):
        logits = model(data)
        loss = loss_fn(logits, label)
        return loss, logits
    
    grad_fn = ms.value_and_grad(forward_fn, None, optimizer.parameters)
    (loss, _), grads = grad_fn(data, label)
    optimizer(grads)
    return loss

2.3 报错信息

[ERROR] DEVICE(123456,ffffffff): RuntimeError: Memory allocation failed at /home/mindspore/mindspore/ccsrc/backend/common/mem_reuse/mem_dynamic_allocator.cc:125
Total memory requested: 15.8 GB
Available memory: 3.2 GB
Batch size: 64
Current epoch: 3, step: 1250

3 根因分析

内存碎片积累：MindSpore使用动态内存管理，随着训练步骤增加，频繁的内存分配和释放会产生显存碎片，导致实际可用连续内存减少。
中间激活值累积：Transformer模型在训练过程中会保存中间激活值用于反向传播，随着训练深入，这些激活值可能累积占用大量内存。
HCCL通信内存占用：在分布式训练中，每个HCCL通信组默认占用200MB内存，随着训练进程运行，通信开销可能导致内存不足。

4 解决方案

1、检查并清理占用进程

# 检查是否有其他进程占用NPU
ps -ef | grep python
# 如果有其他MindSpore进程在运行，请先终止它们

2、优化HCCL通信内存

# 设置HCCL通信缓冲区大小，减少内存占用
export HCCL_BUFFSIZE=1024  # 根据实际情况调整，单位MB

3、配置MindSpore内存管理参数

import mindspore as ms

# 设置最大设备内存，为HCCL预留空间
ms.set_context(max_device_memory="28GB")  # 从30GB降低到28GB

# 启用内存复用优化
ms.runtime.set_memory(optimize_level='O0')  # O0:动态内存复用，O1:静态内存复用

# 或者使用更细粒度的内存池配置
ms.runtime.set_memory(init_size='10GB', increase_size='2GB', max_size='28GB')

4、使用混合精度减少内存占用

# 启用梯度检查点（激活值重计算）
from mindspore.nn.transformer import RecomputeConfig
recompute_config = RecomputeConfig()
recompute_config.recompute = True
recompute_config.recompute_slice_activation = True

# 在自定义Transformer中应用
class CustomTransformer(nn.Cell):
    def __init__(self, ...):
        super().__init__()
        self.recompute = recompute_config.recompute
        # ... 其他初始化

# 使用混合精度训练减少内存占用
from mindspore import amp
net_with_loss = amp.build_train_network(
    train_net, 
    optimizer=optimizer,
    loss_scale_manager=loss_scaler,
    amp_level='O2'  # 自动混合精度
)

5、使用数据预取和缓存优化

# 确保数据加载器正确释放内存
def create_dataset(batch_size=32):  # 进一步降低batch size
    dataset = ...
    # 使用数据预取和缓存优化
    dataset = dataset.prefetch(buffer_size=4)
    dataset = dataset.cache()
    return dataset
MindSpore训练自定义Transformer模型时报错"RuntimeError: Memory allocation failed"