DeepSpeed之ZeRO系列:将显存优化进行到底

前言

目前训练超大规模语言模型主要有两条技术路线:TPU + XLA + TensorFlow/JAX 和 GPU + PyTorch + Megatron-LM + DeepSpeed。前者由Google主导,由于TPU和自家云平台GCP深度绑定,对于非Googler来说, 只可远观而不可把玩,后者背后则有NVIDIA、Meta、MS大厂加持,社区氛围活跃,也更受到群众欢迎。

上面提到的DeepSpeed的核心是ZeRO(Zero Redundancy Optimizer),简单来说,它是一种显存优化的数据并行(data parallelism, DP)方案。而“优化“这个话题又永无止境,在过去两年DeepSpeed团队发表了三篇ZeRO相关的论文,提出了去除冗余参数、引入CPU和内存、引入NVMe等方法,从始至终都围绕着一个目标:将显存优化进行到底。

ZeRO: 一种去除冗余的数据并行方案

ZeRO: Memory Optimizations Toward Training Trillion Parameter Models 发表在SC 20,DeepSpeed项目最初就是论文中ZeRO方法的官方实现。

背景

如今训练大模型离不开各种分布式并行策略,常用的并行

猜你喜欢

转载自blog.csdn.net/u013250861/article/details/131248014