PatrickStar：基于分块内存管理的大规模语言模型并行训练

最新进展

请参阅 CHANGE_LOG.md。

认识 PatrickStar

预训练模型（PTM）正成为自然语言处理研究和工业应用的热点。然而，训练这些模型需要巨大的硬件资源，使得只有少数人工智能领域的从业者能够进行。现在，PatrickStar 将让每个人都能轻松训练 PTM！

显存不足错误（OOM）是每一位训练 PTM 的工程师的噩梦。我们通常不得不增加 GPU 数量来存储模型参数以避免此类错误。PatrickStar 为这一问题提供了一种更好的解决方案。通过 异构训练（DeepSpeed Zero Stage 3 也采用了该技术），PatrickStar 能够充分利用 CPU 和 GPU 的显存，从而用更少的 GPU 来训练更大的模型。

系统设计

Patrick 的核心思想如下：在训练过程中，非模型数据（主要是激活值）是动态变化的，而现有的异构训练方案却将模型数据 静态地 分配到 CPU 和 GPU 上。为了更好地利用 GPU，PatrickStar 提出了一种基于分块内存管理模块的 动态 内存调度机制。PatrickStar 的内存管理可以将除当前计算部分之外的所有内容卸载到 CPU 上，从而节省 GPU 显存。此外，分块内存管理在扩展到多 GPU 时，对于集体通信也非常高效。有关 PatrickStar 的设计理念，请参阅论文和 这篇文档。

结果

实验中，PatrickStar v0.4.3 在微信数据中心节点上，使用 8 块 Tesla V100 GPU 和 240GB 显存，成功训练了一个拥有 180 亿（18B）参数的模型，其网络拓扑结构如 此处 所示。PatrickStar 比 DeepSpeed 大了两倍多，且在相同规模的模型上性能也更优。图中的 pstar 表示 PatrickStar v0.4.3 的性能，而 deeps 则表示使用官方示例 DeepSpeed 示例 中默认开启激活优化的 ZeRO3 阶段的 DeepSpeed v0.4.3 的性能。

我们还在 A100 SuperPod 的单个节点上评估了 PatrickStar v0.4.3。它可以在 8 块 A100 GPU 和 1TB CPU 内存的配置下训练一个 680 亿参数的模型，这比 DeepSpeed v0.5.7 大了 6 倍以上。除了模型规模之外，PatrickStar 的效率也远超 DeepSpeed。基准测试脚本见 这里。

我们在微信 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 上。以下是一个 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": {  # loss scaler 参数
        "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()

我们使用的 config 格式与 DeepSpeed 配置 JSON 相同，主要包括优化器参数、损失缩放器参数以及一些 PatrickStar 特有的配置。

有关上述示例的详细说明，请参阅指南 这里。

更多示例请查看 这里。

一个快速入门的基准测试脚本见 这里。该脚本使用随机生成的数据运行，因此无需准备真实数据。它还演示了 PatrickStar 的所有优化技术。有关运行基准测试的更多优化技巧，请参阅 优化选项。

许可证

BSD 3-Clause 许可证

引用我们

@article{fang2021patrickstar,
  title={PatrickStar: 基于分块内存管理的预训练模型并行训练},
  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={基于分块动态内存管理的预训练模型并行训练},
  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

由微信 AI 团队和腾讯 NLP Oteam 提供支持。

PatrickStar 快速上手指南

PatrickStar 是一款基于分块内存管理（Chunk-based Memory Management）的大语言模型并行训练工具。它通过动态调度 CPU 和 GPU 显存，显著降低显存占用，使开发者能够使用更少的显卡训练更大的模型（例如在 8 张 A100 上训练 68B 模型，或在 32 张 GPU 上训练 GPT-3 175B）。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux
• 编译器：GCC 版本 7 或更高（必需）
• 深度学习框架：PyTorch
• 推荐镜像：为了方便配置环境，推荐使用 NVIDIA NGC 提供的 PyTorch 镜像（已预装兼容的驱动和库）：

  docker pull nvcr.io/nvidia/pytorch:21.06-py3
  

安装步骤

您可以直接通过源码进行安装。建议在国内网络环境下配置 pip 国内镜像源以加速下载。

# 使用国内镜像源安装（推荐）
pip install . -i https://pypi.tuna.tsinghua.edu.cn/simple

注意：如果在构建过程中遇到编译错误，请再次确认 gcc --version 是否 >= 7。

基本使用

PatrickStar 基于 PyTorch 构建，迁移现有项目非常简单。其配置格式与 DeepSpeed 兼容，主要包含优化器、混合精度训练参数及 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()

快速测试

如果您想立即验证安装并查看性能基准，可以运行官方提供的快速基准测试脚本（该脚本使用随机生成的数据，无需准备真实数据集）：

./examples/run_transformers.sh

更多详细的优化技巧和完整示例，请参考项目目录下的 examples 文件夹及 doc/optimization_options.md 文档。