PatrickStar 是一款腾讯开发的分布式深度学习训练工具，它的设计目标是支持以 GPT、Bert 为代表的超大预训练模型训练。
用法
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, } 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 特定配置的参数。
引用我们
@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}
}
