Почему inference в torch.compile тормозит на первом батче и как прогревать

[sabaza]

Поставил torch.compile на сервинг, и первый запрос после старта пода летит 14 секунд вместо обычных 40мс. Дальше всё норм, но в k8s с readinessProbe это боль, под считается живым и в него сразу прилетает трафик, который висит. Torch 2.5.1, режим reduce-overhead. Кто как лечит этот первый батч?

[joseph23]

так это же компиляция графа в рантайме, оно при первом проходе jit-ит ядра. ничего удивительного. прогревай фиктивным прогоном на старте до того как откроешь healthcheck

[tolkien]

прогревай фиктивным прогоном на старте до того как откроешь healthcheck

это работает только если у тебя один фиксированный размер входа. у меня динамические длины последовательностей, и torch.compile рекомпилит граф на каждую новую форму. прогрел на 128 токенов, пришёл запрос на 200 и снова 9 секунд колом. так что одним warmup не отделаешься

[lostangel]

У нас та же тема была на T4 в Селектеле. Решили в три захода и оно реально работает.

Первое, mark_dynamic на ось длины. torch._dynamo.mark_dynamic(input_ids, 1) перед компиляцией, тогда дайнамо не плодит граф под каждую форму, а собирает один с символьной осью. Рекомпиляции почти ушли, осталась пара на крайних бакетах.

Второе, бакетинг входов. Паддим до ближайшей степени из набора 64/128/256/512, прогреваем все четыре на старте контейнера. Да, чуть лишних вычислений на коротких, зато формы стабильные и граф один на бакет.

Третье, кэш компиляции между рестартами. TORCHINDUCTOR_CACHE_DIR на примонтированный PVC, плюс с torch 2.4 есть mega-cache, дамп через torch.compiler.save_cache_artifacts. После рестарта пода inductor не пересобирает ядра с нуля, холодный старт упал с 14с до примерно 2с. И главное, readinessProbe вешай на отдельный эндпоинт, который отвечает 200 только после того как все бакеты прогреты. Тогда трафик не приходит в непрогретый под.

По итогу p99 на старте перестал скакать, и автоскейл перестал плодить полудохлые поды.

[corvet]

@joseph23, max-autotune вообще не трогай для сервинга если важен холодный старт, он перебирает конфиги ядер и компиляция растягивается до минуты. reduce-overhead норм компромисс

[delphin]

readinessProbe вешай на отдельный эндпоинт, который отвечает 200 только после того как все бакеты прогреты

вот это ключевое что все пропускают. у меня HPA на rps плодил поды в пик, они вставали в роутинг непрогретыми, p99 улетал в небо, HPA видел рост латенси и плодил ещё больше. петля. развязал именно отдельным проб-эндпоинтом

[dannii]

а зачем вообще compile если можно torchscript
лол

[valru]

@anonymous torchscript半 deprecated по факту, его уже еле поддерживают и на новых моделях он падает на каждом втором кастомном слое. compile это замена. так что нет