Megatron重计算

背景与挑战

在大模型的训练流程中，传统的实践要求存储前向传播阶段产生的激活值，以供后续反向传播过程中的梯度计算使用。这一需求导致了激活值保存数量随模型深度线性增长的现象，显著加剧了对硬件内存资源的压力。

解决方案

为应对上述挑战，提出了重计算策略。具体而言，在前向传播与损失函数计算阶段，即时释放不再需要的激活值内存空间，仅在反向传播时根据需要重新计算激活值。此方法通过有效缩短激活值的生命周期，显著减轻了内存负担，提升了整体的资源利用效率。

使用场景

使用方法

• 选择性重计算

  脚本中添加如下配置：

  --recompute-activations    # 开启选择性重计算      

• 完全重计算
  脚本中添加如下配置：

  --recompute-granularity full    # 开启完全重计算
  --recompute-method uniform     # 配置具体重计算方式

  或

  --recompute-granularity full    # 开启完全重计算
  --recompute-method block     # 配置具体重计算方式

  --recompute-method可配置参数uniform或block，任选其一：

  • --recompute-method uniform：将Transformer层均匀划分组（每组大小--recompute-num-layers），按组存储输入和激活值。
  • --recompute-method block：将前--recompute-num-layers个transformer层重计算，剩余层不进行重计算。

同时配置--recompute-activations和--recompute-granularity full生效选择性重计算。

使用效果

采用重计算策略后，系统能够实现以下显著改进：

• 通过避免长时间保留大量中间计算结果，大幅降低了对内存的需求。
• 尽管重计算激活值会带来额外的计算成本，但综合考虑，这种方法在保持模型训练质量的同时，实现了更高效的资源管理。