Meta首颗训练芯片,正式亮相

半导体行业观察·2026年08月18日 12:26
Meta发布首款DLRM训练AI芯片MTIA 300,性能媲美GPU。

在今年三月,Meta披露了MTIA 300等一系列芯片的计划。日前,他们终于揭开了路线图首颗芯片的神秘面纱。

据介绍,MTIA 300是 Meta 首款针对深度学习推荐模型(DLRM:Deep Learning Recommendation Models ) 优化的 AI 训练芯片。与 GenAI 训练不同,DLRM 训练虽然所需的浮点运算能力不高,但需要大容量内存、高网络带宽和频繁的集体通信;这种组合通常会导致加速器利用率低下。MTIA 300 通过三项关键创新应对这些挑战:

内置 12×800 Gbps RDMA 网卡的网卡芯片组,确保高网络性能并消除 PCIe 开销;

专用消息引擎用于集体卸载,在仅占用三分之一芯片面积的情况下,提供与计算引擎相当的吞吐量;

近内存计算,将消息引擎放置在靠近 HBM、缓存和 I/O 的位置,并配备专用硬件来加速基于归约的集体操作,例如 ReduceScatter。

虽然 MTIA 300 的设计初衷是面向 DLRM,但这些广泛适用的设计原则也适用于后续针对 GenAI 模型优化的 MTIA 系列芯片(MTIA 400、450 和 500)。

据我们所知,MTIA 300 是首款内置网卡芯片和通用集体卸载引擎的加速器,能够实现高效的纵向扩展和横向扩展通信。此外,MTIA 300 的中等浮点运算能力和更高的 HBM 带宽及容量最初是针对训练 DLRM 模型而优化的,但这些特性也使其适用于 GenAI 推理。性能评估表明,MTIA 300 在性能上可与同代 GPU 相媲美,同时还具有成本优势。

简介

人工智能工作负载在整个行业持续快速增长。在 Meta,Facebook 和 Instagram 等产品依赖深度学习推荐模型 (DLRM) 来提供个性化内容,包括广告、短视频和好友动态。这种增长推动了 Meta 自主研发人工智能芯片。我们之前推出了 MTIA 1 和 MTIA 2i(又名 MTIA 100 和 200),这两款芯片均针对 DLRM 推理进行了优化,其中 MTIA 2i 目前已部署数十万颗芯片。

本文将介绍这一发展历程的下一步:Meta 的首款训练芯片 MTIA 300。

Meta 支持两大类人工智能工作负载:GenAI 和 DLRM。DLRM 训练与 GenAI 训练有所不同:它需要的浮点运算次数 (FLOPS) 更少,但需要更大的 HBM 容量、更高的网络带宽和更频繁的通信。这种组合通常会导致加速器利用率较低。与通用 GPU 不同,MTIA 300 针对 DLRM 进行了优化,其设计有意与 GPU 架构有所不同,以体现这一重点。

为了阐明 MTIA 300 针对深度学习模型 (DLRM) 的设计动机,我们首先总结了工作负载特性。DLRM 通常将多层感知器 (MLP:multi-layer perceptrons ) 应用于密集特征(例如用户年龄),并将嵌入表应用于稀疏类别特征(例如帖子 ID),并通过密集交互层连接。密集组件需要高浮点运算能力 (FLOPS)(尽管远低于 GenAI),而稀疏组件需要不定期地访问超大型嵌入表,并且通常受限于内存或指令。由于嵌入表通常超过单个加速器的内存容量(有时占模型参数的 99% 以上),DLRM 在训练中采用混合并行:密集层使用数据并行,而稀疏层按表或按行对嵌入表进行分片。这种方法能够在不牺牲密集层计算效率的情况下训练大型模型。

针对不同模型组件的数据并行和模型并行相结合,引入了复杂的通信模式。在数据并行方面,每个加速器接收一个本地批次,并执行 AllReduce 操作以同步梯度。模型并行则通过 AllToAllv 集合来交换特征,并在前向和反向传播过程中重新分配结果。此外,许多 DLRM 算法对密集组件使用分布式 Shampoo 优化器,该优化器在优化阶段会添加 AllGather 操作。高效执行这些集合对于实现高性能至关重要。

由于 DLRM 对通信数据路径提出了更高的要求,MTIA 300 集成了以下特性来应对这一挑战:

• 内置网卡芯片:MTIA 300 采用芯片架构,嵌入两个网卡芯片,共包含 12 个高度优化的 800Gbps RDMA 网卡。内置网卡避免了加速器和网卡之间的 PCIe 开销,并且网卡可以灵活地用于纵向或横向扩展网络。

协同卸载:在 GPU 中,计算引擎和主机 CPU 处理协同操作,这通常效率低下。相比之下,MTIA 300 使用专用消息引擎,在仅占用三分之一芯片面积的情况下,即可为这些操作提供与计算引擎相同的通信吞吐量。

近内存计算:计算引擎和消息引擎共享片上网络,用于内存和 I/O 访问。为了避免高带宽集合操作造成拥塞,消息引擎被放置在芯片边缘,靠近 HBM、缓存和 I/O 口。消息引擎的近内存计算逻辑块为所有基于归约的集合操作(包括 Reduce、AllReduce 和 ReduceScatter)提供高吞吐量。

除了概述 MTIA 300 独特的硬件特性和整体架构外,我们还描述了软件堆栈,重点介绍了 collectives 库,该库提供了一个熟悉的界面,同时有效地利用了 MTIA 的专用消息引擎和内置网卡,以实现高效通信。

虽然已有许多GPU和AI ASIC芯片被报道,但我们的经验表明,DLRM存在一些独特的需求,而之前的研究并未涉及这些需求。据我们所知,MTIA 300是首款内置网卡芯片和通用集体卸载引擎的加速器,它避免了使用计算引擎进行集体计算的低效性,并实现了灵活的网络配置,可同时满足纵向扩展和横向扩展的需求。相比之下,尽管TPU的稀疏核心也能卸载对嵌入表的远程访问,但它专用于非RDMA、非交换式环面网络,并且缺乏通用的集体库接口,这限制了其在其他行业加速器中的应用,而这些加速器通常基于RDMA和类似的集体库接口构建。

尽管 MTIA 300 的显著特征——内置网卡芯片、集体卸载和近内存计算——最初是为 DLRM 训练而设计的,但这些设计原则仍然具有广泛的适用性,并已被后续针对 GenAI 模型优化的 MTIA 系列产品所采用。MTIA 300 的研发始于数年前,旨在与 H100 和 H200 GPU 竞争,而其后续产品 MTIA 400 则旨在与 GB300 GPU 一较高下。此外,即将推出的 MTIA 450 和 500 的目标是在 GenAI 推理性能方面超越未来的 GPU,达到业界领先水平。

MTIA 300架构

图 1 展示了 MTIA 300 的概览,它采用了芯片组架构。为了突出其独特之处,我们首先将 MTIA 300 与 MTIA-2i 进行比较,然后再与 GPU 进行比较。

A. MTIA 300 与 MTIA-2i 和 GPU 的比较

表 I 比较了 MTIA 300 和 MTIA-2i 的规格。MTIA-2i 专为推理而设计,而 MTIA 300 专为训练而设计,MTIA 300 引入了多项改进:面积增大约 3 倍,热设计功耗 (TDP) 提高约 10 倍,采用液冷(而非风冷),使用 HBM3E(而非 LPDDR),FP8 计算(而非 INT8),BF16 FLOPS 提升超过 3 倍,采用 2.5D CoWoS 封装,计算芯片尺寸与光罩尺寸相同,支持 RoCE 的网络芯片,以及用于卸载集体通信的消息引擎。值得注意的是,FP32 的 SIMD 计算能力提升超过 6 倍,使得 GEMM 与 SIMD 的比例达到 16:1,而 MTIA-2i 的比例为 32:1;考虑到训练中非GEMM计算(例如,基于表格的批量嵌入前向/后向迭代以及优化器)的需求不断增长且种类繁多,这一增长幅度相当可观。

与通用GPU不同,MTIA 300针对DLRM进行了优化。表II重点介绍了它与H100 GPU的区别。值得注意的是,MTIA 300弱化了峰值浮点运算性能,而着重提升了HBM带宽和网络性能,其特点是内置网卡芯片和专用的硬件支持,用于集体通信卸载。

B. MTIA 300 计算芯片

图 2 展示了 MTIA 300 计算芯片的架构。它由一个 12×6 的处理单元 (PE) 网格组成,用于执行计算任务,并包含 16 个消息引擎 (ME),用于执行集体操作。在东西两侧,SRAM 存储体可用作末级缓存 (LLC) 或末级暂存区 (LLS)。每一侧都连接到 3 个 12 层高的 HBM3E 堆栈。处理单元之间以及处理单元与片上和片外存储器之间通过网状互连连接。

片上网络 (NoC):NoC 是一个二维路由器网状结构,用于连接主网格内的处理单元和消息引擎。它还将计算芯片连接到控制和主机接口模块,以及位于南北两侧的芯片接口 IP 模块,这些模块连接到网络芯片。NoC 提供数据、控制、实用程序(例如,寄存器访问和调试)、同步和归约通道。为了提升性能和可扩展性,我们引入了集群路由器,将六个PE节点本地连接起来,从而降低了总跳数延迟。与MTIA-2i不同,计算芯片不使用内存交叉开关;而是由片上网络(NoC)处理存储体选择路由。它采用Lrouting算法——先沿一个维度(例如X轴)移动,再沿另一个维度(Y轴)移动——将流量均匀地分配到整个网格上,并使用虚拟通道来避免死锁。

主机接口:MTIA 300 提供高性能主机接口,支持 PCIe、DMA 和安全启动处理器。它包含用于主机管理计算和网络芯片的接口,以及调试接口。

控制核心:这是一个 RISC-V 四路 SMP 核心,负责协调 PE 和 ME 之间的执行。它包含相关的上下文 RAM、邮箱寄存器和 MSI-X 中断。

冗余:为了提高光刻胶限制芯片的良率,计算芯片包含一行冗余的 PE。由于 PE 占用面积和分布式内存最大,并且考虑到内存的东西向组织结构和片上网络 (NoC) 的路由路径,添加冗余行是最简单的解决方案。每个 PE 列可以通过替换冗余行中对应的 PE 来容忍一个故障 PE。此配置在启动时完成,对软件透明,并且不会影响 NoC 的性能。

C. 处理单元 (PE)

图 3 显示了处理单元 (PE) 的内部架构。每个 PE 包含两个 RISC-V 内核、用于加速计算和数据传输的固定功能单元,以及通过内存桥连接所有组件的内部存储器。

内存桥 (MB:Memory Bridge ):内存桥通过内部片上网络 (NoC) 提供 PE 中所有组件之间的数据和配置连接。它还包含中断控制器、机器定时器和调试/跟踪模块等外设。

本地存储器:每个 PE 包含 512 KB 的快速本地存储器 (LS)。LS 由软件管理,并被分割成大小可控的循环缓冲区 (CB)。

RISC-V 内核:RISC-V 内核执行应用程序代码,并向固定功能单元的命令处理器发出命令。MTIA 300 具有两个 64 B 宽的矢量内核,可提供额外的 SIMD 吞吐量和对称编程模型,允许相同的代码序列在任一内核上运行。MTIA 使用异步数据流执行模型。程序员编写一个内核,为固定功能单元生成一系列自定义指令,在解决依赖关系的过程中进行数据移动和计算。

内存布局单元 (MLU):MLU 执行内存布局变换,包括转置、重塑、切片和连接。

点积引擎 (DPE:Dot Product Engine):DPE 执行通用矩阵乘法 (GEMM) 运算,用于训练的前向和反向传播。它处理两个输入张量:第一个张量被读取并缓存到 DPE 中,而第二个张量从 LS 流式传输,并与第一个张量的所有行进行点积运算。DPE 包含两个 32×64B×32 的乘加 (MAC) 模块,每个 PE 的总吞吐量为 7.82 TFLOPS,支持 FP16/BF16 输入和 FP32 输出。它还支持 FP8 输入(S1E4M3 或 S1E5M2 格式)和 TF32 输入,这对于某些需要更高精度的排名和推荐应用场景非常有用。

归约引擎(RE:Reduction Engine):RE 存储来自 DPE 的中间矩阵乘法结果,并通过专用归约网络执行 PE 间的归约操作。它可以接收并累积结果,然后再将其转发给下一个 PE 或 SIMD 引擎进行进一步处理。

SIMD引擎(SFU):SFU支持量化、逐元素运算和非线性函数。它由一个执行流水线组成,包含浮点ALU和查找表(LUT),用于近似计算非线性函数。SFU可以从正则表达式(RE)接收输入,也可以直接从逻辑序列(LS)读取数据。为了进行训练,我们移除了INT8,并添加了FP8以及FP16、BF16和FP32。SIMD宽度从每个周期32个元素增加到128个元素,以实现16:1的GEMM:SIMD比率(MTIA-2i的一半),这反映了计算中很大一部分都消耗在了非GEMM运算上。两个RISC-V内核提供了额外的SIMD吞吐量。

基于DLRM训练的需求,我们对SFU进行了多项改进,包括提高高精度数据类型非线性运算的吞吐量,以及增加对最小值/最大值、钳位和随机舍入的支持。 MTIA 300 还支持硬件加速的基数排序,以加速反向嵌入操作。在前向传播过程中,稀疏偏移量和索引被打包,使得单个输出索引映射到输入的一个连续子集。在反向传播过程中,稀疏索引必须排序,使得一个连续子集映射到单个嵌入表索引。基数排序从 LS 中获取元素,通过分桶对其进行排序,创建直方图,并将分桶后的元素存储在内存中,从而加速反向嵌入操作。

命令处理器 (CP:Command Processor):该处理器负责处理 RISC-V 内核中自定义指令在固定功能单元之间的执行,包括调度和依赖关系检查。CP 仲裁 RISC-V 内核和固定功能单元之间的 LS 访问。它还为程序员提供循环缓冲区 (CB) 抽象,并管理依赖关系跟踪,以确保生产者和消费者正确使用 CB。

Fabric 接口 (FI:Fabric Interface):FI 是一个 DMA 引擎,用于在 PE 本地内存和片上或片外内存之间通过片上网络 (NoC) 传输数据。它还强制执行数据包分片和漏桶流量整形,以平滑流量并减少拥塞。

FI 和命令处理器的两项增强功能提供了更强大的数据移动抽象。首先,MTIA 300 支持用于张量切片的字节对齐 DMA,从而消除了布局转换的软件开销。其次,它增加了用于分散和聚集操作的硬件加速索引 DMA 传输。命令处理器使用 LS 中的索引列表生成读取或写入序列,这对于嵌入表查找尤其有用。

D. 消息引擎 (ME:Message Engine)

由于 DLRM 训练对通信基础设施要求很高,ME 的设计旨在解决 GPU 的限制并实现以下目标:

避免主机参与数据路径:通过将网卡集成到 MTIA 300 封装中,我们避免了使用 PCIe 数据路径。此外,通过主机 CPU 管理 1.2 TB/s 的 I/O 会减慢工作提交和完成队列的处理速度,并占用大量主机核心。因此,我们将这些任务卸载到 ME。

从 PE 卸载集体操作:在 GPU 中,计算核心负责处理集体归约操作,这可能造成面积效率低下。ME 仅使用 PE 核心三分之一的面积即可提供类似的归约带宽,从而最大限度地提高集体操作的面积效率。

减少片上网络 (NoC) 的争用:计算核心上的高带宽集体通信也会因流量过大而给 NoC 带来压力。将 ME 放置在 PE 网格的边缘可以最大限度地减少跨网格拥塞(图 4)。

图 5 展示了为解决这些限制而设计的 ME 架构。它由三个主要功能模块组成。

CPU-M 和外设:ME 包含一个标量 RISC-V 内核(CPU-M),类似于 PE 的矢量内核,以及 256 KB 的上下文 SRAM。一个重要的特性是每个 ME 都拥有一个共享的大型完成队列 (CQ),这消除了轮询多个队列的需要,并防止了 CQ 溢出。

网卡接口:由于 MTIA 300 封装内包含 12 个独立的 RDMA 网卡,我们希望避免 ME 在软件中管理大量的门铃地址。这由网卡接口处理,该接口接收单个 FIFO 队列中的工作请求 (WR),并将其分发到相应网卡上的正确门铃。

近内存计算 (NMC:Near Memory Compute):NMC 是一个归约模块,每个周期可处理 128 字节的归约或 DMA 数据,如果所有模块同时运行,则每个周期的处理能力会降至 96 字节。ME 可提供高达 2.8 TB/s 的归约处理能力,是 I/O 带宽 (1.2 TB/s) 的两倍多。所有基于归约的集体操作(包括 Reduce、AllReduce 和 ReduceScatter)都会用到该模块。

E. MTIA 300 网络芯片

为了避免 PCIe 的开销,我们将基于第三方网卡 IP 的 RDMA 网卡直接集成到 MTIA 300 封装中,作为网络芯片,如图 4 所示。每个网络芯片包含六个定制的 800 Gbps (100 GB/s) RDMA IP 模块,每个芯片可提供 600 GB/s 的吞吐量。我们采用芯片间接口和 112G SerDes 来实现高带宽密度。定制的 RDMA IP 模块进行了如下优化。

Express Doorbells:为了最大限度地减少事务提交延迟,我们引入了“Express Doorbells”机制,该机制使用工作请求 (WR) 本身作为门铃写入操作,避免了额外的 HBM 环形缓冲区读取(每次事务 800 ns)。借助快速门铃,每个 IP 模块最多可支持 24,576 个未完成的工作请求,这些请求分布在 1,024 个队列对 (QP) 中。

移除 QP 缓存:我们移除了 QP 缓存,因为它会占用大量的芯片面积。这会将每个网卡芯片组的活动 QP 数量限制在 1100 个,远远超过我们的工作负载需求(通常只有几百个进程)。

简化数据包处理流水线:典型的网卡支持许多功能,例如虚拟交换和 TC 卸载(例如 cls flower),而这些功能我们并不需要。移除这些功能可以简化数据包处理流水线。

AXI 转向标签(AXI steering tag):通过支持自定义转向标签,我们可以利用计算芯片组中针对不同流量类型的独立缓存分区等功能。

系统架构:机架和网络

A. 机架系统

随着模型快速演进,我们致力于提供灵活的系统架构,以适应突发变更,弥补芯片开发流程的不足——芯片开发需要提前很久做出决策。为此,我们设计了采用独立计算刀片和网络刀片的 MTIA 训练系统。每个 MTIA 训练机箱包含 16 个计算刀片插槽和 6 个网络刀片插槽,所有插槽均通过一对电缆背板连接。并非所有网络刀片都需要安装。计算刀片采用垂直放置,以最大限度地减小电缆背板的尺寸,因为电缆背板历来对加工精度要求较高。

我们设计了两种类型的网络刀片,分别用于纵向扩展(scale-up)和横向扩展(scale-out)。纵向扩展刀片采用延迟和功耗更低的 ASIC 芯片,多个机架可以组合成更大的纵向扩展域。横向扩展刀片支持解耦调度架构,该架构通过使用数据包喷射来避免热链路,同时确保数据包按序交付,并通过架构级端到端信用机制提供可靠性。每个计算刀片可提供 200 GB/s 的 I/O 性能,刀片数量可根据所选的纵向扩展或横向扩展方案进行配置。两个刀片均采用液冷散热。

计算刀片包含一个 CPU、512 GB 内存和一个 MTIA 300 加速器。1:1 的映射关系简化了节点结构,使其能够应对不可预测的 CPU 需求,并避免 PCIe 争用。CPU 和加速器均采用液冷散热。

B. 网络设计

MTIA 300 的网络架构如图 7 所示。横向扩展网络提供 200 GB/s 的带宽,而纵向扩展网络目前每个加速器可提供 800 GB/s 的带宽,并可选配最高 1 TB/s 的带宽。这种方案可根据需求优化功耗和成本。我们使用一个包含 16 个节点(一个机架)的纵向扩展域和一个包含 4,096 个节点的一级横向扩展域,并可根据需要添加二级交换网络,将横向扩展域扩展到 16,000 个节点或更多。

软件栈

图 8 展示了 MTIA 300 的软件栈,它提供原生的 PyTorch 体验。该软件栈支持多种 PyTorch 库,例如 FSDP2、DTensor、TorchRec 和 XFormers,可快速开发训练和推理解决方案。它采用最新的 PyTorch 技术栈,并使用 TorchDynamo 进行图捕获,使用 TorchInductor 生成 Triton 代码。它同时支持 eager 模式和图模式,提供 CUDA 兼容的运行时 API,并允许使用 C++ 或 Triton 编写内核。我们还利用编码代理实现内核的自动化生成。

MTIA 编译器作为定制的 PyTorch 后端构建,使用 TorchDynamo 追踪正向图,并通过 AOTAutograd 生成反向图。它应用了 MTIA 优化的算子分解,并支持手写模式融合和通过 TorchInductor 实现的编译器驱动融合。它实现了一套针对训练工作负载的内存优化方案。例如,其图调度器使用启发式算法(例如整数线性规划)来降低每次训练迭代中的内存峰值压力。此外,编译器执行激活重物化(rematerialization),以实现更大的有效批次大小并提高 HBM 利用率。

与 GPU 的一个显著区别在于,MTIA 能够在单个图中捕获计算和集体操作。通过 `torch.export` 或 `torch.compile` 追踪的集体操作会与计算运算符一起编译成一个整体图。这降低了子图启动开销,提高了效率和确定性。计算和通信之间的依赖关系通过信号量进行管理。目前,MTIA 完全支持静态形状的集体操作,并且正在开发动态形状的集体操作以及具有动态发送/接收计数的设备驻留 AllToAll 功能。

与 GPU 相比,MTIA 300 具有用于通信卸载的消息引擎 (ME) 和用于集成网卡的网络芯片,从而无需 PCIe 流量或主机/PE 参与即可实现高效通信。为了利用这些硬件特性,我们的集体通信库 HCCL(参见 IV-B 节)使用 MTIA 流接口提交所有通信操作。在内部,HCCL 使用 RDMA 动词来驱动网络芯片。

为了管理数据流顺序,控制核心(CPU-C)在MTIA-2i的基础上进行了扩展,用于调度PE网格和ME的工作。工作以数据包的形式到达控制核心,这些数据包包含计算和/或通信任务;核心会在调度前确保满足依赖关系。用于通信的工作数据包可能包含多个子图,HCCL会将这些子图映射到ME以进行并行执行。每个ME可以同时运行多个子图;因此,16个ME可以并行处理多个子图。完成后,ME会将状态报告给控制核心,以便释放后续工作的阻塞。

A. 集体图处理

图 9 展示了将子图分派到 CPU-M 的示例。子图表示为工作队列条目 (WQE) 数组,每个 WQE 描述几种类型的操作之一,例如:SEND、RECV、WRITE、WAIT、SET 和 REDUCE。某些操作直接映射到 RDMA 工作请求(即 SEND、RECV 和 WRITE),具有相似的字段和语义(例如,队列对 ID、本地地址和目标地址、长度、lkey、rkey 等)。SET WQE 将值写入本地内存(HBM 或缓存),而 WAIT WQE 会暂停,直到满足内存位置上的比较器条件(例如,等待地址 0xabcdef > 10)。REDUCE WQE 执行求和运算 S = A+B,其中 S 可以与 A 或 B 重叠,或者可以选择性地执行内存复制,其作用类似于 DMA 引擎。

WQE 还包含一些内部字段以辅助处理。流控制字段定义了 WQE 之间的顺序,从而支持常见的通信模式,例如环、递归倍增和有序树。以下是可用于指定顺序的字段:

wqe_sync:在指定的前一个 WQE(从当前 WQE 向前计数)完成之前,不要发出此 WQE。

fence:在当前 WQE 完成之前,不要发出任何其他 WQE。

rx_sync:在发出此 WQE 之前,等待所有未完成的接收 WQE 完成。

sync:在发出此 WQE 之前,等待所有先前的 WQE 完成。

近内存计算 (NMC) 引擎(第二部分 D 节)能够高效地执行归约操作。由于其靠近缓存和 HBM,因此无需使用 PE 网格即可执行 AllReduce 或 ReduceScatter 的关键操作,从而最大限度地减少通信和计算操作之间的资源共享。

综上所述,图 10 展示了一个简化的 AllReduce 环形算法,该算法分为 ReduceScatter 阶段和 AllGather 阶段,并转换为具有依赖关系的 WQE。从下往上读取时,第一个接收和发送操作没有依赖关系,可以并行执行。后续的加法操作依赖于前一个接收操作,从而解除下一个接收和发送操作的阻塞。此过程重复进行,直到所有数据在本地完成归约(ReduceScatter 阶段)。最后三个步骤描述了 AllGather 阶段,其中每个步骤都依赖于前一个步骤进行数据移动。

B. 集体通信库 (HCCL)

HCCL 构建使用网络芯片和 NMC 引擎的工作包和子图。它使用 RDMA 动词来管理控制路径。例如,要在两个对等体之间创建队列对 (QP),它会调用 ibv_create_qp,然后调用 ibv_modify_qp 将它们分别移动到“准备接收”和“准备发送”状态,这与任何标准的 RDMA 工作流程相同。创建这些 QP 还涉及通过 NIC 接口将它们映射到 Express Doorbell。工作包 (WC) 错误由 ME 在线检测。未触发 WC 的错误通过 ibv_get_async_event 捕获。

由于缺乏硬件 QP 缓存(参见第二部分 E),每个 NIC 支持的可供 HCCL 访问的活动 QP 数量有限 (1,088)。 HCCL 通过仅在需要时连接队列对 (QP)、设计集体算法以最大限度地减少未使用的 QP 以及在通信器内重用 QP 来管理此问题。由于我们灵活地使用网卡 (NIC),我们拥有 12 个 NIC,共计 13,056 个队列对,这些队列对可以在纵向扩展和横向扩展的通信域之间进行拆分或共享。

对于数据路径,HCCL 提供了一个类似于其他集体通信库的 API,包括常用的集体通信方式(例如 AlltoAll、AllReduce、ReduceScatter、AllGather)以及使用通信器定义参与设备组的点对点通信(即发送和接收)。此 API 主要通过 PyTorch 分布式接口(包括 PyTorch Distributed [8] 和 torchcomms [35])公开。这些接口的后端将 PyTorch 语义(例如张量和进程组)转换为通信库语义,即连续缓冲区和带有关联资源的通信器。

当用户调用 HCCL 通信 API 时,该库会准备将通信任务卸载到设备。与许多通信由主机驱动的库不同,HCCL 在工作到达设备后便不再参与其中。它会根据待处理的工作、系统拓扑、通信类型和其他参数预先确定最高效的算法和通信通道,并选择相应的算法和资源。然后,该库会将这些参数转换为工作包、子图和 WQE,供设备执行。

并行性体现在多个层面。单个任务的工作包在同一流中执行,并在需要时强制执行算法阶段之间的顺序。在工作包内部,子图逻辑上并行执行,但硬件可用性可能会将部分子图排队,直到所需的 ME 空闲。子图还可以使用如上所述的 WAIT 和 SET 等语义来定义彼此之间的依赖关系。WQE 按顺序发出,但仅在通过流控制字段指定时才会阻塞。这些机制共同使 HCCL 能够最大限度地提高并行性,高效利用网络带宽和缩减容量。图 11 可视化了这一流程。

HCCL 维护一个线程来跟踪未完成工作的状态,以便清理不再需要的资源并监控错误。该线程还允许 HCCL 根据所使用的接口(例如 torchcomms 接口中的 TorchWork 对象)向应用程序发出操作完成的信号。

评估:

对 MTIA 300 进行恰当的评估需要将其置于快速发展的 AI 加速器领域中进行考察。MTIA 300 的研发始于数年前,旨在与 H100 和 H200 GPU 竞争。其后续产品 MTIA 400 旨在与 GB300 GPU 抗衡,而 MTIA 450 和 500 则致力于在 GenAI 推理方面实现业界领先的性能,以超越未来的 GPU。在此背景下,表 III 对 MTIA 300、H100 和 H200 进行了比较。我们使用微基准测试和生产规模的 DLRM 训练工作负载来评估 MTIA 300。此外,尽管 MTIA 300 针对 DLRM 进行了优化,但其卓越的 HBM 容量和带宽使其也适用于 GenAI 推理,我们将在本节中对此进行评估。

A. 计算操作

首先,我们使用几个微基准测试来评估计算性能。

嵌入性能:嵌入算子,特别是表格批处理嵌入,是DLRM稀疏部分的重要计算组件。如图12所示,我们对来自生产工作负载的各种形状和输入分布的嵌入算子性能进行了评估,结果表明MTIA 300具有强大的性能:与H100和H200相比,MTIA 300在前向操作上分别实现了2.0倍和1.6倍的加速,在后向操作上分别实现了2.1倍和1.6倍的加速(几何平均值)。这种高性能得益于MTIA 300的高内存和缓存带宽,以及其专用的功能单元,例如基数排序,这些功能单元可以加速嵌入算子。我们注意到,嵌入性能与 HBM 带宽并非线性增长,因此与 H100 相比,速度提升并未达到 2.5 倍。当输入数据倾斜(大多数索引指向同一特征)时,嵌入运算符通常会受限于缓存或指令,而不是内存带宽。

内存带宽:图 13 展示了在 MTIA 300、H100 和 H200 上使用简单的 BF16 加法内核,针对不同张量大小进行逐元素加法运算的内存带宽。结果表明,对于大数据量,MTIA 300 的运算速度可达 5.57 TB/s(达到其峰值 HBM 带宽的 91%),超过了 H100 的 2.26 TB/s(94%)和 H200 的 4.40 TB/s(92%)。虽然 MTIA 300 在大内存操作方面展现出相当的 HBM 效率和更高的带宽,但其在细粒度内核方面性能较低。我们将在第六节讨论这种单内核延迟。

GEMM 性能:GEMM 是 DLRM 中密集层的重要组成部分。图 14 展示了 BF16-GEMM 在生产环境中遇到的各种大小下的性能。结果表明,由于其高内存带宽,MTIA 300 在处理较小、内存密集型的 GEMM 时表现良好。然而,对于较大的矩阵,MTIA 300 的单位性能较低,受限于其峰值 FLOPS。当运算强度超过 400 字节/FLOPS 时,H100 的效率达到 63%,而 MTIA 300 在所测试的形状下效率为 59%,尽管 MTIA 300 的效率高于 H200 的 54%。由于 MTIA 300 在我们的微基准测试中对特定形状的效率超过 90%,因此其较低的效率是由于我们的 GEMM 库仍需进一步优化。

B. 集体操作

本节比较了三种在生产训练模型中广泛使用的集体操作:AllGather、AllReduce 和 AllToAll。图 15 显示了它们在 MTIA 300 和 H100 上,不同加速器数量下的归一化延迟。在我们的模型使用的消息大小范围内(图 15 中的“时间百分比”所示),MTIA 300 对所有三种集体操作的性能普遍更优。当使用 16 个或更多加速器或消息大小超过 16 MB 时,这种加速效果尤为显著,我们认为这归因于 MTIA 300 更大的扩展域和 2.2 倍更高的扩展带宽。对于小消息,目前使用 NCCL 的 H100 的性能往往优于使用 HCCL 的 MTIA 300。我们尚未针对小消息对软件栈进行全面优化,因为它们目前在我们积极运行于 MTIA 300 上的训练工作负载中仅占实际运行时间的一小部分。

C. 计算与集体操作的重叠

MTIA 300 采用的硬件架构能够卸载集体操作,从而有效降低并发执行期间对计算内核性能的影响。为了评估计算内核和集体操作之间的交互,我们开发了一个微基准测试,该测试运行 1000 个大小为 4K × 4K × 4K 的 TF32 GEMM,同时在 16 个加速器上执行各种具有代表性消息大小的集体操作。图 16 显示了并行执行计算和通信时的效率;100% 的效率表示平台达到的性能与仅运行计算或仅运行通信操作时的性能相同。MTIA 300 对两者都保持了高效率,表明两者之间的干扰极小。这得益于专用的消息引擎和近内存计算单元,它们独立于主计算引擎处理集体操作。相比之下,当集体操作和计算并发执行时,H100 会面临流式多处理器资源的争用,导致性能显著下降。 MTIA 300 将集体操作与计算解耦的能力对 DLRM 来说尤其有利,因为 DLRM 依赖于频繁的集体操作进行分布式训练。

D. 生产环境训练工作负载

MTIA 300 在生产环境训练工作负载上展现出了令人瞩目的性能。本节将考察 MTIA 300 在训练 DLRM 模型 时的性能,该模型包含约 1500 亿个参数(其中 99% 为稀疏分量)。训练单个样本大约需要 30 亿次浮点运算。我们使用 TorchRec [14] 实现了该模型,并在 MTIA 300 和 H100 上分别使用 PyTorch 的图编译器 TorchInductor [2] 对其进行了完全编译,以最大限度地提高性能。该模型使用分布式 Shampoo [优化器,并通过分布式数据并行方案进行并行化。

1) 集体训练:为了比较集体训练的执行时间,我们使用 40 个加速器配置了模型,本地批处理大小为 6144。图 17 展示了每次训练迭代的消息统计信息和延迟。如第一张图所示,AllReduce 和 AllGather 操作处理大量数据,传入消息大小分别为 1.6 GB 和 2.1 GB。然而,35 个 AllToAllv 操作带来了独特的性能挑战;它们涉及稀疏数据分布,消息大小变化很大,从 1 KB 到 1 GB 不等。第二张图展示了 MTIA 300 在大规模 AllToAll 和 AllReduce 操作方面优于 H100 的性能。总体而言,MTIA 300 的通信性能比 H100 高出 3.9 倍。

2) 端到端训练性能:虽然为 GPU 实现的模型可以直接在 MTIA 300 上运行,但当模型与架构协同设计时,性能才能最大化。下面,我们将重点介绍三种协同设计策略,这些策略共同使 MTIA 300 在 DLRM 模型上实现了比 H100 高 1.42 倍的性能/总拥有成本 (Perf/TCO)。

Shampoo 算法的 CPU 卸载:矩阵特征值分解算子是 Shampoo 优化器预处理步骤中计算密集型的组成部分。虽然 H100 使用 cuSOLVER 库进行此操作,但由于性能与精度之间的权衡,在 MTIA 300 上实现数值精确且高效的版本较为困难。然而,MTIA 300 1:1 的主机与加速器比例允许将这些计算卸载到主机 CPU,从而确保足够的数值精度。相比之下,使用 1:8 的比例(H100 系统的典型比例)会导致 7.8% 的性能损失,这凸显了 MTIA 300 均衡架构的优势。

禁用量化通信:由于该模型最初针对 H100 进行了优化,因此默认启用逐行 FP8 量化通信以减少数据量。然而,在 MTIA 300 上,此过程依赖于低效的 RISC-V 操作,而非原生支持。通过禁用 FP8 量化并利用 MTIA 300 的高网络带宽,我们可以避免这些资源密集型操作,从而实现 4.4% 的性能提升。

更大的训练批次:在 MTIA 300 上增加本地批次大小可以优化分布式训练性能。通过增大本地批次大小,我们可以在保持全局批次大小不变的情况下减少训练器的数量,从而提高内核粒度并最大限度地减少通信开销。表 IV 比较了不同本地批次大小下的性能/总拥有成本 (Perf/TCO)。MTIA 300 强大的 HBM 容量支持更大的批次大小,例如 10,240,这比 H100 使用的 6,144 基线提高了 2% 的性能/总拥有成本。利用 MTIA 300 的内存来增加本地批处理大小,是提高训练效率的有效策略。

未来的优化机会:我们旨在通过多项针对性的优化来进一步提升性能。除了改进 GEMM 和小规模集体操作的内核之外,我们还在开发针对 MTIA 300 独特硬件的定制策略。这包括一种专门针对 MTIA 300 性能特性的 TorchRec 分片策略,该策略利用其独特的内核、通信和内存权衡来减少负载不均衡和迭代时间。此外,我们正在探索在同一网格内进行内核共置(而不是在 12×6 网格上顺序执行),以提高 PE 的利用率。这些优化措施有望共同提升 MTIA 300 的效率和吞吐量。

E. 大型语言模型推理

MTIA 300 最初是为训练深度学习资源模型 (DL RM) 而优化的,但它同样适用于大型语言模型 (LLM)。本研究使用 H200 作为参考平台,而非 H100,因为 H100 的 HBM 容量不足以运行评估中使用的 DeepSeek-R1 模型。我们使用开源的 LLM 推理基准测试工具 InferenceMax [5],它运行在 vLLM 推理runtime上。注意力机制和键值缓存 (KV-cache) 使用 BF16 精度存储和计算,而专家混合模型 (MoE) 的计算则使用 FP8 精度;这些精度在 MTIA 300 和 H200 上均原生支持。我们关注在线短提示场景:每个请求的输入和输出长度均独立且均匀地从 [0.8×1024, 1024] 个词元中随机抽取。这些测试涵盖了一系列批次大小(并发级别从 4 到 256),并测量了吞吐量(token/s)和延迟。

由于单个加速器无法承载完整的模型,我们评估了两个平台上的 8 加速器配置(8 个 MTIA 300 和 8 个 H200)。我们考虑了两种分片策略:TP8-TP8(在 8 个加速器上对注意力机制和 MoE 都应用张量并行)和 DP8-EP8(注意力组件采用数据并行,MoE 采用专家并行)。在 TP8-TP8 中,每一层的权重矩阵都被分片到所有 8 个设备上;每个设备计算每个操作的一个切片,并在每个并行区域结束后通过 AllReduce 进行同步。在 DP8-EP8 中,密集层(注意力机制和共享的 MLP)被复制到所有 8 个设备上,每个设备处理一个独立的请求批次;而 MoE 部分则被划分,使得每个设备拥有 1/8 的专家模型,token通过 AllToAll 通信在设备之间路由。在本评估中,我们使用混合批处理配置了服务器。

图 18 展示了 MTIA 300 和 H200 的性能,其中 x 轴表示每个请求的客户端端到端延迟,y 轴表示每个加速器的总token吞吐量。结果表明,在 InferenceMax 基准测试中,MTIA 300 的整体性能优于 H200。在此配置下,执行时间主要集中在解码阶段,这使得 MTIA 300 能够利用其高 HBM 带宽,并在高并发级别(超过 64)下超越 H200。然而,性能差距并非在所有操作点上都一致;在低并发级别下差距较小。这是由于 MTIA 300 的小批量通信开销高于 H200 通过 NVLink 实现的快速节点内通信,以及某些内核在批处理/token维度之外并未完全并行化,导致 MTIA 300 上的 PE 内核利用率不足。

挑战与局限性

尽管 MTIA 300 针对 DLRM 训练进行了优化,但它仍存在一些局限性,我们将在后续的 MTIA 版本中加以解决。尤其重要的是,平衡特化和泛化能力对于适应不断演进的模型仍然至关重要。除了支持 DLRM 之外,我们针对 MTIA 400、450 和 500 的路线图还扩展了对 LLM 工作负载的支持,并特别优先考虑 LLM 推理性能。

计算密集型工作负载:现代 GenAI 模型是计算密集型的,需要大规模、低精度的矩阵运算。这种趋势正开始影响 DLRM,它们正在采用更大的密集组件和 Transformer 架构,进一步增加了计算需求。因此,我们设计了具有更高 FLOPS 的 MTIA 400 来应对这些挑战。

低精度数据类型:在 MTIA 300 上支持新的数据类型极具挑战性,尤其是在模型采用自定义缩放格式时,例如 MX4 和 NVFP4 的行级或块级缩放。虽然 MTIA 300 可以处理标准的硬件加速类型转换,但这些新兴数据类型通常需要特殊的缩放方式,而 MTIA 300 本身并不支持这些缩放方式,因此必须回退到 RISC-V 执行。这种回退会降低性能,并使使用这些高级数据表示的模型的集成和优化变得更加复杂。为了解决这一限制,我们在 MTIA 400 中引入了对 MX4 的原生支持。

数值精度:硬件可互换性——即期望在不同平台上获得一致的结果——至关重要。在 MTIA 300 和现有系统之间实现数值一致性极具挑战性:浮点运算(精度、舍入和支持的数据类型)以及内核或运算符实现(算法选择和运算顺序)的差异可能会导致收敛问题,而这些问题可能在经过数小时甚至数天的训练后才会出现。我们开发了数值调试工具来应对这些挑战,但这些工具仍需时间完善。

Eager模式:Eager模式便于模型开发,尤其适用于调试和快速原型设计,尽管生产环境中通常使用图模式。然而,Eager模式面临着显著的性能挑战。与NVIDIA高度优化的低延迟系统不同,MTIA 300的Eager模式会产生大量的主机端开销——这些开销源于Python解释器、动态调度和设备主机通信——而这些开销无法随着加速器硬件的提升而扩展。因此,随着加速器的改进,主机成为瓶颈,凸显了对更好的Eager模式支持的必要性。

训练时间:优化MTIA 300的训练时间需要权衡各种复杂的因素。MTIA 300在处理大批量训练时表现出色,这得益于其丰富的HBM和高效的计算能力。然而,为了保持模型收敛,必须控制全局批次大小,这通常会导致局部批次较小。虽然一些用户倾向于通过运行更多训练器并使用更小的批次来加快训练速度,但这可能会低估 HBM 的性能,并增加细粒度算子带来的主机运行时开销。挑战在于如何协调这些相互冲突的需求。

模型启用和算子编写:训练 DL RM 所需的算子和形状数量级比推理要多。有些算子对性能非常敏感,值得开发人员投入大量精力以实现最佳性能,而另一些则是“尾部”算子,则不然。这与 LLM 形成对比,LLM 的算子和形状范围相对明确。此外,排序工程师会进行大量的实验,导致模型频繁迭代。

我们启用的每个新模型快照都需要新的算子、形状和图变换。除了在 TorchInductor 中添加支持和编写手动内核之外,我们还需要大量的基础设施来收集生产环境跟踪信息、合成运算符和形状测试,以及维护 CI/CD 环境,以提高覆盖率和开发速度。为此,我们已在利用编码代理实现自动化内核生成方面取得了初步成功。

相关工作

由于人工智能工作负载的激增,包括谷歌、亚马逊、华为、IBM、微软、阿里巴巴、百度、OpenAI、特斯拉、腾讯和Meta在内的各大IT公司自主研发人工智能芯片的趋势已显著增强。此外,许多初创公司,例如Cerebras、Groq、SambaNova、Cambricon和Rivos,也提供人工智能芯片,作为成熟厂商GPU的替代方案。

与Meta之前的推理芯片相比,MTIA 300(我们的首款训练芯片)引入了三个显著特点:(1)内置网卡芯片组;(2)用于集体卸载的专用消息引擎;(3)近内存计算,用于加速基于归约的集体计算。据我们所知,这些特点在现有的人工智能加速器中是独一无二的。

尽管谷歌 TPU 的稀疏核心可以卸载通信任务,但它专用于非 RDMA、非交换式环面网络,并且缺乏通用的集体库接口。这限制了稀疏核心技术在其他行业加速器中的应用,而这些加速器通常基于 RDMA 和类似的集体库接口构建。

大多数定制 AI ASIC 专注于图模式下的推理,缺乏对即时模式的一流支持,而即时模式对于模型开发者使用 PyTorch 至关重要。相比之下,MTIA 提供了一个原生的、以 PyTorch 为核心的软件生态系统,支持 TorchDynamo、TorchInductor、Triton 以及即时模式和图模式,并提供类似 CUDA 的运行时 API,简化了模型移植。值得注意的是,虽然 MTIA 和其他芯片都采用了数据流架构(计算在其依赖关系解析过程中进行),但 MTIA 是唯一提供原生 PyTorch 体验的芯片。最后,虽然现代 AI ASIC 主要针对 GenAI 模型进行了优化,但 MTIA 300 针对 DLRM 进行了优化。

结论

我们介绍了 MTIA 300 的设计和评估,这是我们首款自主研发的用于 DLRM 训练的 AI 芯片。凭借其首创的集成网卡和集体卸载引擎,我们证明,与 GPU 相比,MTIA 300 在推荐模型训练工作负载方面具有更高的效率。展望未来,我们正在加速 Meta 下一代 AI 芯片 的开发,该芯片将满足 LLM 训练和推理工作负载以及推荐任务日益增长的需求。我们预计,向自主研发 AI 芯片的转变将推动新的模型协同设计机遇,并带来下一波模型创新浪潮。

本文来自微信公众号“半导体行业观察”(ID:icbank),作者:meta,36氪经授权发布。

+1
4

好文章,需要你的鼓励

参与评论
评论千万条,友善第一条
后参与讨论
提交评论0/1000

最有深度的半导体新媒体,实讯、专业、原创、深度。

36氪AI测评

选靠谱AI,看真实评测
查看
36氪AI测评官方交流社区
加入

36氪项目推荐

咨询项目审核和入驻
联系
36氪项目推荐订阅号
关注

下一篇

七夕景区送新婚免票,实为抢占长期消费入口。

1小时前

36氪APP让一部分人先看到未来
36氪
鲸准
氪空间

推送和解读前沿、有料的科技创投资讯

一级市场金融信息和系统服务提供商

聚焦全球优秀创业者,项目融资率接近97%,领跑行业