自适应选择重计算

背景与挑战

Megatron重计算技术是优化内存使用的关键手段，通过避免冗余数据存储以减少内存占用。然而，传统的重计算策略往往预设固定模式，未能动态响应实际内存需求，限制了内存资源的高效利用。

解决方案

为了实现NPU内存资源的最大化利用并显著提升模型训练效率，我们引入了自适应选择重计算特性。这一创新机制能够智能调整训练过程中的内存分配，依据当前内存状况动态选择最佳的重计算策略。

自适应选择重计算特性由三大关键组件构成：

重计算策略搜索：根据实时内存状态，智能筛选最适宜的重计算方案。
SwapManager功能集成：确保关键张量能在CPU与NPU间无缝迁移，预防因内存溢出而导致的训练中断。
内存管理：适配Ascend Extension for PyTorch的NPU PluggableAllocator接口，实现OOM（Out of Memory）异常的及时响应与处理。

自动选择重计算策略流程如下图所示。

图1 自动选择重计算策略流程

图2 SwapManager管理流程

使用场景

特别适用于训练阶段，尤其是当全重计算模式下NPU内存仍有较大剩余空间，追求更优训练性能的场景。

使用方法

环境配置
 设置如下环境变量：
```
export ADAPTIVE_RECOMPUTING=1
```
脚本调整
在训练脚本中加入如下参数：
```
--adaptive-recompute-device-swap
```
自适应选择性重计算与配置兼容性：
- 重计算方法与粒度配置失效
   启用自适应选择性重计算功能后，原有的完全重计算相关配置选项“--recompute-method”与“--recompute-granularity”将自动失效。这意味着系统将根据内部算法动态决定重计算策略，无需手动指定重计算方法或粒度。
- 特性兼容性限制
   当前版本中，自适应选择性重计算功能与以下特性存在兼容性限制：
  完全重计算（Full Recomputation）
  激活函数重计算（Activation Recomputation）
  
  这些特性与自适应选择性重计算的机制存在潜在冲突，可能导致性能下降或系统不稳定。在启用自适应选择性重计算时，请确保上述特性处于禁用状态，以维护系统的稳定运行与优化效果。
可选参数
- 如下参数用于手动指定训练内存大小，范围为0至设备最大内存，单位为MB。
```
--adaptive-recompute-device-size   N 
```
  低于0的值视为无效，超出范围则自动采用设备最大内存。设定过低的内存可能使性能逼近全重计算水平。如遇内存溢出（OOM，Out of Memory）情况，系统将提示相应错误。此时，用户需重新设定一个更为合理的内存分配值，进而重启模型训练流程。对于寻求性能最优化的高级用户，我们推荐采用二分查找算法（Binary Search Algorithm）来精确确定最优内存配置。通过迭代缩小测试范围，逐步逼近最适宜的内存阈值，从而在保障训练稳定性的同时，最大化利用可用资源。然而，鉴于该方法对用户的专业知识有一定要求，初学者或不熟悉此特性的用户应谨慎操作，避免不当配置带来的训练中断风险或效率低下。
- 如下参数用于设置停止收集训练信息的步数，默认为10步，推荐值大于5。
```
--adaptive-recompute-profiling-step    N
```
  非标准设置（小于5或者大于总步数的十分之一）仅产生警告信息，不影响训练进程及结果。

鉴于“自适应选择重计算”与“内存碎片优化”两项特性均涉及对PyTorch内存管理机制的深度调整，同时启用这两项功能可能引发操作上的不兼容性。为确保软件架构的稳定性和安全性，MindSpeed已预先部署了断言检测机制（assertion check），以防止此类冲突状况的发生。我们建议用户根据具体的应用场景和资源需求，谨慎选择其中一项特性进行启用，以避免潜在的系统异常。

使用效果

相较于全重计算模式，Llama2-7B模型训练场景下性能提升约16.29%，Llama2-13B模型训练场景下性能提升约12.05%，显著优化了训练效率与资源利用。

父主题： 内存资源优化