AI-Hypercomputer/maxtext

Что это за инструмент

MaxText — это высокопроизводительная библиотека для обучения больших языковых моделей (LLM) на базе JAX, оптимизированная для Google Cloud TPUs и GPU.

Зачем нужен

Инструмент позволяет эффективно проводить pre-training и post-training (включая SFT и RL-методы вроде GRPO/GSPO) моделей от одного хоста до кластеров с десятками тысяч чипов. Он обеспечивает высокую утилизацию вычислительных ресурсов (MFU) и скорость генерации токенов, оставаясь при этом простым в использовании благодаря автоматической оптимизации со стороны JAX и XLA.

Что можно реализовать

• Pre-training и тонкая настройка (SFT) популярных LLM, таких как Gemma, Llama, DeepSeek, Qwen и Mistral.
• Проведение исследований с использованием методов обучения с подкреплением, включая GRPO и GSPO.
• Запуск обучения на масштабируемых кластерах (до десятков тысяч чипов) с высокой эффективностью использования FLOPs.
• Оценка моделей с помощью фреймворков lm-eval, evalchemy и пользовательских бенчмарков.
• Эксперименты с новыми архитектурами, такими как Qwen3.5, Kimi-K2 и DeepSeek-V3.2.

Ключевые возможности

• Поддержка широкого спектра моделей: Gemma, Llama, DeepSeek, Qwen, Mistral, Kimi-K2, Qwen3.5.
• Высокая производительность (MFU) и масштабируемость от одного хоста до крупных кластеров без необходимости ручной оптимизации.
• Встроенная поддержка передовых методов post-training: Supervised Fine-Tuning (SFT), GRPO, GSPO.
• Наличие готовых конфигураций для новейших моделей (например, Gemma 4, DeepSeek-V3.2, Kimi-K2-Thinking).
• Возможность работы в decoupled mode без зависимости от инфраструктуры GCP.

👤 Кому подойдёт: Исследователи в области ИИ, ML-инженеры и разработчики, занимающиеся обучением и тонкой настройкой больших языковых моделей на GPU/TPU.