§Контекст главы

Тринадцатая глава книги Fregly посвящена практическому профилированию, тюнингу и масштабированию рабочих нагрузок PyTorch на GPU. Глава охватывает полный цикл: от инструментов профилирования (PyTorch Profiler, Nsight Systems, Nsight Compute, Linux perf) и методологии выявления узких мест до конкретных техник оптимизации (torch.compile, CUDA-графы, тюнинг аллокатора памяти, activation checkpointing, оптимизация конвейера данных) и масштабирования на несколько GPU и узлов (DDP, FSDP, tensor parallel, pipeline parallel). Завершается глава обсуждением непрерывного бенчмаркинга и защиты производительности от регрессий.

INVTX-маркеры и инструменты профилирования

NVTX-маркеры вставляются в код через CUDA C++, PyTorch или любую C++/Python-библиотеку, поддерживающую GPU NVIDIA. Многие библиотеки уже вставляют NVTX-маркеры за разработчика для критических участков кода, таких как train_step, forward, backward, optimizer_step. Но их можно добавлять и самостоятельно с помощью torch.profiler.record_function() и torch.cuda.nvtx.range_push().

IIПрофилирование PyTorch: пример MoE

В главе профилирование демонстрируется на примере модели трансформера со смесью экспертов (MoE). MoE — это LLM с несколькими экспертными слоями, где каждый эксперт представляет собой сеть прямого распространения. Маршрутизация токенов к экспертам управляется системой гейтинга. Архитектура богата на оптимизационные возможности: множество matmul среднего размера, кастомные ядра маршрутизации (dispatch/combine) и динамическое распределение токенов.

Использование PyTorch Profiler

Сначала настраивается модель и входные данные. Используется Hugging Face Transformers для загрузки модели и токенизатора, модель перемещается на GPU, подготавливается небольшой батч:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "..."
tokenizer = AutoTokenizer.from_pretrained(model_name)
device = torch.device("cuda")
model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, fused=True)

batch_size = 4
input_texts = ["MoEs are great."] * batch_size
enc = tokenizer(input_texts, return_tensors="pt", padding=True, truncation=True)
input_ids = enc.input_ids.to(device)
attention_mask = enc.attention_mask.to(device)
labels = input_ids.clone()

Чтобы не захватывать разовые затраты на настройку, модель прогревается несколькими итерациями перед профилированием. Это подготовит модель к анализу: скомпилирует JIT-ядра, заполнит кэши.

for _ in range(5):
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)

Теперь профилируется одна итерация обучения, высокоуровневые регионы помечаются через record_function и NVTX-диапазоны:

from torch import profiler

with profiler.profile(
    activities=[profiler.ProfilerActivity.CPU, profiler.ProfilerActivity.CUDA],
    record_shapes=True,
    profile_memory=True,
    with_stack=True,
    with_flops=True
) as prof:
    with profiler.record_function("train_step"):
        torch.cuda.nvtx.range_push("forward")
        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
            outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
            loss = outputs.loss
        torch.cuda.nvtx.range_pop()

        torch.cuda.nvtx.range_push("backward")
        loss.backward()
        torch.cuda.nvtx.range_push("optimizer_step")
        optimizer.step()
        torch.cuda.nvtx.range_pop()
        optimizer.zero_grad()
        torch.cuda.nvtx.range_pop()

После выполнения можно изучить результаты на уровне операций, вызвав prof.key_averages().table(). Запрашивается топ-10 операций, отсортированных по self_cuda_time_total:

print(
    prof.key_averages()
    .table(
        sort_by="self_cuda_time_total",
        row_limit=10,
        fields=["self_cuda_time_total", "calls"]
    )
)

Системное профилирование с Nsight Systems и NVTX

Вставленные NVTX-маркеры делают анализ временной шкалы в Nsight Systems простым. Чтобы агрегировать метрики по фазам, код профилируется через nsys со сводкой на основе NVTX:

nsys profile \
    --output=profile \
    --stats=true \
    -t cuda,nvtx \
    python train.py

nsys stats --report=nvtx_gpu_proj_sum profile.nsys-rep

Механизм проекции NVTX — это не точная привязка ядра к диапазону, а статистическая оценка. CUPTI фиксирует, когда CPU вызвал push/pop NVTX-маркера, и какие GPU-ядра были запущены в этом временном окне на хосте.

IIIRoofline-анализ GEMM-ядер

Для глубокого анализа экспертного matmul вызывается CLI-профайлер Nsight Compute (ncu), нацеленный на конкретные GEMM-ядра по имени:

ncu \
    --target-processes all \
    --kernel-name-regex "matmul" \
    --metrics \
        gpu__time_duration.avg, \
        gpu__dram_throughput.avg.pct_of_peak_sustained_elapsed, \
        lts__throughput.avg.pct_of_peak_sustained_elapsed, \
        sm__sass_thread_inst_executed_op_fp32_pred_on.sum, \
        sm__warps_active.avg.pct_of_peak_sustained_active \
    --csv full \
    -o matmul_roofline_report \
    python train.py

В базовом запуске GEMM-ядро достигает лишь 50% пиковых FLOPS, 70% пиковой пропускной способности памяти и умеренной 60% SM occupancy. Это недостаточно для полного скрытия задержек памяти. Метрики указывают на memory-bound ядро. После оптимизации (слияние ядер, рост арифметической интенсивности, снижение перемещения памяти) пиковые FLOPS растут до 85%, пропускная способность памяти падает до 40%, а occupancy растёт до 80%. Ядро сдвигается из memory-bound в compute-bound.

IVCPU-профилирование через Linux perf

Для более целостной картины того, где тратится время — на хосте и на устройстве — применяется Linux perf для анализа циклов CPU, кэш-промахов, branch-misses и т. д. Затем эти данные используются для серии оптимизаций.

perf stat -e \
    cycles,instructions,cache-misses,branch-misses \
    python train.py

perf record -F 2000 -g --call-graph dwarf -o perf.data \
    python train.py

perf report --stdio -n -g -i perf.data

На основе профиля выявляются следующие узкие места и меры их устранения:

• Избыточные накладные расходы Python (45% в py::forward) — использовать torch.compile для устранения накладных расходов интерпретатора и слияния Python-операций в оптимизированный CUDA-код.
• Крупное горячее место matmul (20.5%) — PyTorch Compiler или кастомное CUDA C++-ядро для обхода Python и прямого использования оптимизированного CUDA-кода.
• Простои загрузки данных (10.2%) — увеличить num_workers DataLoader, включить persistent_workers=True.
• Накладные расходы синхронизации градиентов (8.7%) — увеличить размер бакета градиентов в DDP с 25 МБ до 50 МБ.
• Узкое место I/O (5.3% в read) — pinned memory, non-blocking копии, контейнерные форматы (Arrow, WebDataset, TFRecord).

VPyTorch Compiler (torch.compile)

Использовать компилятор с настройками по умолчанию просто: model = torch.compile(model). Под капотом TorchDynamo трассирует Python-код, AOT Autograd захватывает backward-проход, а TorchInductor, опираясь на OpenAI Triton для генерации GPU-ядер, автоматически создаёт эффективные fused-ядра.

compiled_model = torch.compile(model, mode="reduce-overhead")

В этом случае eager-режим занимает примерно 248 мс на итерацию, а скомпилированный — около 173 мс. Ускорение примерно на 30% в основном от объединения множества мелких GPU-ядер.

Профилирование и отладка компилятора

Когда torch.compile не даёт ожидаемого ускорения, используются инструменты диагностики. torch._dynamo.explain(model) выводит отчёт о graph breaks, причинах их возникновения и о том, какие части модели не были захвачены TorchDynamo. Переменная TORCH_LOGS="+dynamo,+inductor" включает подробное логирование компилятора.

import torch._dynamo as dynamo

def toy_example(a, b):
    x = a / (torch.abs(a) + 1)
    print("woo")
    if b.sum() < 0:
        b = -b
    return x * b

explanation = dynamo.explain(toy_example)(torch.randn(10), torch.randn(10))
print(explanation)

Отчёт покажет количество графов, количество graph breaks и причины их возникновения. Для минимизации graph breaks используется torch.where() вместо Python-условий, а также torch.compiler.set_stance("fail_on_recompile") для принудительного выброса ошибки при обнаружении graph break.

VIОптимизированные механизмы внимания

Трансформерные модели значительную часть времени проводят в механизмах внимания. PyTorch предоставляет несколько уровней оптимизации:

• Scaled Dot-Product Attention (SDPA) — torch.nn.functional.scaled_dot_product_attention автоматически использует самое быстрое доступное ядро (FlashAttention, memory-efficient или cuDNN) в зависимости от форм тензоров и типов данных.
• FlexAttention — компиляторный подход для кастомных паттернов разреженности (block-sparse, sliding-window). Позволяет описывать маску как функцию, а компилятор генерирует оптимизированные ядра.
• FlexDecoding — оптимизация фазы декодирования для LLM. Компилирует специализированное ядро для декодирования (Q_len=1) с учётом растущего KV-кэша.
• Context Parallel — шардирование внимания по длине последовательности для обработки сверхдлинных контекстов.

VIIКонкурентность с CUDA-потоками

CUDA-потоки обеспечивают конкурентность и перекрытие операций на GPU. По умолчанию PyTorch планирует все операции на стандартном потоке устройства (stream 0) последовательно. Однако многие задачи независимы, и GPU может выполнять их параллельно с несколькими потоками.

device = 'cuda'
transfer_stream = torch.cuda.Stream(device=device)
compute_stream = torch.cuda.default_stream(device=device)

dataloader_iter = iter(dataloader)

first_batch = next(dataloader_iter, None)
if first_batch:
    with torch.cuda.stream(transfer_stream):
        next_inputs, next_labels = (
            first_batch[0].to(device, non_blocking=True),
            first_batch[1].to(device, non_blocking=True),
        )

for _ in range(len(dataloader)):
    compute_stream.wait_stream(transfer_stream)
    inputs, labels = next_inputs, next_labels

    batch = next(dataloader_iter, None)
    if batch:
        with torch.cuda.stream(transfer_stream):
            next_inputs, next_labels = (
                batch[0].to(device, non_blocking=True),
                batch[1].to(device, non_blocking=True),
            )

    with torch.cuda.stream(compute_stream):
        outputs = model(inputs)
        loss = loss_fn(outputs, labels)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)

Синхронизация потоков через события

event = torch.cuda.Event(enable_timing=False)

with torch.cuda.stream(stream1):
    kernel_launch(...)
    event.record()

stream2.wait_event(event)

with torch.cuda.stream(stream2):
    other_kernel_launch(...)

VIIICUDA-графы для снижения накладных расходов

CUDA-графы устраняют накладные расходы на запуск на итерацию, снижают накладные расходы CPU и ликвидируют крошечные простои между ядрами. PyTorch предоставляет API torch.cuda.CUDAGraph для захвата и воспроизведения CUDA-графов.

g = torch.cuda.CUDAGraph()
capture_stream = torch.cuda.Stream()

static_input = torch.randn(batch_shape, device='cuda')
static_output = torch.empty(output_shape, device='cuda')

with torch.cuda.stream(capture_stream):
    tmp = model(static_input)
    static_output.copy_(tmp)

capture_stream.synchronize()

with torch.cuda.graph(g, stream=capture_stream):
    tmp = model(static_input)
    static_output.copy_(tmp)

capture_stream.synchronize()

CUDA Graph Trees используются torch.compile в режиме reduce-overhead для компиляции и кэширования отдельных статических графов под каждую форму входа. Рекомендуется поддерживать формы входов максимально согласованными между итерациями.

IXТюнинг CUDA Memory Allocator

PyTorch использует кэширующий аллокатор памяти для CUDA. По умолчанию он адаптируется к паттернам аллокации, разбивая и переиспользуя блоки GPU-памяти. Ключевые параметры настройки через переменные окружения:

• PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:N — максимальный размер блока, который аллокатор может разделять. Указывает аллокатору сохранять большие свободные блоки (до 256 МБ) вместо постоянного дробления на мелкие фрагменты.
• PYTORCH_CUDA_ALLOC_CONF=roundup_power2_divisions:[N:M,...] — группировка запросов в фиксированные бакеты. Например, в диапазоне 512–1024 МБ с 512:2 запрос в 600 МБ округляется до 768 МБ. Это снижает фрагментацию и увеличивает переиспользование.
• PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync — асинхронный аллокатор CUDA. Устраняет глобальные синхронизации при аллокации и освобождении памяти.

export PYTORCH_CUDA_ALLOC_CONF=\
max_split_size_mb:256,\
roundup_power2_divisions:[256:1,512:2,1024:4,:8],\
backend:cudaMallocAsync

XActivation Checkpointing для экономии памяти

Для очень больших моделей activation checkpointing (gradient checkpointing) необходим для управления памятью. Вместо сохранения всех активаций forward-прохода сохраняются только входы каждого блока. При backward-проходе активации пересчитываются заново.

Не обязательно чекпоинтить всё. Можно сфокусироваться на самых больших слоях — обычно это блоки трансформера, которые держат массу активаций, но не меньшие слои вроде layer norm и эмбеддингов.

import torch
from torch.utils.checkpoint import checkpoint

class TransformerBlock(torch.nn.Module):
    def forward(self, x):
        return checkpoint(self._forward, x, use_reentrant=False)

    def _forward(self, x):
        return self.attention(self.mlp(x))

XIОптимизация входного конвейера данных

Частая причина неэффективности — GPU простаивают в ожидании данных. DataLoader PyTorch поддерживает запуск нескольких воркеров для параллельной загрузки и предобработки данных, пока обучается модель.

• pin_memory=True — передача через page-locked память. Ускоряет копирование CPU→GPU.
• non_blocking=True и persistent_workers=True — для скрытия копирования и избежания пересоздания воркеров между эпохами.
• prefetch_factor — сколько батчей каждый воркер предзагружает заранее. По умолчанию 2.

loader = DataLoader(
    dataset,
    batch_size=B,
    num_workers=8,
    pin_memory=True,
    persistent_workers=True,
    prefetch_factor=4,
)

XIIFSDP: автоматический checkpointing и offloading

FSDP в PyTorch — стратегия распределённого параллелизма, шардирующая параметры модели, градиенты и активации между GPU. Технически FSDP реализует стратегию ZeRO Stage-3 для шардирования состояний модели между GPU.

import torch
import torch.nn as nn
import torch.distributed as dist
from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP,
    CPUOffload,
    ShardingStrategy,
    MixedPrecision
)
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy

dist.init_process_group("nccl")
torch.cuda.set_device(dist.get_rank() % torch.cuda.device_count())

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(4096, 4096),
            nn.ReLU(),
            nn.Linear(4096, 4096),
        )

    def forward(self, x):
        return self.layers(x)

model = MyModel().cuda()

auto_wrap_policy = transformer_auto_wrap_policy(
    model,
    min_num_params=1e8,
)

fsdp_model = FSDP(
    model,
    auto_wrap_policy=auto_wrap_policy,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
    use_orig_params=True,
    cpu_offload=CPUOffload(offload_params=True, pin_memory=True),
)

Когда сеть медленная — или доступно всего несколько GPU — можно использовать ShardingStrategy.SHARD_GRAD_OP (ZeRO Stage-2) для шардирования только градиентов и состояния оптимизатора. На суперчиповых системах Grace-Blackwell доступны SuperOffload и GraceAdam.

XIIIМасштабирование с PyTorch Distributed

DDP с torch.compile

При использовании DistributedDataParallel, синхронизирующего градиенты между GPU через all-reduce, PyTorch автоматически создаёт graph breaks в точках синхронизации. Дизайн PyTorch — компилировать backward-вычисления каждого бакета как отдельный граф, чтобы между этими графами выполнять all-reduce.

FSDP с torch.compile

Рекомендуется оборачивать каждый блок трансформера — а не просто отдельный низкоуровневый слой — в собственный FSDP-модуль с use_orig_params=True. Так TorchDynamo вставит graph break на каждой границе шарда, где нужна коммуникация.

Tensor и pipeline parallel с torch.compile

Tensor parallel и pipeline parallel ортогональны torch.compile. Компилятор PyTorch в основном фокусируется на оптимизации вычислений внутри каждого сегмента. Рекомендуется использовать torch.compile для оптимизации вычислений внутри каждого слоя или бакета.

TorchTitan, AsyncTP, AutoParallel и SimpleFSDP

TorchTitan — набор референсных реализаций для крупномасштабного обучения. SimpleFSDP перереализует FSDP дружественным к torch.compile способом через DTensor. В экспериментах TorchTitan SimpleFSDP снизил потребление памяти до ~28%, а пропускную способность улучшил до ~69% по сравнению с традиционным FSDP2 eager-путём.

XIVМного-GPU профилирование с HTA

При масштабировании на множество GPU управление трассировками может стать громоздким. HTA от Meta помогает объединять и визуализировать многорабочие трассировки. HTA поглощает JSON-трассировки от torch.profiler с каждой GPU/ранга и представляет единый временной график.

С HTA можно увидеть трассировки всех 8 GPU, выровненные по времени. NVTX-маркеры с каждого ранга выравниваются и визуализируются. Так можно заметить, что ранг 0 входит в backward-проход в момент T, а ранг 1 — только в T+1, возможно, потому что ранг 1 ждал ранг 0 в all-reduce.

from torch.profiler import profile, schedule, tensorboard_trace_handler

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=schedule(wait=1, warmup=1, active=3, repeat=1),
    on_trace_ready=tensorboard_trace_handler('./log'),
    record_shapes=True,
    profile_memory=True,
    with_stack=True
) as prof:
    for step, batch in enumerate(dataloader):
        train_step(batch)
        prof.step()

XVCI и бенчмаркинг

После применения всех оптимизаций критически важно поддерживать их и непрерывно проверять на регрессии производительности. Следует настроить простую CI регрессии производительности с TorchBench и GitHub Actions для автоматического отлова замедлений.

- name: Run MoE benchmark
  run: |
    torchbench run --model moe --iters 10 --batch-size 4 --json results.json

- name: Compare throughput
  run: |
    python scripts/compare_perf.py baseline.json results.json

XVIПроверка прогнозов

XVIIДополнительные замечания

XVIIIКлючевые положения

1. NVTX-маркеры — единый язык профилирования. PyTorch Profiler, Nsight Systems и Nsight Compute понимают NVTX. Это делает кросс-инструментальный анализ сквозным.
2. Профилирование выявляет узкие места. В MoE-модели главное узкое место — matmul в экспертных FFN, за ним — dispatch и combine.
3. Roofline-анализ показывает memory-bound vs compute-bound. Базовая GEMM-ядро MoE — memory bound (50% FLOPS, 70% памяти). После оптимизаций — compute bound (85% FLOPS, 40% памяти).
4. Linux perf дополняет GPU-профилирование. 45% накладных расходов Python, 8.7% — на синхронизацию градиентов, 10.2% — на загрузку данных.
5. torch.compile даёт ~30% ускорения. Слияние мелких GPU-ядер и устранение накладных расходов Python. Режим reduce-overhead использует CUDA Graphs.
6. CUDA-графы снижают задержку итерации на 25%. Устраняют накладные расходы запуска ядер. CUDA Graph Trees кэшируют графы под каждую форму входа.
7. Аллокатор памяти тюнится через max_split_size_mb, roundup_power2_divisions и backend:cudaMallocAsync.
8. Activation checkpointing обменивает вычисления на память. Селективный подход позволяет выбрать оптимальный компромисс.
9. DataLoader: три ключевых параметра — pin_memory=True, persistent_workers=True, prefetch_factor.
10. FSDP реализует ZeRO Stage-3. Шардирование на уровне блока трансформера даёт максимальную эффективность. SimpleFSDP снижает память до ~28% и увеличивает пропускную способность до ~69%.
11. Graph breaks в DDP и FSDP ожидаемы. Компиляция каждого бакета отдельно — дизайн, а не баг.
12. HTA агрегирует трассировки всех рангов. Позволяет выявлять дисбаланс и синхронизационные задержки в multi-GPU обучении.
13. CI с TorchBench выявляет регрессии. Устойчивое падение пропускной способности ≥5% на трёх прогонах — сигнал для расследования.
14. PyTorch HUD — публичный дашборд. Ночные бенчмарки и отслеживание регрессий по коммитам.

Источники: Chris Fregly, «AI Systems Performance Engineering» (O'Reilly, 2026), Глава 13 — цитаты приводятся в ограниченном объёме в целях обзора и критического анализа. Официальная документация PyTorch (torch.profiler, torch.compile, FSDP, CUDA Graphs); NVIDIA Nsight Systems/Compute documentation; Holistic Trace Analysis (Meta); MLPerf Training/Inference v5.0; TorchBench; PyTorch HUD.