MindSpore训练自定义Transformer模型时报错"RuntimeError: Memory allocation failed"

[chengxiaoli (Cheng_li)]

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 根因分析

• 此处由用户填写。

4 解决方案

• 此处由用户填写。
• 包含文字方案和最终脚本代码 请将正确的脚本打包并上传附件。