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

PatrickStar 项目介绍

认识 PatrickStar

在自然语言处理领域，预训练模型(PTM)已经成为学术研究和工业应用的热点。然而，训练这些模型需要巨大的硬件资源，使得它们仅对AI领域中的少部分人群可用。为了改变这一现状，PatrickStar旨在让PTM训练变得人人可及。

在训练PTM的过程中，内存溢出错误(Out-of-memory error, OOM)常常是工程师们的噩梦。为了解决这个问题，通常我们需要引入更多的GPU来存储模型参数。然而，PatrickStar提供了一个更优的解决方案。通过异构训练方法（DeepSpeed Zero Stage 3也使用类似技术），PatrickStar能够充分利用CPU和GPU的内存，使用户能够用更少的GPU训练更大的模型。

系统设计

PatrickStar的核心思想是：在训练过程中，非模型数据（主要是激活值）会不断变化，而当前的异构训练方案常常是静态地将模型数据分割到CPU和GPU上。为了更好地利用GPU资源，PatrickStar提出了一种基于块的动态内存管理模块，进行动态内存调配。PatrickStar的内存管理支持卸载除当前计算部分以外的所有模型数据到CPU，以节省GPU资源。此外，基于块的内存管理在扩展到多个GPU时，对于集体通信非常高效。

项目成果

在实验中，PatrickStar v0.4.3能够在微信数据中心的节点上使用8个Tesla V100 GPU和240GB GPU内存训练一个具有180亿（18B）参数的模型，比DeepSpeed的能力强两倍。此外，PatrickStar即使在相同大小的模型上也表现优异。

我们还在NVIDIA的A100 SuperPod单节点上评估了PatrickStar v0.4.3，它能在1TB CPU内存下，用8个A100 GPU训练68B的模型，比DeepSpeed v0.5.7处理的规模大6倍以上。更重要的是，除了模型规模之外，PatrickStar在效率上远胜于DeepSpeed。

在多个SuperPod节点上，PatrickStar成功在32个GPU上训练GPT3-175B模型。这是首次在如此小的GPU集群上运行GPT3，而微软则使用了10,000个V100来进行相同的工作。

我们还用PatrickStar训练了CLUE-GPT2模型，其损失和准确率曲线表现优异。

安装方法

使用以下命令即可安装PatrickStar：

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,
        },
    },
    ...
}

def model_func():
    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()

PatrickStar的配置格式与DeepSpeed配置JSON相同，主要包括优化器参数、损失标量以及一些PatrickStar特有的配置。

详细的使用说明可参阅指导文档。

开源协议

本项目基于BSD 3-Clause开源协议。

联系我们

如果您有任何问题或需要支持，请联系以下邮箱：

本项目由微信AI团队和腾讯NLP Oteam共同支持。