Ray 加速 PyTorch Lightning 模型訓練與 GPU 資源管理

– 12 分鐘閱讀

PyTorch Lightning 簡化了 PyTorch 模型訓練流程，而 Ray 則提供了強大的分散式計算能力。結合兩者可以更有效率地訓練大規模模型，例如用於葡萄酒品質預測的模型。PyTorch Lightning 的 Trainer 類別封裝了訓練迴圈，方便管理訓練步驟和最佳化策略。透過 Ray 的分散式外掛，例如 RayPlugin，可以將 PyTorch Lightning 的訓練過程平行化到多個節點和 GPU 上，進一步提升訓練速度。然而，分散式訓練會引入額外的開銷，例如遠端呼叫和資料傳輸，因此需要仔細評估效能提升與開銷之間的平衡。除了 RayPlugin，Ray 還提供了 HorovodRayPlugin 和 RayShardedPlugin 等外掛，分別整合了 Horovod 和 FairScale 的分片 DDP 機制，提供更多彈性的分散式訓練選項。Ray 最初是為強化學習（RL）設計的，其 RLlib 函式庫提供了豐富的演算法和工具，方便開發和佈署 RL 應用。RL 的核心概念包括環境、狀態、獎勵、策略和價值，RLlib 則提供了高度分散式的 RL 工作負載支援。

在 Ray 中使用 GPU 進行加速時，需要考慮 GPU 資源的有效管理和釋放。由於許多高層次的函式庫不會主動釋放 GPU 資源，因此建議使用 max_calls=1 引數來強制重啟工作節點，確保資源釋放。對於需要重複使用 GPU 資料的場景，可以考慮使用長壽命的 Actor，但需要注意資源鎖定的問題。Ray 的 ML 函式庫，例如 Ray Train 和 Ray Tune，也提供了 GPU 相關的組態選項，方便使用者控制 GPU 的使用。在叢集組態中，建議明確定義不同型別的節點和其對應的硬體資源，以便 Ray 的自動擴充套件功能可以根據需求選擇合適的節點。最後，由於 Ray 本身不支援 CPU 備援模式，可以透過手動檢查 GPU 可用性，並根據結果選擇使用 GPU 或 CPU 執行任務，實作簡單的 CPU 備援機制。除了 GPU 之外，Ray 也支援其他型別的加速器，例如 TPU 和 FPGA，可以根據實際需求選擇合適的硬體加速方案。

使用 Ray 加速 PyTorch Lightning 的模型訓練

在現代機器學習中，訓練大規模模型的效率和效能是關鍵。PyTorch Lightning 和 Ray 是兩個強大的工具，它們可以幫助我們更高效地進行模型訓練。本文將探討如何使用 PyTorch Lightning 和 Ray 來加速葡萄酒品質預測模型的訓練過程。

PyTorch Lightning 的基本使用

PyTorch Lightning 是一個根據 PyTorch 的高層次框架，旨在簡化模型訓練的過程。它提供了一個 Trainer 類別，內部實作了訓練迴圈，使得所有必要的最佳化都可以在訓練迴圈中進行。

以下是使用 PyTorch Lightning 訓練葡萄酒品質預測模型的基本範例：

from pytorch_lightning import Trainer
# 訓練步驟設定
trainer = Trainer(max_steps=1000)
# 開始訓練
trainer.fit(model, train_dl)

內容解密：

這段程式碼展示瞭如何使用 PyTorch Lightning 進行模型訓練。首先，我們建立了一個 Trainer 物件，並設定最大訓練步數為 1000。然後，我們呼叫 trainer.fit 方法來開始訓練過程，並將模型和訓練資料集（train_dl）作為引數傳遞。

分散式訓練的實作

PyTorch 和 PyTorch Lightning 都支援分散式訓練，可以透過 Joblib 或 Horovod 來實作。然而，Ray 提供了一個更簡單且強大的方式來進行分散式訓練。

Ray 的 Distributed PyTorch Lightning Training 函式庫為 PyTorch Lightning 提供了分散式訓練的外掛，這些外掛允許我們在不改變任何程式碼的情況下，將訓練過程平行化到多核心或多節點、多 GPU 叢集中。

以下是如何使用 Ray 外掛來實作分散式訓練的範例：

from pytorch_lightning import Trainer
from ray.train.lightning import RayPlugin

plugin = RayPlugin(num_workers=6)
trainer = Trainer(max_steps=1000, plugins=[plugin])
trainer.fit(model, train_dl)
print(f"Build model in {time.time() - start}")
print(model)

內容解密：

這段程式碼展示瞭如何使用 Ray 外掛來實作分散式訓練。首先，我們匯入 RayPlugin 模組，並建立一個 RayPlugin 物件，指定工作節點數量為 6。然後，我們在 Trainer 中加入這個外掛，並啟動訓練過程。這樣可以讓我們的模型在多個工作節點上進行平行化訓練。

三種方法的效能比較

使用不同的方法來進行模型訓練時，執行時間會有所不同。以下是三種方法的效能比較：

  • PyTorch：16.6 秒
  • PyTorch Lightning：8.2 秒
  • 分散式 PyTorch Lightning with Ray：25.2 秒

可以看到，分散式 PyTorch Lightning with Ray 的執行時間較長，這是因為分散式系統中的遠端呼叫會帶來額外的開銷。

其他分散式外掛

除了 RayPlugin 外，Ray 函式庫還提供了其他兩種外掛：

  • HorovodRayPlugin：整合 Horovod 作為分散式訓練協定。
  • RayShardedPlugin：整合 FairScale 提供分片 DDP 機制。這種方式可以顯著降低大規模模型的記憶體使用量。

強化學習與 Ray

Ray 最早是作為強化學習（Reinforcement Learning, RL）平台開發的。RL 是機器學習的一種技術，透過代理在互動環境中透過試錯學習。

RL 的關鍵組成部分包括環境、狀態、獎勵、政策、價值等。Ray 的 RLlib 提供了高度分散式的 RL 工作負載支援。

此圖示解說：

此圖示展示了強化學習中的主要組成部分及其關係。環境中代理處於特定狀態下，根據當前狀態選擇行動並獲得獎勵反饋；同時根據政策決定最佳行動以最大化未來累積獎勵；價值則表示未來可能獲得的獎勵。

RLlib 的應用

RLlib 是建立在 Ray 上的一個強化學習函式庫，提供了多種高度分散式演算法、政策和損失函式。它支援多種應用場景，包括但不限於自動駕駛、遊戲 AI 和機器人控制等領域。

此圖示解說：

此圖示展示了 RLlib 的核心組成部分及其功能支援範圍：包含政策（Policy）、演算法（Algorithms）、損失函式（Loss Functions）和預設模型（Default Models）。

GPU與加速器在Ray中的使用

Ray主要專注於水平擴充套件，但有時候使用特殊的加速器如GPU，比單純增加更多「普通」的計算節點更為經濟且快速。GPU非常適合向量化運算，即同時對資料塊進行相同的操作。機器學習（ML）和更廣泛的線性代數都是主要應用場景之一，深度學習特別適合向量化處理。

然而，GPU資源通常比CPU資源昂貴，Ray的架構使得只在必要時才請求GPU資源變得相對簡單。要充分利用GPU，必須使用專門的函式庫，因為這些函式庫涉及直接記憶體存取，其結果並不總是可序列化的。在GPU計算領域，NVIDIA和AMD是兩個主要選擇，並且它們有不同的函式庫來進行整合。

GPU適合什麼樣的問題？

並非所有問題都適合GPU加速。GPU特別適契約時對大量資料點執行相同計算的問題。如果一個問題非常適合向量化，那麼GPU很可能也適合它。

以下是一些常見的受益於GPU加速的問題：

  • 機器學習（ML）
  • 線性代數
  • 物理模擬
  • 圖形處理（這裡沒有驚喜）

然而，GPU並不適合分支密集型非向量化工作流程，或是資料複製成本等同於或高於計算成本的工作流程。

基礎元件

使用GPU涉及額外的開銷，類別似於分散式任務的開銷（雖然速度稍快）。這些開銷來自於資料序列化以及通訊，儘管CPU與GPU之間的連線通常比網路連線更快。與Ray中的分散式任務不同的是，GPU沒有Python直譯器。高層次工具通常會生成或呼叫原生GPU程式碼。CUDA和Radeon Open Compute（ROCm）是與NVIDIA和AMD互動的兩個事實標準低層次函式庫。

NVIDIA首先發布了CUDA，並迅速在許多高層次函式庫和工具中獲得了認可，包括TensorFlow。AMD的ROCm起步較慢，並未見到相同程度的採用。一些高層次工具，包括PyTorch，現在已經整合了ROCm支援，但許多其他工具需要使用特定分叉版本的ROCm，如TensorFlow（tensorflow-rocm）或LAPACK（rocSOLVER）。

確保基礎元件正確可能會出乎意料地困難。例如，根據我們的經驗，在Linux4Tegra上構建具有Ray功能的NVIDIA GPU Docker容器花了幾天時間。ROCm和CUDA函式庫有特定版本支援特定硬體，同樣地，您可能希望使用的一些高層次程式只支援某些版本。如果您在Kubernetes或類別似容器化平台上執行時，可以從預構建容器如NVIDIA的CUDA映像或AMD的ROCm映像開始作為基礎。

高層次函式庫

除非您有特別需求，否則最好使用高層次函式庫來為您生成GPU程式碼，例如基本線性代數子程式（BLAS）、TensorFlow或Numba。應該嘗試在您正在使用的基礎容器或機器映像中安裝這些函式庫，因為它們通常涉及大量編譯時間。

一些函式庫如Numba會動態重寫您的Python程式碼。要讓Numba操作您的程式碼，您需要為函式新增裝飾器（例如@numba.jit）。不幸的是，numba.jit和其他動態重寫功能在Ray中並不直接支援。相反地，如果您使用這樣的一個函式庫時，只需將呼叫進行包裝即可。

此圖示

@startuml
skinparam backgroundColor #FEFEFE
skinparam componentStyle rectangle

title Ray 加速 PyTorch Lightning 模型訓練與 GPU 資源管理

package "機器學習流程" {
    package "資料處理" {
        component [資料收集] as collect
        component [資料清洗] as clean
        component [特徵工程] as feature
    }

    package "模型訓練" {
        component [模型選擇] as select
        component [超參數調優] as tune
        component [交叉驗證] as cv
    }

    package "評估部署" {
        component [模型評估] as eval
        component [模型部署] as deploy
        component [監控維護] as monitor
    }
}

collect --> clean : 原始資料
clean --> feature : 乾淨資料
feature --> select : 特徵向量
select --> tune : 基礎模型
tune --> cv : 最佳參數
cv --> eval : 訓練模型
eval --> deploy : 驗證模型
deploy --> monitor : 生產模型

note right of feature
  特徵工程包含：
  - 特徵選擇
  - 特徵轉換
  - 降維處理
end note

note right of eval
  評估指標：
  - 準確率/召回率
  - F1 Score
  - AUC-ROC
end note

@enduml
from numba import cuda, float32

# CUDA kernel
@cuda.jit
def mul_two(io_array):
    pos = cuda.grid(1)
    if pos < io_array.size:
        io_array[pos] *= 2 # do the computation

@ray.remote(num_gpus=1)
def remote_mul(input_array):
    # This implicitly transfers the array into the GPU and back, which is not free
    return mul_two(input_array)

#### 內容解密：
上述程式碼展示了一個簡單的CUDA範例以及如何在Ray中呼叫該範例。

首先定義了一個簡單的CUDA核函式`mul_two`。這個函式接受一個輸入陣列`io_array`並將其每個元素乘以2。
接著定義了一個Ray遠端函式`remote_mul`。這個函式接受一個輸入陣列並將其轉移到GPU上執行前面定義好的CUDA核函式。
遠端函式`remote_mul`使用裝飾器`@ray.remote(num_gpus=1)`來指示Ray在具有1個GPU資源的節點上執行該函式。

資源管理

您可以透過新增num_gpus引數到ray.remote裝飾器中來請求GPU資源，類別似於請求記憶體和CPU資源。與Ray中的其他資源（包括記憶體）一樣，GPU不保證可用性且Ray不會自動清理資源。儘管Ray不會自動清理記憶體資源給你， Python會做一些記憶體清理工作（但還是需要小心避免記憶體洩漏），因此相比記憶體洩漏更容易出現GPU洩漏。

使用GPU與加速器的Ray設計模式

在使用Ray進行大規模資料處理或機器學習時，GPU資源的管理與釋放是一個關鍵問題。許多高層次的函式庫在執行完畢後不會主動釋放GPU資源，這可能會導致資源浪費。玄貓將探討如何在Ray中有效地管理和釋放GPU資源，並介紹一些實用的設計模式。

釋放GPU資源

大多數高層次的Python函式庫在完成工作後不會主動釋放GPU資源，這意味著GPU可能會被佔用而無法被其他任務使用。為瞭解決這個問題，可以透過設定Ray的max_calls引數來強制釋放GPU資源。具體來說，在使用@ray.remote裝飾器時，可以新增max_calls=1來確保每次呼叫完成後都會重啟Python虛擬機器，從而釋放GPU資源。

# 請求全額GPU資源
@ray.remote(num_gpus=1, max_calls=1)
def do_serious_work():
    # 處理工作內容
    pass

內容解密：

  • 此程式碼片段展示瞭如何在Ray中請求並釋放GPU資源。
  • @ray.remote(num_gpus=1, max_calls=1)這行程式碼表示每次呼叫do_serious_work函式後，Ray都會重啟該函式所屬的工作節點，從而釋放GPU資源。
  • max_calls=1這個引數確保每次呼叫完成後都會重啟Python虛擬機器，從而避免GPU資源被長期佔用。

長壽命Actor

雖然重啟工作節點可以確保GPU資源被釋放，但這樣做的副作用是無法重用已經存在於GPU中的資料。為瞭解決這個問題，可以使用長壽命的Actor來代替普通的函式呼叫。這樣雖然可以保留資料，但也需要鎖定相應的資源。

Ray ML函式庫組態

Ray內建的ML函式庫也可以組態為使用GPU資源。例如，在Ray Train中，需要在Trainer建構函式中設定use_gpu=True來啟用GPU訓練。此外，Ray Tune提供了更靈活的資源請求方式，可以在tune.run中指定每個試驗所需的資源。

# 組態Ray Train使用GPU
trainer = Trainer(use_gpu=True, num_workers=4)

# 組態Ray Tune使用兩個CPU和一個GPU
tune.run(trainable, num_samples=10, resources_per_trial={"cpu": 2, "gpu": 1})

內容解密：

  • trainer = Trainer(use_gpu=True, num_workers=4)這行程式碼表示建立一個Trainer物件，並組態其使用GPU進行訓練。
  • tune.run(trainable, num_samples=10, resources_per_trial={"cpu": 2, "gpu": 1})這行程式碼表示在Ray Tune中執行一個試驗集合，並為每個試驗組態兩個CPU和一個GPU。

自動擴充套件與多種節點

Ray的自動擴充套件功能可以根據請求的資源型別選擇不同型別的節點。這在處理GPU資源時尤為重要，因為GPU通常比其他資源更加昂貴且數量有限。玄貓建議在叢集組態中明確定義不同型別的節點型別和其對應的硬體資源。

podTypes:
  rayGPUWorkerType:
    memory: 10Gi
    maxWorkers: 4
    minWorkers: 1
    CPU: 1
    rayResources:
      CPU: 1
      GPU: 1
      memory: 1000000000
    nodeSelector:
      node.kubernetes.io/gpu: gpu

內容解密：

  • podTypes:部分定義了不同型別的Pod（容器）。
  • rayGPUWorkerType:表示一種具有特定硬體組態的Pod型別。
  • memory: 10Gi、maxWorkers: 4、minWorkers: 1分別表示該Pod型別所需的記憶體大小、最大工作者數量和最小工作者數量。
  • CPU: 1、rayResources:部分則詳細說明瞭該Pod型別所需的CPU和其他Ray資源。
  • nodeSelector:部分指定了該Pod應佈署到哪些具有特定標籤（如具有GPU）的節點上。

CPU備援設計模式

大多數高層次函式庫都支援CPU備援模式，即當沒有可用的GPU時自動切換到CPU執行。然而，Ray本身並沒有內建這樣的功能。玄貓介紹了一種簡單但有效的方法來實作CPU備援：首先嘗試請求一個具有GPU的遠端任務，如果無法成功則回落到CPU執行。

# 檢查是否有可用的GPU
@ray.remote(num_gpus=1)
def do_i_have_gpus():
    return True

# 授予4分鐘時間檢查是否能取得到GPU
futures = [do_i_have_gpus.remote()]
ready_futures, rest_futures = ray.wait(futures, timeout=240)

resources = {"num_cpus": 1}
if ready_futures:
    resources["num_gpus"] = 1

# 根據可用資源選擇執行任務
@ray.remote(**resources)
def optional_gpu_task():
    pass

內容解密：

  • do_i_have_gpus()函式嘗試請求一個具有GPU的遠端任務。
  • ray.wait(futures, timeout=240)這行程式碼授予4分鐘時間來檢查是否能取得到可用的GPU。
  • 根據檢查結果選擇是否將任務組態為使用CPU或GPU。

其他非GPU加速器

除了以NVIDIA GPU為代表之外，各種特殊處理器如TPUs、FPGAs以及視訊記憶體技術等也是加速計算中的關鍵技術。這些技術同樣適用於提升計算效能。例如Numba可利用特殊CPU特性進行加速、TensorFlow可利用TPUs等高效硬體進行高效計算；此外Non-Volatile Memory Express (NVMe)驅動器也是提升I/O效能的一大利器。以上技術與策略均適用於其他硬體加速技術之設計考量。

玄貓 BlackCat

玄貓 BlackCat

技術愛好者，專注於分享程式開發、雲端技術與 AI 應用的心得體會。