模型并行性能够促进视觉任务的性能。但是目前，还没有一个标准库可以让我们像采用混合精度等其他 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) # forward
- loss = ts.nn.functional.parallel_cross_entropy(x, y) # parallel loss function
- loss.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 resnet
- model = 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 ranks
- state_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 ranks
- state_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）一起混合使用。
- # gradscaler
- scaler = 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 step
- scaler.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 ZeroRedundancyOptimizer
- if 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 来确保同一模型并行组中的模块具有相同的训练数据。
