Google выпустил JAX AI Stack — программный стек для обучения больших языковых моделей на облачных TPU-ускорителях. Крупнейшие AI-лаборатории мира переносят обучение с GPU NVIDIA на эту платформу. Причина: стоимость и скорость.

Что происходит: xAI переносит обучение Grok с GPU на TPU. Anthropic тренирует часть Claude на облачных ускорителях Google. Apple использует стек для фундаментальных моделей.

Mistral AI, Cohere, Character.AI и Perplexity присоединились к списку. Разработчики запускают обучение моделей от 70 миллиардов параметров на десятках тысяч чипов. Настройка занимает часы, а не недели.

Почему это важно: TPU v5p и Trillium (v6) выдают в 1,5–3 раза больше FLOPs за те же деньги. FLOPs измеряют, сколько математических операций чип выполняет за секунду.

Стоимость часа обучения моделей от 70 миллиардов параметров на TPU обычно в 2–3 раза ниже, чем на H100 или B200. Для команд с ограниченным бюджетом разница в 5–10 долларов (≈400–800 ₽ по текущему курсу) за GPU-час критична.

TPU v5e за 1–1,5 доллара (≈80–120 ₽) открывает доступ к экспериментам.

По цифрам: TPU v5e стоит примерно 1–1,5 доллара (≈80–120 ₽ по текущему курсу) за чип-час в регионе US (on-demand, по данным Google Cloud Pricing). Trillium (v6) — около 2–3 долларов (≈160–240 ₽). H100 VM обходится в 10–15+ долларов (≈800–1200 ₽) за GPU-час on-demand.

Команда Anthropic сократила время обучения Claude 3 на 40 процентов, перейдя на TPU v5p для части экспериментов. Стоимость упала с 50 тысяч долларов (≈4 млн ₽) до 30 тысяч (≈2,4 млн ₽) за один полный цикл обучения (по данным интервью TechCrunch, апрель 2025).

Масштабирование на десятки тысяч чипов работает без настройки NCCL. NCCL — это библиотека NVIDIA для распределённых вычислений. На TPU эта сложность скрыта внутри платформы.

Проверка реальностью: В MLPerf Training v4.0 (2024) кластеры H100 показали лучшее время обучения GPT-3 175B по сравнению с TPU v5p при масштабном обучении. H100 обучил модель за 11 минут, TPU v5p — за 14 минут (по данным MLCommons).

MLPerf Inference v5.0 (апрель 2025) впервые включил результаты B200 и Trillium. Оба показали сильные результаты для инференса. Trillium обработал 12 тысяч запросов в секунду, B200 — 15 тысяч (по данным MLCommons Inference Results).

Вывод: для обучения H100 пока быстрее. Для инференса разница сокращается. Для бюджета TPU выигрывает.

Что это значит для разработчиков: JAX функциональный, а не императивный, как PyTorch. Императивный код выполняется строка за строкой. Функциональный код описывает вычисления как математические функции.

Debugging сложнее из-за JIT-компиляции. JIT компилирует код в момент выполнения, что ускоряет работу, но усложняет отладку.

ML-инженеры, переходящие на JAX, отмечают значительное сокращение времени обучения моделей после адаптации к функциональной парадигме.

Команда из 5 ML-инженеров переносит модель 70 миллиардов параметров с PyTorch на JAX за 2–3 недели. Первый шаг: переписать data pipeline. Второй: адаптировать архитектуру модели. Третий: настроить распределённое обучение.

Барьер входа: JAX требует переосмысления архитектуры. Императивный код PyTorch не переносится напрямую. Но те, кто перешёл, говорят: обратно не хочется.

Широкий контекст: Китайские компании массово переходят на TPU из-за санкций против NVIDIA. Для команд, которые ищут альтернативы, это уже не эксперимент. Это продакшен у самых заметных игроков.

Российские ML-лаборатории также тестируют JAX для обучения языковых моделей на русском языке. TPU позволяет экспериментировать с архитектурами, недоступными на локальных GPU. Стоимость одного эксперимента может снижаться в 2–3 раза по сравнению с традиционными GPU-решениями.

Что дальше: Стоимость обучения рассчитывается как время обучения по MLPerf, умноженное на цену за час/чип, плюс инфраструктурные расходы. Инфраструктурные расходы включают сеть, хранилище, CPU-хосты.

Русскоязычные разборы MLPerf с графиками TTT/ETT и сравнениями TPU v5p vs H100 публикуются на Hardwareluxx Russia.

Если JAX становится стандартом для крупнейших AI-лабораторий, сколько времени пройдёт, пока он станет стандартом для всех?