Migliora le prestazioni e la resilienza dell'allenamento su Runtime AI

Man mano che un lavoro si espande verso più GPU, aumenta la probabilità di guasti hardware e software. Questa pagina tratta strategie per rendere le tue corse di allenamento più veloci e più resistenti ai guasti:

• Carica i dati in modo efficiente così le GPU non restano inattive.
• Modello di checkpoint e stato dell'ottimizzatore in modo efficiente verso volumi del Catalogo Unity.
• Recupero automatico dopo un'interruzione.
• Salva lo stato della pipeline dei dati in modo che, quando un’esecuzione viene ripresa, l’addestramento continui sui dati corretti.

Con questi schemi, eseguire il checkpoint del modello costa poco, quindi puoi eseguire checkpoint frequenti, riprendere l’elaborazione con costi contenuti e migliorare l’efficienza di calcolo delle tue GPU.

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

Carica i dati in modo efficiente per minimizzare il tempo inattivo della GPU

Un passaggio di addestramento dovrebbe sovrapporsi al calcolo GPU con la preparazione dei dati per il passaggio successivo. Su AI Runtime, tutto l'accesso ai dati passa attraverso Unity Catalog. Per i dataset basati su file nei volumi del Catalogo Unity, si usa databricks.air.data.UCVolumeDataset, che copia ogni file dal supporto FUSE a una cache locale veloce al primo accesso e fornisce il percorso locale memorizzato.

Abbinalo a databricks.air.data.DataLoader, una sottoclasse drop-in del PyTorch DataLoader ottimizzata per AI Runtime I/O, che recupera e memorizza file contemporaneamente mentre la GPU calcola.

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

import databricks.air.data

dataset = databricks.air.data.UCVolumeDataset("/Volumes/my-catalog/my-schema/my-volume/data")

loader = databricks.air.data.DataLoader(
    dataset,
    batch_size=64,
)

for batch in loader:
    local_paths = batch  # open these immediately; see the caching note below
    ...

Decodifica i file racchiudendo databricks.air.data.UCVolumeDataset in un secondo IterableDataset che consuma il flusso del percorso. Il wrapper riceve percorsi locali già memorizzati nella cache, quindi il parsing non tocca mai il montaggio FUSE:

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: databricks.air.data.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 = databricks.air.data.UCVolumeDataset("/Volumes/my-catalog/my-schema/my-volume/images")
dataset = ImageDataset(path_dataset)
loader = databricks.air.data.DataLoader(dataset, batch_size=64)

Due requisiti per la scalabilità orizzontale:

• Usare sempre databricks.air.data.DataLoader per l'addestramento multi-epoca. Impone persistent_workers=True quando num_workers > 0, così il tracker di espulsione della cache in memoria di ogni worker persiste tra le epoche. La versione standard di PyTorch DataLoader ricrea i processi worker a ogni epoca per impostazione predefinita, il che provoca una perdita nella directory della cache condivisa finché questa non si riempie.
• Tutti i livelli devono superare lo stesso num_workers. databricks.air.data.UCVolumeDataset Partiziona i file usando un passo globale attraverso world_size × num_workers gli slot. Valori non corrispondenti causano duplicazioni o salti tra i file di fascia.

Quando torch.distributed viene inizializzato, databricks.air.data.UCVolumeDataset rileva il rank al momento dell'iterazione e partiziona automaticamente i file tra i rank, quindi non è necessario alcun DistributedSampler per i dati di volume basati su file.

Checkpoint con Checkpoint Distribuito (DCP)

Usa il PyTorch Distributed Checkpoint (DCP) invece di torch.save. Ogni rango scrive in parallelo il proprio shard in una directory di checkpoint, utilizzando l’intera larghezza di banda I/O complessiva ed evitando il picco di memoria causato dalla raccolta dell’intero stato in un unico rango. DCP archivia anche i metadati globali dei tensori, quindi un checkpoint salvato con un certo numero di GPU può essere ripristinato con un numero diverso di GPU.

Su AI Runtime, databricks.air.data.UCVolumeWriter e databricks.air.data.UCVolumeReader sono i backend di archiviazione DCP. Fanno transitare tutto l'I/O attraverso una directory locale veloce (/tmp, basata su NVMe sui nodi GPU AIR) e caricano su o scaricano da un volume di Unity Catalog, operazione più rapida che scrivere gli shard direttamente sul mount FUSE.

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

Il DCP vale la pena di essere utilizzato anche per l'addestramento puramente data-parallelo (DDP), dove i pesi vengono replicati tra i ranghi. DCP scrive una singola copia deduplicata dei pesi replicati pur catturando lo stato unico di ogni rango (posizione dati e stato RNG, trattato sotto), ed è la stessa API necessaria se successivamente passi a FSDP o parallelismo tensoriale.

Salva in modo asincrono

Un salvataggio sincrono blocca l'allenamento finché i byte non sono durevoli nel volume. Per un checkpoint grande, quello è il tempo di inattività della GPU. dcp.async_save copia lo stato in un buffer di staging (veloce) e poi lo carica in background mentre l'addestramento continua. Poiché ogni checkpoint richiede quasi zero tempo GPU, puoi permetterti di eseguirli molto più spesso, il che limita la quantità di lavoro persa dopo un'interruzione.

I salvataggi asincroni richiedono un backend CPU sul gruppo di processi, quindi inizializzalo con entrambi gloo e nccl:

import torch.distributed as dist
import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import get_state_dict
import databricks.air.data

dist.init_process_group(backend="cpu:gloo,cuda:nccl")

checkpoint_future = None

def save_async(step, model, optimizer):
    global checkpoint_future
    # Ensure the previous async save finished before starting a new one.
    if checkpoint_future is not None:
        checkpoint_future.result()

    model_sd, optim_sd = get_state_dict(model, optimizer)
    state_dict = {"model": model_sd, "optim": optim_sd, "step": step}
    writer = databricks.air.data.UCVolumeWriter(f"/Volumes/my-catalog/my-schema/my-volume/checkpoints/step_{step}")
    checkpoint_future = dcp.async_save(state_dict, storage_writer=writer)

Recupera automaticamente dall'ultimo checkpoint valido

Un'esecuzione può essere interrotta durante il salvataggio, lasciando una directory di checkpoint parziale. databricks.air.data.UCVolumeWriter pubblica il .metadata file nel volume solo dopo che i file di dati dello shard hanno terminato di caricarsi, quindi la presenza di .metadata è un segnale affidabile che un salvataggio è stato completato. Usalo per selezionare il checkpoint valido più recente al riavvio.

import os

def find_latest_valid(checkpoint_root):
    """Return the newest checkpoint directory that finished writing, or None."""
    candidates = sorted(
        (d for d in os.listdir(checkpoint_root) if d.startswith("step_")),
        key=lambda d: int(d.split("_")[1]),
        reverse=True,
    )
    for name in candidates:
        path = os.path.join(checkpoint_root, name)
        if os.path.exists(os.path.join(path, ".metadata")):  # save completed
            return path
    return None  # nothing valid; start fresh

Un ciclo di addestramento resiliente seleziona l'ultimo checkpoint valido, ripristina da esso e i checkpoint frequentemente. L'intervallo tra i checkpoint limita la quantità di lavoro persa dopo un'interruzione, quindi salvataggi asincroni frequenti ed economici mantengono ridotta la quantità di ricalcolo:

CHECKPOINT_EVERY = 100

latest = find_latest_valid("/Volumes/my-catalog/my-schema/my-volume/checkpoints")
start_step = 0
if latest is not None:
    model_sd, optim_sd = get_state_dict(model, optimizer)
    state = {"model": model_sd, "optim": optim_sd, "step": 0}
    dcp.load(state, storage_reader=databricks.air.data.UCVolumeReader(latest))
    set_state_dict(
        model,
        optimizer,
        model_state_dict=state["model"],
        optim_state_dict=state["optim"],
    )
    start_step = state["step"]

for step in range(start_step, total_steps):
    train_step(...)
    if step % CHECKPOINT_EVERY == 0:
        save_async(step, model, optimizer)  # inexpensive, so run it often

Questo modello di checkpoint a loop e stato dell'ottimizzatore. Non ripristina ancora la posizione della pipeline di dati, che viene trattata nella sezione successiva.

Crea un checkpoint della pipeline di dati

Un checkpoint del modello salva lo stato del modello e dell'ottimizzatore, ma non la posizione della tua pipeline di dati nel dataset. Supponiamo di ripristinare il modello al passo 1.900 ma il dataloader si riavvia dall'inizio del dataset. L’esecuzione ripresa viene nuovamente addestrata su esempi già visti in questa epoca di addestramento e salta gli esempi vicini al punto di interruzione, alterando silenziosamente la distribuzione dei dati senza generare alcun errore.

Per riprendere con i dati corretti, traccia la tua posizione nel dataset come parte del tuo stato di addestramento e ripristinala nel curriculum. Ci sono quattro punti da considerare:

Tenere traccia di un campione o di un offset di shard

Registra fino a che punto sei avanzato nell'epoca usando un indice globale di campioni, un conteggio dei lotti o una lista degli ID degli shard consumati nel codice di stato del checkpoint, e salta a quella posizione quando riprendi. Questo mantiene la posizione dei dati sotto il tuo controllo diretto ed esplicito, invece di affidarsi a un dataloader per serializzarne lo stato interno.

# Include the data position in the checkpoint state dict:
state_dict = {
    "model": model_sd,
    "optim": optim_sd,
    "step": step,
    "epoch": epoch,
    "samples_seen": samples_seen,  # your own counter, advanced each batch
}

Per un dataset in stile mappa con un campionatore deterministico, si saltano i lotti già consumati in questa epoca al momento della ripresa. Poiché l'ordine del campionatore è deterministico per un determinato (seed, epoch) (vedi Rendere deterministica la pipeline), l'avanzamento rapido riproduce esattamente la posizione:

resume_batch = state["samples_seen"] // batch_size

for epoch in range(start_epoch, num_epochs):
    sampler.set_epoch(epoch)
    for batch_idx, batch in enumerate(loader):
        # On the resumed epoch only, skip batches already processed.
        if epoch == start_epoch and batch_idx < resume_batch:
            continue
        train_step(batch)
        samples_seen += batch_size

Per un dataset sharded e in streaming, tieni invece traccia dell'insieme degli shard completati e fornisci all'esecuzione ripresa solo gli shard rimanenti. Questo evita di dover rieseguire tutti i batch di un'intera epoca solo per tornare al punto di interruzione:

# Filter the shard list down to work not yet done, then build the loader from it.
remaining = [s for s in all_shards if s not in state["completed_shards"]]
dataset = ShardDataset(remaining)

Salva lo stato interno di un dataset personalizzato

Se definisci un tuo dataset, forniscigli i metodi per serializzare e ripristinare la sua posizione e includi tale stato nel checkpoint. Questo mantiene la logica di ripresa accanto a quella di iterazione e il dataset sa esattamente cosa deve far avanzare rapidamente (lo shard corrente, l'offset al suo interno, il contenuto del buffer di mescolamento e così via), invece di doverla far ricostruire al ciclo di addestramento a partire da un contatore esterno.

from torch.utils.data import IterableDataset

class ResumableShardDataset(IterableDataset):
    """A streaming dataset that can checkpoint and restore its own position."""

    def __init__(self, shards):
        self._shards = shards
        self._shard_idx = 0      # position advanced during __iter__
        self._offset = 0

    def state_dict(self):
        return {"shard_idx": self._shard_idx, "offset": self._offset}

    def load_state_dict(self, state):
        self._shard_idx = state["shard_idx"]
        self._offset = state["offset"]

    def __iter__(self):
        for i in range(self._shard_idx, len(self._shards)):
            self._shard_idx = i
            for j, example in enumerate(self._read_shard(self._shards[i])):
                if j < self._offset:
                    continue  # skip examples already consumed from this shard
                self._offset = j + 1
                yield example
            self._offset = 0

# Save and restore the dataset position with the rest of the checkpoint state.
state_dict["dataset"] = dataset.state_dict()
# On resume:
dataset.load_state_dict(state["dataset"])

Riavviare da un limite di epoca

Se lo skip-ahead non è praticabile, crea un checkpoint solo alla fine di ogni epoca e riprendi all'inizio dell'epoca successiva. L’esecuzione vede quindi ogni esempio esattamente una volta per epoca, nonostante l’interruzione, al costo di perdere fino a un’epoca di avanzamento per ogni guasto. Questo è più semplice quando le epoche sono brevi rispetto al tasso di guasto.

Rendere la pipeline deterministica

Ciascuna delle due strategie riprende correttamente dai dati corretti solo se la pipeline è riproducibile a partire dallo stato salvato. Il rimescolamento e la data augmentation si basano su generatori di numeri casuali (RNG), quindi inizializzateli con un seed e includetene lo stato nel checkpoint. Altrimenti, dopo il riavvio l'ordine di mescolamento e di aumento dei dati non corrisponde a quello precedente all'interruzione e un offset di avanzamento punta ai campioni sbagliati.

Per riprendere a metà pipeline anziché solo al confine tra epoche, gli RNG usati per il rimescolamento e l’augmentazione devono anch’essi essere soggetti a checkpoint. Il seeding da solo riproduce la sequenza dall'inizio, ma non il punto in cui sei stato interrotto. Usa oggetti RNG il cui stato interno possa essere serializzato e salva un checkpoint di tale stato, in modo che ogni RNG riprenda esattamente dal punto in cui si era interrotto. Affidarsi a un RNG globale che viene reinizializzato solo all'inizio dell'epoca fa sì che vengano riprodotte le stesse estrazioni dall'inizio dell'epoca, che non corrispondono più a un offset di avanzamento a metà epoca.

Semina tutte le fonti di casualità:

import random
import numpy as np
import torch

def seed_everything(seed: int):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

Salva e ripristina lo stato RNG insieme al modello affinché le sequenze di aumento e mescolare continuino senza soluzione di continuità:

# Save
state_dict["rng"] = {
    "python": random.getstate(),
    "numpy": np.random.get_state(),
    "torch": torch.get_rng_state(),
    "cuda": torch.cuda.get_rng_state_all(),
}

# Load
rng = state["rng"]
random.setstate(rng["python"])
np.random.set_state(rng["numpy"])
torch.set_rng_state(rng["torch"])
torch.cuda.set_rng_state_all(rng["cuda"])

Se usi un DistributedSampler invece di un dataloader con stato, chiama sampler.set_epoch(epoch) all'inizio di ogni epoca. Poiché la mescolatura è una funzione deterministica di (seed, epoch), il ripristino del contatore delle epoche riproduce la permutazione esatta:

for epoch in range(start_epoch, num_epochs):
    sampler.set_epoch(epoch)  # deterministic reshuffle per epoch
    for batch in loader:
        ...

Pagine correlate

• Carica dati su AI Runtime: caricamento di dati tabulari e non strutturati tramite Unity Catalog.
• Tracciamento e osservabilità degli esperimenti: tracciamento degli esperimenti MLflow, visualizzazione dei log e monitoraggio delle risorse GPU.
• Addestramento distribuito nei notebook: il decoratore @distributed e l'addestramento multi-GPU.