FSDP2, GPU profilleme ve mobil hesaplama

Üretim ortamlarında yapay zeka mühendisliği, dağıtık çoklu GPU yürütmesini, donanım düzeyinde profillemeyi ve platformlar arası mobil/uç çalışma zamanlarını bir araya getirir. Aşağıda bu üç alan ayrıntılı olarak ele alınmaktadır.

1. PyTorch FSDP2 ile Çoklu GPU Üzerinde Dağıtık Eğitim

PyTorch'un modern FSDP2 (fully_shard) sistemi, DTensor kullanarak model parametrelerini, gradyanları ve optimizasyon durumlarını parametre bazında GPU'lar arasında parçalar; böylece iletişimi hesaplamayla örtüştürürken bellekteki çoğaltmayı ortadan kaldırır.

Aşağıda eksiksiz bir çoklu GPU eğitim betiği yer almaktadır. torchrun --nproc_per_node=4 train_fsdp.py kullanarak 4 GPU üzerinde başlatın.

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

• Arka Planda Olanlar: FSDP, her katmanın ileri geçişinden önce yalnızca o bloğun tüm bfloat16 ağırlıklarını yeniden oluşturmak için NVLink üzerinden asenkron bir NCCL All-Gather işlemi başlatır, hesaplamayı yapar ve parçalanmamış ağırlıkları hemen atar. loss.backward() sırasında gradyanların ortalaması alınır ve Reduce-Scatter aracılığıyla yeniden GPU'lar arasında bölüştürülür.

2. Bir PyTorch Modelini Profilleme: PyTorch Profiler & NVIDIA Nsight Compute (ncu)

GPU kodunu optimize etmek, bir çekirdeğin bellek bant genişliğiyle mi sınırlı olduğunu (HBM okuma/yazmalarını beklerken durakladığını), yoksa hesaplamayla mı sınırlı olduğunu (Tensor Core FLOPS kapasitesini doldurduğunu) bilmeyi gerektirir.

Adım A: torch.profiler ile Makro Düzeyde İzleme

CPU başlatma boşluklarını, yavaş veri yükleyicilerini ve en çok zaman alan CUDA çekirdeklerini belirlemek için PyTorch Profiler kullanın.

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

Adım B: NVIDIA Nsight Compute (ncu) ile Mikro Düzeyde Roofline Analizi

torch.profiler yavaş bir çekirdek belirledikten sonra, donanım performans sayaçlarını doğrudan Streaming Multiprocessor (SM) birimleri üzerinde incelemek için terminalden ncu çalıştırın:

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

kernel_report.ncu-rep içindeki Speed of Light (SOL) ve Roofline grafiğini incelerken iki sayıya bakın:

3. Uç & Mobil GPU Yürütmesi: Apple Metal & Vulkan Compute

Sinir ağlarını akıllı telefonlara, dizüstü bilgisayarlara ve gömülü cihazlara dağıtmak, NVIDIA sürücüleri veya devasa güç bütçeleri olmadan çalışmayı gerektirir; bunun yerine Birleşik Bellek Mimarileri (UMA) ile Apple Metal ve Vulkan Compute kullanılır.

• Birleşik Bellek Mimarisi (Sıfır Kopyalı Yürütme).

• Tensörleri PCIe üzerinden CPU RAM'inden VRAM'e kopyalamayı gerektiren ayrık NVIDIA GPU'larından farklı olarak Apple Silicon (M-series / A-series) ve mobil SoC'ler (Qualcomm Snapdragon, ARM Mali), CPU, GPU ve NPU arasında tek bir fiziksel LPDDR5/LPDDR5X bellek havuzunu paylaşır.
• Apple MLX, llama.cpp ve ExecuTorch gibi çerçeveler model ağırlıklarını doğrudan paylaşılan adres alanına (MTLResourceStorageModeShared) eşleyerek ana sistemden cihaza aktarım gecikmesini tamamen ortadan kaldırır.

• Apple Metal & Metal Performance Shaders (MPS).

• macOS ve iOS'ta CUDA'nın yerini Metal alır. Geliştiriciler, C++14 türevi Metal Shading Language (MSL) ile hesaplama gölgelendiricileri yazar; iş parçacığı gruplarını Apple GPU çekirdeklerine ve Simdgroup Matrix komutlarına (Apple'ın NVIDIA Tensor Core birimlerinin karşılığı) dağıtır.
• PyTorch'ta device="mps" ayarı, tensör işlemlerini Apple'ın MPSGraph derleyicisine yönlendirir; bu derleyici işlemleri dinamik olarak birleştirir ve GPU yürütme birimleri arasında işi zamanlar.

• Platformlar Arası Vulkan Compute & SPIR-V.

• Android, Windows, Linux ve AMD/Intel/Qualcomm GPU'larında Vulkan Compute, evrensel ve üreticiden bağımsız GPU arka ucu olarak kullanılır.
• Hesaplama çekirdekleri (GLSL veya HLSL ile yazılan ya da MLIR/IREE'den derlenen) SPIR-V bayt koduna derlenir. Vulkan çekirdekleri, Cooperative Matrix uzantılarını (VK_KHR_cooperative_matrix) kullanarak tek bir taşınabilir ikili dosyayla Qualcomm Adreno, ARM Mali, AMD RDNA ve Intel Arc GPU'larındaki donanımsal matris çarpma birimlerine doğrudan erişir.