AI ランタイムでのトレーニングのパフォーマンスと回復性を向上させる

備考

プレビュー

この機能は パブリック プレビュー段階です。

ジョブがより多くの GPU にスケールするにつれて、ハードウェアおよびソフトウェアの障害が発生する確率が高まります。このページでは、トレーニングランをより高速かつフォールトトレラントにするための戦略について説明します:

• GPUがアイドル状態にならないようにデータを効率的に読み込む。
• モデルとオプティマイザーの状態を Unity Catalog ボリュームに効率的にチェックポイントする。
• 中断後に自動的に回復する。
• データパイプラインをチェックポイントして、再開されたランが正しいデータでトレーニングすることを継続できるようにします。

これらのパターンを使用すると、モデルのチェックポイント処理は低コストになるため、頻繁にチェックポイントを作成し、安価に再開し、GPUの有効なコンピュートを向上させることができます。

注記

serverless_gpu.data.UCVolumeDataset、serverless_gpu.data.DataLoader、serverless_gpu.data.UCVolumeWriter、および serverless_gpu.data.UCVolumeReader には GPU 環境 5 以上（Serverless GPU Python API 0.5.16 以上）が必要です。

GPU のアイドル時間を最小限に抑えるためにデータを効率的に読み込む

トレーニングステップでは、GPU コンピュートと次のステップのためのデータ準備をオーバーラップさせる必要があります。AI ランタイム では、すべてのデータアクセスは Unity Catalog を経由します。Unity Catalog ボリューム内のファイルベースのデータセットには serverless_gpu.data.UCVolumeDataset を使用します。これは、初回アクセス時に FUSE マウントから高速なローカルキャッシュに各ファイルをコピーし、キャッシュされたローカルパスを生成します。

GPU コンピュート中にファイルを並行してフェッチおよびキャッシュするように調整された、Serverless GPU I/O 用の PyTorch DataLoader のドロップインサブクラスである serverless_gpu.data.DataLoader と組み合わせて使用してください。

Python

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

警告

serverless_gpu.data.UCVolumeDataset によって生成されるパスは エフェメラル （一時的）です。空きディスク容量がthreshold（defaultはキャッシュファイルシステムの10％、環境変数 SGC_FSLAYER_MIN_FREE_DISK_BYTES で上書き可能）を下回ると、キャッシュは最も最近downloadされていないファイルから削除するため、次のアイテムを取得するとすぐにパスが削除される可能性があります。同じループの反復内で、それを開く、デコードする、またはコピーします。後で再度開くために、返されたパスをリストや辞書に保存しないでください。

パスストリームを消費する 2 番目の IterableDataset で serverless_gpu.data.UCVolumeDataset をラップして、ファイルをデコードします。ラッパーは既にキャッシュされたローカルパスを受け取るため、解析が FUSE マウントに触れることはありません：

Python

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)

スケールアウト時の2つの要件：

• マルチエポックトレーニングには常に serverless_gpu.data.DataLoader を使用してください。 num_workers > 0 の際に persistent_workers=True を強制するため、各ワーカーのメモリ内キャッシュ削除トラッカーはエポックをまたいで保持されます。標準の PyTorch DataLoader は default でエポックごとにワーカーを再フォークするため、共有キャッシュディレクトリが一杯になるまでリークが発生します。
• すべてのランクは同じnum_workersを渡す必要があります。 serverless_gpu.data.UCVolumeDatasetは、world_size × num_workersスロットにわたるグローバルなストライドを使用してファイルをパーティション分割します。値が一致しないと、ファイルがランク間で重複またはスキップされます。

torch.distributedが初期化されると、serverless_gpu.data.UCVolumeDatasetは反復時にランクを読み取り、ファイルをランク間で自動的にパーティション分割するため、ファイルベースのボリューム データにDistributedSamplerは不要です。

分散チェックポイント (DCP) によるチェックポイント

PyTorchのDistributed Checkpoint（DCP）をtorch.saveではなく使用します。各ランクは、それぞれのシャードをチェックポイント ディレクトリに並列で書き込み、集約I/O帯域幅全体を使用し、すべての状態を1つのランクに集約する際のメモリスパイクを回避します。DCPはグローバルなテンソル メタデータも保存するため、あるGPU数で保存されたチェックポイントを別のGPU数で再開できます。

AI ランタイムでは、serverless_gpu.data.UCVolumeWriterとserverless_gpu.data.UCVolumeReaderがDCPのストレージバックエンドです。これらはすべてのI/Oを高速なローカルディレクトリ（/tmp、AIR GPUノード上のNVMeバックアップ）でステージングし、Unity Catalogボリュームに対してuploadまたはdownloadを行います。これは、FUSEマウントに直接シャードを書き込むよりも高速です。

Python

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 は、重みがランク間で複製される純粋なデータ並列（DDP）トレーニングであっても使用する価値があります。DCP は、各ランク固有の状態（データ位置と RNG 状態。後述）を保持しつつ、複製された重みの重複排除されたコピーを 1 つ書き込みます。これは、後に FSDP やテンソル並列処理に移行する場合に必要な API と同じです。

非同期で保存する

同期保存は、バイトがボリューム内に永続化されるまでトレーニングをブロックします。大規模なチェックポイントの場合、それは GPU のアイドル時間となります。dcp.async_save は状態をステージングバッファにコピーし（高速）、トレーニングの継続中にバックグラウンドでuploadします。各チェックポイントは GPU 時間をほとんど消費しないため、より頻繁にチェックポイントを作成できます。これにより、中断後の損失作業を制限できます。

非同期保存にはプロセスグループ上のCPUバックエンドが必要なため、glooとncclの両方で初期化してください：

Python

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)

最新の有効なチェックポイントから自動的に回復する

ランは保存の途中で中断される可能性があり、その場合、チェックポイントディレクトリが不完全な状態で残ります。serverless_gpu.data.UCVolumeWriter は、シャードデータファイルの upload が完了した後にのみ .metadata ファイルをボリュームに公開するため、.metadata の存在は保存が完了したことを示す信頼できるシグナルとなります。再起動時に最新の有効なチェックポイントを選択するために使用します。

Python

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

回復力のあるトレーニングループは、最新の有効なチェックポイントを選択してそこから復元し、頻繁にチェックポイントを作成します。チェックポイントの間隔は中断後の損失作業を制限するため、頻繁で低コストな非同期保存により再計算を最小限に抑えられます：

Python

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

このループは、モデルとオプティマイザーの状態をチェックポイントします。データパイプラインの位置はまだ復元されません。これについては次のセクションで説明します。

データパイプラインをチェックポイント化する

モデルのチェックポイントはモデルとオプティマイザーの状態をキャプチャしますが、データセット内のデータパイプラインの位置はキャプチャしません。ステップ1,900でモデルを復元したと仮定しますが、データローダーはデータセットの最初から再開されます。再開されたランは、このエポックですでに確認済みの例で再トレーニングを行い、中断ポイント付近の例をスキップするため、エラーなしでデータ分布に暗黙的なバイアスがかかります。

正しいデータで再開するには、データセット内の位置を独自のトレーニング状態の一部として追跡し、再開時にそれを復元します。考慮すべき点は 4 つあります：

サンプルまたはシャードのオフセットを追跡する

チェックポイント状態辞書内のグローバルサンプルインデックス、バッチカウント、または消費されたシャードIDのリストを使用してエポックの進行状況を記録し、再開時にその位置までスキップします。これにより、データローダーが内部状態をシリアル化することに依存するのではなく、データ位置を明示的に制御下に置くことができます。

Python

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

確定的サンプラーを使用するマップスタイルのデータセットの場合、再開時にこのエポックで既に消費されたバッチをスキップします。サンプラーの順序は特定の (seed, epoch) に対して確定的であるため（パイプラインを確定的（deterministic）にするを参照）、早送りによって正確な位置が再現されます：

Python

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

シャード化されたストリーミングデータセットの場合は、代わりに完了したシャードのセットを追跡し、再開されたランには残りのシャードのみを渡します。これにより、中断ポイントに到達するためだけにエポック全体のバッチを再実行することを回避できます：

Python

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

カスタムデータセットの内部状態をチェックポイントする

独自のデータセットを書く場合は、その位置をシリアル化および復元するメソッドを持たせ、その状態をチェックポイントに取り込みます。これにより、再開ロジックを反復ロジックの近くに保ち、データセットは高速に進めるために必要なもの（現在のシャード、その内部のオフセット、シャッフルバッファーの内容など）を外部カウンターからトレーニング ループが再構築するのではなく、正確に把握できます。

Python

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

Python

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

エポック境界から再開

スキップアヘッドが実用的でない場合は、エポックの境界でのみチェックポイントを作成し、次のエポックの起動時に再開します。その場合、ランは中断を挟んでエポックごとに各例を正確に1回ずつ処理しますが、障害ごとに最大1エポック分の進捗が失われる可能性があります。これは、失敗率に対してエポックが短い場合に最も簡単です。

パイプラインを決定的にする

どちらの戦略も、保存された状態からパイプラインを再現できる場合にのみ、正しいデータで再開されます。シャッフルと拡張は RNG から抽出されるため、シードを設定し、その状態をチェックポイントに保持してください。そうでない場合、再起動後のシャッフルと拡張の順序が中断前の順序と一致しなくなり、スキップアヘッドオフセットが誤ったサンプルを指すことになります。

エポック境界だけでなくパイプラインの途中から再開するには、シャッフルと拡張を駆動する RNG 自体がチェックポイント可能である必要があります。シード設定だけでは起動からのシーケンスは再現されますが、中断した時点からの再現はできません。内部状態をシリアル化できる RNG オブジェクトを使用し、その状態をチェックポイントすることで、各 RNG が中断した正確な位置から継続できるようにします。エポック起動時にのみ再シードされるグローバル RNG に依存すると、エポックの最初から同じドローが再生されてしまい、エポック途中のスキップアヘッドオフセットと一致しなくなります。

すべてのランダム性のソースをシードします:

Python

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)

モデルと一緒に RNG 状態を保存および復元し、拡張およびシャッフルシーケンスがシームレスに継続されるようにします：

Python

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

ステートフルなデータローダーの代わりに DistributedSampler を使用する場合は、各エポックの起動時に sampler.set_epoch(epoch) を呼び出してください。シャッフルは (seed, epoch) の決定論的な関数であるため、エポックカウンターを復元することで正確な順列が再現されます：

Python

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

注記

データパイプラインの正確性を確保するには、データの順序と拡張ストリームが再現可能である必要があり、これは上記のシード設定と RNG チェックポイントによって提供されます。通常、ビット単位で同一のフォワードパスは必要ありません。torch.use_deterministic_algorithms(True) は決定論的なカーネルを強制しますが、throughputが低下する可能性があり、すべての操作をカバーしているわけではありません。

関連ページ

• AI ランタイム でデータを読み込む: Unity Catalog を介した表形式データおよび非構造化データの読み込み。
• エクスペリメントのトラッキングと可観測性：MLflow エクスペリメントのトラッキング、Logs の表示、GPU リソースのモニタリング。
• ノートブックでの分散トレーニング: @distributed デコレーターとマルチ GPU トレーニング。