大型模型的张量并行性
一种将单个神经网络层内的数学拆分到多个 GPU 上的方法,这样对于一个设备来说太大的模型仍然可以运行。
概述
It matters because frontier models have hundreds of billions of parameters that no single GPU can hold or compute fast enough alone.
深入探讨
张量并行性(也称为层内模型并行性)在 GPU 上分割各个权重矩阵,而不是将整个层放在单独的设备上。在 Transformer 中,大矩阵乘法(注意力投影和前馈 MLP)被拆分:例如,MLP 的第一个权重矩阵按列分区,第二个权重矩阵按行分区,因此每个 GPU 计算一个切片,然后单个 all-reduce 组合结果。注意力分散在多个头中,每个 GPU 处理一个子集。由于每个 GPU 同时执行每一层的一部分,因此张量并行性会减少每个 GPU 的内存并加快计算速度,但它需要每层 GPU 之间进行频繁的高带宽通信。这就是为什么它通常被限制在由 NVLink 连接的节点内,并与管道和数据并行性相结合,以实现非常大的训练和服务作业。
技术洞察
由 Megatron-LM 推广的技巧是选择分区尺寸,从而最大限度地减少通信。按列拆分第一个 MLP 矩阵可让每个 GPU 在本地应用非线性,无需同步;逐行拆分第二行意味着输出只需要一次全归约即可对部分结果求和。因此,每一层大约会产生两个全归约(向前)和两个(向后)。由于这些集合发生在每一层,因此延迟占主导地位,因此张量并行性存在于 NVLink 等快速节点内链接背后,而不是较慢的节点间网络背后。
战略影响
成本与预算
多年来,架构决策决定着性能和运营成本。
更清晰的判决
技术教育帮助团队选择正确的堆栈,而不仅仅是最新的堆栈。
质量控制
更好的工程选择可以减少生产中的可靠性事故。
大型模型张量并行的未来
张量并行性仍然是基础性的,但越来越多地融入“3D并行性”(张量+管道+数据),并与混合专家模型的专家并行性相结合。 Megatron-LM、DeepSpeed 和 vLLM 等框架可自动执行分片。随着 GPU 互连(NVLink、NVSwitch)和光学结构变得更快,节点边界限制放宽,允许更广泛的张量并行组。期望更智能的自动并行化,选择分片尺寸和组大小,以最大限度地减少给定集群拓扑的通信。
现实世界的实施
使用 Megatron-LM 将每一层的权重矩阵分片到一个 NVLink 连接节点中的 8 个 GPU 上,从而训练 175B 参数模型。
在 vLLM 中使用 tensor_parallel_size=4 提供 70B 参数的聊天模型,以便权重适合四个 GPU 并实时响应。
将 Transformer 注意力头拆分到 GPU 上,以便每个设备计算一个子集,然后连接下一层的输出。
结合节点内的张量并行性和跨节点的管道并行性,在大型 GPU 集群上训练万亿参数模型。
风险与防护栏
优化一项基准测试可以隐藏更广泛的系统弱点。
基础设施和维护成本常常被低估。
随着系统变得更加复杂,安全性和可观察性差距可能会扩大。
实施路线图
在实施之前定义延迟、质量和成本目标。
在实际负载和数据条件下进行基准测试。
仪器监控错误、漂移和用户影响。
在扩展之前准备回滚和事件响应路径。
不断探索
Free newsletter
Get the daily AI briefing
Three verified AI stories every weekday morning, written in plain English. Free forever, no ads.
One email each weekday. Unsubscribe in one click. We never sell or share your address.
Test yourself
Take the Tensor Parallelism for Large Models quiz
Instant feedback on every answer, and a shareable certificate with a verifiable ID once you pass a course.
Support free AI education. AI Understanding is a 501(c)(3) nonprofit — no ads, no paywall, ever. Make a donation
常见问题
What is Tensor Parallelism for Large Models?
一种将单个神经网络层内的数学拆分到多个 GPU 上的方法,这样对于一个设备来说太大的模型仍然可以运行。这很重要,因为前沿模型拥有数千亿个参数,没有任何一个 GPU 能够单独保存或计算得足够快。
GPU 之间的张量并行性有何不同?
张量(层内)并行性将层内的权重矩阵进行分片,因此每个 GPU 都计算同一层的一部分,这与管道并行性不同,管道并行性将整个层放置在不同的 GPU 上。
在威震天式 MLP 分割中,两个权重矩阵如何划分以最小化通信?
按列分割第一个矩阵可以让每个 GPU 在本地应用非线性;将第二个按行拆分意味着单个全归约对部分输出求和,从而最大限度地减少同步。
为什么张量并行性通常保留在单个节点内?
每层都会触发集体通信(全归约),因此对延迟敏感的流量需要快速的节点内链接(如 NVLink),而不是较慢的节点间网络。
在张量并行下,注意力机制通常是如何并行的?
注意力头是独立的,因此它们分布在 GPU 上——每个设备计算一些头,然后组合输出。
在张量并行中,什么集体操作通常组合部分结果?
All-reduce 对每个 GPU 上计算的部分输出求和,以便它们就该层的结果达成一致;在前向和后向传递中,每层都会发生多次这种情况。