最高花費1700萬美元，這是租卡訓練谷歌5400億參數PaLM的成本

機器之心報道

機器之心編輯部

假設我們普通人想用雲計算來訓練一個 PaLM，我們需要準備多少錢？一位網友算出的結果是：900~1700 萬美元。

從去年開始，谷歌人工智能主管 JeffDean 就開始了「谷歌下一代人工智能架構」——Pathways 的預告。與之前為數千個任務訓練數千個模型的方法不同，新架構的願景是訓練一個模型做成千上萬件事情。

一年之後，Pathways 系統論文終于亮相，Jeff Dean 所在的團隊還公布了用它訓練的一個大型語言模型——PaLM。實驗表明，PaLM 在多語言任務和代碼生成方面具有強大的能力，可以出色地完成笑話解讀、bug 修複、從表情符号中猜電影等語言、代碼任務。

PaLM 解讀笑話示例。

PaLM 是一個隻有解碼器的密集 Transformer 模型，參數量達到了 5400 億。為了訓練這個模型，谷歌動用了 6144 塊 TPU，讓 Pathways 在兩個 Cloud TPU v4 Pods 上訓練 PaLM。這是名副其實的「鈔能力」。

驚歎之餘，有人可能想問：假設我們普通人（不像谷歌那樣擁有大量 TPU）想用雲計算來訓練一個 PaLM，我們需要準備多少錢？一位網友算出的結果是：900~1700 萬美元。

我們一起來看一下他是怎麼算的。

論文裡的可用資訊

随着參數的增加，ML 模型的計算成本也在飙升。谷歌曾總結過，自 2010 年以來，ML 模型的訓練計算量驚人地增長了 100 億，平均每 6 個月就翻一番。如今，PaLM 站上了 C 位。

論文顯示，PaLM 的最終訓練運作消耗的算力是 2.56e24 FLOPs。

論文還提到，PaLM-540B 在 6144 塊 TPU v4 晶片上訓練了 1200 小時，在 3072 塊 TPU v4 晶片上訓練了 336 小時，包括一些停機時間（downtime）和重複步驟。

是以，PaLM-540B 的訓練總共花費：

2.56e24 FLOPs；

8404992 個 TPUv4 chip-hour（每個晶片包含 2 個核，約合 16809984 個 TPUv4core-hour）；

大約 64 天。

此外，在 TPU 使用率方面，PaLM-540B 的訓練使用了 rematerialization，因為帶有 rematerialization 的可行批大小可實作更高的訓練吞吐量。不考慮 rematerialization 成本，最終模型在沒有自注意力機制的情況下 FLOP 使用率為 45.7%，有自注意力為 46.2%。PaLM 的分析計算硬體 FLOPs 使用率（包括 rematerialization FLOPs）為 57.8%。

估算結果

現在有兩種方法可以估計 PaLM-540B 的成本：

1、最後訓練運作使用 2.56×10 FLOPs

可以通過租用 TPU 執行個體來估計每個 flops 的成本（假設使用率為 57.8%）。

從其他雲提供商（例如使用 NVIDIA A100 的雲提供商）擷取每個 FLOP 的成本，然後估計總成本。

2、使用 8404992 個 TPUv4 chip-hour 這一資料

方法 2 似乎更加準确，但遺憾的是，作者沒有拿到有關 TPUv4 的租用價格資料（需要咨詢銷售代表）。

是以，他選擇了第一種方法，并按三種方式分别計算了一下。

通過 Google Cloud 租用 TPUv3

我們可以以每小時 32 美元的價格租用一個 32 核的 TPUv3 pod。約合 1 美元一個 TPUcore-hour。

考慮上述 16809984 個 TPUv4core-hour，如果用 TPUv3 替換 TPUv4，則花費約 1700 萬美元。

不過這裡的估算存在兩個問題：

一方面，TPUv3 的性能不如 TPUv4，是以我們需要更多的時間或更多的 TPU。按這個思路計算，實際花費應該高于這個數字；

另一方面，我們經常看到，雖然平台提供的硬體更好了，但價格可能不怎麼變化。是以，如果 Google Cloud 給出的 TPUv4 與 TPUv3 的價格大緻相同，那麼這種估計就是公平的。但如果二者價格相差較大，那麼實際成本也會有較大差異。

按每 FLOP 的花銷算

我們知道，TPUv3 晶片為 bfloat16 提供大約每秒 123 TFLOP（TFLOPS）的算力。當然，這隻是表格裡顯示的峰值性能。

由于不同情況下，硬體的使用率存在差異，是以實際的 TFLOPS 資料往往低于峰值資料。前面說過，PaLM 的 FLOPs 使用率達到了驚人的 57.8%。與之前的諸多模型相比（如下圖），這是一項新的記錄。

是以，作者假設 PaLM 在 TPU v3 上訓練時硬體使用率達到 50%：

按照這個算法，我們每一美元可以買到 221 PFLOPs。考慮到最後一次訓練要用 2.56×10^24FLOPs 的算力，我們的最終花費大約是 1160 萬美元：

按租用 NVIDIA 顯示卡算

兩年前，有人給 GPT-3 算過一筆賬，發現如果使用當時市場上價格最低的 GPU 雲（使用 Lambda GPU 執行個體）來訓練 GPT-3，花費最低為 460 萬美元。

圖源：https://lambdalabs.com/blog/demystifying-gpt-3/

如果隻考慮 PaLM 的訓練計算量是 GPT-3 的 10 倍這一事實，PaLM 的最終訓練成本應為 4600 萬美元左右。

但是，文章裡的資料畢竟是兩年前的，而且用的是 Tesla V100。現在的 NVIDIA A100 性能（Tensor 性能）已經提升了一個數量級。

是以，如果按硬體性能提高到原來的 10 倍，使用率是 50% 來計算，PaLM 的訓練成本大概是 920 萬美元左右。

結論

作者用三種方法估計了 PaLM 的最終訓練成本，結果分别為 1700 萬美元、1160 萬美元和 920 萬美元左右。

但需要注意的是：

1、谷歌并不需要花那麼多錢，他們擁有硬體。這裡是假設終端消費者因租用 TPUv3 pod 訓練 PaLM 而向 Google Cloud 支付的錢；

2、如果租用時間比較長，你可以拿到折扣（1 年 37% 的折扣）；

3、作者沒有 TPUv4 的價格資料，是以使用了 TPUv3 的。

4、這裡假設你知道如何高效利用 TPUv3 pod，将使用率提到 50%，這一使用率非常驚人；

5、這裡隻讨論最後一次訓練的成本，不包括其他困難且費錢的工作，如工程、研究、測試等。

參考連結：https://blog.heim.xyz/palm-training-cost/