PatrickStar:通过基于块的内存管理实现大型语言模型的并行训练
最新进展
请参见CHANGE_LOG.md。
认识PatrickStar
预训练模型(PTM)已成为NLP研究和行业应用的热点。然而,PTM的训练需要巨大的硬件资源,这使得它只对AI社区中的少数人可用。现在,PatrickStar将使PTM训练对每个人都可用!
内存不足错误(OOM)是每个训练PTM的工程师的噩梦。我们常常不得不引入更多的GPU来存储模型参数以防止此类错误。PatrickStar为此类问题带来了更好的解决方案。通过异构训练(DeepSpeed Zero Stage 3也使用了它),PatrickStar可以充分利用CPU和GPU内存,从而可以使用更少的GPU来训练更大的模型。
系统设计
PatrickStar的想法是这样的。非模型数据(主要是激活)在训练过程中是变化的,但当前的异构训练解决方案是静态地将模型数据拆分到CPU和GPU中。为了更好地利用GPU,PatrickStar提出了一种动态内存调度,借助基于块的内存管理模块。PatrickStar的内存管理支持将除当前计算部分以外的所有模型数据卸载到CPU,从而节省GPU。此外,基于块的内存管理在扩展到多GPU时对于集体通信也是高效的。 有关PatrickStar背后的思想,请参见论文和此文档。
实验结果
在实验中,PatrickStar v0.4.3能够在WeChat数据中心节点的8xTesla V100 GPU和240GB GPU内存上训练一个180亿(18B)参数的模型,其网络拓扑如此处所示。PatrickStar的模型规模超过DeepSpeed两倍以上。对于相同规模的模型,PatrickStar的性能也优于DeepSpeed。pstar表示PatrickStar v0.4.3。deeps表示使用官方示例DeepSpeed example的Zero3阶段并默认启用激活优化的DeepSpeed v0.4.3的性能。
我们还在A100 SuperPod的单个节点上评估了PatrickStar v0.4.3。它可以在8xA100和1TB CPU内存上训练68B模型,规模是DeepSpeed v0.5.7的6倍以上。除了模型规模外,PatrickStar的效率也远高于DeepSpeed。基准测试脚本在此处。
WeChat AI数据中心和NVIDIA SuperPod的详细基准测试结果发布在这个Google文档中。
将PatrickStar扩展到SuperPod上的多台机器(节点)。 我们成功在32个GPU上训练了GPT3-175B。就我们所知,这是第一个在如此小的GPU集群上运行GPT3的工作。 微软使用了10,000个V100来预训练GPT3。 现在你可以在32个A100 GPU上微调它,甚至可以预训练你自己的GPT3,太棒了!
我们还使用PatrickStar训练了CLUE-GPT2模型,损失和准确率曲线如下所示:
安装
pip install .
请注意,PatrickStar要求gcc版本为7或更高版本。你也可以使用NVIDIA NGC镜像,以下镜像经过测试:
docker pull nvcr.io/nvidia/pytorch:21.06-py3
使用方法
PatrickStar基于PyTorch,便于迁移pytorch项目。以下是PatrickStar的一个示例:
from patrickstar.runtime import initialize_engine
config = {
"optimizer": {
"type": "Adam",
"params": {
"lr": 0.001,
"betas": (0.9, 0.999),
"eps": 1e-6,
"weight_decay": 0,
"use_hybrid_adam": True,
},
},
"fp16": { # 损失缩放器参数
"enabled": True,
"loss_scale": 0,
"initial_scale_power": 2 ** 3,
"loss_scale_window": 1000,
"hysteresis": 2,
"min_loss_scale": 1,
},
"default_chunk_size": 64 * 1024 * 1024,
"release_after_init": True,
"use_cpu_embedding": False,
"client": {
"mem_tracer": {
"use_async_mem_monitor": args.with_async_mem_monitor,
}
},
}
def model_func():
# MyModel是torch.nn.Module的派生类
return MyModel(...)
model, optimizer = initialize_engine(model_func=model_func, local_rank=0, config=config)
...
for data in dataloader:
optimizer.zero_grad()
loss = model(data)
model.backward(loss)
optimizer.step()
我们使用与DeepSpeed配置JSON相同的config
格式,主要包括优化器参数、损失缩放器参数和一些PatrickStar特定的配置。
关于上述示例的详细说明,请查看此处的指南。
更多示例请查看此处。
一个快速启动的基准测试脚本在此处。它使用随机生成的数据执行,因此你不需要准备真实数据。它还展示了PatrickStar的所有优化技术。有关运行基准测试的更多优化技巧,请参见优化选项。
许可证
BSD 3-Clause许可证
引用我们
@article{fang2021patrickstar,
title={PatrickStar: Parallel Training of Pre-trained Models via a Chunk-based Memory Management},
author={Fang, Jiarui and Yu, Yang and Zhu, Zilin and Li, Shenggui and You, Yang and Zhou, Jie},
journal={arXiv preprint arXiv:2108.05818},
year={2021}
}
@article{fang2022parallel,
title={Parallel Training of Pre-Trained Models via Chunk-Based Dynamic Memory Management},
author={Fang, Jiarui and Zhu, Zilin and Li, Shenggui and Su, Hui and Yu, Yang and Zhou, Jie and You, Yang},
journal={IEEE Transactions on Parallel and Distributed Systems},
volume={34},
number={1},
pages={304--315},
year={2022},
publisher={IEEE}
}
联系我们
{jiaruifang, zilinzhu, josephyu}@tencent.com
由腾讯NLP团队WeChat AI Team提供技术支持。