欢迎关注我的公众号 [极智视界],获取我的更多经验分享
大家好,我是极智视界,本文来介绍一下 Colossal-AI高效异构内存管理系统。
邀您加入我的知识星球「极智视界」,星球内有超多好玩的项目实战源码下载,链接:https://t.zsxq.com/0aiNxERDq
首先需要了解一下异构内存中的数据移动,由于 GPU 的内存容量有限,一般没有办法直接容下大模型,这样的话可以使用 CPU 的内存以及 SSD 硬盘来共同存储大模型。在训练深度学习模型的时候,可以在运行时适当地把数据在不同的设备中进行交换,也就是所谓的 data swapping,理想情况下只需要将需要计算的数据保存在 GPU 中,从而可以训练更加大的模型。
微软在自家的大模型加速框架 deepspeed 中提出了 ZeRO (Zero Redundancy Optimizer) 的异构内存管理解决方案。在数据并行的训练中,若不管三七二十一直接加载到每个 GPU,则会存在最大的内存冗余。deepspeed 发现了这个缺点,将模型的参数、梯度和优化器状态进行切分并存储在不同的 G