女性门户网站源码,外贸网络推广哪个好,网页微信版官方,织梦cms同步wordpress演讲人#xff1a;林伟#xff0c;阿里云研究员#xff0c;阿里云人工智能平台 PAI 技术负责人 本文旨在探讨阿里云 TorchAcc#xff0c;这是一个基于 PyTorch/XLA 的大模型分布式训练框架。
过去十年 AI 领域的显著进步#xff0c;关键在于训练技术的革新和模型规模的快…演讲人林伟阿里云研究员阿里云人工智能平台 PAI 技术负责人 本文旨在探讨阿里云 TorchAcc这是一个基于 PyTorch/XLA 的大模型分布式训练框架。
过去十年 AI 领域的显著进步关键在于训练技术的革新和模型规模的快速攀升。尽管大模型展现了堪比人类的理解力但其训练却对算力提出了极高的要求。唯有配备充足的计算资源方能在海量数据上有效训练大模型确保其在有限时间内实现优质收敛。 图片来源于 GTC 2024大会China AI Day 线上专场的演讲《TorchAcc基于TorchXLA的分布式训练框架》
根据上图左侧图表显示过去五年大模型规模的增长态势尤为突出平均每两年大小翻 15 倍而对于 Transformer 为代表的语言模型以及多模态模型而言其规模膨胀速度更加惊人每隔两年以 750 倍剧增。对比之下右侧图表揭示了一个明显的矛盾点不论是单个 GPU 的计算能力抑或是 GPU 显存容量的发展速度都无法跟上模型规模如此急剧的扩张步伐。这一现实状况直接催生了对分布式训练的迫切需求。分布式训练不再局限于以往单纯的数据并行模式而是在此基础上更加重视并采取模型并行策略以弥补单个计算单元算力与存储提升速度相对于模型规模增长的滞后性。
在分布式训练实践中开发人员普遍认同构建模型并行的分布式训练系统相比数据并行更为复杂。数据并行从分布式角度来看其逻辑相对直接和简洁因为每个计算节点执行的任务本质上是对等且一致的。在这种情况下只需在训练过程末尾插入 AllReduce 步骤将各个工作节点worker独立计算出的梯度差异累加整合然后求平均值并将最终梯度结果广播至所有参与工作的节点用以同步更新全局模型参数。
这类简单的分布式训练范式确实呈现出类似单机计算的特点主要涉及全局梯度同步的 AllReduce。然而步入大模型时代由于模型规模过大已无法容纳于单个 GPU 之内我们就必须采用模型并行策略其开发难度也就陡然上升了。
原因是模型并行需要根据模型的规模和结构来决定如何恰当地“分割”模型即将其分割为多个可以平衡计算负载的模块。在不同的分割策略下模型在各个节点上算子的算法实现方式会发生变化同时不同分割方法还会引起节点间通信原语的差异需要精心选择最优分割方案以及配套的通信原语。
在模型分割完成后接下来的任务就是选用适合的通信原语并精细地调度各个算子及其相关的通信操作力求最大化计算与网络通信的重叠overlap以充分发挥底层计算资源的效率。正是由于存在多种可能的分割选项与调度决策寻求最优模型并行策略的复杂性明显高于数据并行对开发者的技巧和经验提出了更高的要求。 图片来源于 GTC 2024 大会 China AI Day 线上专场的演讲《TorchAcc基于 TorchXLA 的分布式训练框架》
本文将围绕四个核心方面展开。首个议题是如何在 TorchAcc 中实现多样化的并行策略涵盖了常规的数据并行以及当下备受关注的 FSDPFully Sharded Data Parallel又称 ZeRO (Zero Redundancy Optimizer)) 。此外还包括了模型并行的各种形态诸如算子并行即 Tensor Parallelism以及流水线并行Pipeline Parallelism等。
TorchAcc 的一大亮点在于其能够自动探寻并有机整合各类并行策略并为用户提供高度自动化的分布式策略配置方案与此同时为了满足高级开发者的定制化需求TorchAcc 还提供了半自动化的控制接口允许用户介入并调整自动探索并行策略的过程从而在兼顾灵活性的同时最大程度地提升训练效率和资源利用率。
通过上述方式TorchAcc 有效地助力算法开发者将精力集中于模型自身的结构设计、训练方法的优化以及追求模型收敛性能的提升上而非花费精力在分布式训练的具体实现细节。TorchAcc 将智能化地协助开发者探寻并实现最佳的分布式训练方案从而显著提升计算资源利用效率和算法迭代效率。
其次模型并行技术的必要性是因为大模型尺寸超出单个 GPU 显存容量的限制。显存容量对于模型训练至关重要如何打破显存瓶颈对于提升分布式训练的整体效率来说至关重要。因此TorchAcc 提供了一种显存智能分配器通过对显存资源的精细化调度与地址分配策略最大限度地提高模型并行训练时的效率确保模型能充分利用现有的显存地址空间。
再者随着模型结构日益复杂且规模不断增大用户对计算资源的需求也在持续攀升因此进一步优化模型在训练过程中的计算密集度及减少访存开销也非常关键。
最后考虑到当前数据中心基础设施的发展趋势大模型训练对网络条件的要求日渐严苛。现代数据中心服务器间的互联带宽已达到 TB 级别以满足大规模模型并行训练对高速数据交换的需求。然而模型并行所带来的复杂通信模式与高频次的数据交互亦会对整体训练效率构成挑战。因此如何有效利用网络带宽减少通信过程在迭代计算中占据的时间比例也就成了训练效率提升的另一重要因素。
在具体实现上TorchAcc 通过一系列技术手段成功地将用户在前端无论是基于 PyTorch 还是 TensorFlow 构建的模型训练过程转化为统一的中间表示层Model IR的 graph。其中对于 TensorFlow 而言因其自身就是一种计算图模型转化过程相对直接而对于 PyTorch我们采用了符号式追踪symbolic tracing以及 LazyTensor 等技术捕获计算图进而转化为 IR Graph。 图片来源于 GTC 2024 大会 China AI Day 线上专场的演讲《TorchAcc基于 TorchXLA 的分布式训练框架》
基于中间表示层IR Graph的构建TorchAcc 实施了一系列多元化的优化策略涵盖计算优化、存储优化、通信优化以及分布式策略优化IR Graph 以各类组合并反复执行这些优化的 Pass 后最终得到一个最优的执行 Plan。然后交由底层 Backend 执行以实现模型训练性能的最大化提升。
通过这一整套方案TorchAcc 在多个模型的分布式训练场景中表现出了显著的性能优势。部分模型的训练过程得以实现高达 3 倍的性能提速充分证明了 TorchAcc 在解决分布式训练难题上的高效性和实用性。 图片来源于 GTC 2024 大会 China AI Day 线上专场的演讲《TorchAcc基于 TorchXLA 的分布式训练框架》
这张图片主要展示了 TorchAcc 的框架总体架构。TorchAcc 以 Pytorch/XLA 为基础并 TorchAcc 依托于 OpenXLA构建了一套大模型训练加速框架。TorchAcc 在处理使用不同前端构建的模型时会灵活采用适宜的图捕获技术如 Symbolic Trace 和 LazyTensor进而生成两种不同层级的图表示FX Graph 和 HLO Graph。其中FX Graph 位于较高抽象层次而 HLO Graph 则更为底层。
基于捕获到的模型计算图TorchAcc 即可进一步展开了四类优化工作即前文提及的计算优化、存储优化、通信优化以及分布式策略优化。 图片来源于 GTC 2024 大会 China AI Day 线上专场的演讲《TorchAcc基于 TorchXLA 的分布式训练框架》
在分布式策略优化层面TorchAcc 支持业界广泛使用的各种并行策略并能够灵活地结合这些策略对给定模型进行有效的并行化处理。具体而言对于数据并行 DPData Parallelism、流水并行 PPPipeline Parallelism以及 FSDPFully Sharded Data Parallel, 也称为 ZeRO这三种分布式策略其实现和优化都是在 FX Graph 这一较高抽象层次上完成的。
选择在 FX Graph 层面对并行策略进行操作的原因在于这一层级所包含的关于计算图结构和操作的信息已足够丰富足以支撑开发人员设计出适应不同并行策略的优化方案。相较于在更低层的 HLO Graph 上直接进行优化由于 FX Graph 具有更高的抽象性和概括性在这一层面上进行优化的成本通常较低更容易实施高效且针对性强的分布式策略调整。
以流水并行作为例子系统能够自动检测 FX Graph 层级上的不同阶段并确定合适的分割点从而有效地将模型分割为多个连续执行的阶段实现流水线并行化。在此过程中我们可以利用 FX Graph 提供的详细计算结构信息来进行智能分割。
至于 Tensor Parallelism 张量并行和 Sequence Parallelism 序列并行这两种更为复杂的并行策略它们要求更为细致精确的信息以便进行决策。为了实现这一点系统需要对前向传播和反向传播的整个计算图的执行计划来进行分析。这时的工作主要在 HLO 这一低级别表示层面上进行。
通过利用 PyTorch/XLA 提供的 mark sharding 接口系统能够在模型参数上添加相应的拆分标记然后将这些拆分信息传递给 OpenXLA 的 SPMD 优化 Pass进而触发计算图的拆分、优化、推导和重写过程最终实现自动的 Tensor Parallelism 和 Sequence Parallelism 功能。 图片来源于 GTC 2024 大会 China AI Day 线上专场的演讲《TorchAcc基于 TorchXLA 的分布式训练框架》
在算子优化层面TorchAcc 引入 FlashAttention 技术来提升 Attention 模块的执行效率。首先通过 XLA 的 custom call 功能将 FlashAttention 的实现无缝地融入到了 OpenXLA 编译器和运行时框架中。这意味着 FlashAttention 可以直接在 XLA 内核层级被执行从而充分利用硬件加速能力。
在整合过程中要处理好在 PyTorch 与 XLA 之间 Tensor 数据的传递问题确保在两个系统间转换时的数据一致性与性能优化同时还要妥善处理 FlashAttention内部参数传递等细节问题保证在并行计算和优化的过程中这些关键参数能够正确且高效地应用到计算中进一步提升模型在执行注意力机制部分的运算速度和资源利用率。
为了用户能便捷地使用 FlashAttention 优化功能我们提供了两种接口用户也可以直接通过 Python 接口调用预先写好的 FlashAttention 算子第三种方法是用户可以使用我们在 OpenXLA 上写好的 Pattern Match Pass该 Pass 能够自动识别计算图中的 Attention Block并将这部分计算结构提取出来替换为FlashAttention 的 custom call。这样设计的优势在于既能充分利用 XLA 原本就十分出色的 Kernel fusion 等算子优化功能又能结合 FlashAttention 带来的先进计算优化技术。 图片来源于 GTC 2024 大会 China AI Day 线上专场的演讲《TorchAcc基于 TorchXLA 的分布式训练框架》
在 Llama 2-7B 模型的性能测试中我们能够明显观察到上述计算优化带来的效果。通过利用 XLA 自身的优化技术尤其是 kernel fusion我们将大量的访存密集型算子做了有效合并从而大幅减少其数量在叠加 FlashAttention 后优化性能进一步提升。 图片来源于 GTC 2024 大会 China AI Day 线上专场的演讲《TorchAcc基于 TorchXLA 的分布式训练框架》
在通信优化层面我们主要完成了三项核心任务以提升分布式训练效率首先我们合并了一些零散的 collective 通讯算子通过减少算子数量来降低通讯开销和调度复杂度其次我们将合并的 collective 通讯算子移至独立的 CUDA Stream 上执行这样一来就能够异步实现计算与通讯的重叠执行。最后我们充分利用了 OpenXLA 的 Latency Hiding Scheduler 功能对通讯算子的调度进行了精细优化使其尽早启动和执行从而增强通讯与计算之间的重叠效果。 图片来源于 GTC 2024 大会 China AI Day 线上专场的演讲《TorchAcc基于 TorchXLA 的分布式训练框架》
通过在 Llama2 -7B 模型上进行的端到端多机性能测试我们发现应用了通讯优化策略后在 128 张 GPU 卡上进行分布式训练优化后的加速比从原来的 88 提升到了 116通过 timeline 图我们也可以直观地看到优化后的通讯算子更加有序并且能够更好地和计算重叠执行。 图片来源于 GTC 2024 大会 China AI Day 线上专场的演讲《TorchAcc基于 TorchXLA 的分布式训练框架》
本文最后一个章节绍 TorchAcc 的显存优化功能该功能通过优化计算图中算子的执行顺序以及 Tensor 在显存中的地址分配来降低显存开销。
如图举例说明假设有一个包含四个算子 V0、V1、V2、V3 的计算图如果不控制算子执行顺序如左图所示按照 V0-V1-V2-V3 的顺序执行若每个 Tensor 按照默认方式进行显存地址申请则可能出现如 B 图左半部分所示的情况即显存容量不足以容纳所有 Tensor导致 out of memory 错误。
然而如果我们能够预判并精细管理内存分配即在分配地址时预知后续执行的算子序列即可如 B 图右半部分所示进行更优的显存布局使得整体计算可在有限显存内顺利完成。更进一步通过精确控制执行顺序比如按照 V0-V2-V1-V3 的方式执行可以进一步压缩显存需求至原始需求的 70% 左右。
这一理念是基于 XLA 中间表示层已有的 scheduler 和 buffer 管理机制我们在此基础上提出了更先进的显存优化方法。目前业界存在多种优化显存分配的方法如启发式算法、约束求解等但这些方法往往难以兼顾时效性和高效性在实际生产环境的集群中应用时可能存在局限性。 图片来源于 GTC 2024 大会 China AI Day 线上专场的演讲《TorchAcc基于 TorchXLA 的分布式训练框架》
在训练场景中实现有效且高效的显存优化是一项极具挑战的任务原因主要包括以下几个方面 NP-Hard 问题本质由于模型的规模、算子的种类繁多以及算子间显存分配的复杂性显存优化问题成为一个典型的 NP-hard 问题即找到全局最优解在计算上通常是不可行的。 算子执行灵活性训练过程中前向传播、反向传播和权重更新等操作具有很高的灵活性特别是在权重更新方面梯度产生后随时可以被用于权重更新但不同的执行时机会影响显存的申请和释放增加了优化难度。 显存复用复杂性在训练过程中前向和反向传播可以通过复用显存减少重新计算但 Tensor 生命周期的多样性和尺寸的变化使得显存复用变得极为复杂这对启发式算法等传统优化手段构成了严峻挑战。
为了解决上述难题我们采取了一种分治策略 Memory-aware Weight Update Scheduler引入了显存感知的权重更新调度器它会根据梯度产生的时机、使用的优化器类型以及当前显存资源状况选择合适的权重更新时间点避免即时更新加重显存压力特别是对于复杂的优化器如 Adam需考虑动量和其他变量的存储。 Graph 分割与局部优化将大计算图根据关键节点 (memory insensitive operator) 分割成多个内存无关性的子图子图间执行顺序固定而子图内部的执行顺序则可以多样化。通过这种方式可以将复杂的全局线性规划问题分解成多个局部问题在子图范围内采用高效的优化方法如线性规划求解最优执行顺序。
通过上述分治策略最终我们能够聚合这些子图的求解结果这也就是我们提出的 ROAM (Reorder Operators and Arrange Tensors Address to Reduce Memory Usage) 这一内存优化探索方式。
上述方法可以成功实现对显存优化问题的高效处理。实验结果显示与原生 PyTorch、启发式算法以及 Facebook 近期基于整数线性规划的优化方法等 baseline 相比ROAM 分别节省了约 16%、13% 和 27% 的显存开销且在优化时长和可扩展性方面表现出色证实了这种方法的有效性。 图片来源于 GTC 2024 大会 China AI Day 线上专场的演讲《TorchAcc基于 TorchXLA 的分布式训练框架》 图片来源于 GTC 2024 大会 China AI Day 线上专场的演讲《TorchAcc基于 TorchXLA 的分布式训练框架》
从另一个维度衡量效果我们考察了算法求解的时间开销。实验证明在常见的深度学习场景中我们的优化算法能够在短短几分钟内得出优化结果。从右图所示对比中可以看出相较于 Facebook 最近提出的 MODeL一种基于线性规划的优化方法我们的方法在求解时间上实现了显著的缩减。原因在于MODeL 在处理大规模图时并未对其进行有效分割而我们的方法通过引入 memory-aware weight update scheduler 和子图划分策略有效地降低了优化问题的空间复杂度从而提高了求解效率。
综上所述TorchAcc 在显存优化、计算优化、通信优化以及并行策略优化等方面均取得显著成效全方位提升了分布式训练的效率与性能。 以上内容来源于 GTC 2024 大会 China AI Day 线上中文演讲专场。扫描图片二维码或登录大会官网即可观看演讲视频并可下载讲义。
