训练大模型也不怕,轻量级TorchShard库减少GPU内存消耗,API与PyTorch相同

第一时间获取价值内容

来自:机器之心

训练大模型时,如何优雅地减少 GPU 内存消耗?你不妨试试这个 TorchShard 库,兼具模型并行与数据并行等特点,还具有与 PyTorch 相同的 API 设计。

模型并行性能够促进视觉任务的性能。但是目前,还没有一个标准库可以让我们像采用混合精度等其他 SOTA 技术那样轻松地采用模型并行性。
最近,马里兰大学帕克分校计算机科学系的研究者 Kaiyu Yue 开源了一个工具TorchShard,这是一个轻量级的引擎,用于将 PyTorch 张量切片成并行的 shard。当模型拥有大量的线性层(例如 BERT、GPT)或者很多类(数百万)时,TorchShard 可以减少 GPU 内存并扩展训练规模,它具有与 PyTorch 相同的 API 设计。
项目地址:https://github.com/KaiyuYue/torchshard
BERT 和 GPT 等超大模型正在成为 NLP 领域应用中的趋势。然而训练这种大模型面临内存限制的问题,为了解决这个难题,研究者使用 Megatron-LM 和 PyTorch-Lightning 模型并行性扩大训练。其中,Megatron-LM 只专注于大规模训练语言模型,而 PyTorch-Lightning 仅基于 sharded 优化器状态和梯度,如 DeepSpeed。
在计算机视觉任务中,我们会在训练基于 Transformer、MLP 模型或在数百万个类中训练模型时遇到同样的问题。TorchShard 的目标是:
  • 建立一个标准的 PyTorch 扩展库,用于使用模型并行性进行扩展训练;

  • 以一种简单、自然的方式使用 PyTorch。

TorchShard 是对模型并行单元(mpu)的彻底重写,是 Megatron-LM 核心。最重要的是,TorchShard 具有与 PyTorch 相同的 API 设计,这意味着所有的子类和子函数都保持与 PyTorch 相同。例如,如果你想让原来的线性层 torch.nn. linear 是并行的,只需将 torch 变成 ts,并调用带有 dim 参数的子类 nn.ParallelLinear,如下所示:
import torchshard as ts
ts.init_process_group(group_size=2) # init parallel groups
m = torch.nn.Sequential( torch.nn.Linear(20, 30, bias=True), ts.nn.ParallelLinear(30, 30, bias=True, dim=None), # equal to nn.Linear() ts.nn.ParallelLinear(30, 30, bias=True, dim=0), # parallel in row dimension ts.nn.ParallelLinear(30, 30, bias=True, dim=1), # parallel in column dimension).cuda()
x = m(x) # forwardloss = ts.nn.functional.parallel_cross_entropy(x, y) # parallel loss functionloss.backward() # backward
torch.save( ts.collect_state_dict(m, m.state_dict()), 'm.pt') # save model state
除此之外,TorchShard 还支持与 DDP 一起使用时的各种特性,保存和加载 shard checkpoints,初始化 shard 参数,以及跨多台机器和 GPU 处理张量。具体如下:
  • torchshard 包含必要的功能和操作,如 torch 包;

  • torchshard.nn 包含图形的基本构建块,如 torch.nn 包;

  • torchshard.nn.functional 包含 torchshard.nn 的相应功能操作,如 torch.nn.functional 包;

  • torchshard.distributed 包含处理分布式张量和组的基本功能,如 torch.distributed 包更容易使用。

如何开始 TorchShard?
安装要求:Python 版本 3.6 以上(含)以及 PyTorch 版本 1.9.0 以上(含)。通过 pip 安装 TorchShard 库:
pip install torchshard
这里以 ImageNet 上训练 ResNet-50 为例,展示仅需几行代码就能在项目中使用 TorchShard。通常 ResNet-50 设计范式包含两部分:卷积块和全连接层,如下图 1 所示。一般来说,由于大量的类依赖于数据集,最后的线性层比卷积块有更多的参数。所以我们切片最后一个线性层来检查其最大尺寸。
图 1:DDP 以及 DDP + TorchShard 前向训练流。
在上图 1 中,左边展示了传统的 DDP 训练范式。假设我们有两个等级,DDP 将强制每个等级有重复的模型参数。然而,TorchShard 会将层级参数切片到不同的等级,从而减少整个 GPU 内存。现在向 ImageNet 官方训练脚本添加一些代码,修改后的版本已经成为 TorchShard 项目的一部分。
首先将 torchshard import 进来:
import torchshard as ts
然后需要初始化模型并行的进程组,就像初始化 DDP 进程组的方法一样。只需要设置一个功能参数来告诉 torchshard 应该从目标层中切片出多少个 shard。
ts.distributed.init_process_group(group_size=args.world_size)
接下来将模型转换为并行版本,其中可以直接将整个模型输入到转换辅助函数中,无需特殊处理。
import resnetmodel = resnet.__dict__[args.arch](pretrained=args.pretrained)ts.nn.ParallelLinear.convert_parallel_linear( model, dim=args.model_parallel_dim)print('=> paralleling model'{}''.format(args.arch))
此外,不要忘记损失函数 torchshard.nn.ParallelCrossEntropy ,该损失函数可以根据输入张量在原始 PyTorch 版本和并行版本之间切换运行模式。例如,如果输入张量是由 torchshard 并行层产生的,torchshard.nn.ParallelCrossEntropy 将以并行方式计算损失值。
criterion = ts.nn.ParallelCrossEntropyLoss().cuda(args.gpu)
当模型并行模式(TorchShard)和数据并行模式(DDP)一起工作时,我们需要处理并行层的输入。每个等级中的参数和训练数据都不同。因此,我们在 ResNet forward 中的并行线性层之前收集输入张量。
x = ts.distributed.gather(x, dim=0) # gather input along the dim of batch size x = self.fc(x)
同样地,我们在计算损失值之前收集目标张量。
output = model(images)if args.enable_model_parallel:    target = ts.distributed.gather(target, dim=0)loss = criterion(output, target)
最后,使用 TorchShard 函数保存和加载 checkpoints 非常简单。TorchShard 提供了名为 torchshard.collect_state_dict 基本函数用于保存 checkpoints,torchshard.relocate_state_dict 用于加载 checkpoints。
保存检查点:
state_dict = model.state_dict()# collect states across all ranksstate_dict = ts.collect_state_dict(model, state_dict)if ts.distributed.get_rank() == 0: torch.save(state_dict, 'resnet50.pt') # save as before
加载检查点:
if ts.distributed.get_rank() == 0:     state_dict = torch.load('resnet50.pt')# relocate state_dict() for all ranksstate_dict = ts.relocate_state_dict(model, state_dict)model.load_state_dict(state_dict) # load as before
现在我们已经完成了在 ImageNet 上为 shard 训练添加代码, 然后可以通过增加类的数量来扩展它,即最后一个线性层的输出特征维度。训练脚本可以在 torchshard/project/imagenet 中找到。下图展示了在 8 个 NVIDIA TITAN-XP (12196 MiB) GPU 、类数 ≤ 1000000 上和 16 个 GPU 、类数为 2000000 上训练 ResNet-50 扩展能力。
图 2:在不同并行策略下使用标准 ResNet 训练设置(即输入大小 224 和批量大小 256)的 GPU 内存成本。
使用 AMP 与 ZeRO
TorchShard 以简单自然的 PyTorch 方式与其他技术(例如自动混合精度 AMP 以及 ZeRO)一起混合使用。
# gradscalerscaler = torch.cuda.amp.GradScaler(enabled=args.enable_amp_mode)

with torch.cuda.amp.autocast(enabled=args.enable_amp_mode): # compute output output = model(images) if args.enable_model_parallel: target = ts.distributed.gather(target, dim=0) loss = criterion(output, target)

# compute gradient and do SGD stepscaler.scale(loss).backward()scaler.step(optimizer)scaler.update()optimizer.zero_grad()
图 3:在不同并行策略以及 AMP 下,使用标准的 ResNet 训练设置时(输入尺寸 224,batch 大小 256),使用 GPU 内存的成本。
ZeRO 是 DeepSpeed 的核心,与 PyTorch >= 1.9.0 一起使用。如果你想测试一个函数,请安装最新版本的脚本来运行,代码如下:
from torch.distributed.optim import ZeroRedundancyOptimizerif args.enable_zero_optim:    print('=> using ZeroRedundancyOptimizer')    optimizer = torch.distributed.optim.ZeroRedundancyOptimizer(        model.parameters(),        optimizer_class=torch.optim.SGD,        lr=args.lr,        momentum=args.momentum,        weight_decay=args.weight_decay)else:    optimizer = torch.optim.SGD(model.parameters(), args.lr,                                momentum=args.momentum,                                weight_decay=args.weight_decay)
图 4:在不同的并行策略和 ZeRO 优化器下,在标准 ResNet 训练设置(输入大小 224 和批大小 256)的 GPU 内存成本。
此外,TorchShard 还提供了基本的 Python API 以及和相应的模板文件,以简化自定义并行层的实现。
研究者将持续开发 TorchShard,如 TorchShard 下一个特性是新的数据采样器 torchshard.utils.data.DistributedGroupSampler,它的命名遵循 torch.utils.data.DistributedSampler。该采样器旨在帮助用户构建 M-way 数据并行、N-way 模型并行,使得其就像 DDP 中的 DistributedSampler 一样简单。用户唯一要做的就是设置模型并行组号,然后 DistributedGroupSampler 来确保同一模型并行组中的模块具有相同的训练数据。
原文链接:https://medium.com/pytorch/torchshard-a31fcbfdc354

-结束-

👆 长按识别,即可关注

(0)

相关推荐