转载:【AI系统】张量并行

news/2024/12/12 18:20:57/文章来源:https://www.cnblogs.com/xueaigc/p/18603156

在大模型的训练中,单个设备往往无法满足计算和存储需求,因此需要借助分布式训练技术。其中,模型并行(Model Parallelism, MP)是一种重要的方法。模型并行的基本思想是将模型的计算任务拆分到不同的设备上执行,以提高训练效率和处理更大规模的模型。下面将重点介绍模型并行中的张量并行。

朴素张量并行

张量并行广泛应用于分布式训练技术。之前的部分已经解释了如何使用数据并行在多个设备上训练神经网络;这种方法一般来说将相同的模型复制到所有设备,每个设备消耗不同部分的输入数据。虽然这可以显著加速训练过程,但在某些情况下模型过大无法放入单个设备时,这种方法并不奏效。

本文展示了如何通过使用朴素张量并行解决这个问题。与数据并行相反,张量并行将单个模型拆分到不同的设备上,而不是在每个设备上复制整个模型(具体来说,假设一个模型 m 包含 6 层:使用数据并行时,每个设备将拥有这 6 层的副本,而使用张量并行在两个设备上时,每个设备只拥有 3 层)。

image

来看一个简单的张量并行的例子:

from torchvision.models.resnet import ResNet, Bottlenecknum_classes = 1000class ModelParallelResNet50(ResNet):def __init__(self, *args, **kwargs):super(ModelParallelResNet50, self).__init__(Bottleneck, [3, 4, 6, 3], num_classes=num_classes, *args, **kwargs)self.seq1 = nn.Sequential(self.conv1, self.bn1, self.relu,self.maxpool, self.layer1,self.layer2).to('npu:0')self.seq2 = nn.Sequential(self.layer3, self.layer4, self.avgpool,).to('npu:1')self.fc.to('npu:1')def forward(self, x):x = self.seq2(self.seq1(x).to('npu:1'))return self.fc(x.view(x.size(0), -1))

上面的代码展示了如何将 torchvision.models.resnet50() 分解到两个设备,将每个块放置在不同的设备上,并移动输入和中间输出以匹配层设备。思路是继承现有的 ResNet 模块,并在构造过程中将层分配到两个设备。然后,重写 forward 方法,通过移动中间输出连接两个子网络。

朴素张量并行实现解决了模型过大无法放入单个设备的问题。然而,你可能已经注意到,如果模型能够放入单个设备,朴素张量并行将比在单个设备上运行更慢。这是因为在任何时候,只有一个设备在工作,而另一个设备处于空闲状态。当中间输出需要从 npu:0 复制到 npu:1 时,性能会进一步恶化。

image

实际上朴素张量并行实现的执行时间比现有的单设备实现慢 7%。因此,可以得出结论,跨设备复制张量的开销约为 7%。仍有改进的空间,因为知道在整个执行过程中有一个设备是空闲的。一种选择是进一步将每个批次分成流水线的分片,这样当一个分片到达第二个子网络时,下一个分片可以进入第一个子网络。这样,两个连续的分片可以在两个设备上并行运行。

朴素朴素张量并行的优点在于实现相对简单,不需要复杂的通信和同步机制。然而,这种方法的缺点也很明显:如果模型的各部分计算量不均衡,可能会导致某些设备的利用率很低,从而影响整体训练效率。此外,对于依赖较强的模型结构,简单的朴素张量并行也可能难以实现。

张量并行

张量并行(Tensor Parallelism,TP)是一种更细粒度的模型并行方法,它将单层内部的参数和计算任务拆分到不同的设备上执行,这种方法特别适合于具有大量参数的大规模模型。它最初是在 Megatron-LM 论文中提出的,它是一种高效的模型并行技术,可用于训练大型 Transformer 模型。

通过张量并行,可以将矩阵乘法等计算操作的矩阵按行或按列切分,然后在不同设备上并行执行部分计算,最后通过集合通信操作合并结果。张量并行可以分为 MatMul 并行、Transformer 并行、Embedding 并行、Cross Entropy Loss 并行。

其中序列并行(Sequence Parallel,SP)也是张量并行的一种变体,它在序列维度上对 nn.LayerNorm 或 RMSNorm 进行分割,以进一步节省训练过程中的激活内存。当模型变得越来越大时,激活内存就会成为瓶颈,因此在张量并行训练中,通常会将序列并行应用于 LayerNorm 或 RMSNorm 层。

张量并行的主要挑战在于如何切分参数和计算任务,以保证计算的一致性和通信的高效性。例如,在进行矩阵乘法时,必须确保各设备上的部分结果在数学上是一致的。此外,通信开销也是一个重要考虑因素,需要在计算和通信之间找到平衡点,以达到最佳性能。

MatMul 并行

矩阵乘法(MatMul)是深度学习中最常见的操作之一。在张量并行中,可以将矩阵按列或者按行切分,然后在不同设备上并行执行部分计算。以矩阵乘法 $A \times B = C$ 为例,假设将矩阵 $B$ 按列切分成 $B_1$ 和 $B_2$,分别存储在设备 1 和设备 2 上。在这种情况下,设备 1 和设备 2 可以分别计算 $B_1 \times A$ 和 $B_2 \times A$,最终通过合并结果得到 $C$。

image

Transformer 并行

在 Transformer 模型中,主要包括多层感知机(MLP)和自注意力(Self-Attention)模块,它们的计算本质上也是矩阵乘法。对于 MLP 模块,可以将输入矩阵 $X$ 和权重矩阵 $A$ 按列切分,不同设备分别计算一部分乘积,然后合并结果。对于自注意力模块,可以将查询(Query)、键(Key)和值(Value)矩阵按列切分,不同设备分别计算注意力得分和加权求和,最后合并结果。

image

对于多层感知机(MLP),对 A 采用列切割,对 B 采用行切割,在初始时使用函数 f 复制 X,结束时使用函数 g 通过 All-Reduce 汇总 Z,这样设计的原因是,尽量保证各设备上的计算相互独立,减少通信量。对 A 来说,需要做一次 GELU 计算,而 GELU 函数是非线形的,$GeLU(X + Y) \not = GeLU(X) + GeLU(Y)$,对 A 采用列切割,那每块设备就可以继续独立计算了。

对于自注意力(Self-Attention)对三个参数矩阵 Q K V,按照列切割。对线性层 B,按照行切割,切割的方式和 MLP 层基本一致。

需要注意的是在使用 dropout 时两个设备独立计算,第一个 dropout 在初始化时需要用不同的随机种子,这样才等价于对完整的 dropout 做初始化,然后再切割。最后一个 dropout 需要用相同的随机种子,保证一致性。

Embedding 并行

在大型 Transformer 模型中(如:LLM),词嵌入的并行处理是一种有效的技术,可以减轻单个设备的内存负担并提高计算效率,通常有两种主要的切分方式:表切分(Table-wise split)和列切分(Column-wise split)。

image

  • 表切分模式(a)下,每个设备存储部分的嵌入表。例如:每个嵌入表对应一个类别特征,每个设备存储一部分嵌入表。设备 1 存储嵌入表 0,设备 2 存储嵌入表 1,依此类推。在这种模式下,每个设备只处理它存储的嵌入表对应的类别特征。这个方法的优点是每个设备只需处理和存储一部分数据,减少了单个设备的内存负担。

  • 列切分模式(b)下,每个设备存储每个嵌入表的一部分列。例如:将嵌入表按列切分,每个设备存储不同的列范围。设备 1 存储嵌入表 0 的 0 到 63 维度,设备 2 存储嵌入表 0 的 64 到 127 维度,依此类推。在这种模式下,每个设备处理它存储的部分列,并行计算所有嵌入表的部分结果。然后通过 All-Gather 操作将各部分结果汇总,得到完整的嵌入输出。

表切分模式适用于类别特征较多的场景,每个类别特征的嵌入表较小。而列切分模式适用于单个嵌入表较大的场景,每个嵌入表的列数较多。

Cross Entropy Loss 并行

Cross Entropy Loss 并行用于在计算损失函数时节省内存和通信,因为模型输出通常非常大。在 Cross Entropy Loss 并行中,当模型输出在(通常是巨大的)词汇维度上进行分片时,可以高效地计算交叉熵损失,而无需将所有模型输出聚集到每一个设备上。这不仅大大减少了内存消耗,还通过减少通信开销和并行分片计算提高了训练速度。

image

Cross Entropy Loss 并行可以分为以下几步:

  1. 数据拆分:将 logits (input) 按照 vocab 维度进行拆分,同时将不同部分分发到各设备,labels (target) 需要先进行 one hot 操作,然后 scatter 到各个设备上;

  2. input(logits) 最大值同步:input(logits) 需要减去其最大值后求 softmax,All Reduce (Max) 操作保证了获取的是全局最大值,有效防止溢出;

  3. exp sum 与 softmax 计算:exp sum 即 softmax 计算中的分母部分,All Reduce (Max) 操作保证了获取的是全局的和;

  4. 计算 Loss:input (logits) 与 one_hot 相乘并求和,得到 label 位置值 im ,并进行 all_reduce (Sum) 全局同步,最后计算 log softmax 操作并加上负号,得到分布式交叉熵的损失值 loss。

DeviceMesh 实现 TP

可以通过 PyTorch DeviceMesh 进行多维度并行的实现。

PyTorch 的张量并行应用程序接口(PyTorch Tensor Parallel APIs)提供了一套模块级原语,用于为模型的各个层配置分片功能。它利用 PyTorch DTensor 进行分片张量封装,DeviceMesh 抽象进行设备管理和分片。它们分为:

  • ColwiseParallel 和 RowwiseParallel:以列或行方式对 Linear 和 Embedding 层进行分片。

  • SequenceParallel:在 LayerNorm 和 Dropout 上执行分片计算。

  • PrepareModuleInput 和 PrepareModuleOutput:通过正确的通信操作配置模块输入输出分片布局。

由于 Tensor Parallel 会将单个张量分片到一组设备上,因此需要先建立分布式环境。Tensor Parallelism 是一种单程序多数据(SPMD)分片算法,类似于 PyTorch DDP/FSDP,它通常在一台主机内工作。

下面尝试初始化一个 8 NPU 的张量并行:

# run this via torchrun: torchrun --standalone --nproc_per_node=8 ./tp_tutorial.py
import torch.nn.functional as F
from torch.distributed.tensor.parallel import loss_parallel
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor.parallel import (parallelize_module,ColwiseParallel,RowwiseParallel,PrepareModuleInput,SequenceParallel,
)tp_mesh = init_device_mesh("cuda", (8,))layer_tp_plan = {# Now the input and output of SequenceParallel has Shard(1) layouts,# to represent the input/output tensors sharded on the sequence dimension"attention_norm": SequenceParallel(),"attention": PrepareModuleInput(input_layouts=(Shard(1),),desired_input_layouts=(Replicate(),),),"attention.wq": ColwiseParallel(),"attention.wk": ColwiseParallel(),"attention.wv": ColwiseParallel(),"attention.wo": RowwiseParallel(output_layouts=Shard(1)),"ffn_norm": SequenceParallel(),"feed_forward": PrepareModuleInput(input_layouts=(Shard(1),),desired_input_layouts=(Replicate(),),),"feed_forward.w1": ColwiseParallel(),"feed_forward.w2": RowwiseParallel(output_layouts=Shard(1)),"feed_forward.w3": ColwiseParallel(),
}model = parallelize_module(model,tp_mesh,{"tok_embeddings": RowwiseParallel(input_layouts=Replicate(),output_layouts=Shard(1),),"norm": SequenceParallel(),"output": ColwiseParallel(input_layouts=Shard(1),# use DTensor as the outputuse_local_output=False,),},
)pred = model(input_ids)
with loss_parallel():# assuming pred and labels are of the shape [batch, seq, vocab]loss = F.cross_entropy(pred.flatten(0, 1), labels.flatten(0, 1))loss.backward()

这里使用 init_device_mesh 函数初始化设备网格 tp_mesh。这个网格指定了将使用 8 个 NPU 进行并行计算。

定义一个 layer_tp_plan 字典,指定了模型中各层的并行化策略。通过 parallelize_module 函数,可以并行化模型,指定 tok_embeddings 层进行行并行化,设置对输入进行复制,输出为分片布局(非本地输出);且对 norm 层进行序列并行化。

在前向传播过程中,通过并行化的模型计算预测值 pred。在 loss_parallel 上下文中,进行张量并行交叉熵损失计算,并执行反向传播以计算梯度。

如果您想了解更多AI知识,与AI专业人士交流,请立即访问昇腾社区官方网站https://www.hiascend.com/或者深入研读《AI系统:原理与架构》一书,这里汇聚了海量的AI学习资源和实践课程,为您的AI技术成长提供强劲动力。不仅如此,您还有机会投身于全国昇腾AI创新大赛和昇腾AI开发者创享日等盛事,发现AI世界的无限奥秘~
转载自:https://www.cnblogs.com/ZOMI/articles/18562946

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处:http://www.hqwc.cn/news/851506.html

如若内容造成侵权/违法违规/事实不符,请联系编程知识网进行投诉反馈email:809451989@qq.com,一经查实,立即删除!

相关文章

转载:【AI系统】动手实现 PyTorch 微分

这里记录一下使用操作符重载(OO)编程方式的自动微分,其中数学实现模式则是使用反向模式(Reverse Mode),综合起来就叫做反向 OO 实现 AD 啦。 基础知识 下面一起来回顾一下操作符重载和反向模式的一些基本概念,然后一起去尝试着用 Python 去实现 PyTorch 这个 AI 框架中最…

转载:【AI系统】计算图原理

在前面的文章曾经提到过,目前主流的 AI 框架都选择使用计算图来抽象神经网络计算表达,通过通用的数据结构(张量)来理解、表达和执行神经网络模型,通过计算图可以把 AI 系统化的问题形象地表示出来。 本文将会以 AI 概念落地的时候,遇到的一些问题与挑战,因此引出了计算图…

转载:【AI系统】计算图基本介绍

在 AI 框架发展的最近一个阶段,技术上主要以计算图来描述神经网络。前期实践最终催生出了工业级 AI:TensorFlow 和 PyTorch,这一时期同时伴随着如 Chainer、DyNet、CNTK、PaddlePaddle、JAX 等激发了框架设计灵感的诸多实验课程。 TensorFlow 和 PyTorch,特别是 PyTorch 代…

转载:【AI系统】什么是微分

自动微分(Automatic Differentiation,AD)是一种对计算机程序进行高效准确求导的技术,一直被广泛应用于计算流体力学、大气科学、工业设计仿真优化等领域。 近年来,机器学习技术的兴起也驱动着对自动微分技术的研究进入一个新的阶段。随着自动微分和其他微分技术研究的深入…

转载:【AI系统】Auto-Tuning 原理

在硬件平台驱动算子运行需要使用各种优化方式来提高性能,然而传统的手工编写算子库面临各种窘境,衍生出了自动生成高性能算子的的方式,称为自动调优。在本文我们首先分析传统算子库面临的挑战,之后介绍基于 TVM 的业界领先的三个自动调优系统。 高性能算子挑战 DNN 部署的硬…

转载:【AI系统】微分计算模式

上一篇文章简单了解计算机中常用几种微分方式。本文将深入介绍 AI 框架离不开的核心功能:自动微分。 而自动微分则是分为前向微分和后向微分两种实现模式,不同的实现模式有不同的机制和计算逻辑,而无论哪种模式都离不开雅克比矩阵,所以我们也会深入了解一下雅克比矩阵的原理…

转载:【AI系统】AI 框架之争

在前面的内容主要是讲述了 AI 框架在数学上对自动微分进行表达和处理,最后表示称为开发者和应用程序都能很好地去编写深度学习中神经网络的工具和库,整体流程如下所示:除了要回答最核心的数学表示原理以外,实际上 AI 框架还要思考和解决许多问题,如 AI 框架如何对实际的神…

转载:【AI系统】指令和存储优化

除了应用极广的循环优化,在 AI 编译器底层还存在指令和存储这两种不同优化。 指令优化 指令优化依赖于硬件提供的特殊加速计算指令。这些指令,如向量化和张量化,能够显著提高计算密度和执行效率。向量化允许我们并行处理数据,而张量化则进一步扩展了这一概念,通过将数据组…

转载:【AI系统】算子循环优化

在具体硬件执行计算的时候,实际会大量地使用 for 等循环指令不断地去读取不同的数据执行重复的指令(SIMT/SIMD),因此循环优化主要是为了提升数据的局部性或者计算的并行性,从而提升整体算子性能,当然这二者都需要 AI 芯片硬件的支持。 循环优化挑战 数据局部性 数据的局部…

数字组合转字母删除二叉树节点字符串相乘打家劫舍ii无序数组第k大 无序数组前k大两个有序数组合并中文数字转换为整数最大连续子数组和零钱凑数

一、数字串转换为字符串 1-26个数字分别代表26个字符(A-z)输入"12326〞就可以拆分为【1,2,3,2,6】、(12, 3, 2, 6]. [1, 23, 2, 6]【1,23,26】、【12,3,26】等,将每种组合转成成对应字母输出,输出所有可能的结果返回所有可能的转换结果// 将数字串转换成字母串// 将数…

PDFgear强大的PDF转换编辑工具

点击上方蓝字睿共享资源关注我 前言 PDFgear是一个很好用的PDF软件,它不仅能看PDF文件,还能把PDF转换成其他格式的文件,或者编辑它们。而且,它还有很多实用的功能,比如你可以在PDF上签名、填写表单等等。这个软件的界面看起来很舒服,用起来也很简单,不管你是新手还是老手…