使用 Hugging Face 與 Dask 擴展 AI 資料處理規模

Hugging Face 與 Dask 提供一個可擴展的框架，用於處理超出本機記憶體限制的大規模 AI 資料集。透過結合 Dask 的分散式運算能力與 Hugging Face 的 transformers 與 datasets，使用者可以將 AI 任務（例如模型推論與資料過濾）從筆記型電腦上的少量資料列，擴展至多 GPU 雲端叢集上的數億筆資料列。

使用 Dask 進行分散式資料處理

Dask 支援 out-of-core 計算，透過將資料切分為可管理的區塊，使得處理過大而無法全部載入系統記憶體的資料集成為可能。對於熟悉 pandas 的使用者而言，Dask DataFrame 提供類似的 API，簡化了從本地原型開發到大規模生產的過渡。

Key benefits of using Dask for AI data processing include:

• 高效載入： Dask 原生支援 Parquet，這是 Hugging Face datasets 的預設格式，能夠有效進行欄位過濾與壓縮。
• 平行執行： map_partitions 函式允許使用者在較大的 Dask DataFrame 中，對每個 pandas DataFrame 分區平行套用自訂函式（例如模型推論）。
• 分散式寫入： Dask 支援平行將結果寫回 Parquet 格式，可與 Hugging Face dataset repository 整合。

擴展模型推論：從 Pandas 到 Dask

為了示範擴展，Hugging Face 使用了 FineWeb 資料集——包含 15 兆個英文網路資料的 token——以及 FineWeb-Edu 分類器，以辨識具高教育價值的網頁。

使用 Pandas 進行本地原型開發

在小規模（例如 100 列）時，可使用 pandas 執行 FineWeb-Edu 分類器。在配備 GPU 的 M1 Mac 上，此過程約需 10 秒。工作流程使用 Hugging Face 的 pipeline 進行文字分類，且在函式內會動態選擇硬體裝置（CUDA、MPS 或 CPU），以確保程式碼在後續分散時的相容性。

擴展至 2.11 億列

當擴展至 2.11 億列的資料集（屬於 432 GB 磁碟上的爬蟲資料）時，串列處理變得極度緩慢。透過切換至 Dask DataFrame 並從 Hugging Face 懶惰載入資料，任務即可平行化。

在此擴展工作流程中，compute_scores 函式透過 map_partitions 套用。為了最佳化效能，會將 Hugging Face pipeline 中的 batch_size 提升（例如至 768），以更有效利用 GPU 硬體。

雲端多 GPU 平行推論

為了達到最高吞吐量，Dask 可部署於雲端基礎設施。於範例中，使用 Coiled 佈建了 100 個工作者的叢集，採用 AWS g5.xlarge 實例（搭載 NVIDIA A10 Tensor Core GPU）。

基礎設施自動化

• VM 建置： 自動啟動具 GPU 功能的雲端 VM。
• 環境設定： 處理 NVIDIA 驅動程式與 CUDA 執行環境的安裝。
• 套件同步： 將本地 Python 套件與檔案同步至雲端工作者，以確保環境一致性。

效能與資源利用率

處理 2.11 億列資料約耗時 5 小時。監控顯示硬體資源使用效率高，GPU 中位利用率為 100%，記憶體中位使用量為 21.5 GB（在可用的 24 GB GPU 記憶體中）。

潛在 AI 應用案例

• 基因組資料過濾： 從龐大的基因組資料集中挑選特定感興趣的基因。
• 結構化資料抽取： 使用大型語言模型將非結構化文字轉換為結構化資料集。
• 網路資料清理： 清理與過濾來自 Common Crawl 的大規模抓取資料。
• 多模態推論： 使用多模態模型分析大規模音訊、影像或影片資料集。