极智AI

极智AI

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

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

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

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

微软在自家的大模型加速框架 deepspeed 中提出了 ZeRO (Zero Redundancy Optimizer) 的异构内存管理解决方案。在数据并行的训练中，若不管三七二十一直接加载到每个 GPU，则会存在最大的内存冗余。deepspeed 发现了这个缺点，将模型的参数、梯度和优化器