DLA 神经网络的极限训练方法：gradient checkpointing

gradient checkpointing

一般来说，训练的过程需要保存中间结果（不管是GPU还是CPU）。前向传播根据输入(bottom_data)计算输出(top_data)，后向传播由top_diff计算bottom_diff（如果某个变量打开梯度进行训练的话）。top和bottom是包含数据和梯度的两个结构体，整个网络的每层top和bottom在训练的过程中都会保存，这消耗了大量的内存。

如果不保存这些变量，每次传播时重新分配和计算，会大大减少内存的使用量，但是也会使得网络的训练时间无限延长。为了平衡这两个矛盾，论文Training Deep Nets with Sublinear Memory Cost 使用亚线性内存成本训练深度网络：我们提出了一种系统方法来减少深度的内存消耗神经网络训练。具体来说，我们设计了一种成本高昂的算法 O（sqrt（n））内存来训练 n 层网络，只需计算成本每个小批量的额外前向传递。每隔 sqrt(n)保留一个检查点的feature map。

CODE

https://pytorch.org/docs/stable/checkpoint.html

// https://discuss.pytorch.org/t/trying-to-understand-torch-utils-checkpoint/95224
import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
import numpy as np
from tqdm.notebook import tqdm

from torch import optim
import torchvision.models as models
from torch import nn

CHECKPOINT = True
BATCH_SIZE = 32
dev = "cuda:0"

class ImageDataset(Dataset):
    def __init__(self,length = 100000,size = 244):
        self.length = length
        self.size = 244
    def __len__(self):
        return self.length
    def __getitem__(self,idx,display = False):
        return torch.from_numpy(np.random.randn(2,3,self.size,self.size))
train = ImageDataset()
trainloader = DataLoader(
    train,
    batch_size = BATCH_SIZE,
    num_workers = 24,
    pin_memory = True
)

resnet = models.resnet50(pretrained = False)

class MODEL(nn.Module):
    def __init__(self,model):
        super(MODEL,self).__init__()
        self.model = model
        self.LR = nn.Linear(1000,1000)
    def forward(self,x):
        if CHECKPOINT == False:
            o1 = self.model(x[:,0])
            o2 = self.model(x[:,1])
        else:
            o1 = torch.utils.checkpoint.checkpoint(self.model,x[:,0])
            o2 = torch.utils.checkpoint.checkpoint(self.model,x[:,1])
        
        return torch.mean((self.LR(o1)-o2)**2)
    
resnet = MODEL(resnet).to(dev)

optimizer = optim.SGD(resnet.parameters(),lr = .001)

for T in tqdm(trainloader):
    out = torch.mean(resnet(T.float().to(dev)))
    optimizer.zero_grad()
    out.backward()
    optimizer.step()

CG

在这里插入图片描述

https://github.com/merrymercy/dtr-prototype

ZeRO-Offload

https://arxiv.org/pdf/2101.06840.pdf 大规模模型训练一直是少数人的比赛场地需要复杂的模型重构和访问昂贵的 GPU 集群。ZeRO-Offload 通过使几乎每个人都可以访问大型模型训练。它可以训练模型单个 GPU 上超过 13 亿个参数，与 GPU 相比，大小增加了 10 倍流行的框架，如PyTorch，它不需要任何模型就可以做到这一点。从数据科学家改变或牺牲计算效率。 ZeRO-卸载通过卸载数据和计算来实现大型模型训练中央处理器。为了保持计算效率，它旨在最大限度地减少数据移入/移出 GPU，减少 CPU 计算时间，同时最大化内存节省 GPU 成本。因此，ZeRO-Offload可以在单个上实现40 TFlops / GPU。 NVIDIA V100 GPU 用于 10B 参数模型，与单独使用 PyTorch 的 30TF 相比对于 1.4B 参数模型，可以训练而不会耗尽的最大参数模型的记忆。ZeRO-Offload 还设计为在以下情况下在多个 GPU 上进行扩展可用，可在多达 128 个 GPU 上提供近乎线性的加速。此外，它可以与模型并行性协同工作，训练超过 70 亿的模型单个 DGX-2 盒子上的参数，与模型尺寸相比增加了 4.5 倍单独使用模型并行性。通过将计算和内存效率与易于使用，ZeRO-Offload 使大规模模型训练民主化，使其成为即使是数据科学家也可以访问，只需访问一个 GPU。