載入 AI 執行時的資料

資料與模型資產對於大型語言模型（LLM）和視覺語言模型（VLM）的深度學習及後期訓練工作負載至關重要。 透過 AI 執行環境，所有資料與模型資產皆可透過 Unity 目錄存取：

• Unity 目錄卷：主要用於大型資料集及非結構化檔案，包括圖片、音訊與文字。
• Unity 目錄資料表：用於結構化及表格資料，透過 Spark Connect 存取。

您的磁碟區和資料表必須註冊於 Unity Catalog，且可供您的使用者或服務主體存取。

適用於非結構化資料的 Unity Catalog 磁碟區

Unity 目錄卷提供任何格式的非表格資料的受控存取，包括結構化、半結構化及非結構化資料。 在 AI 執行時中，磁碟區是存取大型資料集、文字、模型資產及模型檢查點的主要機制。

使用者可以使用熟悉的檔案系統操作，在 Unity 目錄卷中列出、讀取及寫入檔案，類似於處理本地磁碟上的檔案：

import os

dir_path = "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir"
file_path = os.path.join(dir_path, "test_file")

os.makedirs(dir_path, exist_ok=True)

# Write to the file
with open(file_path, "w") as file:
    file.write("Hello, World!")

同樣地，Shell 操作的運作方式也相同：

%sh ls -l /Volumes/<catalog-name>/<schema-name>/<volume-name>
%sh mkdir -p /Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir
%sh touch /Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/test_file

Unity 目錄磁碟區的幾個特點使其非常適合機器學習工作負載：

• 分散式儲存：Unity Catalog 採用分散式儲存，允許 AI 執行時工作負載從筆記本與基於 CLI 工作負載中讀寫資料並建模資產。
• 針對機器學習存取模式優化：底層的儲存與存取路徑針對常見機器學習工作負載進行優化，特別是具有連續讀寫的大型檔案。 這使得 Unity Catalog 非常適合用於訓練資料載入、模型資產載入及模型檢查點撰寫。
• 類檔案系統存取：使用者可透過熟悉的檔案系統操作，在 Unity 目錄卷中列出、讀取及寫入檔案，類似於處理本地磁碟上的檔案。

由於自動背景提交，使用者可以預期對 Unity 目錄卷資料的存取是一致的：

• 寫入：AI 執行時會自動提交寫入，讓存取同一 Unity 目錄卷的其他應用程式與工作負載能看到變更。
• 讀取：AI 執行階段會自動偵測磁碟區中的變更，不需要使用者明確執行任何重新整理或同步作業。

調整音量表現

如前所述，Unity Catalog 磁碟區以分散式儲存體為後端，並已針對大型檔案的循序讀寫進行最佳化。

幾個建議可以幫助你從 AI 執行環境中獲得最佳效能：

• 將資料串接成較大的檔案：若可能，將資料整合成較少且較大的檔案，約每檔 1 GiB 至 10 GiB。 這讓 AI 執行時能夠積極預取資料，並自動達到接近最佳的序列讀取效能。

• 對於小檔案工作負載，請使用本地磁碟：如果你的工作負載包含許多小檔案，建議在處理前先用平行複製將檔案複製到本地磁碟。 這可以減少透過磁碟區反覆存取許多小型檔案所產生的額外負荷。

  # Recommended using parallel copy (256 concurrency in this example, you can tune)
  #
  # This takes only 22 seconds to copy 15,375 150KiB small image files.
  %sh cd /Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ && find . -type f -print0 | xargs -0 -P 256 -I {} cp --parents "{}" /tmp/
  
  
  # !!! Avoid doing this !!!
  #
  # Because the files are copied in serial, this copies the same 15,375 150KiB small image files much more slowly.
  # %sh cp -r /Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/* /tmp
  

• 你可以用 UCVolumeDataset 來處理機器學習的工作負載。 它結合了上述優化，以提供從 Unity 目錄卷中高效存取與載入資料。 如需詳細資訊，請參閱下列各節。

使用 UCVolumeDataset 載入非結構化資料

對於儲存在 Unity 目錄卷中的非結構化資料，如圖片、音訊和文字檔，請使用 UCVolumeDataset 套件中的 serverless_gpu.data 資料。 UCVolumeDataset 是 PyTorch IterableDataset，會在首次存取時將每個檔案從磁碟區複製到快速的本機快取，並傳回快取後的本機檔案路徑。 它處理了你原本會手動實現的效能與分布問題：

• 本地快取。 檔案會在首次存取時從 FUSE 掛載複製到本地快取目錄，之後再從快取中提供，因此多紀元訓練不會重新讀取該卷。
• 自動分割。 當 torch.distributed 初始化時，檔案會先分配到各個 rank，接著再進一步分配給 DataLoader 個工作程序，因此每個 (rank, worker) 配對都會取得互不重疊的資料分片，而不需要額外設定。

UCVolumeDataset 產生原始的本地檔案路徑。 若要將這些檔案解碼為張量，請再用第二個 IterableDataset 將其包住，讓它接收路徑串流並套用你的解析邏輯。 這樣可將 I/O 與解析的職責分開。

from serverless_gpu.data import UCVolumeDataset
from torch.utils.data import IterableDataset
from PIL import Image
import torchvision.transforms.functional as TF

class ImageDataset(IterableDataset):
    """Decodes each cached file path from UCVolumeDataset into a tensor."""

    def __init__(self, path_dataset: UCVolumeDataset):
        self._path_dataset = path_dataset

    def __iter__(self):
        for local_path in self._path_dataset:
            image = Image.open(local_path).convert("RGB")
            yield TF.to_tensor(image)

path_dataset = UCVolumeDataset("/Volumes/catalog/schema/my_volume/images")
dataset = ImageDataset(path_dataset)

包裝器接收已快取的本地路徑，因此解析步驟不會觸及該磁碟區。 你可以串連額外的包裝器來增強、標記化或過濾。

為了最佳效能，建議 UCVolumeDataset 搭配 serverless_gpu.data.DataLoader 而非原廠的 PyTorch DataLoader。 它針對 AI 執行時的 I/O 進行調整，並在 GPU 運算時同時擷取與快取檔案。

體積上的檢查點模型

為了替模型建立檢查點，讓你可以從最新的快照繼續訓練或在發生當機後復原，你可以像使用本機檔案系統一樣使用 Unity Catalog 磁碟區。

Databricks 建議使用分散式檢查點（DCP），以提升單 GPU 與多 GPU 工作負載的效能。 請參閱 Databricks 工程部落格中關於 AI 執行環境的快速容錯 PyTorch 訓練 。

import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict

import serverless_gpu

checkpoint_path = "/Volumes/my-catalog/my-schema/my-volume/checkpoints/step_1000"

# Save
model_sd, optim_sd = get_state_dict(model, optimizer)
state_dict = {"model": model_sd, "optim": optim_sd, "step": 1000}
dcp.async_save(
    state_dict,
    storage_writer=serverless_gpu.data.UCVolumeWriter(checkpoint_path))

# Load
model_sd, optim_sd = get_state_dict(model, optimizer)
state_dict = {"model": model_sd, "optim": optim_sd}
dcp.load(
    state_dict,
    storage_reader=serverless_gpu.data.UCVolumeReader(checkpoint_path))

set_state_dict(
    model,
    optimizer,
    model_state_dict=state_dict["model"],
    optim_state_dict=state_dict["optim"],
)

單體式的 torch.save 做法也可行。

• 對於單GPU模型檢查點，

  # The monolithic torch.save approach for single GPU chip
  
  # Save
  torch.save({"model": model.state_dict(), "opt": optimizer.state_dict()},
             "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ckpt.pt")
  
  
  # Load
  ckpt = torch.load(
      "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ckpt.pt",
      weights_only=True)
  model.load_state_dict(ckpt["model"])
  optimizer.load_state_dict(ckpt["opt"])
  

• 對於透過 torchrun 啟動的分散式訓練，

  # The monolithic torch.save approach for multi-GPU distributed training.
  # This snippet assumes your launcher has already called
  # dist.init_process_group(...).
  
  import os
  import torch.distributed as dist
  
  # Save only on rank 0.
  if dist.get_rank() == 0:
      torch.save({"model": model.state_dict(), "opt": optimizer.state_dict()},
                 "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ckpt.pt")
  
  # Wait for rank 0 to finish writing before any rank reads.
  dist.barrier()
  
  # Load on ALL ranks (map to current rank's local GPU).
  local_rank = int(os.environ["LOCAL_RANK"])
  ckpt = torch.load(
      "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ckpt.pt",
      map_location=f"cuda:{local_rank}",
      weights_only=True)
  model.load_state_dict(ckpt["model"])
  optimizer.load_state_dict(ckpt["opt"])
  

載入表格式資料

使用 Spark Connect 從 Delta 表格載入表格式機器學習資料。

針對單節點訓練，你可以使用 PySpark 方法toPandas()將 Apache Spark DataFrames 轉換成 pandas DataFrame，然後再用 PySpark 方法to_numpy()轉換成 NumPy 格式。

Spark Connect 支援大多數 PySpark API，包括 Spark SQL、Spark 上的 Pandas API、結構化串流以及基於 DataFrame 的 MLlib。 請參閱 PySpark API 參考文件 以了解最新支援的 API。

關於其他限制，請參見 無伺服器運算限制。

使用 Unity Catalog 磁碟區載入大型 Delta 資料表

對於過大而無法使用toPandas()轉換的大型 Delta 資料表，請將資料匯出至 Unity Catalog 卷，然後使用 PyTorch 或 Hugging Face 直接載入。

# Step 1: Export the Delta table to Parquet files in a UC volume
output_path = "/Volumes/catalog/schema/my_volume/training_data"
spark.table("catalog.schema.my_table").write.mode("overwrite").parquet(output_path)

# Step 2: Load the exported data directly using Hugging Face datasets
from datasets import load_dataset

dataset = load_dataset("parquet", data_files="/Volumes/catalog/schema/my_volume/training_data/*.parquet")

此方法避免了訓練期間的 Spark 開銷，且適用於單 GPU 及分散式訓練工作流程。