自适应内存优化

问题分析

在大模型训练中，重计算特性可以有效地减少显存使用，但是策略较为固定，无法最大限度使用显存资源。

为了在最大限度地利用NPU显存的同时，提高模型训练的性能，我们支持通过自动调整训练内存大小来自动选择计算策略。这一特性称为自适应内存优化。

自适应内存优化设计主要包括：自适应无损swap、策略生成、策略搜索、SwapManager功能以及内存管理等几部分。

自适应内存优化策略流程如下图所示：

SwapManager功能需要内存管理适配PTA的NPUPluggableAllocator接口拦截OOM，让SwapManager功能可以介入，流程如下图所示：

该特性主要用于训练场景，如果用户发现开启了全重计算功能后， NPU显存剩余较多，此时若想充分利用显存，从而提高训练性能，可以考虑开启该特性。

在训练脚本中添加--adaptive-memory-optimization

注意：

当前自适应内存优化与全重计算、自适应选择重计算、预取特性swap-attention、 recompute-in-bubble等不兼容。
目前自适应内存优化已能够管理一部分使用torch.autograd.Function修饰的auto_function类
- 在调用auto_function的文件中添加 from mindspeed.core.memory.adaptive_memory.adaptive_memory_function import adapt_mem_func_wrapper
- 将 auto_function.apply(*args) 修改为 adapt_mem_func_wrapper(auto_function, *args)
- 以mindspeed.moe.pipe_experts中的PipeExpert类的调用为例，在mindspeed.moe.moe_layer文件中添加from mindspeed.core.memory.adaptive_memory.adaptive_memory_function import adapt_mem_func_wrapper，将expert_output = PipeExpert.apply(*args)修改为expert_output = adapt_mem_func_wrapper(PipeExpert, *args)

这里的gpt-175B是经过裁剪后的

gpt-175B:

特性	参数	NPU卡数	TFLOPs	收益
adaptive-memory-optimization	seq-length=8192、micro-batch-size=10、global-batch-size=40、TP=8、PP=1、DP=1、CP=1、NL=8、hidden-size=12288	8卡（单机）	165.90	-
全重计算	seq-length=8192、micro-batch-size=10、global-batch-size=40、TP=8、PP=1、DP=1、CP=1、NL=3、hidden-size=12288、recompute-num-layers=3	8卡（单机）	145.93	13.68%

特性	参数	NPU卡数	TFLOPs	收益
adaptive-memory-optimization	seq-length=8192、micro-batch-size=3、global-batch-size=9、TP=2、PP=4、DP=1、CP=1、NL=8、hidden-size=12288	8卡（单机）	76.30	-
全重计算	seq-length=8192、micro-batch-size=3、global-batch-size=9、TP=2、PP=4、DP=1、CP=1、NL=8、hidden-size=12288、recompute-num-layers=1	8卡（单机）	66.50	14.17%

特性	参数	NPU卡数	TFLOPs	收益
adaptive-memory-optimization	seq-length=8192、micro-batch-size=2、global-batch-size=8、TP=2、PP=4、VPP=2、DP=1、CP=1、NL=8、hidden-size=12288	8卡（单机）	86.10	-
全重计算	seq-length=8192、micro-batch-size=2、global-batch-size=8、TP=2、PP=4、VPP=2、DP=1、CP=1、NL=8、hidden-size=12288、recompute-num-layers=1	8卡（单机）	75.10	14.65%