PatrickStar

开源

PatrickStar 为 NLP 提供更大、更快、更绿色的预训练模型,并让人人都能使用 AI。

773 stars0 点赞0 次浏览

工具介绍

PatrickStar 为 NLP 提供更大、更快、更绿色的预训练模型,并让人人都能使用 AI。

PatrickStar: 通过基于块的内存管理并行训练大型语言模型

内存不足错误 (OOM) 是每个训练 PTM 的工程师的噩梦。我们通常需要引入更多的 GPU 来存储模型参数，以防止此类错误。PatrickStar 为此问题带来了更好的解决方案。通过异构训练（DeepSpeed Zero Stage 3 也采用此方法），PatrickStar 可以充分利用 CPU 和 GPU 内存，从而使用更少的 GPU 来训练更大的模型。

系统设计 Patrick 的想法是这样的。非模型数据（主要是激活）在训练过程中会发生变化，但当前的异构训练解决方案是静态地将模型数据分割到 CPU 和 GPU 中。为了更好地利用 GPU，PatrickStar 提出了一种基于块的内存管理模块的动态内存调度方案。PatrickStar 的内存管理支持将模型的所有内容（除了当前的计算部分）卸载到 CPU，以节省 GPU。此外，基于块的内存管理在扩展到多个 GPU 时，集群通信也更加高效。

结果

在实验中,Patrickstar v0.4.3 可以在微信数据中心节点中使用 8 个 Tesla V100 GPU 和 240GB GPU 内存,训练一个 18 亿(18B) 参数模型,其网络拓扑如下所示这里。PatrickStar 的规模比 DeepSpeed 大了两倍。而且 PatrickStar 对于相同规模的模型性能更好。pstar 表示 PatrickStar v0.4.3,deeps 表示 DeepSpeed v0.4.3 使用官方示例 DeepSpeed 示例的性能。