PyTorch FSDP2 பல-GPU பயிற்சி, Profiling (ncu) மற்றும் Edge/Mobile GPU இயக்கம் (Metal & Vulkan)

உற்பத்தி (Production) AI பொறியியல், பரவலாக்கப்பட்ட பல-GPU இயக்கம், hardware-மட்ட profiling, மற்றும் குறுக்கு-தள (cross-platform) mobile/edge runtime-கள் ஆகியவற்றை இணைக்கிறது. கீழே மூன்று துறைகளிலும் ஆழ்ந்த ஆராய்ச்சி (deep dive) உள்ளது.

1. PyTorch FSDP2-உடன் பல-GPU பரவலாக்கப்பட்ட பயிற்சி (Distributed Training)

PyTorch-இன் நவீன FSDP2 (fully_shard), DTensor பயன்படுத்தி model அளவுருக்கள் (parameters), gradient-கள், optimizer நிலைகள் ஆகியவற்றை அளவுரு வாரியாக (per-parameter) GPU-களில் பிரிக்கிறது (shards); தகவல்தொடர்பைக் (communication) கணிதத்துடன் மேற்படிப்பு செய்துகொண்டே (overlapping) நினைவக நகலையும் (duplication) நீக்குகிறது.

கீழே முழுமையான பல-GPU பயிற்சி script உள்ளது. torchrun --nproc_per_node=4 train_fsdp.py பயன்படுத்தி அதை 4 GPU-களில் தொடங்குங்கள் (launch).

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

• உள்ளே என்ன நடக்கிறது (Under the Hood): ஒவ்வொரு layer-இன் forward pass-க்கும் முன், FSDP அந்தத் தொகுதிக்கு மட்டுமான முழு bfloat16 weight-களை மறுகட்டமைக்க NVLink மீது ஒத்திசைவற்ற (asynchronous) NCCL All-Gather-ஐ வெளியிடுகிறது; கணிதத்தை இயக்கி, பிரிக்கப்படாத (unsharded) weight-களை உடனே கைவிடுகிறது. loss.backward()-இன் போது, gradient-கள் சராசரியாக்கப்பட்டு Reduce-Scatter வழியே GPU-களில் மீண்டும் பங்கிடப்படுகின்றன.

2. PyTorch Model-ஐ Profiling செய்தல்: PyTorch Profiler & NVIDIA Nsight Compute (ncu)

GPU code-ஐ மேம்படுத்த (optimising), ஒரு kernel Memory-Bandwidth Bound-ஆ (HBM வாசிப்பு/எழுத்துகளுக்காகக் காத்திருந்து நிறுத்தப்பட்டது) அல்லது Compute Bound-ஆ (Tensor Core FLOPS-ஐ நிரப்புகிறது) என அறிவது அவசியம்.

படி A: torch.profiler-உடன் மேல்-மட்ட (Macro-Level) Trace

CPU launch இடைவெளிகளை (gaps), மெதுவான data loader-களை, மற்றும் அதிக நேரம் எடுக்கும் CUDA kernel-களைக் கண்டறிய PyTorch Profiler-ஐப் பயன்படுத்துங்கள்.

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

படி B: NVIDIA Nsight Compute (ncu)-உடன் நுண்-மட்ட (Micro-Level) Roofline பகுப்பாய்வு

torch.profiler ஒரு மெதுவான kernel-ஐக் கண்டறிந்ததும், Streaming Multiprocessor-களின் (SMs) மீது hardware performance counter-களை நேரடியாக ஆராய terminal-இலிருந்து ncu-ஐ இயக்குங்கள்:

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

kernel_report.ncu-rep-இல் Speed of Light (SOL) மற்றும் Roofline Chart-ஐ ஆராயும்போது, இரண்டு எண்களைப் பாருங்கள்:

3. Edge & Mobile GPU இயக்கம்: Apple Metal & Vulkan Compute

Neural network-களை smartphone-கள், laptop-கள், embedded சாதனங்களில் பயன்படுத்த (deploying), NVIDIA driver-கள் அல்லது பாரிய மின்சார ஒதுக்கீடுகள் (power budgets) இல்லாமல் இயக்க வேண்டும் — அதற்குப் பதிலாக Unified Memory Architectures (UMA) கொண்ட Apple Metal மற்றும் Vulkan Compute-ஐ நம்பியிருக்க வேண்டும்.

• Unified Memory Architecture (Zero-Copy இயக்கம்).

• PCIe வழியே CPU RAM-இலிருந்து VRAM-க்கு tensor-களை நகல் செய்ய வேண்டிய discrete NVIDIA GPU-களைப் போலன்றி, Apple Silicon (M-series / A-series) மற்றும் mobile SoC-கள் (Qualcomm Snapdragon, ARM Mali) CPU, GPU, NPU ஆகியவற்றுக்கிடையே ஒரே இயற்கை LPDDR5/LPDDR5X நினைவகக் குளத்தைப் (memory pool) பகிர்கின்றன.
• Apple MLX, llama.cpp, ExecuTorch போன்ற framework-கள் model weight-களை நேரடியாகப் பகிரப்பட்ட முகவரி வெளியில் (shared address space; MTLResourceStorageModeShared) map செய்து, host-to-device இடமாற்ற latency-ஐ முற்றிலும் நீக்கிவிடுகின்றன.

• Apple Metal & Metal Performance Shaders (MPS).

• macOS மற்றும் iOS-இல், Metal, CUDA-ஐ மாற்றுகிறது (replaces). Developer-கள் Metal Shading Language (MSL)-இல் — ஒரு C++14 வழித்தோன்றல் — compute shader-களை எழுதி, Apple GPU core-களிலும் Simdgroup Matrix அறிவுறுத்தல்களிலும் (NVIDIA Tensor Core-களுக்கு Apple-இன் சமானம்) threadgroup-களை அனுப்புகின்றனர் (dispatch).
• PyTorch-இல், device="mps" என அமைப்பது tensor செயல்பாடுகளை Apple-இன் MPSGraph compiler வழியே செலுத்துகிறது; அது செயல்பாடுகளை இயங்கும் வகையில் இணைத்து (fuses), GPU execution அலகுகளில் வேலையை அட்டவணைப்படுத்துகிறது.

• குறுக்கு-தள (Cross-Platform) Vulkan Compute & SPIR-V.

• Android, Windows, Linux, மற்றும் AMD/Intel/Qualcomm GPU-களில், Vulkan Compute உலகளாவிய vendor-நடுநிலை GPU backend-ஆக செயல்படுகிறது.
• Compute kernel-கள் (GLSL, HLSL-இல் எழுதப்பட்டவை, அல்லது MLIR/IREE-இலிருந்து தொகுக்கப்பட்டவை) SPIR-V bytecode-ஆகத் தொகுக்கப்படுகின்றன. Cooperative Matrix நீட்சிகளைப் (VK_KHR_cooperative_matrix) பயன்படுத்தி, Vulkan kernel-கள் ஒரே எடுத்துச்செல்லக்கூடிய (portable) binary மூலம் Qualcomm Adreno, ARM Mali, AMD RDNA, Intel Arc GPU-களில் hardware matrix-multiply அலகுகளை நேரடியாக அணுகுகின்றன.