模型并行性可以提高视觉任务的性能。但是目前，没有标准库可以让我们像其他SOTA技术（例如混合精度）一样轻松地采用模型并行性。最近，马里兰大学帕克分校计算机科学系研究员KaiyuYue开源了一个工具TorchShard，这是一个用于将PyTorch张量切片为并行分片的轻量级引擎。当模型有大量线性层（如BERT、GPT）或许多类（百万）时，TorchShard可以减少GPU内存和规模训练。它具有与PyTorch相同的API设计。项目地址：https://github.com/KaiyuYue/torchshardBERT和GPT等超大模型正在成为NLP应用的趋势。然而，训练如此大的模型面临着内存限制的问题。为了解决这个问题，研究人员使用Megatron-LM和PyTorch-Lightning模型并行扩展训练。其中，Megatron-LM只专注于语言模型的大规模训练，而PyTorch-Lightning只基于分片优化器状态和梯度，例如DeepSpeed。在计算机视觉任务中，我们在训练基于Transformer、基于MLP的模型，或者训练百万级的模型时，都会遇到同样的问题。TorchShard的目标是：构建一个标准的PyTorch扩展库，用于使用模型并行性进行扩展训练；以简单自然的方式使用PyTorch。TorchShard是对模型并行单元(MPU)的完全重写，它是Megatron-LM的核心。最重要的是，TorchShard具有与PyTorch相同的API设计，这意味着所有子类和子函数都与PyTorch保持相同。比如你想让原来的线性层torch.nn.linear是平行的，就把torch转成ts，调用子类nn.ParallelLinear，参数是dim，像这样：importtorchshardaststs.init_process_group(group_size=2)#initparallelgroupsm=torch.nn.Sequential(torch.nn.Linear(20,30,bias=True),ts.nn.ParallelLinear(30,30,bias=True,dim=None),#equaltonn.Linear()ts.nn.ParallelLinear(30,30,bias=True,dim=0),#parallelinrowdimensionts.nn.ParallelLinear(30,30,bias=True,dim=1),#parallelincolumndimension).cuda()x=m(x)#前向损失=ts.nn.functional.parallel_cross_entropy(x,y)#parallellossfunctionloss.backward()#backwardtorch.save(ts.collect_state_dict(m,m.state_dict()),'m.pt')#savemodelstate除其他外，TorchShard与DDP一起使用时支持各种功能，保存和加载分片检查点，初始化分片参数，以及跨多台机器和GPU处理张量。具体如下：torchshard包含了必要的功能和操作，比如torch包；torchshard.nn包含图形的基本构建块，例如torch.nn包；torchshard.nn.functional包含了torchshard.nn对应的函数操作，比如torch.nn。功能包；torchshard.distributed包含处理分布式张量和组的基本功能，例如torch.distributed包更容易使用。如何启动TorchShard？安装要求：Python3.6版本（含）和PyTorch1.9.0版本（含）。通过pip安装TorchShard库：pipinstalltorchshard这里以在ImageNet上训练ResNet-50为例，说明只需要几行代码就可以在项目中使用TorchShard。通常ResNet-50设计范式由两部分组成：卷积块和全连接层，如下图1所示。通常，由于依赖于数据集的大量类，最后一个线性层的参数比卷积块多。所以我们对最后一个线性层进行切片以检查其最大尺寸。图1：DDP和DDP+TorchShard前向训练流程。在上图1中，传统的DDP训练范式显示在左侧。假设我们有两个类，DDP将为每个类强制执行重复的模型参数。然而，TorchShard将层次参数切片为不同的级别，从而减少了整体GPU内存。现在将一些代码添加到官方的ImageNet训练脚本中，修改后的版本已经是TorchShard项目的一部分。先导入torchshard：importtorchshardasts然后需要初始化model并行进程组，方法同初始化DDP进程组一样。只需要设置一个函数参数，告诉torchshard应该从目标层切出多少个shard。ts.distributed.init_process_group(group_size=args.world_size)接下来将模型转换为并行版本，其中整个模型可以直接输入转换辅助函数而无需特殊处理。importresnetmodel=resnet.__dict__[args.arch](pretrained=args.pretrained)ts.nn.ParallelLinear.convert_parallel_linear(model,dim=args.model_parallel_dim)print("=>parallelingmodel'{}'".format(args.arch)))另外，不要忘记损失函数torchshard.nn.ParallelCrossEntropy，它可以根据输入张量在原始PyTorch版本和并行版本之间切换。例如，如果输入张量是由torchshard并行层产生的，torchshard.nn.ParallelCrossEntropy将并行计算损失值。criterion=ts.nn.ParallelCrossEntropyLoss().cuda(args.gpu)当模型并行模式（TorchShard）和数据并行模式（DDP）一起工作时，我们需要处理并行层的输入。每个类的参数和训练数据都不同。因此，我们在ResNetforward中的平行线性层之前收集输入张量。x=ts.distributed.gather(x,dim=0)#gatherinputalongthedimofbatchsizex=self.fc(x)同样，我们在计算损失值之前收集目标张量。output=model(images)ifargs.enable_model_parallel:target=ts.distributed.gather(target,dim=0)loss=criterion(output,target)最后，使用TorchShard函数保存和加载检查点非常简单。TorchShard提供了一个基本函数torchshard.collect_state_dict用于保存检查点，以及torchshard.relocate_state_dict用于加载检查点。保存检查点：state_dict=model.state_dict()#collectstatesacrossallranksstate_dict=ts.collect_state_dict(model,state_dict)ifts.distributed.get_rank()==0:torch.save(state_dict,'resnet50.pt')#saveasbeforeloadcheckpoint:ifts.distributed.get_rank()==0:state_dict=torch.load('resnet50.pt')#relocatestate_dict()forallranksstate_dict=ts.relocate_state_dict(model,state_dict)model.load_state_dict(state_dict)#loadasbefore现在我们是done为了在ImageNet上添加分片训练的代码，然后可以通过增加类的数量来扩展，也就是最后一个线性层的输出特征维度。训练脚本可以在torchshard/project/imagenet中找到。下图展示了在8个NVIDIATITAN-XP(12196MiB)GPU和1,000,000个类别和16个2,000,000个类别的GPU上训练ResNet-50缩放。图2：在不同并行策略下使用标准ResNet训练设置（即输入大小224和批量大小256）的GPU内存成本。将AMP与ZeROTorchShard结合使用，以一种简单自然的PyTorch方式与其他技术混合，例如自动混合精度AMP和ZeRO。#gradscalerscaler=torch.cuda.amp.GradScaler(enabled=args.enable_amp_mode)withtorch.cuda.amp.autocast(enabled=args.enable_amp_mode):#computeoutputoutput=model(images)ifargs.enable_model_parallel:target=ts.distributed.gather(target,dim=0)loss=criterion(output,target)#computegradientanddoSGDstepscaler.scale(loss).backward()scaler.step(optimizer)scaler.update()optimizer.zero_grad()图3：在不同的并行策略和在AMP下使用标准ResNet训练设置（输入大小224，批量大小256）时使用GPU内存的成本。ZeRO是DeepSpeed的核心，与PyTorch>=1.9.0一起使用。如果要测试某个功能，请安装最新版本的脚本运行，代码如下：fromtorch.distributed.optimimportZeroRedundancyOptimizerifargs.enable_zero_optim:print('=>usingZeroRedundancyOptimizer')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：在不同的并行策略和零优化器下，标准ResNet训练设置（输入大小224和批量大小256）的GPU内存成本。此外，TorchShard还提供了基本的PythonAPI和相应的模板文件，以简化自定义并行层的实现。研究人员将继续开发TorchShard。例如，TorchShard的下一个特性是新的数据采样器torchshard.utils.data.DistributedGroupSampler，其名称遵循torch.utils.data.DistributedSampler。该采样器旨在帮助用户构建M路数据并行、N路模型并行，使其像DDP中的DistributedSampler一样简单。用户唯一需要做的就是设置模型并行组编号，然后DistributedGroupSampler将确保同一模型并行组中的模块具有相同的训练数据。
