在 TPU VM 上執行 AI/機器學習工作負載時的儲存空間最佳做法

如要盡可能提高 TPU VM 上 AI/機器學習工作負載的效能和成本效益,請為工作負載選取並設定合適的儲存空間解決方案。移除 I/O 瓶頸後,即可減少 TPU 加速器閒置的時間,進而縮短訓練時間並降低成本。

本文提供工作負載專屬的儲存空間建議,以及在 TPU VM 上訓練、檢查點、服務和快取資料的最佳做法。套用這些做法前,請先查看 TPU 資料的可用儲存空間選項。本文假設您熟悉 TPU VM,並具備佈建 Cloud Storage 資源的基本經驗。

工作負載專屬指引

下表列出不同工作負載的儲存空間建議 (依偏好順序排列):

工作負載 建議 最佳化和工具 (如適用)
訓練資料集,包括資料準備
  1. 使用 Cloud Storage 地區值區搭配 Rapid Cache 和用戶端調整 (例如預先讀取),以最低成本完成作業。
  2. 使用 Cloud Storage Rapid Bucket,獲得最佳效能和規模。
  3. 如果因為中繼資料並行程度高或檔案較小 (小於 1 MB) 而採用 Lustre (檔案系統儲存空間) 標準,請使用 Managed Lustre
  • 如果是區域 Cloud Storage 值區,請啟用階層命名空間
  • 如要使用單一區域 Cloud Storage bucket 或 Rapid Bucket,請使用 Cloud Storage FUSE 和 gRPC。如果是 GKE,請使用 gcsfusecsi-training 設定檔
  • 如果是 Managed Lustre,建議使用 Dynamic 層級來降低成本,並自動提升效能。
檢查點和強化學習權重
  1. 如要獲得高效能,且延遲時間必須極短 (不到 1 毫秒),例如強化學習的同步檢查點和高速權重傳播,請使用 Managed Lustre,不必進行大量調整。
  2. 使用 Cloud Storage Rapid Bucket,以最高處理量進行非同步和多層級檢查點作業,並將檢查點作業平行處理至儲存空間。Rapid Bucket 提供這些工作負載所需的高處理量。
如果是 Rapid Bucket,請使用 gcsfusecsi-checkpointing 設定檔
模型儲存空間和下載
  1. 使用 Hyperdisk ML 進行小型模型快取,最多可達 2,500 個執行個體。
  2. 如要採用單一儲存空間解決方案,請使用與訓練資料集相同的儲存空間解決方案 (Managed Lustre、Rapid Bucket 或區域值區)。
如要下載模型,請使用 GKE Run:ai Model Streamer 或 Cloud Storage FUSE,並透過 gcsfusecsi-serving 設定檔使用個別掛接點。
鍵/值 (KV) 快取卸載
  1. 一律使用主機 CPU 記憶體 (RAM) 做為主要層級,以釋出 VRAM,且不會增加明顯延遲。
  2. 如果 CPU 記憶體不足,請使用 Managed Lustre 做為平行低延遲的第二層快取階層,以滿足高處理量 (最高 10 TB/秒) 的 KV 快取卸載超低延遲 (不到 1 毫秒) 需求。

Cloud Storage 最佳化

下列各節說明使用 Cloud Storage 和 TPU VM 時,最佳化效能的最佳做法。

啟用階層式命名空間,以最佳化中繼資料

如要提升中繼資料效能,請在為 AI/機器學習工作負載建立地區 bucket 時,啟用階層式命名空間。中繼資料效能是指 Cloud Storage 處理作業的速度,這些作業包括查閱、列出或修改物件路徑和資料夾,而不是讀取或寫入檔案內容本身。

在未啟用階層式命名空間的 bucket 中,資料夾並非實際資源,而是以正斜線 (/) 分隔的物件名稱前置字串表示的模擬資料夾。因此,系統必須掃描所有具有該前置字串的物件,導致列出目錄內容或重新命名目錄等作業緩慢。階層式命名空間提供真正的檔案系統結構,對 AI/機器學習工作負載來說相當重要,原因如下:

  • 以不可分割的方式重新命名目錄:ML 架構會重新命名目錄,以完成檢查點。階層命名空間支援以不可分割的方式重新命名,確保查核點能快速完成。
  • 更高的初始 QPS:相較於未啟用階層命名空間的 bucket,階層命名空間最多可將讀取和寫入作業的初始每秒查詢次數 (QPS) 增加八倍。避免多個 TPU 同時存取儲存空間時發生瓶頸。
  • 提升資料夾層級作業效率:大幅加快在特定目錄中尋找及列出檔案的速度,縮短訓練和資料載入期間的回應時間。

透過 Rapid Bucket 提供的區域 bucket 預設會使用階層命名空間。詳情請參閱「階層式命名空間總覽」。

搭配適當的剖析檔使用 Cloud Storage FUSE

Cloud Storage FUSE 是一種 FUSE 轉接器,可讓您將 bucket 掛接為本機檔案系統。使用 Google Kubernetes Engine 時,建議使用 Cloud Storage FUSE CSI 驅動程式Cloud Storage FUSE 設定檔,自動調整效能。

如要進一步瞭解使用 Cloud Storage FUSE 的最佳做法,請參閱效能調校最佳做法

自訂 TPU VM 開機磁碟

如要自訂 TPU VM 的客體 OS 環境,可以使用開機指令碼或建立自訂映像檔。自訂開機磁碟適用於下列情境:

  • 預先載入軟體和程式庫:安裝特定 ML 架構、依附元件或自訂軟體,縮短 VM 啟動時間,並確保環境一致。
  • 使用非標準的 OS 發行版本:使用 Google 管理清單中未列出的 OS 發行版本或版本。
  • 套用安全性與監控設定:套用自訂安全性設定、安裝監控代理程式,或設定環境變數。

不過,TPU VM 的開機磁碟復原功能有限。您無法卸離或建立開機磁碟的快照以進行離線修復,因此請謹慎進行會影響開機程序的變更。遵循這些最佳做法,即可在自訂 TPU VM 環境時,降低啟動失敗的風險。

自訂開機磁碟時,請謹記以下重要原則:

  • 盡量減少開機磁碟修改次數:盡可能在永久磁碟或 Hyperdisk 磁碟區上安裝應用程式及儲存資料,而非大幅修改開機磁碟。

  • 使用 UUID 掛接:在 /etc/fstab 檔案中新增項目時,請一律使用 UUID 識別磁碟和分割區 (UUID=...),而非 /dev/sdb1 等裝置名稱。自動產生的裝置名稱不保證在重新啟動後仍維持不變。

請按照下列建議操作,降低系統變更時發生開機失敗的風險:

  • 錯誤處理:在指令碼中導入完善的錯誤檢查機制和順暢的失敗模式。將詳細訊息記錄到序列埠控制台和 Cloud Logging,協助進行偵錯。

  • 重要依附元件:修改開機必備檔案 (例如 /etc/fstab 檔案、網路設定或系統啟動載入程式設定) 時,請務必謹慎。語法錯誤或輸入錯誤可能會導致 VM 無法開機。

  • 次要磁碟:如果指令碼依賴次要磁碟,請確保指令碼可處理磁碟可能不存在或連接時間比預期長的情況。除非絕對必要,否則請避免開機程序嚴重依賴次要磁碟掛接。

    以下列舉幾個建議和不建議的/etc/fstab項目,用於掛接次要磁碟:

    • 建議使用UUID=a1b2c3d4-e5f6-7890-1234-567890abcdef /mnt/mydata ext4 defaults,nofail 0 2
    • 不建議使用/dev/sdb1 /mnt/mydata ext4 defaults 0 2

    使用 nofail 選項可防止系統在找不到磁碟時停止運作,但請確保應用程式可以處理無法使用的掛接點。

  • 套件管理:新增第三方存放區時請務必謹慎。請確認這些檔案可信,且與基本 OS 映像檔相容。瞭解您安裝的任何套件的依附元件,以及這些依附元件對系統程式庫的潛在影響。

  • 磁碟空間:監控開機磁碟使用率。如果記錄檔過多或安裝的軟體過大,開機磁碟可能會填滿,導致 VM 無法啟動。

  • 記錄:將應用程式和指令碼設定為詳細記錄到序列埠控制台,因為這是診斷 TPU VM 啟動問題的主要工具。

規劃儲存空間容量

規劃工作負載所需的儲存空間容量,充分運用加速器非常重要。包括儲存容量和檢查點頻寬。

預估儲存空間

您可以根據下列估算值,初步瞭解儲存空間需求:

工作負載類型 資料集儲存空間 檢查點儲存空間
LLM 預先訓練 每個 TPU 2 TB 每個 TPU 200 GB
多模態訓練 每個 TPU 12 TB 每個 TPU 1 TB
推論 每個 TPU 1 TB 每個 TPU 1 GB

預估檢查點頻寬

您可以使用下列公式,估算訓練工作負載所需的最低檢查點頻寬。如果是資料讀取、多次訓練執行,或是訓練和推論,請按比例提高預估頻寬需求。

  1. 檢查點大小參數數量 × 每個參數的位元組數 (FP16 + 最佳化工具狀態的每個參數約為 12 到 16 個位元組)。為最佳化工具狀態和不同精確度新增緩衝區 (約 3 倍)。
  2. 查核點間隔:查核點儲存頻率,例如每 15 分鐘儲存一次。
  3. 所需頻寬檢查點大小 ÷ 檢查點間隔

以下範例說明如何估算 Qwen3-72B 的最低檢查點頻寬:

  1. 檢查點大小:720 億個參數 × 12 個位元組 ≈ 每個檢查點 864 GB。加上緩衝區, 3 × 864 GB ≈ 2.5 TB。
  2. 查核點間隔:2 分鐘 = 120 秒。
  3. 所需頻寬:2.5 TB ÷ 120 秒 ≈ 20 GBps。

參考食譜

如需特定硬體和工作負載的儲存空間設定範例,請參閱下列做法:

配額和頻寬限制

Cloud Storage 和 Compute Engine 產品的頻寬會受到預設配額限制。如果超出配額,輸入和輸出要求可能會受到節流。

如要瞭解 Cloud Storage 配額和如何要求增加配額,請參閱 Cloud Storage 說明文件中的「配額與限制」。如要瞭解 Hyperdisk 和 Persistent Disk 的 Compute Engine 配額,請參閱「磁碟配額」。

後續步驟