torch.compile в PyTorch 2.x — реальное ускорение или только на бумаге?

[lev4334]

Перехожу с PyTorch 2.1 на 2.4, начитался про torch.compile и что он даёт 30-40% ускорение тренировки. Попробовал на своей модели (трансформер для NER, ~110M параметров, обучаю на A100) — получил ускорение 8% и кучу предупреждений про graph breaks. Что я делаю не так? Или это реально только для определённых архитектур работает? Использую просто `model = torch.compile(model)` без параметров.

[pavel_sys]

@db_admin, 8% это нормально для первого запуска с дефолтными настройками. Проблема в graph breaks — каждый break это место где компилятор вынужден вернуться в обычный Python, и весь прирост от компиляции этого сегмента теряется. Запусти `TORCH_LOGS=graph_breaks python train.py` чтобы увидеть где ломается граф. Самые частые причины: data-dependent control flow (if tensor.item() > 0), операции с изменяемыми размерами (dynamic shapes), кастомные CUDA кернелы без torch.library.

[nastya_sys61]

@indie_hacker, Для трансформеров рекомендую `torch.compile(model, mode='max-autotune', fullgraph=True)`. fullgraph=True падает с ошибкой если есть graph breaks — это заставляет тебя их найти и убрать. На чистых трансформерных архитектурах без хитрых трюков это обычно достижимо. У меня на аналогичной модели (~130M) после устранения всех breaks получилось +31% к throughput на A100.

[vlad_rust]

Ещё важно: torch.compile очень любит статичные размеры батчей и длин последовательностей. Если у тебя dynamic padding (батчи разного размера) — compile будет перекомпилировать граф под каждый новый shape, первые несколько батчей будут медленными. Либо используй `dynamic=True` в compile (меньший прирост но без рекомпиляций), либо бакетизируй последовательности по длине чтобы shapes были предсказуемы.

[gitvue9120]

Конкретные цифры из нашей практики на A100 80GB: BERT-large fine-tuning без compile — 142 samples/sec, с compile mode='reduce-overhead' — 168 samples/sec (+18%), с mode='max-autotune' и устранёнными breaks — 198 samples/sec (+39%). Compile с первым запуском тратит 3-5 минут на компиляцию, зато потом быстро. На коротких задачах не стоит.

[omegadata7864]

@rust_lover, Отдельная история с compile + FSDP для multi-GPU — там свои нюансы. Порядок важен: сначала FSDP wrap, потом compile. `model = torch.compile(FSDP(model))` а не наоборот. Иначе compile видит FSDP internals и ломается. В 2.4 это починили частично но всё ещё есть edge cases.

[kirill_backend]

@db_admin, Итого: torch.compile реально работает, но требует вложений — надо профилировать breaks и настраивать. На трансформерах без экзотики можно получить 25-40% на A100. На V100/старом железе прирост меньше. Если модель маленькая или датасет небольшой — время на компиляцию не отобьётся. Стоит включать когда обучение длится часы и более.