作者：Paul Zhang, Zain Huda, Sarunya Pumma, Shintaro Iwasaki, Supadchaya Puangpontip, Benson Ma

我们很高兴宣布 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 功能已被标记为 stable（稳定），具有 PyTorch 风格的 BC（向后兼容）保证，并可在全新的 TorchRec 文档中查看。

FBGEMM

FBGEMM 是一个为 CPU 和 GPU 提供高性能内核的库。自 2018 年以来，FBGEMM 通过将其范围从用于 CPU 推理的性能关键型内核扩展到用于训练和推理的更复杂的稀疏算子——以及最近用于生成式 AI 的算子——来支持 Meta 内部和外部 AI/ML 工作负载的高效执行，这些算子可在 CPU 和 GPU 上运行。

FBGEMM 一直通过其为推荐工作负载提供的后端高性能内核实现（从嵌入包内核到锯齿状张量操作）来赋能 TorchRec。与 TorchRec 一起，我们发布了 FBGEMM 1.0，该版本保证了其核心功能的多个稳定 API 的功能性和向后兼容性，并提供了增强的文档。

性能

DLRM (Deep Learning Recommendation Model) 是 Meta 推荐服务的标准神经网络架构，其中类别特征通过嵌入进行处理，而连续（密集）特征则通过底部多层感知器进行处理。下图描绘了 DLRM 的基本架构，包含密集特征和稀疏特征之间的二阶交互层以及用于生成预测的顶部 MLP。

flow diagram

TorchRec 提供了标准化的模块，并在融合嵌入式查找方面进行了显著优化。EBC 是传统的 PyTorch 嵌入模块实现，包含一组 torch.nn.EmbeddingBags。 FusedEBC 由 FBGEMM 提供支持，用于在嵌入表上执行高性能操作，并结合融合优化器和 UVM 缓存/管理以缓解内存限制，是分片 TorchRec 模块中用于分布式训练和推理的优化版本。以下基准测试表明，与传统的 PyTorch 嵌入模块实现 (EBC) 相比，FusedEBC 带来了巨大的性能提升，并且 FusedEBC 能够通过 UVM 缓存处理比 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 提供了一个 planner 类，可自动生成跨多个 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，它调用 planner 生成分片计划（如上所示），并根据该计划对 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 中的功能。