Профилирование PyTorch: полное руководство по torch.profiler для начинающих

Не верь метрикам — верь профайлеру

Ты запустил обучение. GPU грузится на 30%. Loss ползет, но что-то не так. Ты добавляешь workers в dataloader, меняешь batch size, ставишь amp — ноль реакции. Знакомо? Твоя модель — не ты. Она не скажет, где у нее болит. Придется вскрывать. torch.profiler — это скальпель.

GPU не умеет говорить: «Эй, я простаиваю, потому что dataloader не успевает подавать данные». Он просто ждет. А ты теряешь часы и электричество. Profiler раскладывает весь pipeline по полочкам: сколько времени заняла forward, backward, сколько — загрузка данных, сколько — ожидание коммуникаций в распределенке.

Мы не будем читать сотни страниц документации. Я покажу три реальных сценария, где profiler вытаскивает узкие места, которые не видны глазом. А заодно разберем, где новички наступают на грабли (спойлер: я тоже наступал).

Что внутри torch.profiler?

Это не просто хронометр. Это полноценный трекер операций CUDA, CPU, копирований, ядер, вызовов библиотек — вплоть до отдельных kernel launch. Он показывает:

• CPU wall time — сколько времени процессор тратит на подготовку данных и запуск операций.
• CUDA time — сколько GPU реально считает.
• Self time — чистое время операции без вложенных вызовов.
• Memory usage — выделение и освобождение памяти на GPU, утечки.
• Stack trace — какая строчка кода вызвала тормоза.
• Input shapes — размер тензоров на входе (новое в PyTorch 2.7 — при профилировании с record_shapes=True).

Первое вскрытие: как запустить profiler за 30 секунд

Забудь про декорирование всего подряд. Минимальный профайлинг выглядит так:

import torch
import torch.profiler

def train_step(model, data, target):
    output = model(data)
    loss = torch.nn.functional.cross_entropy(output, target)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ],
    record_shapes=True,
    with_stack=True
) as prof:
    train_step(model, input_tensor, target_tensor)

print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

Этот код покажет топ-10 операций по времени на GPU. Если datascience кунг-фу не работает, иди сюда. Ты увидишь, что не forward жрет время, а какая-нибудь to(device) в цикле, о которой ты забыл.

Продвинутый уровень: Chrome Trace и Flame Graph

Таблицы — это хорошо. Но когда у тебя 200 слоев, глаза замыливаются. Profiler умеет экспортировать трассу в формате Chrome Trace (файл .json). Открой его в chrome://tracing — и ты увидишь временную линию операций, наложения, пустующие полоски.

prof.export_chrome_trace("trace.json")

Что ищешь в трейсе?

• Большие зазоры между CPU и CUDA. Если после cpu_op идет долгая пауза перед cudaLaunchKernel, значит CPU занят чем-то другим (например, загрузкой данных).
• Частые синхронизации. torch.cuda.synchronize() в непредвиденных местах — убийца производительности.
• Переключения контекста. Если dataloader blocks CPU на время копирования — это признак того, что num_workers не хватает или используется pin_memory=False.

Flame graph (с помощью prof.export_stacks() и FlameGraph плагина) даст иерархию вызовов. Ты сразу увидишь, что 40% времени уходит на torch.nn.functional.linear, а не на активации.

Кейс 1: Dataloader — корень зла

Типичная картина: GPU utilization 20%, CPU utilization 100%. Profiler показывает, что DataLoader.__iter__().next() жрет CPU время, а между батчами — длинные провалы. Решение — увеличить num_workers, добавить pin_memory, перейти на torchdata или использовать асинхронную загрузку.

# Проверка: замени train_step на замер dataloader отдельно
dataloader = DataLoader(dataset, batch_size=64, num_workers=4, pin_memory=True)

with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU]) as prof:
    for batch in dataloader:
        pass

print(prof.key_averages().table(row_limit=10))

Кейс 2: Распределенное обучение — коммуникации крадут время

Когда модель обучается на 8 GPU, часть времени тратится на all-reduce градиентов. Profiler показывает, как долго висят nccl:all_reduce или nccl:all_gather. Если эти коллы занимают >30% времени, нужно менять стратегию: переходить на gradient accumulation, увеличивать batch size на GPU, или использовать sharded optimizer (FSDP).

# Пример для DDP:
with torch.profiler.profile(
    activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA],
    profile_memory=True,
    with_stack=True
) as prof:
    for i, (data, target) in enumerate(train_loader):
        if i >= 5: break  # первые 5 батчей для прогрева
        output = model(data)
        loss = loss_fn(output, target)
        loss.backward()
        optimizer.step()

prof.export_chrome_trace("ddp_trace.json")

В трейсе ты увидишь, что после backward запускается torch.distributed.all_reduce. Если он длится слишком долго — попробуй torch.distributed.algorithms.join или бампни bucket_cap_mb в DDP. Более глубокий разбор коммуникаций — в статье Масштабирование обучения нейросетей: production-ready pipeline на PyTorch DDP.

Кейс 3: Профилирование LLM и Torch Compile

LLM (крупные языковые модели) — отдельная боль. Их инференс часто bottleneck'ится на attention, особенно если не использовать FlashAttention или не компилировать модель. В PyTorch 2.7 profiler умеет показывать Triton kernels, которые генерирует torch.compile.

model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-3.2-8B").eval()
model = torch.compile(model, mode="reduce-overhead")

input_ids = torch.randint(0, 32000, (1, 512)).cuda()

with torch.profiler.profile(
    activities=[torch.profiler.ProfilerActivity.CUDA],
    with_stack=True,
    record_shapes=True
) as prof:
    with torch.no_grad():
        output = model(input_ids)

prof.export_chrome_trace("llm_trace.json")

Открываешь trace — видишь, что операция triton_attention занимает 70% времени. Если у тебя не стоит use_flash_attention=True или ты используешь устаревшую библиотеку — самое время обновить transformers и установить flash-attn.

Если после компиляции что-то замедлилось (а такое бывает), profiler покажет, какие операции стали тормозом. Чаще всего проблема в нестабильных shapes или перегруженной рекомпиляции. Подробнее про рекомпиляцию — в статье Почему кастомные CUDA-ядра не дают ускорения в реальном обучении.

Пять граблей, на которые наступают все новички

1 Не прогревать модель перед замером

Если начать профилирование сразу, ты попадешь на CUDA kernel compilation и кеширование. Первый батч всегда медленнее. Правило: сделай хотя бы 2-3 итерации без профилирования, потом стартуй запись.

2 Слишком короткий захват

Если профилировать одну итерацию — статистика шумная. Лучше захватывать 5-10 итераций и смотреть среднее. Используй schedule с wait, warmup, active, repeat.

3 Игнорирование memory profiling

Без profile_memory=True ты не увидишь, какие операций аллоцируют память. Бывает, что случайная .expand() создает гигантский тензор и продавливает лимит CUDA memory.

4 Верить таблице без сортировки

По умолчанию key_averages().table() сортирует по CPU time. А ты смотришь на GPU. Всегда указывай sort_by="cuda_time_total" или self_cuda_time_total.

5 Не учитывать оверхед профайлера

Сам профайлер жрет ресурсы. Абсолютные цифры могут быть завышены на 10-30%. Смотри на относительные доли: какая операция занимает самый большой процент.

Когда profiler не поможет? (и что тогда делать)

Есть вещи, которые profiler не схватывает: системные вызовы ввода-вывода, проблемы сети при распределенке, кэши процессора. Для этого есть другие утилиты:

• NVIDIA Nsight Systems — для глобального взгляда на pipeline (CPU+GPU+коммуникации). Я сравнивал его с profiler в статье NVIDIA Nsight vs PyTorch Profiler.
• traceML — для детекции утечек памяти и простоев dataloader на пролонгированных трейнах.
• NCCL benchmarks — для изолированного теста коммуникаций.

Если у тебя TPU — profiler там другой. Глянь гайд Easy-torch-tpu: обучение PyTorch на TPU — там профайлер XLA.

Сводная таблица параметров profiler

Неочевидный совет напоследок

Не верь profiler'у на слово — он врет, но предсказуемо. Запусти профилирование дважды. Если цифры сильно расходятся (<20%), значит система зашумлена: выключи браузер, музыку, дашборды. А еще лучше — зафиксируй seed и используй torch.set_num_threads(1) для детерминизма. Профилирование — это искусство задавать правильные вопросы. Теперь ты знаешь, как их задавать.

Подписаться на канал