跳转到主要内容
博客

TorchRec 和 FBGEMM 1.0 稳定版发布

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

TorchRec

TorchRec 于 2022 年首次开源，提供创建最先进的个性化模型的通用原语

  • 用于数百个 GPU 分布式训练的简单、优化 API
  • 嵌入的高级分片技术
  • 推荐系统创作中常见的模块
  • 通过 TorchRec 模型量化和分片 API 实现无摩擦分布式推理

此后，TorchRec 显著成熟，Meta 内部许多生产推荐模型广泛采用其进行训练和推理，并新增了功能，例如：可变批处理嵌入、嵌入卸载、零冲突哈希等。此外，TorchRec 在 Meta 之外也有应用，例如 Databricks 的推荐模型和 Twitter 算法。因此，标准 TorchRec 功能已被标记为稳定，具有 PyTorch 风格的 BC 保证，并可在改版后的 TorchRec 文档中查看。

FBGEMM

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

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 模块中用于分布式训练和推理的优化版本。下面的基准测试展示了 FusedEBC 相较于传统 PyTorch 嵌入模块实现 (EBC) 在性能上的巨大提升，以及 FusedEBC 能够处理比 GPU 内存可用容量大得多的嵌入，这得益于 UVM 缓存。

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 功能的教程。