Мы мигрировали с PyTorch на JAX. Инференс ускорился в 3 раза, но команда чуть не развалилась

Наша команда 2 года писала на PyTorch. Все привыкли, был наработан опыт, библиотеки, инструменты. Но когда нам нужно было масштабировать инференс до 10,000 RPS, мы уперлись в потолок. Мы решили попробовать JAX. Идея была простая: XLA-компиляция должна дать нам кратный прирост производительности. Первые 3 месяца были адом. Синтаксис функциональный, отладка — nightmare, документация скудная. Половина команды хотела вернуться обратно. Но мы дожали. Написали собственные утилиты для отладки, переучились думать в терминах функционального программирования. Результат: инференс ускорился в 3 раза, потребление GPU упало на 40%. Мы смогли обслуживать тот же трафик на вдвое меньшем железе. Мораль: технологический стек — это не только код. Это люди, процессы, инструменты. Миграция — это всегда боль. Но если боль оправдана бизнес-ценностью — терпите.