前言
在使用昇腾处理器（如 Ascend 910）进行深度学习训练时，我们经常会遇到一个棘手的问题：显存不足 (OOM)。特别是当我们试图训练大语言模型（LLM）或者使用非常大的 Batch Size 来保证 Batch Normalization 的效果时，单卡的显存往往捉襟见肘。
梯度累积（Gradient Accumulation）是一种在不增加物理显存需求的情况下，通过代码逻辑模拟大 Batch Size 的经典技术。
本文将基于MindSpore 2.x的函数式编程范式，手把手教你如何手动实现梯度累积。
核心原理
梯度累积的原理非常直观：
- Mini-Batch 拆分：假设我们想要的 Batch Size 是 64，但显存只能容纳 16。我们将 64 拆分为 4 个 Step（Micro-Batch = 16）。
- 累积梯度：在前 3 个 Step 中，我们只进行前向计算和反向传播，计算出梯度并将其累加到缓存中，但不更新模型参数。
- 参数更新：在第 4 个 Step，我们将当前的梯度也累加进去，然后利用累加后的总梯度更新优化器，并清空梯度缓存。
这样，从数学等价性上，我们完成了一次 Batch Size = 64 的训练，但显存峰值仅为 Batch Size = 16 时的占用量。
代码实战
在 MindSpore 2.x 中，推荐使用函数式变换（Functional Transformations）来构建训练流程。这种方式比旧版的TrainOneStepCell更加灵活，更易于实现梯度累积逻辑。
1. 环境准备与模型定义
首先，导入必要的库并构建一个简单的网络作为演示。这里我们使用一个简单的线性网络，实际场景中可以替换为 ResNet、BERT 或 Llama 等复杂模型。
import mindspore as ms import mindspore.nn as nn import mindspore.ops as ops from mindspore import Tensor, Parameter import numpy as np # 设置运行模式为 Graph 模式（推荐在 Ascend 上使用此模式以获得最佳性能） ms.set_context(mode=ms.GRAPH_MODE, device_target="Ascend") # 定义一个简单的网络 class SimpleNet(nn.Cell): def __init__(self): super(SimpleNet, self).__init__() self.fc = nn.Dense(10, 1) def construct(self, x): return self.fc(x) # 定义损失函数和优化器 net = SimpleNet() loss_fn = nn.MSELoss() optimizer = nn.SGD(net.trainable_params(), learning_rate=0.01)2. 构造梯度累积的核心逻辑
这是本文的重点。我们需要手动维护梯度的累加。
关键点说明：
- Hyper-parameter：
accumulation_steps定义累积步数。 - Loss Scale：为了保证梯度均值正确，计算出的 Loss 通常需要除以
accumulation_steps。 - Parameter 容器：我们需要创建一个与网络参数形状一致的容器来存储累积的梯度。
# 定义累积步数 accumulation_steps = 4 # 1. 创建梯度累积容器 (Accumulated Gradients) # 克隆一份网络参数结构，并将值初始化为0，用于存储累加的梯度 accumulate_grads = optimizer.parameters.clone(prefix="accumulate_grads", init='zeros') # 定义前向计算函数 def forward_fn(data, label): logits = net(data) loss = loss_fn(logits, label) # 关键：Loss 需要根据累积步数进行归一化 return loss / accumulation_steps # 获取梯度计算函数 (value_and_grad) grad_fn = ms.value_and_grad(forward_fn, None, optimizer.parameters) # 定义单步训练逻辑 # 使用 @ms.jit 装饰器加速，编译为静态图 @ms.jit def train_step(data, label, current_step): # 1. 计算当前 micro-batch 的 loss 和梯度 loss, grads = grad_fn(data, label) # 2. 将当前梯度累加到 accumulate_grads 中 # 使用 HyperMap 或 zip 遍历并相加 loss = ops.depend(loss, optimizer.map(ops.assign_add, accumulate_grads, grads)) # 3. 判断是否达到累积步数 if (current_step + 1) % accumulation_steps == 0: # 使用累积后的梯度更新参数 optimizer(accumulate_grads) # 更新完毕后，必须清空梯度累积容器，为下一轮做准备 optimizer.map(ops.assign, accumulate_grads, ops.zeros_like(accumulate_grads)) return loss3. 模拟训练循环
接下来我们构造虚拟数据，运行 Training Loop 来验证逻辑。
# 模拟数据集 def data_generator(): for i in range(20): # 模拟 Batch Size = 16 的数据 data = Tensor(np.random.randn(16, 10).astype(np.float32)) label = Tensor(np.random.randn(16, 1).astype(np.float32)) yield data, label # 开始训练 print(f"Start Training with Gradient Accumulation (Steps={accumulation_steps})...") step_counter = 0 for data, label in data_generator(): loss = train_step(data, label, Tensor(step_counter, ms.int32)) # 打印日志 if (step_counter + 1) % accumulation_steps == 0: print(f"Step {step_counter + 1}: Weights Updated. Current Loss: {loss.asnumpy()}") else: print(f"Step {step_counter + 1}: Gradients Accumulated.") step_counter += 1 print("Training Finished.")性能与注意事项
在昇腾 NPU 上实施梯度累积时，有几点最佳实践需要注意：
- 图模式编译（JIT）：务必使用
@ms.jit装饰train_step函数。由于梯度累积涉及较多的控制流（if 判断），MindSpore 的静态图编译器能够对此进行优化，减少 Host-Device 交互开销。 - Sink Mode（下沉模式）：如果使用
Model.train接口配合Dataset.sink，实现梯度累积会更复杂（通常需要自定义TrainOneStepCell）。上述的自定义循环方式在调试和逻辑控制上更为直观，适合科研和复杂模型开发。 - BatchNorm 的影响：梯度累积虽然模拟了大 Batch 的梯度更新，但 Batch Normalization (BN)层依然是基于 Micro-Batch（小批次）计算统计量的。如果 Micro-Batch 过小（例如 1 或 2），BN 层可能会导致模型难以收敛。这种情况下，建议使用
Group Normalization或Layer Normalization替代 BN，或者冻结 BN 层的统计量。
