Hugging Face + JuiceFS:多使用者多節點環境下提升模型載入效率

JuiceFS發表於2024-09-29

Hugging Face 的 Transformers 是一個功能強大的機器學習框架,提供了一系列 API 和工具,用於預訓練模型的下載和訓練。為了避免重複下載,提高訓練效率,Transformers 會自動下載和快取模型的權重、詞表等資源,預設儲存在 ~/.cache/huggingface/hub 目錄下。這個快取資料的機制。

但是,當有多使用者或多節點處理相同或關聯任務時,每個裝置都需要重複下載相同的模型和資料集,這勢必造成管理難度增大和浪費網路資源的問題。

要解決這個問題,可以將 Hugging Face 的快取資料目錄設定在共享儲存上,讓每個需要該資源的使用者能夠共享使用同一份資料。

在共享儲存的選擇上,如果裝置不多且都在本地,則可以考慮使用 Samba 或 NFS 共享。如果計算資源分佈在不同的雲或不同地區的機房,則需要採用效能和一致性都更有保證的分散式檔案系統,JuiceFS 就是一個非常適合的方案。

利用 JuiceFS 分散式、多端共享、強一致性等特性,不同計算節點間可以高效共享和遷移訓練資源,免於重複準備相同資料,從而顯著最佳化資源使用和儲存管理,提高整個 AI 模型的訓練效率。

JuiceFS 架構

JuiceFS 是開源的雲原生分散式檔案系統,採用了資料與後設資料分離儲存的技術架構,以物件儲存作為底層儲存來儲存資料,以鍵值儲存或關係型資料庫作為後設資料引擎來儲存檔案的後設資料。這些計算資源可以自行搭建,也可以在雲平臺上購買,因此 JuiceFS 是很容易搭建和使用的。

底層資料儲存

在底層儲存方面,JuiceFS 支援市面上幾乎所有主流的公有云物件儲存服務,比如 Amazon S3、Google Cloud Storage、阿里雲 OSS等,也支援 MinIO、Ceph 等私有部署的物件儲存。

後設資料引擎

在後設資料引擎方面,JuiceFS 支援 Redis、MySQL、PostgreSQL 等多種資料庫,另外,也可選購 JuiceFS 官方的雲服務版本,它採用 Juicedata 官方自研的高效能分散式後設資料引擎,可以滿足更高效能要求的場景需求。

JuiceFS 客戶端

JuiceFS 分為開源版和雲服務版,它們所使用的客戶端不同,但使用方法基本一致。本文以開源版為例進行介紹。

JuiceFS 社群版提供了跨平臺的客戶端,支援 Linux、macOS、Windows 等作業系統,可以在各種環境中使用。

當你基於這些計算資源搭建了一個 JuiceFS 檔案系統以後,就可以透過 JuiceFS 客戶端提供的 API 或者 FUSE 介面來訪問 JuiceFS 檔案系統,實現檔案的讀寫、後設資料的查詢等操作。

對於 Hugging Face 的快取資料儲存,可以將 JuiceFS 掛載到 ~/.cache/huggingface/ 目錄下,這樣 Hugging Face 相關的資料就可以儲存到 JuiceFS 中。另外,也可以透過設定環境變數自定義 Hugging Face 快取目錄位置,指向 JuiceFS 掛載的目錄。

接下來就展開介紹如何建立 JuiceFS 檔案系統,以及將 JuiceFS 用於 Hugging Face 的快取目錄的兩種方法。

建立 JuiceFS 檔案系統

假設已經準備好了以下物件儲存和後設資料引擎:

  • 物件儲存 Bucket:https://jfs.xxx.com
  • 物件儲存 Access Key:your-access
  • 物件儲存 Secret Key:your-secret
  • Redis 資料庫:your-redis.xxx.com:6379
  • Redis 密碼:redis-password

其中,物件儲存相關資訊只在建立 JuiceFS 檔案系統時使用一次,它們會被寫入後設資料引擎,在之後的使用中只需要後設資料引擎的地址和密碼。

安裝 JuiceFS 客戶端

對於 Linux 和 macOS 系統,可以透過以下命令安裝 JuiceFS 客戶端:

curl -sSL https://d.juicefs.com/install | sh -

Windows 系統建議在 WSL 2 的 Linux 環境中使用 JuiceFS 客戶端,另外也可以自行下載預編譯的 JuiceFS 客戶端,詳情請參考 JuiceFS 官方文件。

建立 JuiceFS 檔案系統

使用 format 命令建立 JuiceFS 檔案系統:

juicefs format \
  --storage s3 \
  --bucket https://jfs.xxx.com \
  --access-key your-access \
  --secret-key your-secret \
  "redis://:redis-password@your-redis.xxx.com:6379" \
  hf-jfs

其中,

  • hf-jfs 是 JuiceFS 檔案系統的名稱,可以自定義;
  • --storage s3 指定物件儲存型別為 S3,可以參考官方文件瞭解更多物件儲存型別;
  • 後設資料引擎 URL 以 redis:// 開頭,緊接著是 Redis 的使用者名稱和密碼,使用 : 分隔,然後是 Redis 的地-址和埠號,使用 @ 分隔。建議使用引號包裹整個 URL。

將 JuiceFS 預先掛載到 Hugging Face 快取目錄

如果還沒有安裝 Hugging Face Transformers,則可以預先建立資料快取目錄,並將 JuiceFS 掛載到該目錄。

# 建立 Hugging Face 快取目錄
mkdir -p ~/.cache/huggingface

# 掛載 JuiceFS 到 Hugging Face 快取目錄
juicefs mount -d "redis://:redis-password@your-redis.xxx.com:6379" ~/.cache/huggingface

緊接著再安裝 Hugging Face Transformers,它會自動將資料快取到 JuiceFS 中,例如:

pip install transformers datasets evaluate accelerate

由於不同硬體和環境涉及的包會有所不同,請根據實際情況進行安裝配置,這裡不做展開。

透過環境變數指定 Hugging Face 快取目錄

另一種方法是透過環境變數指定 Hugging Face 快取目錄,這樣可以在不修改程式碼的情況下,將 Hugging Face 快取目錄指向 JuiceFS 掛載的目錄。

比如,JuiceFS 的掛載目錄是 /mnt/jfs:

juicefs mount -d "redis://:redis-password@your-redis.xxx.com:6379" /mnt/jfs

透過 HUGGINGFACE_HUB_CACHE 或 TRANSFORMERS_CACHE 環境變數指定 Hugging Face 快取目錄:

export HUGGINGFACE_HUB_CACHE=/mnt/jfs

這樣,Hugging Face Transformers 就會將資料快取到 JuiceFS 掛載的目錄中。

隨處使用 Hugging Face 快取資料

得益於 JuiceFS 分散式多端共享儲存的特性,使用者只需將 Hugging Face 的快取目錄設定到 JuiceFS 掛載點,並完成首次模型資源的下載即可。

隨後,在任何需要該資源的節點上掛載 JuiceFS,選擇上述兩種方法之一進行設定,即可複用快取資料。JuiceFS 採用 “close-to-open” 機制,確保在多個節點上共享讀寫相同資料時的資料一致性。

需要注意的是,從 Hugging Face 下載模型資源的速度會受到網路環境的影響,不同國家和地區、不同網路線路的速度可能有所不同。此外,在使用 JuiceFS 共享資料快取目錄時,也可能面臨類似的網路延遲問題。為了有效改善速度,建議儘量提升頻寬,並降低工作節點與 JuiceFS 底層物件儲存以及訪問後設資料引擎之間的網路延遲。

總結

本文介紹了 JuiceFS 作為 Hugging Face Transformers 的資料快取目錄的兩種方法,分別是將 JuiceFS 預先掛載到 Hugging Face 快取目錄,以及透過環境變數指定 Hugging Face 快取目錄。這兩種方法都可以實現 AI 訓練資源在多個節點間共享複用。

如果你有多個節點需要訪問相同的 Hugging Face 快取資料,或者希望在不同環境中載入相同的訓練資源,JuiceFS 也許是一個非常理想的選擇。

希望本文介紹的內容能夠對你在 AI 模型訓練的過程中提供一定的幫助,如果有相關問題歡迎加入 JuiceFS 微信群與社群使用者共同交流討論。

相關文章