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

• TorchScale - A Library of Foundation Architectures
• Implementation of “BitNet: Scaling 1-bit Transformers for Large Language Models” in pytorch
• 图片速览 BitNet: 1-bit LLM
  
• 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。

梯度累积

        训练时大的batch一般能得到更稳定的训练效果，梯度累积训练方法是一种用于训练深度神经网络的技术，旨在减少显存需求并提高训练效果。在传统的训练方法中，模型的参数是通过单个批次（batch）的数据计算得到的梯度平均值进行更新。但在梯度累积训练中，模型的参数更新是通过多个批次的梯度累积得到的。

以下是梯度累积训练的基本步骤：

1. 设置梯度累积步数（accumulation steps），它决定了要累积多少个批次的梯度。

2. 初始化模型的参数。

3. 对于每个训练批次（batch）：

   • 使用当前批次的数据进行前向传播计算损失。
   • 对损失进行反向传播计算梯度。
   • 累积当前批次的梯度到之前的梯度值上。

4. 当累积达到设置的步数时，将累积的梯度应用于模型参数的更新：

   • 通过将累积的梯度平均化来获得参数的更新值。
   • 使用更新值来更新模型的参数。

5. 重复步骤3和4，直到完成所有的训练批次。

梯度累积训练的主要优势在于能够降低每个批次所需的显存量，允许在具有有限显存的硬件上训练更大的模型。此外，梯度累积还可以改善模型的收敛性，提高模型的性能和泛化能力。

import torch
import torch.optim as optim

# 模型定义
class SimpleModel(torch.nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = torch.nn.Linear(10, 1)

    def forward(self, x):
        return self.fc(x)

# 创建模型、优化器和损失函数
model = SimpleModel()
optimizer = optim.SGD(model.parameters(), lr=0.01)
criterion = torch.nn.MSELoss()

# 梯度累积步数
accumulation_steps = 4

# 训练数据和目标数据
input_data = torch.rand((8, 10))  # 8个样本，每个样本10个特征
target_data = torch.rand((8, 1))   # 对应的目标数据

# 模型训练
for epoch in range(100):
    # Forward pass
    output = model(input_data)
    loss = criterion(output, target_data)

    # Backward pass (梯度计算)
    loss.backward()

    # 每隔 accumulation_steps 步更新一次梯度
    if (epoch + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

    # 打印损失
    if (epoch + 1) % 10 == 0:
        print(f'Epoch [{epoch + 1}/100], Loss: {loss.item():.4f}')

# 评估模型
with torch.no_grad():
    model.eval()
    test_input = torch.rand((4, 10))  # 用于测试的输入数据
    predictions = model(test_input)
    print("Model predictions:", predictions)

• LST: Ladder Side-Tuning for Parameter and Memory Efficient Transfer Learning
  最近在各种领域中采用了对下游任务的大型预训练模型进行微调。但是，更新大型预训练模型的整个参数集的成本很高。尽管最近提出的参数高效迁移学习 （PETL） 技术允许在预训练的骨干网络中为新任务更新一小部分参数（例如仅使用 2% 的参数），但它们只能将训练内存需求降低多达 30%。这是因为可训练参数的梯度计算仍然需要通过大型预训练骨干模型进行反向传播。为了解决这个问题，我们提出了梯形侧调 （LST），这是一种新的 PETL 技术，可以大幅减少训练记忆需求。与在骨干网络中插入额外参数的现有参数高效方法不同，我们训练了一个梯形图侧网络，这是一个小型且独立的网络，它通过来自骨干网络的快捷连接（称为梯形图）将中间激活作为输入并进行预测。LST 的内存要求明显低于以前的方法，因为它不需要通过骨干网络进行反向传播，而只需要通过侧网络和梯形连接。我们使用各种模型（T5 和 CLIP-T5）在 NLP （GLUE） 和视觉和语言 （VQA、GQA、NLVR2 、MSCOCO） 任务上评估我们的方法。LST 节省了 69% 的内存成本来微调整个网络，而其他方法在类似参数使用中仅节省了 26%（因此，节省了 2.7 倍的内存）。此外，LST 在低内存状态下实现了比 Adapter 和 LoRA 更高的精度。为了进一步展示这种更好的内存效率的优势，我们还将 LST 应用于更大的 T5 模型，与完全微调和其他 PETL 方法相比，获得了更好的 GLUE 性能。精度与效率的权衡也适用于VL任务。