FSDP2, GPU පැතිකඩ විශ්ලේෂණය සහ ජංගම පරිගණනය

නිෂ්පාදන භාවිතයේ AI ඉංජිනේරුකරණය බෙදාහැරුණු බහු-GPU ක්‍රියාත්මක කිරීම, දෘඪාංග මට්ටමේ පැතිකඩ විශ්ලේෂණය සහ විවිධ වේදිකා සඳහා ජංගම/එජ් ධාවන පරිසර සම්බන්ධ කරයි. පහත දැක්වෙන්නේ මෙම ක්ෂේත්‍ර තුනම පිළිබඳ සවිස්තර විමසුමකි.

1. PyTorch FSDP2 සමඟ බහු-GPU බෙදාහැරුණු පුහුණුව

PyTorch හි නවීන FSDP2 (fully_shard), DTensor භාවිතයෙන් මාදිලි පරාමිති, අනුක්‍රමණ සහ ප්‍රශස්තකාරක තත්ත්ව එක් එක් පරාමිතිය අනුව GPU අතර කොටස් කරයි. එමඟින් මතක අනුපිටපත් ඉවත් කරන අතර සන්නිවේදනය සහ ගණනය සමගාමීව සිදු කරයි.

සම්පූර්ණ බහු-GPU පුහුණු ස්ක්‍රිප්ටයක් පහත දැක්වේ. torchrun --nproc_per_node=4 train_fsdp.py භාවිතයෙන් එය GPU 4 ක් පුරා දියත් කරන්න.

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

• අභ්‍යන්තරයේ සිදුවන්නේ කුමක්ද: එක් එක් ස්තරයේ ඉදිරි ගමන් වාරයට පෙර, එම බ්ලොක් එක සඳහා පමණක් සම්පූර්ණ bfloat16 බර නැවත ගොඩනැගීමට FSDP විසින් NVLink හරහා අසමමුහුර්ත NCCL All-Gather එකක් නිකුත් කරයි, ගණනය සිදු කරයි, ඉන්පසු කොටස් නොකළ බර වහාම ඉවත් කරයි. 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 ප්‍රස්තාරය පරීක්ෂා කරන විට අගයන් දෙකක් බලන්න:

3. එජ් සහ ජංගම GPU ක්‍රියාත්මක කිරීම: Apple Metal & Vulkan Compute

ස්මාර්ට්ෆෝන්, ලැප්ටොප් සහ කාවැද්දූ උපාංග මත ස්නායුක ජාල යෙදවීමට NVIDIA ධාවක හෝ විශාල බලශක්ති ප්‍රතිපාදන නොමැතිව ධාවනය කළ යුතුය. ඒ වෙනුවට Unified Memory Architectures (UMA) සමඟ Apple Metal සහ Vulkan Compute මත රඳා පවතී.

• ඒකාබද්ධ මතක වාස්තු විද්‍යාව (පිටපත් රහිත ක්‍රියාත්මක කිරීම).

• CPU RAM සිට VRAM වෙත PCIe හරහා ටෙන්සර් පිටපත් කිරීමට අවශ්‍ය වන වෙන් වූ 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 මත CUDA වෙනුවට Metal භාවිත වේ. සංවර්ධකයන් 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 පුරා දෘඪාංග න්‍යාස ගුණන ඒකක සෘජුවම භාවිත කරයි.