FSDP2، GPU پروفائلنگ اور موبائل کمپیوٹنگ

پروڈکشن AI انجینئرنگ تقسیم شدہ ملٹی GPU اجرا، ہارڈویئر کی سطح کی پروفائلنگ، اور مختلف پلیٹ فارموں پر موبائل/ایج رن ٹائمز کو یکجا کرتی ہے۔ ذیل میں ان تینوں شعبوں کا تفصیلی جائزہ پیش ہے۔

1. PyTorch FSDP2 کے ساتھ ملٹی GPU تقسیم شدہ تربیت

PyTorch کا جدید FSDP2 (fully_shard)، DTensor استعمال کرتے ہوئے ماڈل کے پیرامیٹرز، گریڈینٹس اور آپٹیمائزر کی حالتوں کو ہر پیرامیٹر کی سطح پر GPUs میں تقسیم کرتا ہے، جس سے میموری کی نقل ختم ہوتی ہے اور مواصلات و حساب بیک وقت انجام پاتے ہیں۔

ذیل میں ملٹی GPU تربیت کا ایک مکمل اسکرپٹ ہے۔ اسے torchrun --nproc_per_node=4 train_fsdp.py استعمال کرکے 4 GPUs پر چلائیں۔

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

• اندر کیا ہوتا ہے: ہر تہہ کے فارورڈ پاس سے پہلے، FSDP صرف اسی بلاک کے مکمل bfloat16 ویٹس دوبارہ بنانے کے لیے NVLink پر ایک غیر ہم وقت NCCL All-Gather جاری کرتا ہے، حساب انجام دیتا ہے اور غیر منقسم ویٹس کو فوراً خارج کر دیتا ہے۔ loss.backward() کے دوران گریڈینٹس کا اوسط لیا جاتا ہے اور Reduce-Scatter کے ذریعے انہیں دوبارہ GPUs میں تقسیم کیا جاتا ہے۔

2. PyTorch ماڈل کی پروفائلنگ: PyTorch Profiler & NVIDIA Nsight Compute (ncu)

GPU کوڈ کو بہتر بنانے کے لیے جاننا ضروری ہے کہ کرنل میموری بینڈوڈتھ سے محدود ہے (HBM کی پڑھائی/لکھائی کے انتظار میں رکا ہوا) یا حساب سے محدود ہے (Tensor Core FLOPS کی پوری گنجائش استعمال کر رہا ہے)۔

مرحلہ A: torch.profiler کے ساتھ بڑی سطح کا ٹریس

CPU لانچ کے وقفوں، سست ڈیٹا لوڈرز اور سب سے زیادہ وقت لینے والے CUDA کرنلز کی شناخت کے لیے 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) کے ساتھ چھوٹی سطح کا Roofline تجزیہ

جب torch.profiler کسی سست کرنل کی شناخت کر لے، تو Streaming Multiprocessors (SMs) پر ہارڈویئر کی کارکردگی کے کاؤنٹرز کا براہ راست معائنہ کرنے کے لیے ٹرمینل سے ncu چلائیں:

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

kernel_report.ncu-rep کے اندر Speed of Light (SOL) اور Roofline چارٹ کا معائنہ کرتے وقت دو اعداد دیکھیں:

3. ایج اور موبائل GPU اجرا: Apple Metal & Vulkan Compute

اسمارٹ فونز، لیپ ٹاپس اور ایمبیڈڈ آلات پر نیورل نیٹ ورکس تعینات کرنے کے لیے NVIDIA ڈرائیورز یا بہت زیادہ توانائی کے بجٹ کے بغیر چلانا ضروری ہے، جس کے لیے ان کی جگہ Unified Memory Architectures (UMA) کے ساتھ Apple Metal اور Vulkan Compute پر انحصار کیا جاتا ہے۔

• متحدہ میموری آرکیٹیکچر (زیرو کاپی اجرا)۔

• علیحدہ NVIDIA GPUs کے برعکس، جنہیں CPU RAM سے VRAM تک PCIe پر ٹینسرز کی نقل بنانا پڑتی ہے، Apple Silicon (M-series / A-series) اور موبائل SoCs (Qualcomm Snapdragon، ARM Mali) میں CPU، GPU اور NPU ایک ہی طبعی LPDDR5/LPDDR5X میموری پول مشترکہ طور پر استعمال کرتے ہیں۔
• Apple MLX، llama.cpp اور ExecuTorch جیسے فریم ورک ماڈل کے ویٹس کو براہ راست مشترکہ ایڈریس اسپیس (MTLResourceStorageModeShared) میں میپ کرتے ہیں، جس سے ہوسٹ سے ڈیوائس منتقلی کی تاخیر مکمل طور پر ختم ہو جاتی ہے۔

• Apple Metal & Metal Performance Shaders (MPS)۔

• macOS اور iOS پر Metal، CUDA کی جگہ لیتا ہے۔ ڈویلپرز Metal Shading Language (MSL) میں کمپیوٹ شیڈرز لکھتے ہیں، جو C++14 سے اخذ کردہ زبان ہے، اور تھریڈ گروپس کو Apple GPU کورز اور Simdgroup Matrix ہدایات (NVIDIA Tensor Cores کا Apple متبادل) پر بھیجتے ہیں۔
• PyTorch میں device="mps" سیٹ کرنے سے ٹینسر آپریشنز Apple کے MPSGraph کمپائلر کے ذریعے چلتے ہیں، جو متحرک طور پر آپریشنز کو یکجا کرتا ہے اور GPU کی ایگزیکیوشن یونٹس پر کام کا شیڈول بناتا ہے۔

• مختلف پلیٹ فارموں کے لیے Vulkan Compute & SPIR-V۔

• Android، Windows، Linux اور AMD/Intel/Qualcomm GPUs پر Vulkan Compute ایک آفاقی، وینڈر سے غیر وابستہ GPU بیک اینڈ ہے۔
• کمپیوٹ کرنلز (جو GLSL، HLSL میں لکھے جاتے ہیں یا MLIR/IREE سے کمپائل ہوتے ہیں) SPIR-V بائٹ کوڈ میں کمپائل کیے جاتے ہیں۔ Cooperative Matrix ایکسٹینشنز (VK_KHR_cooperative_matrix) استعمال کرتے ہوئے Vulkan کرنلز ایک ہی قابل منتقلی بائنری کے ذریعے Qualcomm Adreno، ARM Mali، AMD RDNA اور Intel Arc GPUs میں ہارڈویئر میٹرکس ضرب یونٹس کو براہ راست استعمال کرتے ہیں۔