自适应选择重计算

问题分析

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

解决方案

为了实现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）
    • 重计算独立调度流水线（Recompute Independent Pipeline, RIPipe）
    • 激活函数重计算（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或者大于总步数的十分之一）仅产生警告信息，不影响训练进程及结果。

使用效果

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

注意事项

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