FSDP2, profilage GPU et calcul mobile

L'ingénierie IA de production relie l'exécution distribuée multi-GPU, le profilage au niveau matériel et les runtimes mobiles/edge multiplateformes. Vous trouverez ci-dessous une analyse approfondie de ces trois domaines.

1. Entraînement distribué multi-GPU avec PyTorch FSDP2

Le FSDP2 moderne de PyTorch (fully_shard) répartit les paramètres du modèle, les gradients et les états de l'optimiseur sur les GPU, paramètre par paramètre, à l'aide de DTensor, éliminant la duplication de mémoire tout en recouvrant la communication par le calcul.

Vous trouverez ci-dessous un script d'entraînement multi-GPU complet. Lancez-le sur 4 GPU avec 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()

• Fonctionnement interne : avant le passage forward de chaque couche, FSDP lance un All-Gather NCCL asynchrone sur NVLink pour reconstruire les poids bfloat16 complets de ce seul bloc, exécute les calculs, puis supprime immédiatement les poids non répartis. Pendant loss.backward(), les gradients sont moyennés et repartitionnés entre les GPU via Reduce-Scatter.

2. Profilage d'un modèle PyTorch : PyTorch Profiler & NVIDIA Nsight Compute (ncu)

L'optimisation du code GPU exige de savoir si un kernel est limité par la bande passante mémoire (bloqué en attente de lectures/écritures HBM) ou limité par le calcul (saturant les FLOPS des Tensor Cores).

Étape A : trace au niveau macro avec torch.profiler

Utilisez PyTorch Profiler pour repérer les intervalles de lancement CPU, les chargeurs de données lents et les kernels CUDA les plus coûteux en temps.

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

Étape B : analyse Roofline au niveau micro avec NVIDIA Nsight Compute (ncu)

Une fois que torch.profiler a identifié un kernel lent, exécutez ncu depuis le terminal pour inspecter directement les compteurs de performance matériels sur les Streaming Multiprocessors (SM) :

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

Lors de l'inspection du Speed of Light (SOL) et du graphique Roofline dans kernel_report.ncu-rep, examinez deux nombres :

3. Exécution GPU edge et mobile : Apple Metal & Vulkan Compute

Déployer des réseaux de neurones sur des smartphones, ordinateurs portables et appareils embarqués exige de fonctionner sans pilotes NVIDIA ni budgets énergétiques massifs, en s'appuyant plutôt sur Apple Metal et Vulkan Compute avec des architectures de mémoire unifiée (UMA).

• Architecture de mémoire unifiée (exécution zéro copie).

• Contrairement aux GPU NVIDIA discrets qui exigent de copier les tenseurs via PCIe de la RAM CPU vers la VRAM, Apple Silicon (séries M / A) et les SoC mobiles (Qualcomm Snapdragon, ARM Mali) partagent un seul pool de mémoire physique LPDDR5/LPDDR5X entre CPU, GPU et NPU.
• Des frameworks comme Apple MLX, llama.cpp et ExecuTorch mappent les poids du modèle directement dans l'espace d'adressage partagé (MTLResourceStorageModeShared), éliminant complètement la latence de transfert hôte-vers-appareil.

• Apple Metal & Metal Performance Shaders (MPS).

• Sur macOS et iOS, Metal remplace CUDA. Les développeurs écrivent des compute shaders en Metal Shading Language (MSL), un dérivé de C++14, en répartissant les threadgroups sur les cœurs GPU Apple et les instructions Simdgroup Matrix (l'équivalent Apple des NVIDIA Tensor Cores).
• Dans PyTorch, définir device="mps" achemine les opérations sur tenseurs via le compilateur MPSGraph d'Apple, qui fusionne dynamiquement les opérations et planifie le travail sur les unités d'exécution du GPU.

• Vulkan Compute multiplateforme & SPIR-V.

• Sur Android, Windows, Linux et les GPU AMD/Intel/Qualcomm, Vulkan Compute sert de backend GPU universel et indépendant du fournisseur.
• Les kernels de calcul (écrits en GLSL, HLSL ou compilés depuis MLIR/IREE) sont compilés en bytecode SPIR-V. Grâce aux extensions Cooperative Matrix (VK_KHR_cooperative_matrix), les kernels Vulkan exploitent directement les unités matérielles de multiplication matricielle des GPU Qualcomm Adreno, ARM Mali, AMD RDNA et Intel Arc avec un seul binaire portable.