我們很高興宣佈推出 TorchRec，這是一個用於推薦系統 (RecSys) 的 PyTorch 領域庫。這個新庫提供了常見的稀疏性和並行性原語，使研究人員能夠構建最先進的個性化模型並將其部署到生產環境中。
我們是如何走到這一步的？
推薦系統 (RecSys) 在當今生產部署的 AI 中佔據了很大一部分，但你可能從 GitHub 上看不出來。與視覺和自然語言處理等領域不同，RecSys 中正在進行的創新和開發大部分都在公司內部進行。對於研究這些技術的學術研究人員或構建個性化使用者體驗的公司來說，這個領域遠未實現民主化。此外，RecSys 作為一個領域，主要透過對稀疏和/或順序事件進行學習模型來定義，這與其他 AI 領域有很大的重疊。許多技術都是可轉移的，特別是對於擴充套件和分散式執行。全球對 AI 的大部分投資都用於開發這些 RecSys 技術，因此將它們隔離開來阻礙了這項投資流入更廣泛的 AI 領域。
到 2020 年中期，PyTorch 團隊收到了很多反饋，認為開源 PyTorch 生態系統中還沒有一個大規模的生產級推薦系統包。當我們試圖找到一個好的解決方案時，Meta 的一群工程師希望將 Meta 的生產 RecSys 堆疊作為一個 PyTorch 領域庫貢獻出來，並致力於圍繞它發展一個生態系統。這似乎是一個有益於 RecSys 領域的研究人員和公司的好主意。因此，我們從 Meta 的堆疊開始，模組化並設計了一個完全可擴充套件的程式碼庫，可適用於各種推薦用例。我們的目標是從 Meta 的軟體堆疊中提取關鍵構建塊，以同時實現創造性探索和規模化。經過近兩年的時間，Meta 內部進行了一系列基準測試、遷移和測試，我們很高興最終能與 RecSys 社群一起踏上這段旅程。我們希望這個軟體包能夠開啟整個 RecSys 行業的對話和協作，並以 Meta 作為第一個重要的貢獻者。
介紹 TorchRec
TorchRec 包含一個可擴充套件的低階建模基礎以及功能豐富的開箱即用模組。我們最初的目標是“雙塔” ([1], [2]) 架構，它們有獨立的子模組來學習候選專案和查詢或上下文的表示。輸入訊號可以是浮點“稠密”特徵或需要訓練大型嵌入表的高基數分類“稀疏”特徵的混合。此類架構的有效訓練涉及結合複製計算“稠密”部分的資料並行和將大型嵌入表分割槽到許多節點上的模型並行。
具體來說，該庫包括：
- 建模原語，例如嵌入包（embedding bags）和鋸齒張量（jagged tensors），使用混合資料並行和模型並行，可以輕鬆構建大型、高效能的多裝置/多節點模型。
- 由 FBGEMM 提供支援的最佳化 RecSys 核心，包括對稀疏和量化操作的支援。
- 一個分片器（sharder），可以採用多種不同策略對嵌入表進行分割槽，包括資料並行、按表、按行、按表按行和按列分片。
- 一個規劃器（planner），可以自動為模型生成最佳化的分片計劃。
- 流水線（Pipelining），用於重疊資料載入裝置傳輸（複製到 GPU）、裝置間通訊（input_dist）和計算（前向、後向），以提高效能。
- GPU 推理支援。
- 推薦系統的常用模組，例如模型和公共資料集（Criteo 和 Movielens）。
為了展示此工具的靈活性，讓我們看一下以下程式碼片段，該片段取自我們的 DLRM 事件預測示例
# Specify the sparse embedding layers
eb_configs = [
EmbeddingBagConfig(
name=f"t_{feature_name}",
embedding_dim=64,
num_embeddings=100_000,
feature_names=[feature_name],
)
for feature_idx, feature_name in enumerate(DEFAULT_CAT_NAMES)
]
# Import and instantiate the model with the embedding configuration
# The "meta" device indicates lazy instantiation, with no memory allocated
train_model = DLRM(
embedding_bag_collection=EmbeddingBagCollection(
tables=eb_configs, device=torch.device("meta")
),
dense_in_features=len(DEFAULT_INT_NAMES),
dense_arch_layer_sizes=[512, 256, 64],
over_arch_layer_sizes=[512, 512, 256, 1],
dense_device=device,
)
# Distribute the model over many devices, just as one would with DDP.
model = DistributedModelParallel(
module=train_model,
device=device,
)
optimizer = torch.optim.SGD(params, lr=args.learning_rate)
# Optimize the model in a standard loop just as you would any other model!
# Or, you can use the pipeliner to synchronize communication and compute
for epoch in range(epochs):
# Train
擴充套件效能
TorchRec 擁有用於大規模推薦 AI 的最先進基礎設施，為 Meta 執行著一些最大的模型。它曾用於訓練一個 1.25 萬億引數模型，並於 1 月投入生產，以及一個 3 萬億引數模型，該模型將很快投入生產。這應該很好地表明 PyTorch 完全能夠解決業界最大規模的 RecSys 問題。我們從社群中的許多人那裡聽說，分片嵌入是一個痛點。TorchRec 很好地解決了這個問題。不幸的是，很難用公共資料集提供大規模基準測試，因為大多數開源基準測試都太小，無法顯示大規模效能。
展望未來
開源和開放技術具有普遍的優勢。Meta 正在向 PyTorch 社群播種一個最先進的 RecSys 包，希望有更多人參與到它的構建中，從而實現新的研究並幫助許多公司。TorchRec 背後的團隊計劃無限期地繼續這個專案，不斷發展 TorchRec 以滿足 RecSys 社群的需求，歡迎新的貢獻者，並繼續為 Meta 的個性化服務提供支援。我們很高興能踏上這段旅程，並期待著貢獻、想法和反饋！
參考文獻
[1] 抽樣偏差校正的神經網路建模用於大規模語料庫專案推薦
[2] DLRM：一種先進的開源深度學習推薦模型