PatrickStar下载 - PatrickStar源代码下载

PatrickStar

Ai源码

v0.4.6

下载

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

遇见帕特里克·斯塔

预训练模型（PTM）正在成为自然语言处理研究和行业应用的热点。然而，PTM 的训练需要大量的硬件资源，因此只有 AI 社区中的一小部分人可以使用。现在， PatrickStar 将为所有人提供 PTM 培训！

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

系统设计

帕特里克的想法是这样的。非模型数据（主要是激活）在训练过程中会发生变化，但当前的异构训练解决方案是将模型数据静态地分割到CPU和GPU上。为了更好地使用 GPU，PatrickStar 提出借助基于块的内存管理模块进行动态内存调度。 PatrickStar的内存管理支持将模型当前计算部分以外的所有内容卸载到CPU以节省GPU。此外，当扩展到多个 GPU 时，基于块的内存管理对于集体通信非常有效。请参阅论文和本文档了解 PatrickStar 背后的想法。

结果

在实验中，Patrickstar v0.4.3能够在微信数据中心节点上使用8xTesla V100 GPU和240GB GPU内存训练180亿（18B）参数的模型，其网络拓扑如下。 PatrickStar 的大小是 DeepSpeed 的两倍多。对于相同尺寸的模型，PatrickStar 的性能也更好。 pstar 是 PatrickStar v0.4.3。 Deeps 使用官方示例 DeepSpeed 示例 Zero3 阶段表示 DeepSpeed v0.4.3 的性能，默认情况下打开激活优化。

alt perf

我们还在 A100 SuperPod 的单个节点上评估了 PatrickStar v0.4.3。它可以在具有 1TB CPU 内存的 8xA100 上训练 68B 模型，比 DeepSpeed v0.5.7 大 6 倍以上。除了模型规模之外，PatrickStar 比 DeepSpeed 效率更高。基准测试脚本在这里。

alt perf

微信人工智能数据中心和 NVIDIA SuperPod 的详细基准测试结果发布在这个 Google Doc 上。

将 PatrickStar 扩展到 SuperPod 上的多台机器（节点）。我们成功在 32 GPU 上训练 GPT3-175B。据我们所知，在如此小的GPU集群上运行GPT3还是第一个作品。微软使用了 10,000 个 V100 来支持 GPT3。现在您可以对其进行微调，甚至可以在 32 A100 GPU 上预训练您自己的模型，太棒了！

alt perf

我们还用PatrickStar训练了CLUE-GPT2模型，损失和准确率曲线如下所示：

线索-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" : {  # loss scaler params
        "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 is a derived class for 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 条款许可证

引用我们

 @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}
}