由张保罗、扎因·胡达、萨鲁娜·普姆玛、石田伸太郎、素帕查亚·庞蓬蒂普、马本森等著

我们很高兴宣布 TorchRec 和 FBGEMM 的稳定版 1.0 发布。TorchRec 是 PyTorch 原生的推荐系统库，由 FBGEMM（Facebook GEneral Matrix Multiplication）的高效、低级内核提供支持。

TorchRec

TorchRec 最初于 2022 年开源，为创建最先进的个性化模型提供常用原语：

  • 简单、优化的 API，支持跨数百个 GPU 的分布式训练
  • 高级分片技术，用于嵌入
  • 常用于推荐系统编写的模块
  • 无摩擦的分布式推理路径，提供 TorchRec 模型的量化和分片 API

自那时起，TorchRec 已显著成熟，在许多 Meta 生产推荐模型中得到了广泛内部采用，用于训练和推理，同时引入了新功能，如：可变批处理嵌入、嵌入卸载、零冲突哈希等。此外，TorchRec 在 Meta 之外也有影响力，例如在 Databricks 的推荐模型和 Twitter 算法中。因此，标准 TorchRec 功能已被标记为稳定，并提供了 PyTorch 风格的 BC 保证，可在重新设计的 TorchRec 文档中看到。

FBGEMM

FBGEMM 是一个为 CPU 和 GPU 提供高性能内核的库。自 2018 年以来，FBGEMM 通过扩展其范围，从针对 CPU 推理的性能关键内核到更复杂的稀疏算子，支持了 Meta 内部和外部 AI/ML 工作的高效执行——最近还扩展到了 CPU 和 GPU 上的生成式 AI——从而支持了 Meta 内部和外部 AI/ML 工作的高效执行。

FBGEMM 通过其高性能内核实现支持 TorchRec，涵盖了从嵌入包内核到锯齿形张量操作的各种推荐工作负载。与 TorchRec 一起，我们发布了 FBGEMM 1.0 版本，该版本保证了几个稳定 API 的功能和向后兼容性，并提供了增强的文档。

性能

DLRM（深度学习推荐模型）是 Meta 推荐系统中的标准神经网络架构，其中分类特征通过嵌入进行处理，而连续（密集）特征则通过底部多层感知器进行处理。以下图展示了 DLRM 的基本架构，包括密集特征和稀疏特征之间的二阶交互层以及用于生成预测的顶部 MLP。

flow diagram

TorchRec 提供了融合嵌入查找的标准化模块和显著的优化。EBC 是一个传统的 PyTorch 嵌入模块实现，包含由 FBGEMM 提供支持的 torch.nn.EmbeddingBags. FusedEBC，它通过融合优化器和 UVM 缓存/管理来缓解内存约束，是 sharded TorchRec 模块中用于分布式训练和推理的优化版本。下面的基准测试展示了 FusedEBC 相较于传统 PyTorch 嵌入模块实现（EBC）的巨大性能提升，以及 FusedEBC 能够处理比 GPU 内存中可用的更大的嵌入的能力。

performance chart

TorchRec 数据类型

TorchRec 提供了标准的数据类型和模块，便于处理分布式嵌入。以下是一个通过 TorchRec 设置嵌入表集合的简单示例：

from torchrec import EmbeddingBagCollection
from torchrec import KeyedJaggedTensor
from torchrec import JaggedTensor

ebc = torchrec.EmbeddingBagCollection(
    device="cpu",
    tables=[
        torchrec.EmbeddingBagConfig(
            name="product_table",
            embedding_dim=64,
            num_embeddings=4096,
            feature_names=["product"],
            pooling=torchrec.PoolingType.SUM,
        ),
        torchrec.EmbeddingBagConfig(
            name="user_table",
            embedding_dim=64,
            num_embeddings=4096,
            feature_names=["user"],
            pooling=torchrec.PoolingType.SUM,
        )
    ]
)

product_jt = JaggedTensor(
    values=torch.tensor([1, 2, 1, 5]), lengths=torch.tensor([3, 1])
)
user_jt = JaggedTensor(values=torch.tensor([2, 3, 4, 1]), lengths=torch.tensor([2, 2]))

kjt = KeyedJaggedTensor.from_jt_dict({"product": product_jt, "user": user_jt})

print("Call EmbeddingBagCollection Forward: ", ebc(kjt))

分片

TorchRec 提供了一个规划类，可以自动生成跨多个 GPU 的优化分片计划。在这里，我们展示了如何在两个 GPU 之间生成分片计划：

from torchrec.distributed.planner import EmbeddingShardingPlanner, Topology

planner = EmbeddingShardingPlanner(
    topology=Topology(
        world_size=2,
        compute_device="cuda",
    )
)

plan = planner.collective_plan(ebc, [sharder], pg)

print(f"Sharding Plan generated: {plan}")

模型并行

TorchRec 的主要分布式训练 API 是 DistributedModelParallel，它调用规划器生成分片计划（如上所示），并根据该计划分片 TorchRec 模块。我们展示了如何使用 DistributedModelParallel 对我们的 EmbeddingBagCollection 进行分片嵌入分布式训练：

model = torchrec.distributed.DistributedModelParallel(ebc, device=torch.device("cuda"))

推理

TorchRec 为模型分布式推理提供了简单的量化分片嵌入 API。用法如下：

from torchrec.inference.modules import (
    quantize_inference_model,
    shard_quant_model,
)
quant_model = quantize_inference_model(ebc)
sharded_model, _ = shard_quant_model(
    quant_model, compute_device=device, sharding_device=device
)

结论

TorchRec 和 FBGEMM 现已稳定，针对大规模推荐系统进行了优化。

搭建 TorchRec 和 FBGEMM，请参阅入门指南。

我们还推荐全面、端到端的教程，用于介绍 TorchRec 和 FBGEMM 的功能。