FSDP2、GPU プロファイリング、モバイルコンピューティング

本番環境の AI エンジニアリングは、分散マルチ GPU 実行、ハードウェアレベルのプロファイリング、クロスプラットフォームのモバイル/エッジランタイムを橋渡しします。以下では、この 3 つの領域すべてを深く掘り下げます。

1. PyTorch FSDP2 によるマルチ GPU 分散トレーニング

PyTorch の最新の FSDP2（fully_shard）は、DTensor を使用してモデルのパラメータ、勾配、オプティマイザの状態をパラメータ単位で複数の GPU にシャーディングし、メモリの重複を排除しながら通信と計算をオーバーラップさせます。

以下は完全なマルチ GPU トレーニングスクリプトです。torchrun --nproc_per_node=4 train_fsdp.py を使って 4 つの GPU で起動します。

import os
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.distributed._composable.fsdp import fully_shard, MixedPrecisionPolicy

class TransformerBlock(nn.Module):
    def __init__(self, dim=2048, n_heads=16):
        super().__init__()
        self.norm1 = nn.RMSNorm(dim)
        self.attn = nn.MultiheadAttention(dim, n_heads, batch_first=True)
        self.norm2 = nn.RMSNorm(dim)
        self.ffn = nn.Sequential(
            nn.Linear(dim, dim * 4, bias=False),
            nn.SiLU(),
            nn.Linear(dim * 4, dim, bias=False),
        )

    def forward(self, x):
        h = self.norm1(x)
        attn_out, _ = self.attn(h, h, h, need_weights=False)
        x = x + attn_out
        return x + self.ffn(self.norm2(x))

class MiniLLM(nn.Module):
    def __init__(self, vocab_size=32000, dim=2048, n_layers=12):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, dim)
        self.layers = nn.ModuleList([TransformerBlock(dim) for _ in range(n_layers)])
        self.head = nn.Linear(dim, vocab_size, bias=False)

    def forward(self, idx):
        x = self.embed(idx)
        for layer in self.layers:
            x = layer(x)
        return self.head(x)

def main():
    # 1. Initialize NCCL process group
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)

    # 2. Build model on meta device or standard CUDA, then shard layer-by-layer
    model = MiniLLM().to(device)
    mp_policy = MixedPrecisionPolicy(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.float32,
    )

    # Shard each TransformerBlock independently to prefetch weights on the fly
    for layer in model.layers:
        fully_shard(layer, mp_policy=mp_policy)
    fully_shard(model, mp_policy=mp_policy)

    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

    # 3. Distributed Training Step (All-Gather -> Compute -> Reduce-Scatter)
    for step in range(5):
        tokens = torch.randint(0, 32000, (8, 1024), device=device)
        targets = torch.randint(0, 32000, (8, 1024), device=device)

        optimizer.zero_grad()
        logits = model(tokens)
        loss = nn.functional.cross_entropy(logits.view(-1, 32000), targets.view(-1))
        loss.backward()
        optimizer.step()

        if local_rank == 0:
            mem_gb = torch.cuda.max_memory_allocated(device) / 1e9
            print(f"Step {step} | Loss: {loss.item():.4f} | Peak VRAM/GPU: {mem_gb:.2f} GB")

    dist.destroy_process_group()

if __name__ == "__main__":
    main()

• 内部的な動作：各レイヤーのフォワードパスの前に、FSDP は NVLink 経由で非同期の NCCL All-Gather を発行し、そのブロック専用の完全な bfloat16 重みを再構築して計算を実行し、直後にシャードされていない重みを破棄します。loss.backward() の間、勾配は Reduce-Scatter によって GPU 間で平均化され、再び分割されます。

2. PyTorch モデルのプロファイリング：PyTorch Profiler と NVIDIA Nsight Compute（ncu）

GPU コードを最適化するには、カーネルがメモリ帯域幅バウンド（HBM の読み取り/書き込み待ちで停滞）なのか、コンピュートバウンド（Tensor Core の FLOPS が飽和）なのかを知る必要があります。

ステップ A：torch.profiler によるマクロレベルのトレース

PyTorch Profiler を使用して、CPU の起動ギャップ、遅いデータローダー、最も時間のかかる CUDA カーネルを特定します。

import torch
from torch.profiler import profile, ProfilerActivity

model = torch.nn.Linear(4096, 4096, device="cuda", dtype=torch.float16)
x = torch.randn(1024, 4096, device="cuda", dtype=torch.float16)

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    record_shapes=True,
    profile_memory=True,
) as prof:
    for _ in range(5):
        out = model(x)
        out.sum().backward()
    torch.cuda.synchronize()

print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
prof.export_chrome_trace("trace.json")  # Open in ui.perfetto.dev

ステップ B：NVIDIA Nsight Compute（ncu）によるミクロレベルの Roofline 分析

torch.profiler が遅いカーネルを特定したら、ターミナルから ncu を実行して、ストリーミングマルチプロセッサ（SM）上のハードウェアパフォーマンスカウンタを直接調べます：

ncu --set full --target-processes all -o kernel_report python my_script.py

kernel_report.ncu-rep 内で Speed of Light（SOL）と Roofline チャートを確認する際は、2 つの数値に注目します：

3. エッジ＆モバイル GPU 実行：Apple Metal と Vulkan Compute

ニューラルネットワークをスマートフォン、ノート PC、組み込みデバイスにデプロイするには、NVIDIA ドライバや大規模な電力予算なしで実行する必要があります。代わりに、統合メモリアーキテクチャ（UMA）を備えた Apple Metal と Vulkan Compute に依存します。

• 統合メモリアーキテクチャ（ゼロコピー実行）。

• CPU RAM から VRAM へ PCIe 経由でテンソルをコピーする必要があるディスクリート NVIDIA GPU とは異なり、Apple Silicon（M シリーズ / A シリーズ）とモバイル SoC（Qualcomm Snapdragon、ARM Mali）は、CPU、GPU、NPU の間で 1 つの物理 LPDDR5/LPDDR5X メモリプールを共有します。
• Apple MLX、llama.cpp、ExecuTorch などのフレームワークは、モデルの重みを共有アドレス空間（MTLResourceStorageModeShared）に直接マッピングし、ホストからデバイスへの転送レイテンシを完全に排除します。

• Apple Metal と Metal Performance Shaders（MPS）。

• macOS と iOS では、Metal が CUDA に取って代わります。開発者は Metal Shading Language（MSL、C++14 の派生言語）でコンピュートシェーダーを記述し、Apple GPU コアと Simdgroup Matrix 命令（NVIDIA Tensor Core に相当する Apple の機能）に threadgroup をディスパッチします。
• PyTorch では、device="mps" を設定すると、テンソル演算が Apple の MPSGraph コンパイラを通じてルーティングされ、演算が動的に融合され、GPU 実行ユニット間で作業がスケジュールされます。

• クロスプラットフォーム Vulkan Compute と SPIR-V。

• Android、Windows、Linux、および AMD/Intel/Qualcomm GPU 全体で、Vulkan Compute はユニバーサルでベンダーニュートラルな GPU バックエンドとして機能します。
• コンピュートカーネル（GLSL、HLSL で記述、または MLIR/IREE からコンパイル）は SPIR-V バイトコードにコンパイルされます。Cooperative Matrix 拡張（VK_KHR_cooperative_matrix）を使用すると、Vulkan カーネルは単一のポータブルバイナリで、Qualcomm Adreno、ARM Mali、AMD RDNA、Intel Arc GPU のハードウェア行列乗算ユニットを直接利用できます。