FSDP2, профилирование GPU и мобильные вычисления

Инженерия ИИ в производственной среде объединяет распределенное выполнение на нескольких GPU, профилирование на уровне оборудования и кроссплатформенные среды выполнения для мобильных и периферийных устройств. Ниже подробно рассматриваются все три области.

1. Распределенное обучение на нескольких GPU с PyTorch FSDP2

Современная система FSDP2 (fully_shard) в PyTorch с помощью DTensor распределяет параметры модели, градиенты и состояния оптимизатора между GPU на уровне отдельных параметров, устраняя дублирование памяти и совмещая обмен данными с вычислениями.

Ниже приведен полный скрипт обучения на нескольких GPU. Запустите его на 4 GPU с помощью 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()

• Что происходит внутри: перед прямым проходом каждого слоя FSDP запускает асинхронную операцию NCCL All-Gather по NVLink, чтобы восстановить полные веса bfloat16 только для этого блока, выполняет вычисления и сразу удаляет неразделенные веса. Во время loss.backward() градиенты усредняются и снова распределяются между GPU через Reduce-Scatter.

2. Профилирование модели PyTorch: PyTorch Profiler & NVIDIA Nsight Compute (ncu)

Для оптимизации кода GPU необходимо знать, ограничено ли ядро пропускной способностью памяти (простаивает в ожидании чтения/записи HBM) или вычислениями (полностью использует FLOPS тензорных ядер).

Шаг A: трассировка на макроуровне с torch.profiler

Используйте PyTorch Profiler для выявления пауз между запусками со стороны CPU, медленных загрузчиков данных и ядер CUDA, занимающих больше всего времени.

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: анализ Roofline на микроуровне с NVIDIA Nsight Compute (ncu)

Когда torch.profiler выявит медленное ядро, запустите ncu из терминала, чтобы исследовать аппаратные счетчики производительности непосредственно на потоковых мультипроцессорах (SM):

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

Изучая Speed of Light (SOL) и диаграмму Roofline в kernel_report.ncu-rep, обратите внимание на два показателя:

3. Выполнение на периферийных & мобильных GPU: Apple Metal & Vulkan Compute

Развертывание нейронных сетей на смартфонах, ноутбуках и встроенных устройствах требует работы без драйверов NVIDIA и огромных энергетических бюджетов; вместо них используются Apple Metal и Vulkan Compute с архитектурами объединенной памяти (UMA).

• Архитектура объединенной памяти (выполнение без копирования).

• В отличие от дискретных GPU NVIDIA, которым требуется копирование тензоров по PCIe из оперативной памяти CPU в VRAM, Apple Silicon (M-series / A-series) и мобильные SoC (Qualcomm Snapdragon, ARM Mali) используют единый физический пул памяти LPDDR5/LPDDR5X, общий для CPU, GPU и NPU.
• Такие фреймворки, как Apple MLX, llama.cpp и ExecuTorch, отображают веса модели непосредственно в общее адресное пространство (MTLResourceStorageModeShared), полностью устраняя задержку передачи с хоста на устройство.

• Apple Metal & Metal Performance Shaders (MPS).

• В macOS и iOS Metal заменяет CUDA. Разработчики пишут вычислительные шейдеры на Metal Shading Language (MSL), производном от C++14, распределяя группы потоков по ядрам GPU Apple и используя инструкции Simdgroup Matrix (аналог тензорных ядер NVIDIA у Apple).
• В PyTorch установка device="mps" направляет операции с тензорами через компилятор MPSGraph от Apple, который динамически объединяет операции и планирует работу исполнительных блоков GPU.

• Кроссплатформенные Vulkan Compute & SPIR-V.

• В Android, Windows, Linux и на GPU AMD/Intel/Qualcomm Vulkan Compute выступает универсальным, независимым от поставщика бэкендом GPU.
• Вычислительные ядра (написанные на GLSL, HLSL или скомпилированные из MLIR/IREE) компилируются в байт-код SPIR-V. С помощью расширений Cooperative Matrix (VK_KHR_cooperative_matrix) ядра Vulkan напрямую обращаются к аппаратным блокам матричного умножения GPU Qualcomm Adreno, ARM Mali, AMD RDNA и Intel Arc, используя один переносимый двоичный файл.