FSDP2, phân tích hiệu năng GPU và điện toán di động

Kỹ thuật AI trong sản xuất kết nối việc thực thi đa GPU phân tán, phân tích hiệu năng ở mức phần cứng, và các runtime di động/edge đa nền tảng. Dưới đây là phần phân tích sâu về cả ba lĩnh vực.

1. Huấn luyện phân tán đa GPU với PyTorch FSDP2

FSDP2 (fully_shard) hiện đại của PyTorch phân mảnh các tham số mô hình, gradient và trạng thái optimizer trên nhiều GPU theo từng tham số bằng DTensor, loại bỏ việc nhân bản bộ nhớ đồng thời chồng lấp giao tiếp với tính toán.

Dưới đây là một script huấn luyện đa GPU hoàn chỉnh. Chạy nó trên 4 GPU bằng 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()

• Điều gì xảy ra bên trong: Trước mỗi forward pass của một layer, FSDP phát một lệnh NCCL All-Gather bất đồng bộ qua NVLink để tái tạo toàn bộ trọng số bfloat16 riêng cho block đó, chạy phép tính, rồi lập tức hủy các trọng số đã ghép lại. Trong loss.backward(), các gradient được lấy trung bình và phân chia lại giữa các GPU qua Reduce-Scatter.

2. Phân tích hiệu năng một mô hình PyTorch: PyTorch Profiler & NVIDIA Nsight Compute (ncu)

Tối ưu hóa code GPU đòi hỏi phải biết một kernel bị giới hạn bởi băng thông bộ nhớ (Memory-Bandwidth Bound - bị đình trệ chờ đọc/ghi HBM) hay giới hạn bởi tính toán (Compute Bound - bão hòa FLOPS của Tensor Core).

Bước A: Trace mức vĩ mô với torch.profiler

Sử dụng PyTorch Profiler để phát hiện các khoảng trống khởi chạy trên CPU, các data loader chậm và các CUDA kernel tốn nhiều thời gian nhất.

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ước B: Phân tích Roofline mức vi mô với NVIDIA Nsight Compute (ncu)

Khi torch.profiler đã xác định được kernel chậm, chạy ncu từ terminal để kiểm tra trực tiếp các bộ đếm hiệu năng phần cứng trên các Streaming Multiprocessor (SM):

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

Khi kiểm tra Speed of Light (SOL) và Roofline Chart trong kernel_report.ncu-rep, hãy xem hai con số:

3. Thực thi GPU Edge & Di động: Apple Metal & Vulkan Compute

Triển khai mạng nơ-ron lên điện thoại thông minh, laptop và thiết bị nhúng đòi hỏi chạy mà không cần driver NVIDIA hay ngân sách điện năng lớn — thay vào đó dựa vào Apple Metal và Vulkan Compute với Kiến trúc Bộ nhớ Thống nhất (UMA).

• Kiến trúc Bộ nhớ Thống nhất (Thực thi Zero-Copy).

• Khác với các GPU NVIDIA rời cần sao chép tensor qua PCIe từ RAM CPU sang VRAM, Apple Silicon (dòng M / dòng A) và các SoC di động (Qualcomm Snapdragon, ARM Mali) dùng chung một vùng nhớ vật lý LPDDR5/LPDDR5X giữa CPU, GPU và NPU.
• Các framework như Apple MLX, llama.cpp và ExecuTorch ánh xạ trọng số mô hình trực tiếp vào không gian địa chỉ dùng chung (MTLResourceStorageModeShared), loại bỏ hoàn toàn độ trễ truyền host-to-device.

• Apple Metal & Metal Performance Shaders (MPS).

• Trên macOS và iOS, Metal thay thế CUDA. Lập trình viên viết compute shader bằng Metal Shading Language (MSL) - một biến thể của C++14 - phân phối threadgroup trên các nhân GPU của Apple và các lệnh Simdgroup Matrix (tương đương của NVIDIA Tensor Core trên Apple).
• Trong PyTorch, đặt device="mps" định tuyến các phép toán tensor qua trình biên dịch MPSGraph của Apple, trình này tự động hợp nhất các phép toán và lên lịch công việc trên các đơn vị thực thi của GPU.

• Vulkan Compute & SPIR-V đa nền tảng.

• Trên Android, Windows, Linux và các GPU AMD/Intel/Qualcomm, Vulkan Compute đóng vai trò backend GPU phổ quát, trung lập với nhà sản xuất.
• Các compute kernel (viết bằng GLSL, HLSL, hoặc biên dịch từ MLIR/IREE) được biên dịch thành bytecode SPIR-V. Nhờ các phần mở rộng Cooperative Matrix (VK_KHR_cooperative_matrix), các kernel Vulkan truy cập trực tiếp vào các đơn vị nhân ma trận phần cứng trên GPU Qualcomm Adreno, ARM Mali, AMD RDNA và Intel Arc bằng một binary di động duy nhất.