Hugging Face 推出工具查詢模型在 AWS Trainium 上的相容性

背景

Hugging Face 先前推出了 TorchNeuron，這是 AWS Trainium 的原生 PyTorch 後端（torch-neuronx 套件）。透過此後端，Transformers 庫可以在 Trainium 上執行其建模代碼，就像在其他裝置上一樣，理論上 Hugging Face Hub 上的任何 Transformers 模型都有可能運行。然而，實際上有哪一些模型能直接運行仍需驗證。 官方原文・Hugging Face Blog

發生什麼事

Hugging Face 發布了一個名為「torch-neuronx model support Space」的工具，用於回答哪些熱門模型能在 AWS Trainium 上無修改運行。該工具針對 Hub 上最流行的模型提供查詢功能，並給出讓使用者自行運行的腳本。 官方原文・Hugging Face Blog

重點內容

該工具預設顯示「Selected view」，這是從 Hub 下載量和趨勢分數最高的模型中挑選出的清單。每一行顯示測試結果、模型分佈在多少個 Trainium 核心上、包含編譯在內的實際運行時間（wall time）以及峰值裝置記憶體。使用者可以切換到「Trending」或「Downloads」視圖，輸入部分名稱來查詢特定模型。這些視圖遵循 Hub 自身的排名，一次針對一個任務。如果某個模型被刻意排除，工具會說明原因（通常是未優先處理，有時是真正不相容）。使用者還可以選擇 Neuron 編譯器（neuronx-cc）和 Transformers 的版本進行查詢。 官方原文・Hugging Face Blog

目前僅測試 Transformers 模型，且選擇範圍限於文本生成（text-generation）和圖像文本到文本（image-text-to-text）模型。Hugging Face 計劃擴展至 diffusers、sentence transformers 和其他 Transformers 任務。要自行運行，使用者需選擇模型以查看方法。工具會列出測試該模型所用的 Neuron 堆疊、需額外安裝的套件以及 run.py 腳本。該腳本遵循 TorchNeuron 發布時的配方：使用 from_pretrained 載入模型（必要時分片），然後使用靜態緩存和 torch.compile(backend="neuron") 進行生成。 官方原文・Hugging Face Blog

在 Trainium2 (trn2) 實例上，例如在 AWS Neuron Deep Learning Container 中，安裝套件後執行 python3 run.py 即可。無需使用 torchrun，因為腳本會在模型分片時自行啟動 torch.distributed ranks。首次運行會開始編譯模型的圖形。若要跳過此步驟，可從公開的 NEFF 緩存桶複製測試期間由相同編譯器構建的圖形。工具會提供帶有正確編譯器版本的命令：hf buckets sync hf://buckets/aws-neuron/torch-neuronx-neff-cache/neuronxcc-<version> /tmp/neff_cache。 官方原文・Hugging Face Blog

對使用者的意義

開發者可以使用此工具查詢特定模型是否在 AWS Trainium 上運行，並獲取具體的執行環境配置和腳本。這有助於在部署模型前確認硬體相容性，並利用公開的 NEFF 緩存加速首次編譯過程。所有顯示的內容都來自公開的 aws-neuron/torch-neuronx-model-support 數據集，每個 neuronx-cc 和庫版本對應一個 Parquet 檔案，使用者可以直接查詢該數據集以獲取詳細結果、錯誤訊息及堆疊版本。 官方原文・Hugging Face Blog

限制與未定事項

目前僅測試能容納在 4 個 Trainium2 裝置（最多 16 個核心）內的模型。PASS 結果不代表吞吐量或更長生成的表現。每個測試將模型載入 trn2 實例，分佈在能容納其權重的最少 Trainium 核心數上。大多數模型使用張量並行（將模型權重切分到多個核心上計算的分片方式）分片；沒有張量並行方案的混合專家（Mixture-of-experts，一種只在部分專家網路參與計算的架構）模型僅使用專家並行（將不同專家分配到不同核心的分片方式）。測試生成 32 個 token 並進行檢查：單核心時與同一檢查點的 CPU 運行結果比較；分片時由小型裁判模型檢查文本是否為流暢的英語。 官方原文・Hugging Face Blog