FSDP2, uchanganuzi wa utendaji wa GPU na ukokotoaji kwenye vifaa vya mkononi

Uhandisi wa AI katika mazingira ya uzalishaji unaunganisha utekelezaji uliosambazwa kwenye GPU nyingi, uchanganuzi wa utendaji katika kiwango cha maunzi, na mazingira ya utekelezaji ya vifaa vya mkononi/pembeni yanayofanya kazi kwenye majukwaa tofauti. Hapa chini kuna uchambuzi wa kina wa maeneo yote matatu.

1. Mafunzo Yaliyosambazwa kwenye GPU Nyingi kwa PyTorch FSDP2

FSDP2 ya kisasa ya PyTorch (fully_shard) hugawanya vigezo vya modeli, gradienti na hali za kiboreshaji kwenye GPU kwa kila kigezo kwa kutumia DTensor, ikiondoa urudufu wa kumbukumbu huku ikiendesha mawasiliano sambamba na ukokotoaji.

Hapa chini kuna hati kamili ya mafunzo kwenye GPU nyingi. Ianzishe kwenye GPU 4 kwa kutumia torchrun --nproc_per_node=4 train_fsdp.py.

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

• Kinachotokea Ndani: Kabla ya upitishaji wa mbele wa kila tabaka, FSDP hutuma NCCL All-Gather isiyosubiri kukamilika kupitia NVLink ili kuunda upya uzani kamili wa bfloat16 kwa bloku hiyo pekee, hufanya hesabu, kisha hutupa mara moja uzani ambao haujagawanywa. Wakati wa loss.backward(), wastani wa gradienti huhesabiwa na gradienti hugawanywa tena kwenye GPU kupitia Reduce-Scatter.

2. Kuchanganua Utendaji wa Modeli ya PyTorch: PyTorch Profiler & NVIDIA Nsight Compute (ncu)

Kuboresha msimbo wa GPU kunahitaji kujua ikiwa kerneli inazuiwa na kipimo-data cha kumbukumbu (imesimama ikisubiri usomaji/uandishi wa HBM) au inazuiwa na ukokotoaji (ikifikia kikomo cha FLOPS za Tensor Core).

Hatua A: Ufuatiliaji wa Kiwango cha Jumla kwa torch.profiler

Tumia PyTorch Profiler kubaini mapengo ya uanzishaji wa CPU, vipakiaji data vya polepole na kerneli za CUDA zinazotumia muda mwingi zaidi.

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

Hatua B: Uchambuzi wa Roofline wa Kiwango cha Ndani kwa NVIDIA Nsight Compute (ncu)

Baada ya torch.profiler kubaini kerneli ya polepole, endesha ncu kutoka kwenye terminali ili kukagua vihesabu vya utendaji wa maunzi moja kwa moja kwenye Streaming Multiprocessors (SMs):

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

Unapokagua Speed of Light (SOL) na Roofline Chart ndani ya kernel_report.ncu-rep, angalia namba mbili:

3. Utekelezaji wa GPU kwenye Vifaa vya Pembeni & vya Mkononi: Apple Metal & Vulkan Compute

Kusambaza mitandao ya neva kwenye simu mahiri, kompyuta mpakato na vifaa vilivyopachikwa kunahitaji kuendesha bila viendeshi vya NVIDIA wala matumizi makubwa ya nishati, badala yake kutegemea Apple Metal na Vulkan Compute zenye Usanifu wa Kumbukumbu Iliyounganishwa (UMA).

• Usanifu wa Kumbukumbu Iliyounganishwa (Utekelezaji Bila Kunakili).

• Tofauti na GPU tofauti za NVIDIA zinazohitaji kunakili tensori kupitia PCIe kutoka RAM ya CPU hadi VRAM, Apple Silicon (M-series / A-series) na SoC za vifaa vya mkononi (Qualcomm Snapdragon, ARM Mali) hushiriki hifadhi moja halisi ya kumbukumbu ya LPDDR5/LPDDR5X kati ya CPU, GPU na NPU.
• Mifumo kama Apple MLX, llama.cpp na ExecuTorch huweka uzani wa modeli moja kwa moja kwenye nafasi ya anwani inayoshirikiwa (MTLResourceStorageModeShared), na kuondoa kabisa muda wa kusubiri wa uhamishaji kutoka kwa mwenyeji hadi kwenye kifaa.

• Apple Metal & Metal Performance Shaders (MPS).

• Kwenye macOS na iOS, Metal huchukua nafasi ya CUDA. Wasanidi programu huandika shader za ukokotoaji katika Metal Shading Language (MSL), lugha iliyotokana na C++14, wakituma threadgroup kwenye viini vya GPU vya Apple na maagizo ya Simdgroup Matrix (sawa ya Apple ya NVIDIA Tensor Cores).
• Katika PyTorch, kuweka device="mps" huelekeza operesheni za tensori kupitia kikusanyaji cha MPSGraph cha Apple, ambacho huunganisha operesheni kwa namna inayobadilika na kupanga kazi kwenye vitengo vya utekelezaji vya GPU.

• Vulkan Compute ya Majukwaa Tofauti & SPIR-V.

• Katika Android, Windows, Linux na GPU za AMD/Intel/Qualcomm, Vulkan Compute ni mfumo wa nyuma wa GPU wa jumla usiofungamana na mtengenezaji mmoja.
• Kerneli za ukokotoaji (zilizoandikwa katika GLSL, HLSL, au zilizokusanywa kutoka MLIR/IREE) hukusanywa kuwa bytecode ya SPIR-V. Kwa kutumia viendelezi vya Cooperative Matrix (VK_KHR_cooperative_matrix), kerneli za Vulkan hutumia moja kwa moja vitengo vya maunzi vya kuzidisha matriki katika GPU za Qualcomm Adreno, ARM Mali, AMD RDNA na Intel Arc kwa kutumia faili moja ya binary inayohamishika.