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, NVLink-এর মাধ্যমে একটি অ্যাসিঙ্ক্রোনাস NCCL All-Gather চালায়, যাতে শুধু ওই ব্লকের সম্পূর্ণ bfloat16 ওয়েট পুনর্গঠন করা যায়। তারপর গণনা করে এবং অশার্ডকৃত ওয়েট সঙ্গে সঙ্গে সরিয়ে দেয়। loss.backward()-এর সময় Reduce-Scatter-এর মাধ্যমে গ্রেডিয়েন্টের গড় করা হয় এবং সেগুলি আবার GPU-গুলিতে ভাগ করে দেওয়া হয়।

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 Chart পরীক্ষা করার সময় দুটি সংখ্যা দেখুন:

3. এজ & মোবাইল GPU নির্বাহ: Apple Metal & Vulkan Compute

স্মার্টফোন, ল্যাপটপ এবং এমবেডেড ডিভাইসে নিউরাল নেটওয়ার্ক স্থাপন করতে NVIDIA ড্রাইভার বা বিশাল বিদ্যুৎ বাজেট ছাড়াই চালাতে হয়। এর পরিবর্তে Unified Memory Architectures (UMA) সহ Apple Metal এবং Vulkan Compute-এর ওপর নির্ভর করতে হয়।

• ইউনিফায়েড মেমরি আর্কিটেকচার (জিরো-কপি নির্বাহ)।

• পৃথক NVIDIA GPU-তে PCIe দিয়ে CPU RAM থেকে VRAM-এ টেনসর কপি করতে হয়। এর বিপরীতে 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-র জায়গা নেয়। ডেভেলপাররা C++14 থেকে উদ্ভূত Metal Shading Language (MSL)-এ কম্পিউট শেডার লেখেন, যা Apple GPU কোর এবং Simdgroup Matrix নির্দেশনায় (NVIDIA Tensor Cores-এর 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-র হার্ডওয়্যার ম্যাট্রিক্স-গুণন ইউনিট সরাসরি ব্যবহার করে।