Colossal-AI高效异构内存管理系统介绍

欢迎关注我的公众号 [极智视界]，获取我的更多经验分享

大家好，我是极智视界，本文来介绍一下 Colossal-AI高效异构内存管理系统。

邀您加入我的知识星球「极智视界」，星球内有超多好玩的项目实战源码下载，链接：t.zsxq.com/0aiNxERDq

首先需要了解一下异构内存中的数据移动，由于 GPU 的内存容量有限，一般没有办法直接容下大模型，这样的话可以使用 CPU 的内存以及 SSD 硬盘来共同存储大模型。在训练深度学习模型的时候，可以在运行时适当地把数据在不同的设备中进行交换，也就是所谓的 data swapping，理想情况下只需要将需要计算的数据保存在 GPU 中，从而可以训练更加大的模型。

微软在自家的大模型加速框架 deepspeed 中提出了 ZeRO (Zero Redundancy Optimizer) 的异构内存管理解决方案。在数据并行的训练中，若不管三七二十一直接加载到每个 GPU，则会存在最大的内存冗余。deepspeed 发现了这个缺点，将模型的参数、梯度和优化器状态进行切分并存储在不同的 GPU 上，从而能够消除内存冗余。除此之外，deepspeed 还支持将模型参数卸载到 CPU 和 NVMe 硬盘上，从而使 GPU 能够容纳更大的模型。

来介绍一下 ZeRO，根据切分状态的不同，可以将 ZeRO 分为三个阶段：

• 第一个阶段：只切分优化器状态，以 7.5B 的 ChatGPT 模型为例，内存降了 4 倍；
• 第二个阶段：切分优化器状态和梯度，这个时候大约能够降低 7 倍左右的内存使用；==> 能够做到和 PyTorch DDP 差不多的运行速度；
• 第三个阶段：同时切分参数、优化器状态和梯度，这种情况下大约能够降低 63 倍的内存使用；==> 能够极致地降低内存使用；==> Colossal-AI 的高效异构内存管理系统就是基于这个进行开发的；

将 GPU 所使用的内存分为两类，第一类为模型数据，就是下图中黄色和蓝色的部分，包括模型的权重、梯度还有优化器状态所占用的内存；第二类为非模型数据，就是下图中的绿色的部分，主要是计算中的临时张量所占用的内存。deepspeed 使用一种静态的内存管理策略，它会在张量在 GPU 中完成计算后按照一定的比例卸载到 CPU 或者是 NVMe 硬盘上。但这种策略会导致在训练的不同阶段，GPU 和 CPU 上的内存空间可能都会存在一定的内存浪费，就像如下图所示的红圈圈，所以这个策略还是存在较大的改进空间。

Colossal-AI 的内存管理与 deepspeed 不同，Colossal-AI 的内存管理系统是建立在基于块的内存管理之上。如下图，将一组连续的参数存储在一个块中。

Colossal-AI 将这个内存管理系统称之为 Gemini，它利用块的机制来管理张量的移动。由于多个张量被放进一个块中，因此在进行数据移动时，它的带宽利用率会更高，其中 ChunkMgr 会记录块的移动。

使用一个有限状态机来定义张量在训练过程中的生命周期，这样张量的每一个状态都会对应一个目标设备，状态的转化意味着张量所在的设备的转换。

使用 torch function 来劫持每一个 torch 的算子，在每个算子的前向和反向之前，就是下图的 pre-forward 和 pre-backward，在这个阶段会进行块的切分或合并、数据移动、动态转移和 CUDA 内存采样。然后在每个算子的前向和反向之后，也就是下图中的 post-forward 和 post-backward，在这个阶段会进行块的状态转移和 CUDA 内存采样。

在训练的前几步中，Gemini 会追踪训练过程中的不同算子的内存使用情况，并且会动态地生成一个策略来决定何时应该将张量卸载到 CPU 内存上，当 GPU 内存不足时，它可以将更多的块卸载到 CPU 上，以避免 GPU 计算内存不足。如下图，deepspeed 是采用一种静态内存管理策略，按照一定比例将某一些张量放在 GPU 上，而另一些张量就放在 CPU 内存中，当这个比例不合适的时候，可能就会出现由于 CPU 内存不够而导致无法训练。所以为了使训练效率最高，deepspeed 是需要手动地去调节这个比例的，在使用上就会产生一定的试错成本。而 Gemini 是采用一种动态的内存管理策略，通过内存采样和 chunkmanage，可以计算出每一个算子所对应的非模型数据所占用的内存大小，然后根据总内存大小和现在已经使用的内存大小计算出能够存放参数的模型大小，从而生成一个智能的内存移动策略，可以动态地将张量合理地分配到 CPU 上。

另外，为了能够最大限度地减少内存移动，采用了最佳页面替换算法，也就是非常经典的 OPT 算法。当需要卸载张量到 CPU 上的时候，会优先卸载那些不会再被使用或最后才会被使用的块，这样可以使内存移动的量是最少的。当 CPU 内存不足的时候，可以将更少的块卸载到 CPU 上，也就是说这样更多的块将保留在 GPU 上。如下图左，deepspeed 在 ADAM 优化的过程中，可能会出现 CPU 内存不足的情况，但是 Gemini 可以动态地将这些内存保留在 GPU 上，从而避免了 CPU 内存的溢出，而且还可以加快训练速度，这样就可以充分利用 GPU 和 CPU 的所有内存进行大规模的训练了。

Gemini 还对比了两种数据块的切分方案，如下图左边的非对称切分方案和右边的对称切分方案，最终选择了对称切分方案。来解释一下，非对称切分方案就是每一块 GPU 保存一组完整的块；而对称切分方案就是将每一个数据块均匀地进行切分，然后每块 GPU 保存其中的一份，这样就允许所有 GPU 进行数据移动，可以充分利用 PCIE 带宽，提高 GPU 和 CPU 之间数据移动的速度。

当以块为单位来管理张量时，相同大小和相同数据类型的块可以有效地重用内存，比如可以将不同时间节点的 fp16 参数和 fp16 梯度保存在同一块的内存空间中。例如当某个参数的梯度计算完毕之后，可能训练中它的 fp16 参数将不会再被需要了，这个时候就可以用它的内存空间来存放 fp16 梯度。如下图，参数 P5、P6 完成反向传播之后，它们在本轮计算中就不会再参与计算，那么可以利用这些空间来存放他们对应的梯度，就是下图中的 G5、G6。这样，内存空间重用技术能够进一步支持大规模训练。

下面给出使用 Colossal-AI 训练 ChatGPT 的示例代码，调度十分简单：

from chatgpt.trainer import PPOTrainer
from chatgpt.trainer.strategies import ColossalAIStrategy
from chatgpt.nn import GPTActor, GPTCritic, RewardModel
from copy import deepcopy
from colossalai.nn.optimizer import HybridAdan

strategy = ColossalAIStrategy()

with strategy.model_init_context():
    # init your model here
    # load pretrained gpt2
    actor = GPTActor(pretrained='gpt2')
    critic = GPTCritic()
    initial_model = deepcopy(actor).cuda()
    reward_model = RewardModel(deepcopy(critic.model), deepcopy(critic.value_head)).cuda()
actor_optim = HybridAdan(actor.parameters(), lr=5e-6)
critic_optim = HybridAdan(critic.parameters(), lr=5e-6)

# prepare models and optimizers
(actor, actor_optim), (critic, critic_optim), reward_model, initial_model = strategy.prepare(
    (actor, actor_optim), (critic, critic_optim), reward_model, initial_model)

# load saved model checkpoint after preparing
strategy.load_model(actor, 'actor_checkpoint.pt', strict=False)
# load saved optimizer checkpoint after preparing
strategy.load_optimizer(actor_optim, 'actor_optim_checkpoint.pt')

trainer = PPOTrainer(
    strategy,
    critic,
    reward_model,
    initial_model,
    actor_optim,
    critic_optim,
    ...
)
trainer.fit(dataset, ...)

好了，以上分享了 Colossal-AI 高效异构内存管理系统，以上内容参考了 @Colossal-AI 柳泓鑫的分享，链接：www.bilibili.com/video/BV1pM…，希望我的分享能对你的学习有一点帮助。

 【公众号传送】

《极智AI | Colossal-AI高效异构内存管理系统》

畅享人工智能的科技魅力，让好玩的AI项目不难玩。邀请您加入我的知识星球， 星球内我精心整备了大量好玩的AI项目，皆以工程源码形式开放使用，涵盖人脸、检测、分割、多模态、AIGC、自动驾驶、工业等。不敢说会对你学习有所帮助，但一定非常好玩，并持续更新更加有趣的项目。 t.zsxq.com/0aiNxERDq