Memuat data pada Runtime AI

Penting

Fitur ini ada di Pratinjau Umum.

Data dan aset model sangat penting untuk pembelajaran mendalam dan beban kerja pasca-pelatihan untuk model bahasa besar (LLM) dan model bahasa visi (VLM). Dengan AI Runtime, semua data dan aset model diakses melalui Unity Catalog:

  • Volume Unity Catalog: digunakan terutama untuk dataset besar dan file tidak terstruktur, termasuk gambar, audio, dan teks.
  • Tabel Unity Catalog: digunakan untuk data terstruktur dan tabel, diakses melalui Spark Connect.

Volume dan tabel Anda harus terdaftar di Unity Catalog dan dapat diakses oleh pengguna atau prinsipal layanan Anda.

Volume Unity Catalog untuk data tidak terstruktur

Volume Unity Catalog menyediakan akses terkontrol ke data non-tabel dalam format apa pun, termasuk data terstruktur, semi-terstruktur, dan tidak terstruktur. Dalam AI Runtime, volume adalah mekanisme utama untuk mengakses dataset besar, teks, aset model, dan titik pemeriksaan model.

Pengguna dapat mendaftarkan, membaca, dan menulis file dalam volume Unity Catalog menggunakan operasi sistem file yang sudah dikenal, mirip dengan bekerja dengan file di disk lokal:

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!")

Begitu pula, operasi shell berfungsi dengan cara yang sama:

%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

Beberapa karakteristik volume Unity Catalog membuatnya sangat cocok untuk beban kerja machine learning:

  • Penyimpanan terdistribusi: Unity Catalog didukung oleh penyimpanan terdistribusi, memungkinkan beban kerja AI Runtime membaca dan menulis data serta memodelkan aset di seluruh platform, baik dari notebook maupun beban kerja berbasis CLI.
  • Dioptimalkan untuk pola akses ML: Jalur penyimpanan dan akses yang mendasarinya dioptimalkan untuk beban kerja ML umum, terutama file besar dengan pembacaan dan penulisan berurutan. Hal ini membuat Unity Catalog sangat cocok untuk pelatihan pemuatan data, pemuatan aset model, dan penulisan checkpoint model.
  • Akses mirip sistem file: Pengguna dapat mendaftarkan, membaca, dan menulis file dalam volume Unity Catalog menggunakan operasi sistem file yang sudah dikenal, mirip dengan bekerja dengan file di disk lokal.

Karena adanya commit latar belakang otomatis, pengguna dapat mengharapkan akses yang konsisten ke data volume Unity Catalog:

  • Menulis: AI Runtime secara otomatis mengkomitkan penulisan, sehingga perubahan terlihat oleh aplikasi dan beban kerja lain yang mengakses volume Unity Catalog yang sama.
  • Baca: AI Runtime secara otomatis menangkap perubahan volume tanpa memerlukan operasi penyegaran atau sinkronisasi eksplisit dari pengguna.

Optimalkan kinerja volume

Seperti yang disebutkan, volume Unity Catalog didukung oleh penyimpanan terdistribusi dan dioptimalkan untuk file besar dengan bacaan dan penulisan berurutan.

Beberapa tips dapat membantu Anda mendapatkan performa terbaik dari AI Runtime:

  • Gabungkan data ke dalam file yang lebih besar: Jika memungkinkan, konsolidasikan data menjadi file yang lebih sedikit dan lebih besar, sekitar 1 GiB hingga 10 GiB per file. Hal ini memungkinkan AI Runtime untuk mengambil data secara agresif dan secara otomatis mencapai kinerja membaca berurutan yang hampir optimal.

  • Untuk beban kerja file kecil, gunakan disk lokal: Jika beban kerja Anda melibatkan banyak file kecil, pertimbangkan untuk menyalin file ke disk lokal menggunakan salinan paralel sebelum memprosesnya. Ini dapat mengurangi overhead akibat akses berulang ke banyak file kecil melalui volume tersebut.

    # 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

  • Anda dapat menggunakan UCVolumeDataset untuk beban kerja pembelajaran mesin Anda. Aplikasi ini menggabungkan optimasi yang dijelaskan di atas untuk memberikan akses dan pemuatan data yang efisien dari volume Unity Catalog. Lihat bagian berikut untuk detail selengkapnya.

Muat data tidak terstruktur dengan UCVolumeDataset

Untuk data tidak terstruktur seperti gambar, audio, dan file teks yang disimpan dalam volume Unity Catalog, gunakan UCVolumeDataset dari databricks.air.data modul. UCVolumeDataset adalah PyTorch IterableDataset yang menyalin setiap file dari volume ke cache lokal yang cepat pada akses pertama dan menghasilkan jalur file lokal yang di-cache. Ini mengatasi persoalan kinerja dan distribusi yang jika tidak demikian harus Anda implementasikan secara manual:

  • Cache lokal. File-file disalin dari mount FUSE ke direktori cache lokal saat pertama kali diakses, lalu selanjutnya diambil dari cache, sehingga pelatihan multi-epoch tidak perlu membaca ulang volume.
  • Pemartisian otomatis. Ketika torch.distributed diinisialisasi, berkas dipartisi di antara rank lalu dibagi lagi ke dalam pekerja DataLoader, sehingga setiap pasangan (rank, worker) menerima bagian yang tidak saling tumpang tindih tanpa konfigurasi tambahan.

Nota

UCVolumeDataset dan databricks.air.data.DataLoader berasal dari paket databricks-sdk-air. Pasang dengan fitur data tambahan, yang juga menginstal torch yang kompatibel:

%pip install "databricks-sdk-air[data]"

UCVolumeDataset menghasilkan jalur file lokal mentah. Untuk mendekode file-file tersebut menjadi tensor, bungkus dengan IterableDataset kedua yang menerima aliran jalur dan menerapkan logika parsing Anda. Ini menjaga I/O dan penguraian tetap terpisah.

from databricks.air.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)

Wrapper menerima path lokal yang sudah tersimpan di cache, sehingga tahap parsing tidak pernah mengakses volume. Anda dapat menautkan pembungkus tambahan untuk augmentasi, tokenisasi, atau pemfilteran.

Untuk kinerja optimal, gunakan UCVolumeDataset bersama databricks.air.data.DataLoader alih-alih DataLoader bawaan PyTorch. Perangkat ini disetel untuk AI Runtime I/O dan mengambil serta menyimpan file secara bersamaan saat GPU menghitung.

Model checkpoint pada volume

Untuk memeriksa model Anda agar dapat melanjutkan pelatihan dari snapshot terbaru atau memulihkan dari crash, Anda dapat menggunakan volume Unity Catalog seperti sistem file lokal.

Databricks merekomendasikan penggunaan checkpoint terdistribusi (DCP) untuk kinerja yang lebih baik pada beban kerja GPU tunggal maupun multi-GPU. Lihat pelatihan PyTorch yang cepat dan tahan kesalahan tentang AI Runtime dari blog teknik Databricks.

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

import databricks.air.data

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=databricks.air.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=databricks.air.data.UCVolumeReader(checkpoint_path))

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

Pendekatan monolitik `torch.save` juga dapat digunakan.

  • Untuk checkpoint model GPU tunggal,

    # 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"])

  • Untuk pelatihan terdistribusi yang diluncurkan melalui 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"])

Memuat data tabular

Gunakan Spark Connect untuk memuat data pembelajaran mesin tabular dari tabel Delta.

Untuk pelatihan node tunggal, Anda dapat mengonversi Apache Spark DataFrames menjadi pandas DataFrames menggunakan metode toPandas(), lalu secara opsional mengonversi ke format NumPy menggunakan metode to_numpy().

Nota

Spark Connect menunda analisis dan resolusi nama ke waktu eksekusi, yang dapat mengubah perilaku kode Anda. Lihat Membandingkan Spark Connect dengan Spark Classic.

Spark Connect mendukung sebagian besar API PySpark, termasuk Spark SQL, Pandas API di Spark, Structured Streaming, dan MLlib (berbasis DataFrame). Lihat dokumentasi referensi API PySpark untuk API terbaru yang didukung.

Untuk batasan lain, lihat Batasan komputasi tanpa server.

Muat tabel Delta besar menggunakan volume Unity Catalog

Untuk tabel Delta besar yang terlalu besar untuk dikonversi dengan toPandas(), ekspor data ke volume Unity Catalog dan muat langsung menggunakan PyTorch atau 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")

Pendekatan ini menghindari overhead Spark selama pelatihan dan bekerja dengan baik untuk GPU tunggal dan alur kerja pelatihan terdistribusi.

  • Terakhir diperbarui pada