Načtení dat do AI Runtime

Důležité

Tato funkce je ve verzi Public Preview.

Data a modelové zdroje jsou klíčové pro hluboké učení a posttrénovací pracovní zátěže pro velké jazykové modely (LLM) a modely vizuálního jazyka (VLM). S AI Runtime jsou všechna data a modelová asseta přístupná prostřednictvím Unity Catalog:

• Svazky katalogu Unity: používají se především pro velké datové sady a nestrukturované soubory, včetně obrázků, zvuku a textu.
• Tabulky Unity Catalog: používají se pro strukturovaná a tabulková data, přístupné přes Spark Connect.

Vaše svazky a tabulky musí být registrovány v Unity Catalog a přístupné vašemu uživateli nebo principálu služby.

Svazek katalogu Unity pro nestrukturovaná data

Svazky Unity Catalog poskytují řízený přístup k netabulkárním datům v jakémkoli formátu, včetně strukturovaných, polostrukturovaných i nestrukturovaných dat. V AI Runtime jsou objemy hlavním mechanismem pro přístup k velkým datovým sadám, textům, modelovým assetům a kontrolním bodům modelu.

Uživatelé mohou vypisovat, číst a zapisovat soubory do svazků Unity Catalog pomocí známých operací souborového systému, podobně jako při práci se soubory na lokálním disku:

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

Podobně fungují operace v shellu stejně:

%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ěkolik charakteristik svazků Unity Catalog je činí vhodnými pro úkoly zaměřené na strojové učení:

• Distribuované úložiště: Unity Catalog je podpořen distribuovaným úložištěm, což umožňuje AI Runtime zátěžím číst a zapisovat data a modelovat prostředky napříč platformou, jak z notebooků, tak z CLI založených na zátěžích.
• Optimalizováno pro přístupové vzory ML: Základní úložné a přístupové cesty jsou optimalizovány pro běžné ML zátěže, zejména pro velké soubory se sekvenčními čteními a zápisy. To činí Unity Catalog velmi vhodným pro načítání trénovacích dat, načítání modelových assetů a psaní kontrolních bodů modelu.
• Přístup podobný souborovému systému: Uživatelé mohou vypisovat, číst a zapisovat soubory do svazků Unity Catalog pomocí známých operací souborového systému, podobně jako při práci se soubory na lokálním disku.

Díky automatickým zálohám na pozadí mohou uživatelé očekávat konzistentní přístup k datům svazků Unity Catalog:

• Zápisy: AI Runtime automaticky potvrzuje zápisy, takže změny jsou viditelné pro ostatní aplikace a úlohy, které přistupují ke stejnému svazku Unity Catalog.
• Čte: AI Runtime automaticky zachytává změny hlasitosti bez nutnosti explicitního obnovování nebo synchronizace od uživatele.

Optimalizace výkonu hlasitosti

Jak již bylo zmíněno, svazky Unity Catalog jsou kryty distribuovaným úložištěm a optimalizovány pro velké soubory se sekvenčními čteními a zápisy.

Několik tipů vám může pomoci dosáhnout co nejlepšího výkonu z AI Runtime:

• Spojte data do větších souborů: Pokud je to možné, agregujte data do menšího počtu větších souborů, přibližně 1 až 10 GiB na soubor. To umožňuje AI Runtime agresivně přednačítat data a automaticky dosáhnout téměř optimálního výkonu sekvenčního čtení.

• Pro malé soubory používejte lokální disk: Pokud vaše zátěž zahrnuje mnoho malých souborů, zvažte jejich zkopírování na lokální disk pomocí paralelních kopií před jejich zpracováním. To může snížit zátěž opakovaného přístupu k mnoha malým souborům přes svazek.

  # 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
  

• Můžete použít UCVolumeDataset pro své úlohy strojového učení. Zahrnuje výše popsané optimalizace pro efektivní přístup k datům a načítání z Unity Catalog svazků. Další podrobnosti najdete v následujících částech.

Načíst nestrukturovaná data pomocí UCVolumeDataset

Pro nestrukturovaná data, jako jsou obrázky, audio a textové soubory uložené ve svazcích Unity Catalog, použijte UCVolumeDataset z modulu databricks.air.data . UCVolumeDataset je PyTorch IterableDataset, který při prvním přístupu zkopíruje každý soubor ze svazku do rychlé místní mezipaměti a vrací cestu k místnímu souboru uloženému v mezipaměti. Řeší problémy s výkonem a distribucí, které byste jinak implementovali ručně:

• Místní ukládání do mezipaměti Soubory se při prvním přístupu zkopírují z přípojného bodu FUSE do místního adresáře mezipaměti a poté se načítají z mezipaměti, takže při trénování přes více epoch nedochází k opětovnému čtení svazku.
• Automatické dělení Při inicializaci torch.distributed se soubory rozdělí mezi ranky a poté dále mezi pracovníky DataLoader, takže každá dvojice (rank, worker) obdrží nepřekrývající se část bez nutnosti další konfigurace.

Poznámka:

UCVolumeDataset a databricks.air.data.DataLoader pocházejí z balíčku databricks-sdk-air. Nainstalujte ho s doplňkem data, který také nainstaluje kompatibilní torch:

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

UCVolumeDataset poskytuje nezpracované místní cesty k souborům. Chcete-li tyto soubory dekódovat na tensory, zabalte je do druhého IterableDataset, které zpracovává proud cest a používá vaši logiku parsování. Díky tomu jsou vstupně-výstupní operace a analýzy oddělené.

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)

Obal přijímá již uložené lokální cesty, takže krok parsování nikdy nezasahuje do objemu. Můžete řetězit další obaly pro rozšíření, tokenizaci nebo filtrování.

Pro optimální výkon spárujte UCVolumeDataset s databricks.air.data.DataLoader namísto výchozího PyTorch DataLoader. Je naladěný na AI Runtime I/O a současně načítá a ukládá soubory do cache, zatímco GPU počítá.

Modely kontrolních bodů ve svazcích

Chcete-li ukládat kontrolní body modelu, abyste mohli pokračovat v trénování od posledního snímku nebo se zotavit po pádu, můžete použít svazky Unity Catalog stejně jako místní souborový systém.

Databricks doporučuje použití distribuovaného checkpointu (DCP) pro lepší výkon jak na single-GPU, tak na multi-GPU zátěžích. Přečtěte si Rychlé trénování PyTorch odolné vůči chybám na AI Runtime na technickém blogu 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"],
)

Monolitický přístup torch.save také funguje.

• Pro ukládání kontrolních bodů modelu s jedním 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"])
  

• Pro distribuované školení zahájené přes 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"])
  

Načítání tabulkových dat

Použijte Spark Connect k načtení tabulkových dat strojového učení z tabulek Delta.

Pro trénování s jedním uzlem můžete datové rámce Apache Sparku převést na datové rámce pandas pomocí metody toPandas() a pak volitelně převést na formát NumPy pomocí metody to_numpy().

Poznámka:

Spark Connect odkládá analýzu a řešení názvů na dobu provádění, což může změnit chování vašeho kódu. Viz Porovnání Spark Connect s Klasickým Sparkem.

Spark Connect podporuje většinu rozhraní API PySpark, včetně Spark SQL, rozhraní Pandas API ve Sparku, strukturovaném streamování a knihovně MLlib (založené na datovém rámci). Nejnovější podporovaná rozhraní API najdete v referenční dokumentaci k rozhraní API PySpark .

Další omezení najdete v tématu Omezení výpočetních prostředků bez serveru.

Načítání velkých Delta tabulek pomocí svazků Unity Catalog

U velkých tabulek Delta, které jsou příliš velké na převod s toPandas(), exportujte data do svazku katalogu Unity a načtěte je přímo pomocí PyTorch nebo 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")

Tento přístup eliminuje režijní zátěž softwaru Spark během trénování a dobře funguje jak pro pracovní postupy s jednou GPU, tak pro distribuované tréninkové workflow.