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

[nodice]

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

[sainty]

@nodice, 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.

[coder_vlad]

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

[solidity_nerd]

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

[krayzie]

Конкретные цифры из нашей практики на 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 минут на компиляцию, зато потом быстро. На коротких задачах не стоит.

[davkar]

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

[Omoto]

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