FSDP2, GPU-Profiling und mobiles Computing

Die produktive KI-Entwicklung verbindet verteilte Multi-GPU-Ausführung, Profiling auf Hardware-Ebene und plattformübergreifende Mobile-/Edge-Runtimes. Im Folgenden finden Sie einen tiefgehenden Einblick in alle drei Bereiche.

1. Verteiltes Multi-GPU-Training mit PyTorch FSDP2

PyTorchs modernes FSDP2 (fully_shard) shardet Modellparameter, Gradienten und Optimizer-Zustände parameterweise über GPUs hinweg mithilfe von DTensor, wodurch Speicherduplizierung entfällt und Kommunikation mit Berechnung überlappt wird.

Im Folgenden finden Sie ein vollständiges Multi-GPU-Trainingsskript. Starten Sie es auf 4 GPUs mit torchrun --nproc_per_node=4 train_fsdp.py.

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

• Was im Hintergrund passiert: Vor dem Forward-Pass jeder Schicht stößt FSDP ein asynchrones NCCL-All-Gather über NVLink an, um die vollständigen bfloat16-Gewichte nur für diesen Block zu rekonstruieren, führt die Berechnung aus und verwirft die ungeshardeten Gewichte sofort wieder. Während loss.backward() werden die Gradienten gemittelt und per Reduce-Scatter wieder auf die GPUs verteilt.

2. Profiling eines PyTorch-Modells: PyTorch Profiler & NVIDIA Nsight Compute (ncu)

Die Optimierung von GPU-Code erfordert zu wissen, ob ein Kernel speicherbandbreitenbegrenzt (blockiert durch Warten auf HBM-Lese-/Schreibzugriffe) oder rechenbegrenzt (Sättigung der Tensor-Core-FLOPS) ist.

Schritt A: Makro-Trace mit torch.profiler

Verwenden Sie den PyTorch Profiler, um CPU-Startlücken, langsame Data-Loader und die zeitaufwendigsten CUDA-Kernel zu erkennen.

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

Schritt B: Mikro-Roofline-Analyse mit NVIDIA Nsight Compute (ncu)

Sobald torch.profiler einen langsamen Kernel identifiziert hat, führen Sie ncu im Terminal aus, um die Hardware-Leistungszähler direkt auf den Streaming Multiprocessors (SMs) zu untersuchen:

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

Achten Sie bei der Untersuchung von Speed of Light (SOL) und Roofline-Diagramm in kernel_report.ncu-rep auf zwei Zahlen:

3. Edge- und mobile GPU-Ausführung: Apple Metal & Vulkan Compute

Der Einsatz neuronaler Netze auf Smartphones, Laptops und Embedded-Geräten erfordert den Betrieb ohne NVIDIA-Treiber und ohne massive Energiebudgets – stattdessen kommen Apple Metal und Vulkan Compute mit Unified Memory Architectures (UMA) zum Einsatz.

• Unified Memory Architecture (Zero-Copy-Ausführung).

• Anders als diskrete NVIDIA-GPUs, die das Kopieren von Tensoren über PCIe vom CPU-RAM in den VRAM erfordern, teilen sich Apple Silicon (M-Serie / A-Serie) und mobile SoCs (Qualcomm Snapdragon, ARM Mali) einen physischen LPDDR5/LPDDR5X-Speicherpool zwischen CPU, GPU und NPU.
• Frameworks wie Apple MLX, llama.cpp und ExecuTorch mappen Modellgewichte direkt in den gemeinsamen Adressraum (MTLResourceStorageModeShared) und eliminieren die Host-zu-Gerät-Transferlatenz vollständig.

• Apple Metal & Metal Performance Shaders (MPS).

• Unter macOS und iOS ersetzt Metal CUDA. Entwickler schreiben Compute-Shader in Metal Shading Language (MSL), einem C++14-Derivat, und verteilen Threadgroups auf Apple-GPU-Kerne und Simdgroup-Matrix-Instruktionen (Apples Äquivalent zu NVIDIA Tensor Cores).
• In PyTorch leitet die Einstellung device="mps" Tensoroperationen über Apples MPSGraph-Compiler, der Operationen dynamisch fusioniert und die Arbeit auf die GPU-Ausführungseinheiten verteilt.

• Plattformübergreifendes Vulkan Compute & SPIR-V.

• Auf Android, Windows, Linux und AMD-/Intel-/Qualcomm-GPUs dient Vulkan Compute als universelles, herstellerneutrales GPU-Backend.
• Compute-Kernel (geschrieben in GLSL, HLSL oder kompiliert aus MLIR/IREE) werden zu SPIR-V-Bytecode kompiliert. Über Cooperative-Matrix-Erweiterungen (VK_KHR_cooperative_matrix) greifen Vulkan-Kernel mit einer einzigen portablen Binärdatei direkt auf die Hardware-Matrixmultiplikationseinheiten von Qualcomm-Adreno-, ARM-Mali-, AMD-RDNA- und Intel-Arc-GPUs zu.