PyTorch 功能欠缺,Meta 推出 TorchRec 來補救?

PyTorch 功能欠缺,Meta 推出 TorchRec 來補救?

作者 | Meta AI - Donny Greenberg

編譯 | 張潔 責編 | 屠敏

近日,Meta 宣佈了 TorchRec,一個用於推薦系統的 PyTorch 庫。這個新庫提供了通用的稀疏性和並行性原語,使研究人員能夠構建最先進的個性化模型並將其部署到生產中。

PyTorch 被吐槽,正式釋出推薦系統庫 TorchRec

之所以會有 TorchRec 的出現,是因為 RecSys (推薦系統) 包含很多成熟的 AI 技術,程式設計師一般不會透過GitHub去了解它。與 Vision 和 NLP 等領域不同,RecSys 的大部分開發都是在公司秘密進行的。對於相關研究人員或構建個性化使用者體驗的公司而言,該領域沒有實現通用化。

此外,RecSys 作為一個領域,很大程度上是由順序事件(一組依次執行的指令,通常在螢幕上從上到下或從左到右)上的學習模型定義的,這與 AI 的其他領域有很多的重疊。許多技術是可遷移、可擴充套件和可分散式執行的。

到 2020 年年中,PyTorch 團隊收到了大量的反饋,稱開源 PyTorch 生態系統中還沒有大規模的推薦系統包。當研究人員試圖找到一個好的方法時,Meta 的一組工程師希望將 Meta 生產的 RecSys 堆疊作為 PyTorch 庫貢獻出來,並致力於圍繞它發展一個新的生態系統。這似乎是一個讓 RecSys 領域的研究人員和公司受益的好主意。

從企業的角度來看, Meta 公司開始了軟體堆疊,研究人員也開始模組化並設計一個完全可擴充套件的程式碼庫。目標是從 Meta 的軟體堆疊中提取關鍵模組,方便同時實現創造性擴充套件。經過近兩年的跨 Meta 基準測試和遷移,終於與 RecSys 社群一起帶來了 TorchRec。

介紹靈活的TorchRec

TorchRec 包括一個可擴充套件的低階建模基礎和電池模組。研究人員最初的目標是“雙塔”架構( 大型語料庫專案推薦的取樣偏差校正神經模型和 DLRM:一個先進的、開源的 DL 推薦模型),這些架構具有單獨的子模組,用它們來學習候選專案。

該庫包括:

建模原語。例如嵌入包和鋸齒狀張量,可以使用混合資料並行和模型並行來創作大型、高效能的多裝置或多節點模型。

由 FBGEMM(一個針對伺服器端推理進行了最佳化的高效能核心庫)提供支援的最佳化 RecSys 核心,包括對稀疏和量化操作的支援。

一個分片器,它可以使用各種不同的策略對嵌入表進行分割槽,包括資料並行、逐表、逐行、逐錶行和逐列分片。

可以為模型自動生成最佳化的資料分片。

流水線用重疊資料載入裝置傳輸、裝置間通訊和計算來提高效能。

GPU(圖形處理器)推理的支援。

RecSys 的通用模組,例如模型和公共資料集(Criteo 和 Movielens)。

為了展示此工具的靈活性,來看一下在 DLRM 事件預測示例中提取的程式碼片段:

# Specify the sparse embedding layerseb_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 allocatedtrain_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 computefor epoch in range(epochs): # Train

TorchRec 規模宏大,用於訓練三萬億模型

TorchRec 擁有最先進的大規模 AI 基礎設施,為 Meta 的一些大模型提供支援。它用於訓練一個 1。25 萬億引數的模型,於 1 月投入生產,還有一個即將投入生產的 3 萬億引數模型。這已經很好地表明 PyTorch 完全有能力解決業界規模最大的 RecSys 問題。

研究人員從社群中聽說,分片嵌入是一個痛點。隨後 TorchRec 解決了這一問題。但不幸的是,使用公共資料集提供的大規模基準測試非常具有挑戰性,因為大多數的開源基準測試都太小了,無法大規模地顯示出效能。

Meta 正在為 PyTorch 社群提供最先進的 RecSys 軟體包,同時 TorchRec 背後的團隊計劃持續地進行這個專案,建立 TorchRec 來滿足 RecSys 社群的需求,並繼續推動 Meta 的個性化。

參考連結:https://pytorch。org/blog/introducing-torchrec/

—END—

《新程式設計師001-004》全面上市,對話世界級大師,報道中國IT行業創新創造

頂部