FSDP2, pemrofilan GPU, dan komputasi seluler

Rekayasa AI tingkat produksi menjembatani eksekusi multi-GPU terdistribusi, pemrofilan tingkat perangkat keras, dan runtime seluler/edge lintas platform. Berikut ulasan mendalam tentang ketiga bidang tersebut.

1. Pelatihan Terdistribusi Multi-GPU dengan PyTorch FSDP2

FSDP2 modern milik PyTorch (fully_shard) membagi (shard) parameter model, gradien, dan status optimizer ke seluruh GPU per parameter menggunakan DTensor, menghilangkan duplikasi memori sambil menumpangtindihkan komunikasi dengan komputasi.

Berikut skrip pelatihan multi-GPU lengkap. Jalankan di 4 GPU menggunakan 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()

• Apa yang Terjadi di Balik Layar: Sebelum forward pass setiap layer, FSDP menjalankan NCCL All-Gather asinkron melalui NVLink untuk merekonstruksi bobot bfloat16 penuh hanya untuk blok tersebut, menjalankan perhitungan, lalu langsung membuang bobot yang tidak di-shard. Selama loss.backward(), gradien dirata-ratakan dan dipartisi kembali ke seluruh GPU melalui Reduce-Scatter.

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

Mengoptimalkan kode GPU memerlukan pengetahuan apakah sebuah kernel Terbatas Bandwidth Memori (tertahan menunggu pembacaan/penulisan HBM) atau Terbatas Komputasi (menjenuhkan FLOPS Tensor Core).

Langkah A: Penelusuran Tingkat Makro dengan torch.profiler

Gunakan PyTorch Profiler untuk menemukan jeda peluncuran CPU, data loader yang lambat, dan kernel CUDA yang paling banyak memakan waktu.

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

Langkah B: Analisis Roofline Tingkat Mikro dengan NVIDIA Nsight Compute (ncu)

Setelah torch.profiler mengidentifikasi kernel yang lambat, jalankan ncu dari terminal untuk memeriksa penghitung kinerja perangkat keras langsung pada Streaming Multiprocessor (SM):

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

Saat memeriksa Speed of Light (SOL) dan Roofline Chart di dalam kernel_report.ncu-rep, perhatikan dua angka:

3. Eksekusi GPU Edge & Seluler: Apple Metal & Vulkan Compute

Menerapkan jaringan saraf ke ponsel pintar, laptop, dan perangkat tertanam mengharuskan berjalan tanpa driver NVIDIA atau anggaran daya yang besar—dengan mengandalkan Apple Metal dan Vulkan Compute beserta Unified Memory Architecture (UMA).

• Unified Memory Architecture (Eksekusi Zero-Copy).

• Tidak seperti GPU NVIDIA diskrit yang memerlukan penyalinan tensor melalui PCIe dari RAM CPU ke VRAM, Apple Silicon (seri M / seri A) dan SoC seluler (Qualcomm Snapdragon, ARM Mali) berbagi satu kolam memori fisik LPDDR5/LPDDR5X antara CPU, GPU, dan NPU.
• Framework seperti Apple MLX, llama.cpp, dan ExecuTorch memetakan bobot model langsung ke ruang alamat bersama (MTLResourceStorageModeShared), menghilangkan sepenuhnya latensi transfer host-ke-perangkat.

• Apple Metal & Metal Performance Shaders (MPS).

• Di macOS dan iOS, Metal menggantikan CUDA. Pengembang menulis compute shader dalam Metal Shading Language (MSL)—turunan C++14—yang mengirimkan threadgroup ke seluruh inti GPU Apple dan instruksi Simdgroup Matrix (padanan Tensor Core NVIDIA di Apple).
• Di PyTorch, mengatur device="mps" mengarahkan operasi tensor melalui compiler MPSGraph milik Apple, yang secara dinamis menggabungkan operasi dan menjadwalkan pekerjaan di seluruh unit eksekusi GPU.

• Vulkan Compute Lintas Platform & SPIR-V.

• Di Android, Windows, Linux, dan GPU AMD/Intel/Qualcomm, Vulkan Compute berfungsi sebagai backend GPU universal yang netral vendor.
• Kernel komputasi (ditulis dalam GLSL, HLSL, atau dikompilasi dari MLIR/IREE) dikompilasi menjadi bytecode SPIR-V. Dengan Cooperative Matrix extensions (VK_KHR_cooperative_matrix), kernel Vulkan langsung memanfaatkan unit perkalian matriks perangkat keras di GPU Qualcomm Adreno, ARM Mali, AMD RDNA, dan Intel Arc menggunakan satu biner portabel.