FSDP2, GPU profiling and mobile compute

Production AI engineering bridges distributed multi-GPU execution, hardware-level profiling, and cross-platform mobile/edge runtimes. Below is a deep dive into all three domains.

1. Multi-GPU Distributed Training with PyTorch FSDP2

PyTorch’s modern FSDP2 (fully_shard) shards model parameters, gradients, and optimizer states across GPUs per-parameter using DTensor, eliminating memory duplication while overlapping communication with computation.

Below is a complete multi-GPU training script. Launch it across 4 GPUs using 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()

• What Happens Under the Hood: Before each layer’s forward pass, FSDP issues an asynchronous NCCL All-Gather over NVLink to reconstruct the full bfloat16 weights for just that block, runs the math, and immediately discards the unsharded weights. During loss.backward(), gradients are averaged and partitioned back across GPUs via Reduce-Scatter.

2. Profiling a PyTorch Model: PyTorch Profiler & NVIDIA Nsight Compute (ncu)

Optimising GPU code requires knowing whether a kernel is Memory-Bandwidth Bound (stalled waiting on HBM reads/writes) or Compute Bound (saturating Tensor Core FLOPS).

Step A: Macro-Level Trace with torch.profiler

Use PyTorch Profiler to spot CPU launch gaps, slow data loaders, and top time-consuming CUDA kernels.

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

Step B: Micro-Level Roofline Analysis with NVIDIA Nsight Compute (ncu)

Once torch.profiler identifies a slow kernel, run ncu from the terminal to inspect hardware performance counters directly on the Streaming Multiprocessors (SMs):

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

When inspecting the Speed of Light (SOL) and Roofline Chart inside kernel_report.ncu-rep, look at two numbers:

3. Edge & Mobile GPU Execution: Apple Metal & Vulkan Compute

Deploying neural networks onto smartphones, laptops, and embedded devices requires running without NVIDIA drivers or massive power budgets—relying instead on Apple Metal and Vulkan Compute with Unified Memory Architectures (UMA).

• Unified Memory Architecture (Zero-Copy Execution).

• Unlike discrete NVIDIA GPUs that require copying tensors over PCIe from CPU RAM to VRAM, Apple Silicon (M-series / A-series) and mobile SoCs (Qualcomm Snapdragon, ARM Mali) share one physical LPDDR5/LPDDR5X memory pool between CPU, GPU, and NPU.
• Frameworks like Apple MLX, llama.cpp, and ExecuTorch map model weights directly into shared address space (MTLResourceStorageModeShared), eliminating host-to-device transfer latency completely.

• Apple Metal & Metal Performance Shaders (MPS).

• On macOS and iOS, Metal replaces CUDA. Developers write compute shaders in Metal Shading Language (MSL)—a C++14 derivative—dispatching threadgroups across Apple GPU cores and Simdgroup Matrix instructions (Apple’s equivalent of NVIDIA Tensor Cores).
• In PyTorch, setting device="mps" routes tensor operations through Apple’s MPSGraph compiler, which dynamically fuses operations and schedules work across GPU execution units.

• Cross-Platform Vulkan Compute & SPIR-V.

• Across Android, Windows, Linux, and AMD/Intel/Qualcomm GPUs, Vulkan Compute serves as the universal vendor-neutral GPU backend.
• Compute kernels (written in GLSL, HLSL, or compiled from MLIR/IREE) are compiled into SPIR-V bytecode. Using Cooperative Matrix extensions (VK_KHR_cooperative_matrix), Vulkan kernels tap directly into hardware matrix-multiply units across Qualcomm Adreno, ARM Mali, AMD RDNA, and Intel Arc GPUs using a single portable binary.