PatrickStar: Parallel Training of Large Language Models via a Chunk-based Memory Management
Meeting PatrickStar
Pre-Trained Models (PTM) are becoming the hotspot of both NLP research and industry application. They are models that are trained with massive data and have learned generic features of the language. In pactice, they are fine-tuned for downstream tasks with task-specific datasets. In this way, PTMs have achieved great performances in almost every tasks. However, the training of PTMs requires enormous hardware resources, which makes it only accessible to small portion of people in the AI community. Now, PatrickStar will make PTM training available to everyone!
Out of memory error (OOM) is the nightmare of every engineer training PTMs. To prevent such error, we often have to introduce more GPUs to store the model params. PatrickStar brings a better solution for such problem. With the heterogeneous training (DeepSpeed Zero Stage 3 also uses it), PatrickStar could make full use of both the CPU and GPU memory, so that you could use fewer GPUs to train larger models.
We noticed that the GPU memory usage varies during training, but the current heterogenous training solutions are all statically spliting the model and optimizer states to CPU and GPU. To make better use of the GPU, PatrickStar proposes a dynamic memory scheduling with the help of a chunk-based memory management module. The memory management of PatrickStar supports offloading everything but the current computing part of the model to CPU. This results in training a much larger model within the same hardware environment. In terms of performance, the chunk-based memory management takes advantage of the linear structure of the transformer-based PTMs, so that it will inherently prefetch the upcoming layers to GPUs, resulting in a great performance.
In experiment, Patrickstar is able to train a 12B param model with 8 Tesla V100 GPU and 240GB GPU memory, which is twice as large as the state of art. And the performance of PatrickStar is better for models of the same size as well. The deeps indicates performance of DeepSpeed v0.4.3 using the official example DeepSpeed example zero3 stage with activation optimzations openinig by default.
We've also trained the CLUE-GPT2 model with PatrickStar, the loss and accuracy curve is shown below:
Installation
pip install .
Note that PatrickStar requires gcc of version 7 or higher. You could also use NVIDIA NGC images, the following image is tested:
docker pull nvcr.io/nvidia/pytorch:21.06-py3
Usage
PatrickStar is based on PyTorch, which makes it easy to migrate a pytorch project. Here is a example of 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()
We use the same config format as DeepSpeed configuration JSON, which mainly includes params of optimizer, loss scaler and some PatrickStar specific configuration.
For some detail explanation of the above example, please check the guide here
For more examples, please check here.
Inside PatrickStar
See this doc for the idea behind PatrickStar.
License
BSD 3-Clause License
Cite Us
@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}
}
Contact Us
{jiaruifang, zilinzhu, josephyu}@tencent.com
Powered by WeChat AI Team, Tencent NLP Oteam.
161 Dec 16, 2022
21 Dec 29, 2022
3 Jan 14, 2022
20 Oct 11, 2022
947 Dec 28, 2022
0 Jun 28, 2022
13 Sep 02, 2022
1 Dec 23, 2021
2.9k Dec 31, 2022
1 Apr 28, 2022
4 Feb 23, 2022
69 Dec 25, 2022
7 Mar 12, 2022
55 Nov 28, 2022
7 Jan 15, 2022
63 Oct 30, 2022
84 Dec 20, 2022
10 Dec 14, 2022
45 Oct 29, 2022
28 Nov 10, 2022