内存碎片优化

背景与挑战

在大模型训练阶段，系统需大量内存以承载模型参数、中间计算结果、优化器状态及批量输入数据等关键信息。频繁的内存分配与释放操作可能导致内存碎片化现象，从而降低资源利用率并影响整体性能。

解决方案

通过精细化管理不同生命周期的张量（Tensor），有效缓解内存碎片化问题，提升训练效率。

• 识别与分类：区分长短期生命周期的张量，针对性优化内存管理策略
  • 长生命周期张量：主要包括模型初始化时生成的权重与梯度、前向传播过程中的激活值、未被及时清理的张量（保留至下一迭代步骤）、以及梯度收敛后生成的优化器状态张量。
  • 短生命周期张量：主要包括即时计算张量、辅助计算张量等。
• 隔离与优化：采用多内存池技术，对不同生命周期的张量进行隔离管理
  • 将识别出的长短生命周期张量放入不同的内存池分别管理。
  • 对于长生命周期的大张量，实施精确匹配策略，为其分配与自身尺寸相匹配的内存块，结合多级分配策略，防止相关内存池产生碎片。

使用场景

此优化机制专为训练场景设计，特别适用于内存资源紧张的情况。当用户遇到OOM（Out of Memory，内存不足）错误，且reserved与allocated内存之间存在显著差距（例如，差距超过1GB）时，表明系统内存碎片化严重。启用本特性有助于减少内存碎片，避免内存不足问题，确保训练流程稳定高效。

Tried to allocated 3384.00 MiB (NPU 2; 61.22 GiB total capacity; 53.87 GiB already allocated; 53.87 GiB current activate; 1.59 GiB free; 56.60 GiB reserved in total by PyTorch)

上述输出示例中，reserved与allocated内存之间的差值高达2.73GB，表明内存碎片化程度较高。此时，启动内存碎片优化机制显得尤为重要。

使用方法

--memory-fragmentation

• 使用内存碎片化优化时需注意内存管理与HCCL资源竞争

  在内存资源充足的情况下，该特性倾向于为新的计算任务分配连续的内存空间，而非复用已分配但未完全使用的内存片段。这一策略可能导致在特定情形下与HCCL（Huawei Collective Communication Library，华为集合通信库）发生内存资源竞争。尤其当HCCL面临内存不足时，由于无法有效释放由torch预留的额外闲置空间，可能会触发“HCCL内存不足”的异常错误。

  为规避上述问题，建议通过调用torch_npu.npu.set_per_process_memory_fraction接口，预先设定PyTorch进程所能占用内存的上限比例。例如，设置x=0.94意味着PyTorch最多可占用单张显卡总内存的94%，以此实现对内存资源的有效管控。

  配置示例：

  # MindSpeed/mindspeed/core/memory/memory_fragmentation/memory_recorder.py
  import torch_npu 
  torch_npu.npu.set_per_process_memory_fraction(0.94)

• 内存碎片优化与自适应重计算冲突

  内存碎片优化和自适应选择重计算两个特性均对PyTorch的内存管理机制进行了扩展。为防止潜在的冲突，MindSpeed内置了assert检查机制。

使用效果

启用内存碎片优化后，系统将有效减少内存碎片，提升资源利用效率，确保大模型训练任务的顺利进行。