FSDP2, perfilado de GPU y cómputo móvil

La ingeniería de IA en producción une la ejecución distribuida en varias GPU, el perfilado a nivel de hardware y los entornos de ejecución multiplataforma para móvil y edge. A continuación, un análisis a fondo de los tres ámbitos.

1. Entrenamiento distribuido multi-GPU con PyTorch FSDP2

El FSDP2 moderno de PyTorch (fully_shard) fragmenta los parámetros del modelo, los gradientes y los estados del optimizador entre las GPU, parámetro a parámetro, mediante DTensor. Así elimina la duplicación de memoria y solapa la comunicación con el cálculo.

A continuación se muestra un script completo de entrenamiento multi-GPU. Lánzalo en 4 GPU con 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()

• Qué ocurre internamente: antes del paso hacia delante (forward) de cada capa, FSDP emite un All-Gather asíncrono de NCCL sobre NVLink para reconstruir los pesos completos en bfloat16 solo de ese bloque, ejecuta las operaciones y descarta de inmediato los pesos sin fragmentar. Durante loss.backward(), los gradientes se promedian y se vuelven a repartir entre las GPU mediante Reduce-Scatter.

2. Perfilado de un modelo de PyTorch: PyTorch Profiler y NVIDIA Nsight Compute (ncu)

Optimizar el código de GPU exige saber si un kernel está limitado por ancho de banda de memoria (detenido a la espera de lecturas/escrituras en HBM) o limitado por cómputo (saturando los FLOPS de los Tensor Cores).

Paso A: traza a nivel macro con torch.profiler

Usa PyTorch Profiler para detectar huecos en el lanzamiento desde la CPU, cargadores de datos lentos y los kernels de CUDA que más tiempo consumen.

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

Paso B: análisis roofline a nivel micro con NVIDIA Nsight Compute (ncu)

Cuando torch.profiler identifique un kernel lento, ejecuta ncu desde el terminal para inspeccionar directamente los contadores de rendimiento del hardware en los Streaming Multiprocessors (SM):

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

Al inspeccionar Speed of Light (SOL) y el gráfico Roofline dentro de kernel_report.ncu-rep, fíjate en dos cifras:

3. Ejecución en GPU de borde y móviles: Apple Metal y Vulkan Compute

Desplegar redes neuronales en teléfonos, portátiles y dispositivos embebidos obliga a funcionar sin controladores de NVIDIA ni grandes presupuestos de energía, apoyándose en Apple Metal y Vulkan Compute con arquitecturas de memoria unificada (UMA).

• Arquitectura de memoria unificada (ejecución sin copias).

• A diferencia de las GPU NVIDIA discretas, que exigen copiar tensores por PCIe desde la RAM de la CPU a la VRAM, Apple Silicon (serie M / serie A) y los SoC móviles (Qualcomm Snapdragon, ARM Mali) comparten un único conjunto de memoria física LPDDR5/LPDDR5X entre CPU, GPU y NPU.
• Marcos como Apple MLX, llama.cpp y ExecuTorch asignan los pesos del modelo directamente al espacio de direcciones compartido (MTLResourceStorageModeShared), eliminando por completo la latencia de transferencia de host a dispositivo.

• Apple Metal y Metal Performance Shaders (MPS).

• En macOS e iOS, Metal sustituye a CUDA. Los desarrolladores escriben compute shaders en Metal Shading Language (MSL), un derivado de C++14, y reparten threadgroups entre los núcleos de la GPU de Apple y las instrucciones Simdgroup Matrix (el equivalente de Apple a los Tensor Cores de NVIDIA).
• En PyTorch, al establecer device="mps", las operaciones con tensores pasan por el compilador MPSGraph de Apple, que fusiona operaciones de forma dinámica y planifica el trabajo entre las unidades de ejecución de la GPU.

• Vulkan Compute multiplataforma y SPIR-V.

• En Android, Windows, Linux y GPU de AMD/Intel/Qualcomm, Vulkan Compute funciona como el backend de GPU universal y neutral respecto al fabricante.
• Los kernels de cómputo (escritos en GLSL, HLSL o compilados desde MLIR/IREE) se compilan a bytecode SPIR-V. Con las extensiones Cooperative Matrix (VK_KHR_cooperative_matrix), los kernels de Vulkan acceden directamente a las unidades de multiplicación de matrices del hardware en las GPU Qualcomm Adreno, ARM Mali, AMD RDNA e Intel Arc con un único binario portable.