MTIA 300:把网络通信做到训练芯片里的推荐模型加速器

2026-08-25 47 预计阅读时间: 1 分钟
来源: engineering.fb.com AI 摘要 Original link

Disclaimer: This article is an AI-assisted summary. Read it together with the original source when precision matters. The summary may omit context, version differences, or edge cases and is not official documentation.

预计阅读时间:9 分钟

训练排序与推荐模型时,算力并不是唯一的瓶颈。大规模稀疏特征、嵌入表分片和频繁的集体通信,会让加速器长期等待数据到位。Meta 公布的 MTIA 300 面向这一类训练与推理负载,其关键变化是将 NIC chiplet 与通信卸载引擎直接集成进芯片,并与 HCCL 通信库协同设计。

这意味着,系统不再完全依赖通用 GPU 加上外置网络设备来完成通信路径。对于通信密集的推荐模型训练,芯片、网络和通信运行时开始作为一个整体优化。

为什么推荐模型特别依赖通信

排序和推荐模型通常同时包含稠密网络与大量稀疏嵌入特征。数据并行、模型并行或嵌入表分片后,训练过程会持续发生数据交换,例如:

  • 聚合不同设备上的梯度;
  • 在分片嵌入表之间请求和返回特征向量;
  • 为并行阶段传递激活值与中间结果;
  • 在每个训练 step 中执行 all-reduce、all-to-all 等集体通信。

通用加速器可以执行这些操作,但通信常常要经过主机、外部 NIC、驱动和通信库。随着训练规模扩大,软件栈调度开销、网络排队和算力等待会累积成可见的 step 时间。

MTIA 300 的设计重点就在于缩短这条路径:内置的 NIC chiplet 提供网络能力,通信卸载引擎承担部分通信处理工作,让计算单元尽量持续处理模型算子,而不是参与大量协议和数据搬运。

内置 NIC 与通信卸载改变了什么

从摘要可知,MTIA 300 是 Meta 自研训练和推理加速器家族中,首个为训练排序、推荐模型优化的芯片,并通过内置 NIC chiplet 满足这类训练的通信需求。其价值不只是“网卡离芯片更近”,而是系统协同设计带来的路径重构。

可以从三个层面理解:

  1. 数据路径更短:网络功能集成在加速器设计中,减少传统外接组件和通用路径上的转发、协调成本。
  2. 通信与计算可重叠:由专用通信引擎处理部分通信工作时,计算资源可以更专注于前向、反向和优化器相关计算。
  3. 库与硬件共同定义行为:Meta 同时设计 HCCL 与硬件能力。通信库不只是调用底层网络接口,也可以针对硬件拓扑和卸载能力安排集体通信。

这是一种很明确的工作负载导向设计。对通用 GPU 而言,通用性是重要目标;对推荐训练基础设施而言,若通信模式足够稳定、规模足够大,就有理由把常见瓶颈固化为专用硬件能力。

HCCL 的位置:让硬件能力真正进入训练框架

通信硬件不会自动改善训练吞吐。框架和通信库必须在正确的时间、以正确的拓扑和分片方式发起操作,才能将网络带宽转化为更短的训练 step。

MTIA 300 采用与 HCCL 协同设计的方式,重点在于让通信库理解设备端能力。一个典型训练循环中,应用代码可能只调用一次梯度同步;真正复杂的工作由通信库决定,包括张量如何切块、选择什么算法、是否与反向计算重叠,以及如何利用设备拓扑。

下面用 PyTorch 分布式训练示意这一层抽象。这是通用实践示例,不是 MTIA 300 或 HCCL 的官方 API。它展示了应用侧应如何把通信边界保持清晰,从而让底层通信后端有机会执行优化。

# train_ddp.py
# Run with: torchrun --standalone --nproc_per_node=4 train_ddp.py
import os
import torch
import torch.distributed as dist
from torch import nn
from torch.nn.parallel import DistributedDataParallel as DDP


def main():
    dist.init_process_group(backend="nccl")  # Replace with the platform backend when required.
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)

    model = nn.Sequential(
        nn.Linear(128, 512),
        nn.ReLU(),
        nn.Linear(512, 1),
    ).cuda(local_rank)
    model = DDP(model, device_ids=[local_rank])
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

    for step in range(100):
        features = torch.randn(4096, 128, device=local_rank)
        labels = torch.randn(4096, 1, device=local_rank)

        prediction = model(features)
        loss = nn.functional.mse_loss(prediction, labels)

        optimizer.zero_grad(set_to_none=True)
        loss.backward()  # DDP schedules gradient collective communication here.
        optimizer.step()

        if local_rank == 0 and step % 10 == 0:
            print(f"step={step}, loss={loss.item():.4f}")

    dist.destroy_process_group()


if __name__ == "__main__":
    main()

运行前需要安装与本机 GPU、PyTorch 和通信后端匹配的版本。若迁移到专用训练加速器,通常不应直接假设 nccl 仍是正确后端;应按该平台提供的 PyTorch 集成、设备运行时和通信库替换初始化配置。

可以怎样评估通信是否真是瓶颈

硬件拥有通信卸载能力,不代表每个模型都会获得相同收益。落地时应先测量,再决定是否调整并行策略。一个实用的检查顺序如下:

# 以 PyTorch profiler 为例,先收集训练 trace。
torchrun --standalone --nproc_per_node=4 train_ddp.py

# 在训练脚本中开启 profiler 后,可使用 TensorBoard 查看 trace:
tensorboard --logdir ./runs

需要重点观察的不是单独的网络带宽峰值,而是以下信号:

  • 一个 step 中集体通信占用的总时间;
  • 通信是否与反向计算重叠;
  • GPU 或加速器是否在等待通信完成;
  • all-to-all 是否随 embedding 分片数增加而显著恶化;
  • 小消息、高频同步是否造成大量启动开销。

对于推荐模型,还要区分稠密部分和稀疏部分。稠密网络的梯度同步常以 all-reduce 为主;分片 embedding 的查找与回传可能更接近 all-to-all。两者的热点不同,适合的分片策略和通信优化也不同。

采用这类架构时的工程判断

MTIA 300 体现的是一个趋势:当训练负载的通信模式足够明确时,网络不再只是数据中心外围资源,而会成为加速器架构的一部分。

评估此类平台时,可以围绕四个问题建立基准:

  • 目标模型中,通信时间在单个训练 step 里占多大比例?
  • embedding、特征交互和稠密塔分别采用了什么并行方式?
  • 通信库是否能与现有框架、调度系统和可观测性工具顺利集成?
  • 模型收益是否足以覆盖平台迁移、算子适配和运维体系变化的成本?

对于通信占比较低、模型结构变化频繁的任务,通用加速器的灵活性仍然重要。对于长期稳定、以大规模排序和推荐训练为核心的工作负载,像 MTIA 300 这样将 NIC 与通信卸载纳入芯片设计的方案,提供了更直接的性能优化空间。


相关推荐