Ladda data i AI Runtime

Data- och modelltillgångar är avgörande för djupinlärning och efterträningsarbetsbelastningar för stora språkmodeller (LLM) och vision-language-modeller (VLM). Med AI Runtime nås all data och modelltillgångar via Unity Catalog:

• Unity-katalogvolymer: används främst för stora datamängder och ostrukturerade filer, inklusive bilder, ljud och text.
• Unity-katalogtabeller: används för strukturerad och tabellarär data, åtkomlig via Spark Connect.

Dina volymer och tabeller måste vara registrerade i Unity-katalogen och tillgängliga för din användare eller tjänstehuvudansvarig.

Unity-katalogvolym för ostrukturerad data

Unity-katalogvolymer ger styrd åtkomst till icke-tabellarär data i alla format, inklusive strukturerad, semistrukturerad och ostrukturerad data. I AI Runtime är volymer den primära mekanismen för att komma åt stora datamängder, text, modelltillgångar och modellkontrollpunkter.

Användare kan lista, läsa och skriva filer i Unity Catalog-volymer med hjälp av välbekanta filsystemoperationer, liknande att arbeta med filer på en lokal disk:

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

På liknande sätt fungerar skaloperationer på samma sätt:

%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

Några egenskaper hos Unity-katalogvolymer gör dem väl lämpade för maskininlärningsarbetsbelastningar:

• Distribuerad lagring: Unity Catalog backas upp av distribuerad lagring, vilket gör det möjligt för AI-körningsarbetsbelastningar att läsa och skriva data och modelltillgångar över plattformen, både från anteckningsböcker och CLI-baserade arbetsbelastningar.
• Optimerat för ML-åtkomstmönster: De underliggande lagrings- och åtkomstvägarna är optimerade för vanliga ML-arbetsbelastningar, särskilt stora filer med sekventiella läsningar och skrivningar. Detta gör Unity Catalog väl lämpad för att träna datainläsning, ladda modelltillgångar och skriva av kontrollpunkter för modeller.
• Filsystemliknande åtkomst: Användare kan lista, läsa och skriva filer i Unity-katalogvolymer med välbekanta filsystemoperationer, liknande att arbeta med filer på en lokal disk.

På grund av automatiska bakgrundsincheckningar kan användare förvänta sig konsekvent åtkomst till volymdata i Unity Catalog:

• Skriver: AI Runtime gör automatiskt skrivningar, vilket gör ändringar synliga för andra applikationer och arbetsbelastningar som använder samma Unity Catalog-volym.
• Läser: AI Runtime plockar automatiskt upp ändringar i volymen utan att användaren behöver någon explicit uppdatering eller synkronisering.

Prestanda för volyminställning

Som nämnts backas Unity Catalog-volymer upp av distribuerad lagring och är optimerade för stora filer med sekventiella läsningar och skrivningar.

Några tips kan hjälpa dig att få bästa möjliga prestanda från AI Runtime:

• Sammanfoga data till större filer: När det är möjligt, konsolidera data till färre, större filer, ungefär 1 GiB till 10 GiB per fil. Detta gör att AI Runtime aggressivt kan förhämta data och automatiskt uppnå nästan optimal sekventiell läsprestanda.

• För arbetsbelastningar med små filer, använd lokal disk: Om din arbetsbelastning involverar många små filer, överväg att kopiera filerna till den lokala disken med parallella kopior innan du bearbetar dem. Detta kan minska omkostnaderna vid upprepad åtkomst till många små filer via volymen.

  # 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
  

• Du kan använda UCVolumeDataset för dina maskininlärningsarbetsbelastningar. Den innehåller de optimeringar som beskrivs ovan för att möjliggöra effektiv dataåtkomst och laddning från Unity Catalog-volymer. Mer information finns i följande avsnitt.

Ladda ostrukturerad data med UCVolumeDataset

För ostrukturerad data såsom bilder, ljud och textfiler lagrade i Unity Catalog-volymer, använd UCVolumeDataset från modulen databricks.air.data. UCVolumeDataset är en PyTorch IterableDataset som kopierar varje fil från volymen till en snabb lokal cache vid första åtkomsten och returnerar sökvägen till den lokalt cachelagrade filen. Den hanterar de prestanda- och distributionsproblem som du annars skulle implementera för hand:

• Lokal cachelagring. Filer kopieras från FUSE-monteringen till en lokal cachekatalog vid första åtkomsten och hanteras från cachen efteråt, så träning med flera epoker läser inte volymen igen.
• Automatisk partitionering. När torch.distributed initieras fördelas filerna mellan ranker och delas sedan vidare mellan DataLoader worker-processer, så att varje (rank, worker)-par får en icke-överlappande del utan någon extra konfiguration.

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

UCVolumeDataset ger rådata lokala filsökvägar. Om du vill avkoda dessa filer till tensorer omsluter du dem på en sekund IterableDataset som förbrukar sökvägsströmmen och tillämpar din parsningslogik. Detta håller I/O och parsningsproblem åtskilda.

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)

Wrappern tar emot lokala sökvägar som redan finns i cacheminnet, så tolkningssteget kommer därför aldrig åt volymen. Du kan länka ytterligare omslutningar för förstoring, tokenisering eller filtrering.

För optimal prestanda, använd UCVolumeDataset tillsammans med databricks.air.data.DataLoader i stället för standardversionen av PyTorchs DataLoader. Den är anpassad för AI Runtime I/O och hämtar och cachar filer samtidigt medan GPU:n beräknar.

Kontrollpunktsmodeller på volymer

För att checkpointa din modell så att du kan återuppta träningen från den senaste snapshoten eller återhämta dig från en krasch, kan du använda Unity Catalog-volymer precis som ett lokalt filsystem.

Databricks rekommenderar att använda en distribuerad kontrollpunkt (DCP) för bättre prestanda på både arbetsbelastningar med ett enkelt och flera GPU. Se snabb, feltolerant PyTorch-träning på AI Runtime från Databricks ingenjörsblogg.

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

Den monolitiska torch.save-metoden fungerar också.

• För skapande av kontrollpunkter för en modell med en enda 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"])
  

• För distribuerad träning som startas via 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"])
  

Läsa in tabelldata

Använd Spark Connect för att läsa in tabellbaserade maskininlärningsdata från Delta-tabeller.

För träning med en nod kan du konvertera Apache Spark DataFrames till Pandas DataFrames med pyspark-metodentoPandas() och sedan konvertera till NumPy-format med pyspark-metodento_numpy().

Spark Connect stöder de flesta PySpark-API:er, inklusive Spark SQL, Pandas API på Spark, Structured Streaming och MLlib (DataFrame-baserad). Se referensdokumentationen för PySpark API för de senaste API:erna som stöds.

Andra begränsningar finns i Begränsningar för serverlös beräkning.

Ladda stora Delta-tabeller med hjälp av Unity Catalog-volymer

För stora Delta-tabeller som är för stora för att konvertera med toPandas()exporterar du data till en Unity Catalog-volym och läser in dem direkt med PyTorch eller 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")

Den här metoden undviker Spark-omkostnader under träning och fungerar bra för både arbetsflöden för enskild GPU och distribuerad utbildning.