Improve training performance and resiliency on AI Runtime

Important

This feature is in Public Preview.

As a job scales to more GPUs, the probability of hardware and software failure increases. This page covers strategies to make your training runs faster and more fault-tolerant:

With these patterns, checkpointing your model is inexpensive, so you can checkpoint frequently, resume cheaply, and improve the effective compute of your GPUs.

Note

serverless_gpu.data.UCVolumeDataset, serverless_gpu.data.DataLoader, serverless_gpu.data.UCVolumeWriter, and serverless_gpu.data.UCVolumeReader require GPU environment 5 or above (Serverless GPU Python API 0.5.16 or above).

Load data efficiently to minimize idle GPU time

A training step should overlap GPU compute with data preparation for the next step. On AI Runtime, all data access goes through Unity Catalog. For file-based datasets in Unity Catalog volumes, use serverless_gpu.data.UCVolumeDataset, which copies each file from the FUSE mount to a fast local cache on first access and yields the cached local path.

Pair it with serverless_gpu.data.DataLoader, a drop-in subclass of the PyTorch DataLoader tuned for serverless GPU I/O which fetches and caches files concurrently while the GPU computes.

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

Warning

The path yielded by serverless_gpu.data.UCVolumeDataset is ephemeral. The cache evicts least-recently-downloaded files once free disk drops below a threshold (default 10% of the cache filesystem, overridable with the SGC_FSLAYER_MIN_FREE_DISK_BYTES environment variable), so a path may be deleted as soon as you pull the next item. Open, decode, or copy it in the same loop iteration. Never store a returned path in a list or dict to reopen later.

Decode files by wrapping serverless_gpu.data.UCVolumeDataset in a second IterableDataset that consumes the path stream. The wrapper receives already-cached local paths, so parsing never touches the FUSE mount:

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)

Two requirements when scaling out:

  • Always use serverless_gpu.data.DataLoader for multi-epoch training. It forces persistent_workers=True when num_workers > 0, so each worker's in-memory cache-eviction tracker survives across epochs. The stock PyTorch DataLoader re-forks workers every epoch by default, which leaks the shared cache directory until it fills.
  • All ranks must pass the same num_workers. serverless_gpu.data.UCVolumeDataset partitions files using a global stride across world_size × num_workers slots. Mismatched values cause files to be duplicated or skipped across ranks.

When torch.distributed is initialized, serverless_gpu.data.UCVolumeDataset reads the rank at iteration time and partitions files across ranks automatically, so you do not need a DistributedSampler for file-based volume data.

Checkpoint with Distributed Checkpoint (DCP)

Use PyTorch Distributed Checkpoint (DCP) rather than torch.save. Every rank writes its own shard in parallel into a checkpoint directory, using the full aggregate I/O bandwidth and avoiding the memory spike of gathering all state to one rank. DCP also stores global tensor metadata, so a checkpoint saved on one number of GPUs can be resumed on a different number.

On AI Runtime, serverless_gpu.data.UCVolumeWriter and serverless_gpu.data.UCVolumeReader are the DCP storage backends. They stage all I/O through a fast local directory (/tmp, NVMe-backed on AIR GPU nodes) and upload to or download from a Unity Catalog volume, which is faster than writing shards directly to the FUSE mount.

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 is worth using even for pure data-parallel (DDP) training, where the weights are replicated across ranks. DCP writes a single deduplicated copy of the replicated weights while still capturing each rank's unique state (data position and RNG state, covered below), and it is the same API you need if you later move to FSDP or tensor parallelism.

Save asynchronously

A synchronous save blocks training until the bytes are durable in the volume. For a large checkpoint, that is GPU idle time. dcp.async_save copies state to a staging buffer (fast) and then uploads in the background while training continues. Because each checkpoint costs almost no GPU time, you can afford to checkpoint far more often, which is what bounds lost work after an interruption.

Asynchronous saves require a CPU backend on the process group, so initialize it with both gloo and 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)

Recover automatically from the most recent valid checkpoint

A run can be interrupted mid-save, leaving a partial checkpoint directory. serverless_gpu.data.UCVolumeWriter publishes the .metadata file to the volume only after the shard data files finish uploading, so the presence of .metadata is a reliable signal that a save completed. Use it to select the most recent valid checkpoint on restart.

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

A resilient training loop selects the latest valid checkpoint, restores from it, and checkpoints frequently. The checkpoint interval bounds lost work after an interruption, so frequent, cheap async saves keep recomputation small:

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

This loop checkpoints model and optimizer state. It does not yet restore the data pipeline position, which is covered in the next section.

Checkpoint the data pipeline

A model checkpoint captures model and optimizer state, but not the position of your data pipeline within the dataset. Suppose you restore the model at step 1,900 but the dataloader restarts from the beginning of the dataset. The resumed run re-trains on examples already seen this epoch and skips examples near the interruption point, silently biasing the data distribution with no error.

To resume on the correct data, track your position in the dataset as part of your own training state and restore it on resume. There are four points to consider:

Track a sample or shard offset

Record how far into the epoch you have progressed by using a global sample index, batch count, or list of consumed shard IDs in the checkpoint state dict, and skip ahead to that position when you resume. This keeps the data position under your explicit control rather than relying on a dataloader to serialize its internal state.

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

For a map-style dataset with a deterministic sampler, skip the batches already consumed in this epoch when resuming. Because the sampler order is deterministic for a given (seed, epoch) (see Make the pipeline deterministic), fast-forwarding reproduces the exact position:

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

For a sharded, streaming dataset, track the set of completed shards instead and hand the resumed run only the shards that remain. This avoids replaying an entire epoch's worth of batches just to reach the interruption point:

# 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 a custom dataset's internal state

If you write your own dataset, give it methods to serialize and restore its own position, and fold that state into the checkpoint. This keeps the resume logic next to the iteration logic and the dataset knows exactly what it needs to fast-forward (current shard, offset within it, shuffle buffer contents, and so on) rather than the training loop reconstructing it from an external counter.

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

Restart from an epoch boundary

If skip-ahead is impractical, checkpoint only at epoch boundaries and resume at the start of the next epoch. The run then sees every example exactly once per epoch across the interruption, at the cost of losing up to one epoch of progress per failure. This is simplest when epochs are short relative to the failure rate.

Make the pipeline deterministic

Either strategy only resumes on the correct data if the pipeline is reproducible from saved state. Shuffling and augmentation draw from RNGs, so seed them and carry their state in the checkpoint. Otherwise the shuffle and augmentation order after restart does not match the order before the interruption, and a skip-ahead offset points at the wrong samples.

To resume midway through a pipeline rather than only at an epoch boundary, the RNGs driving shuffling and augmentation must themselves be checkpointable. Seeding alone reproduces the sequence from the start, but not the point you were interrupted at. Use RNG objects whose internal state you can serialize, and checkpoint that state so each RNG continues from exactly where it left off. Relying on a global RNG that is only re-seeded at epoch start replays the same draws from the beginning of the epoch, which no longer matches a mid-epoch skip-ahead offset.

Seed all sources of randomness:

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)

Save and restore RNG state alongside the model so augmentation and shuffle sequences continue seamlessly:

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

If you use a DistributedSampler instead of a stateful dataloader, call sampler.set_epoch(epoch) at the start of each epoch. Because the shuffle is a deterministic function of (seed, epoch), restoring the epoch counter reproduces the exact permutation:

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

Note

For data pipeline correctness you need the data order and augmentation stream to be reproducible, which the seeding and RNG checkpointing above provide. You generally do not need bitwise-identical forward passes; torch.use_deterministic_algorithms(True) forces deterministic kernels but can reduce throughput and does not cover every operation.

  • Last updated on