FairScale 库测试实验（一）-- 大模型训练基础之模型并行

DDP的分布式训练方法采用数据并行方式，相当于通过增大数据的batch来加快训练。但对于大模型（LLM）来说，DDP已经不适用了。因为LLMs的模型本身太大，一块GPU都放不下怎么可能去复制从而实现数据并行呢。所以LLM的训练采用模型并行的方式来训练。

FairScale 是一个用于高性能和大规模训练的 PyTorch 扩展库。该库扩展了基本的 PyTorch 功能，同时添加了新的 SOTA 扩展技术。FairScale 以可组合模块和易于使用的 API 的形式提供最新的分布式训练技术。这些 API 是研究人员工具箱的基本组成部分，因为他们试图用有限的资源扩展模型。（来源官网）

本次熟悉一下其常用的设置。

目录

一、预先准备

二、使用 PIPELINE PARALLEL 进行模型分片

1、官网的教程 

2、实际应用举例

3、模型分片后，数据前向传播流程分析

总结

---

一、预先准备

随便准备一个模型作为例子。以下面模型为例（可以运行），简单的分类任务。

import torch
import torch.nn as nn

import random
import torchvision.datasets as data
import torchvision.transforms as transforms
import torch.optim as optim


class Classifier(nn.Module):
    def __init__(self):
        super(Classifier, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(in_channels=3,out_channels=32, kernel_size=4,stride=2, padding=1),
            nn.ReLU(),
            nn.MaxPool2d((2,2),1),
            nn.Conv2d(in_channels=32,out_channels=64, kernel_size=3, stride=2, padding=1),
            nn.ReLU(),
            nn.MaxPool2d((2,2),1),
            nn.Conv2d(64, 128, 2, 2,1),
            nn.ReLU(),
            nn.MaxPool2d((2,2),1)
        )
        #
        self.classifier = nn.Sequential(

            nn.Linear(3*3*128,4096),
            nn.ReLU(),

            nn.Dropout(p=0.5),
            nn.Linear(4096,2048),
            nn.ReLU(),

            nn.Linear(2048,1024),
            nn.ReLU()
        )
        self.last_layer_input = nn.Sequential(nn.Linear(1024,10),
            nn.Softmax())

    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        x = self.last_layer_input(x)

        return x

if __name__ == '__main__':

    batchSize = 50
    nepoch = 45

    print("Random Seed: 88")
    random.seed(88)
    torch.manual_seed(88)
    device = 'cuda' if torch.cuda.is_available() else 'cpu'

    dataset = data.CIFAR10(root='/root/data/zjx/Datasets/cifar10',  # 这个路径自己改
                         train=True,
                         transform=transforms.Compose([transforms.ToTensor()]),
                         download=True
                         )

    dataloader = torch.utils.data.DataLoader(dataset,
                                             batch_size=batchSize,
                                             shuffle=True)

    Model = Classifier().to(device)
    Cross_entropy = nn.BCELoss().to(device)
    Optimizer = optim.Adam(Model.parameters(), lr=0.00001)

    for epoch in range(nepoch):
        for i, (data, label) in enumerate(dataloader, 0):
            data = data.to(device)
            label_onehot = torch.zeros((data.shape[0], 10)).to(device)
            label_onehot[torch.arange(data.shape[0]), label] = 1
            output = Model(data)
            loss = Cross_entropy(output, label_onehot)
            print('{}/{}: Loss is {}'.format(i, epoch, loss.data))
            Model.zero_grad()
            loss.backward()
            Optimizer.step()

二、使用 PIPELINE PARALLEL 进行模型分片

1、官网的教程 

官网的示例看这里，主要关键点包括切片的语法设置，以及设备对齐。

模型切片的要求格式

import fairscale
import torch
import torch.nn as nn

model = nn.Sequential(
            torch.nn.Linear(10, 10),
            torch.nn.ReLU(),
            torch.nn.Linear(10, 5)