FSDP2, GPU प्रोफाइलिंग आणि मोबाइल कॉम्प्युट

प्रोडक्शन AI अभियांत्रिकी वितरित मल्टी-GPU एक्झिक्युशन, हार्डवेअर-स्तरीय प्रोफाइलिंग आणि क्रॉस-प्लॅटफॉर्म मोबाइल/एज रनटाइम यांना जोडते. खाली तिन्ही क्षेत्रांचा सखोल आढावा दिला आहे.

1. PyTorch FSDP2 सह मल्टी-GPU वितरित ट्रेनिंग

PyTorch चे आधुनिक FSDP2 (fully_shard) DTensor वापरून मॉडेल पॅरामीटर, ग्रेडियंट आणि ऑप्टिमायझर स्टेट प्रति-पॅरामीटर पद्धतीने GPU मध्ये विभागते, त्यामुळे मेमरीची द्विरुक्ती टळते आणि कम्युनिकेशन गणनेसोबत ओव्हरलॅप होते.

खाली संपूर्ण मल्टी-GPU ट्रेनिंग स्क्रिप्ट दिली आहे. ती torchrun --nproc_per_node=4 train_fsdp.py वापरून 4 GPU वर चालवा.

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 द्वारे ती पुन्हा GPU मध्ये विभागली जातात.

2. PyTorch मॉडेलचे प्रोफाइलिंग: PyTorch Profiler & NVIDIA Nsight Compute (ncu)

GPU कोड ऑप्टिमाइझ करण्यासाठी कर्नल Memory-Bandwidth Bound (HBM रीड/राइटची वाट पाहत थांबलेले) आहे की Compute Bound (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 Chart तपासताना, दोन आकडे पाहा:

3. एज आणि मोबाइल GPU एक्झिक्युशन: Apple Metal & Vulkan Compute

स्मार्टफोन, लॅपटॉप आणि एम्बेडेड उपकरणांवर न्यूरल नेटवर्क चालवण्यासाठी NVIDIA ड्रायव्हर किंवा मोठ्या वीज बजेटशिवाय चालावे लागते—त्याऐवजी Unified Memory Architectures (UMA) सह Apple Metal आणि Vulkan Compute वर अवलंबून राहावे लागते.

• Unified Memory Architecture (झिरो-कॉपी एक्झिक्युशन).

• ज्यात टेन्सर PCIe वरून CPU RAM मधून VRAM मध्ये कॉपी करावे लागतात अशा डिस्क्रीट NVIDIA GPU च्या विपरीत, Apple Silicon (M-series / A-series) आणि मोबाइल SoC (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 Core चे Apple समतुल्य) थ्रेडग्रुप डिस्पॅच करतात.
• PyTorch मध्ये, device="mps" सेट केल्यास टेन्सर ऑपरेशन्स Apple च्या MPSGraph कंपायलरमधून जातात, जो गतिशीलपणे ऑपरेशन्स फ्यूज करतो आणि GPU एक्झिक्युशन युनिटवर कामाचे शेड्यूलिंग करतो.

• क्रॉस-प्लॅटफॉर्म Vulkan Compute & SPIR-V.

• Android, Windows, Linux आणि AMD/Intel/Qualcomm GPU वर, Vulkan Compute सार्वत्रिक व्हेंडर-न्यूट्रल GPU बॅकएंड म्हणून काम करते.
• कॉम्प्युट कर्नल (GLSL, HLSL मध्ये लिहिलेले किंवा MLIR/IREE मधून कंपाइल केलेले) SPIR-V बायटकोडमध्ये कंपाइल केले जातात. Cooperative Matrix एक्स्टेन्शन (VK_KHR_cooperative_matrix) वापरून, Vulkan कर्नल एकाच पोर्टेबल बायनरीसह Qualcomm Adreno, ARM Mali, AMD RDNA आणि Intel Arc GPU वरील हार्डवेअर मॅट्रिक्स-गुणाकार युनिटचा थेट वापर करतात.