Verbessern Sie die Trainingsleistung und Resilienz auf AI Runtime

Wenn ein Job auf mehr GPUs skaliert, steigt die Wahrscheinlichkeit von Hardware- und Softwarefehlern. Diese Seite behandelt Strategien, um Ihre Trainingsläufe schneller und fehlertoleranter zu gestalten:

• Lade Daten effizient, damit GPUs nicht im Leerlauf sind.
• Modell- und Optimierungszustände effizient als Checkpoints in Unity Catalog Volumes speichern.
• Erholen Sie sich automatisch nach einer Unterbrechung.
• Setzen Sie einen Checkpoint in der Datenpipeline, damit ein fortgesetzter Lauf das Training mit den korrekten Daten fortsetzt..

Mit diesen Mustern ist das Checkpointen Ihres Modells kostengünstig, sodass Sie häufig Checkpoints durchführen, günstig weitermachen und die effektive Berechnung Ihrer GPUs verbessern können.

Daten effizient laden, um die Leerlaufzeit der GPU zu minimieren

Ein Trainingsschritt sollte die GPU-Berechnung mit der Datenvorbereitung für den nächsten Schritt überschneiden. Bei AI Runtime läuft der gesamte Datenzugriff über den Unity Catalog. Für dateibasierte Datensätze in Unity Catalog-Volumes verwenden Sie serverless_gpu.data.UCVolumeDataset, das jede Datei beim ersten Zugriff vom FUSE-Mount in einen schnellen lokalen Cache kopiert und den zwischengespeicherten lokalen Pfad zurückgibt.

Kombinieren Sie es mit serverless_gpu.data.DataLoader, einer Drop-in-Unterklasse des PyTorch DataLoader , die für serverlose GPU-I/O optimiert ist und Dateien gleichzeitig abruft und cachet, während die GPU berechnet.

import serverless_gpu.data

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

loader = serverless_gpu.data.DataLoader(
    dataset,
    batch_size=64,
)

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

Dekodieren Sie Dateien, indem Sie serverless_gpu.data.UCVolumeDataset in ein zweites IterableDataset einbetten, das den Pfadstrom verbraucht. Der Wrapper erhält bereits zwischengespeicherte lokale Pfade, sodass das Parsen nie auf den FUSE-Einhängepunkt zugreift:

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

Zwei Anforderungen beim Skalieren:

• Verwenden Sie immer serverless_gpu.data.DataLoader für das Multi-Epoch-Training. Es erzwingt persistent_workers=True wenn num_workers > 0, sodass der Cache-Evition-Tracker jedes Arbeiters im Speicher über Epochen hinweg überlebt. Das Stock-PyTorch DataLoader erzeugt standardmäßig bei jeder Epoche neue Worker-Knoten, wodurch das gemeinsam genutzte Cache-Verzeichnis so lange ungenutzt bleibt, bis es voll ist.
• Alle Ränge müssen die gleiche num_workers bestehen. serverless_gpu.data.UCVolumeDataset Partitioniert Dateien mithilfe eines globalen Schrittes über Slots hinweg world_size × num_workers . Nicht übereinstimmende Werte führen dazu, dass Dateien zwischen den Rängen dupliziert oder übersprungen werden.

Wenn torch.distributed initialisiert wird, liest serverless_gpu.data.UCVolumeDataset den Rang zum Iterationszeitpunkt aus und partitioniert Dateien automatisch auf die Ränge, sodass Sie für dateibasierte Volumendaten kein DistributedSampler benötigen.

Checkpoint mit verteiltem Checkpoint (DCP)

Verwenden Sie PyTorch Distributed Checkpoint (DCP) anstelle von torch.save. Jeder Rang schreibt parallel seinen eigenen Shard in ein Checkpoint-Verzeichnis, wobei die gesamte aggregierte I/O-Bandbreite ausgeschöpft und ein sprunghafter Speicheranstieg beim Sammeln des gesamten Zustands auf einem Rang vermieden wird. DCP speichert außerdem globale Tensor-Metadaten, sodass ein Checkpoint, der auf einer bestimmten Anzahl von GPUs gespeichert ist, auf einer anderen Zahl fortgesetzt werden kann.

In AI Runtime sind serverless_gpu.data.UCVolumeWriter und serverless_gpu.data.UCVolumeReader die DCP-Speicher-Backends. Sie führen sämtliche E/A über ein schnelles lokales Verzeichnis aus (/tmp, auf AIR-GPU-Knoten durch NVMe unterstützt) und laden Daten in ein Unity Catalog-Volume hoch bzw. daraus herunter, was schneller ist, als Shards direkt auf den FUSE-Einhängepunkt zu schreiben.

import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict
import serverless_gpu.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=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"],
)

DCP lohnt sich auch für reines datenparalleles (DDP)-Training, bei dem die Gewichte über die Ränge hinweg repliziert werden. DCP schreibt eine einzelne deduplizierte Kopie der replizierten Gewichte, während es weiterhin den einzigartigen Zustand jedes Ranges erfasst (Datenposition und RNG-Zustand, siehe unten), und es ist dieselbe API, die Sie benötigen, wenn Sie später zu FSDP oder Tensorparallelität wechseln.

Asynchron speichern

Ein synchroner Speichervorgang blockiert das Training, bis die Bytes dauerhaft auf dem Datenträger gespeichert sind. Bei einem großen Checkpoint ist das die GPU-Leerlaufzeit. dcp.async_save Kopiert den Zustand in einen Staging-Puffer (schnell) und lädt dann im Hintergrund hoch, während das Training weiterläuft. Da jeder Checkpoint so gut wie keine GPU-Zeit kostet, kannst du es dir leisten, deutlich häufiger Checkpoints zu erstellen, wodurch der Arbeitsverlust nach einer Unterbrechung begrenzt wird.

Asynchrone Speichervorgänge erfordern ein CPU-Backend für die Prozessgruppe, daher initialisieren Sie sie mit sowohl gloo als auch nccl:

import torch.distributed as dist
import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import get_state_dict
import serverless_gpu.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 = serverless_gpu.data.UCVolumeWriter(f"/Volumes/my-catalog/my-schema/my-volume/checkpoints/step_{step}")
    checkpoint_future = dcp.async_save(state_dict, storage_writer=writer)

Automatisch vom letzten gültigen Prüfpunkt wiederherstellen

Ein Lauf kann während des Speichervorgangs unterbrochen werden, sodass ein unvollständiges Checkpoint-Verzeichnis zurückbleibt. serverless_gpu.data.UCVolumeWriter schreibt die .metadata-Datei erst dann in das Volume, nachdem die Datendateien der Shards vollständig hochgeladen wurden, sodass das Vorhandensein von .metadata ein zuverlässiges Zeichen dafür ist, dass ein Speichervorgang abgeschlossen wurde. Nutze es, um beim Neustart den aktuellsten gültigen Kontrollpunkt auszuwählen.

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

Eine robuste Trainingsschleife wählt den neuesten gültigen Checkpoint aus, wird daraus wiederhergestellt und erstellt häufig Checkpoints. Das Checkpoint-Intervall begrenzt den Arbeitsverlust nach einer Unterbrechung, sodass häufige, kostengünstige asynchrone Sicherungen den Neuberechnungsaufwand gering halten:

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=serverless_gpu.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

Diese Schleife speichert den Modell- und Optimiererzustand als Checkpoint. Die Position der Datenpipeline wird noch nicht wiederhergestellt, die im nächsten Abschnitt behandelt wird.

Checkpoint für die Datenpipeline erstellen

Ein Modell-Checkpoint erfasst den Zustand des Modells und des Optimierers, aber nicht die Position Ihrer Datenpipeline im Datensatz. Angenommen, du stellst das Modell bei Schritt 1.900 wieder her, aber der Dataloader startet vom Anfang des Datensatzes neu. Die fortgesetzte Ausführung trainiert die Daten anhand der in dieser Epoche bereits gesehenen Beispiele neu und überspringt Beispiele in der Nähe des Unterbrechungspunkts, wodurch die Datenverteilung stillschweigend und ohne Fehler verzerrt wird.

Um mit den korrekten Daten fortzufahren, verfolgen Sie Ihre Position im Datensatz als Teil Ihres eigenen Trainingsstatus und stellen Sie diese beim Fortfahren wieder her. Es gibt vier Punkte zu beachten:

Einen Sample- oder Shard-Offset nachverfolgen

Notieren Sie, wie weit Sie in der Epoche fortgeschritten sind, indem Sie einen globalen Sample-Index, Batch-Count oder Liste der verbrauchten Shard-IDs im Checkpoint-Status-Dict verwenden, und springen Sie zu dieser Position vor, wenn Sie fortsetzen. Dadurch bleibt die Datenposition unter deiner expliziten Kontrolle, anstatt dich auf einen Dataloader zu verlassen, der den internen Zustand serialisiert.

# 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
}

Für einen Datensatz im Map-Stil mit einem deterministischen Sampler überspringen Sie beim Wiederaufnehmen die in dieser Epoche bereits verarbeiteten Batches. Da die Abtastreihenfolge für ein gegebenes (seed, epoch) deterministisch ist (siehe Die Pipeline deterministisch gestalten), wird durch Vorspulen die exakte Position wiederhergestellt:

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

Für einen in Shards aufgeteilten Streaming-Datensatz erfassen Sie stattdessen, welche Shards bereits abgeschlossen sind, und übergeben Sie beim Fortsetzen des Laufs nur die noch verbleibenden Shards. Dadurch wird vermieden, sämtliche Batches einer ganzen Epoche erneut abzuspielen, nur um die Unterbrechungsstelle zu erreichen:

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

Checkpoint den internen Zustand eines benutzerdefinierten Datensatzes

Wenn du deinen eigenen Datensatz schreibst, gib ihm Methoden, um seine eigene Position zu serialisieren und wiederherzustellen, und integriere diesen Zustand in den Checkpoint. Dadurch bleibt die Logik für die Fortsetzung neben der Iterationslogik, und der Datensatz weiß genau, was er für den Schnellvorlauf benötigt (aktueller Shard, Offset darin, Inhalt des Shuffle-Puffers usw.), anstatt dass die Trainingsschleife dies anhand eines externen Zählers rekonstruiert.

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

Neu starten ab einer Epochengrenze

Wenn das Vorspringen unpraktisch ist, kontrollieren Sie nur an Epochengrenzen und setzen Sie zu Beginn der nächsten Epoche weiter. Die Ausführung verarbeitet dann jedes Beispiel über die Unterbrechung hinweg genau einmal pro Epoche, wobei pro Ausfall der Fortschritt von bis zu einer Epoche verloren gehen kann. Dies ist am einfachsten, wenn die Epochen im Verhältnis zur Ausfallrate kurz sind.

Machen Sie die Pipeline deterministisch

Beide Strategien setzen nur dann mit den korrekten Daten auf, wenn die Pipeline aus dem gespeicherten Zustand reproduzierbar ist. Mischen und Anreichern greifen auf Zufallszahlengeneratoren zurück, daher sollten diese initialisiert und ihr Zustand im Kontrollpunkt gespeichert werden. Andernfalls stimmt die Misch- und Augmentierungsreihenfolge nach dem Neustart nicht mit der Reihenfolge vor der Unterbrechung überein, und ein Skip-Ahead-Offset verweist auf die falschen Samples.

Um mitten im Verlauf einer Pipeline fortzusetzen und nicht nur an einer Epochen­grenze, müssen die RNGs, die das Mischen und die Augmentierung steuern, selbst checkpointfähig sein. Allein das Seeding reproduziert die Sequenz von Anfang an, aber nicht den Punkt, an dem du unterbrochen wurdest. Nutze RNG-Objekte, deren inneren Zustand du serialisieren kannst, und checkpoint diesen Zustand, damit jeder RNG genau dort weiterläuft, wo er aufgehört hat. Die Verwendung eines globalen Zufallszahlengenerators, der nur zu Beginn einer Epoche neu initialisiert wird, führt dazu, dass dieselben Ziehungen vom Anfang der Epoche wiederholt werden, was nicht mehr mit einem in der Mitte der Epoche vorgenommenen Zeitsprung übereinstimmt.

Alle Zufallsquellen initialisieren:

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)

Speichern und stellen Sie den RNG-Status zusammen mit dem Modell wieder her, sodass Augmentierungs- und Mischsequenzen nahtlos ablaufen:

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

Wenn Sie einen DistributedSampler anstelle eines zustandsbehafteten Datenladers verwenden, rufen Sie sampler.set_epoch(epoch) zu Beginn jeder Epoche auf. Da das Mischen eine deterministische Funktion von (seed, epoch) ist, ergibt die Wiederherstellung des Epochenzählers erneut genau dieselbe Permutation:

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

Verwandte Seiten

• Daten auf der KI-Laufzeit laden: Laden von tabellarischen und unstrukturierten Daten über den Unity-Katalog.
• Experimentverfolgung und Beobachtbarkeit: MLflow-Experimentverfolgung, Logging und GPU-Ressourcenüberwachung.
• Verteiltes Training in Notebooks: der @distributed-Decorator und Multi-GPU-Training.