训练排序与推荐模型时,算力并不是唯一的瓶颈。大规模稀疏特征、嵌入表分片和频繁的集体通信,会让加速器长期等待数据到位。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 满足这类训练的通信需求。其价值不只是“网卡离芯片更近”,而是系统协同设计带来的路径重构。
可以从三个层面理解:
- 数据路径更短:网络功能集成在加速器设计中,减少传统外接组件和通用路径上的转发、协调成本。
- 通信与计算可重叠:由专用通信引擎处理部分通信工作时,计算资源可以更专注于前向、反向和优化器相关计算。
- 库与硬件共同定义行为: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 与通信卸载纳入芯片设计的方案,提供了更直接的性能优化空间。