AIKraft Agents
ВойтиТарифы и покупка

AI-Hypercomputer/maxtext

A simple, performant, and scalable Jax LLM!

Инференс⭐ 2 457Pythonпоследний релиз: maxtext-v0.2.5
Открыть на GitHub ↗Сайт проекта ↗

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

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.

Релизы

maxtext-v0.2.5major
🕐 1 мес назад · релиз на GitHub ↗

Релиз MaxText v0.2.5: поддержка DeepSeek-V3/V4, Qwen3.5, Cosmos 3, мультимодальных моделей и GRPO-обучения.

  • Added: Добавлена поддержка моделей DeepSeek-V3/V4, Qwen3.5, Qwen3-Next, Cosmos 3 и Weaver.
  • Added: Реализована инфраструктура для обучения с подкреплением (GRPO) и квантованного Multi-Token Prediction.
  • Added: Внедрена мультимодальная SFT-пайплайн на базе Grain и поддержка вывода Qwen3-VL через vLLM.
  • Added: Оптимизирован JAX Flash Attention для длинных последовательностей и добавлена контекстная параллельность (CP).
  • Added: Мигрированы скрипты DiLoCo в Cluster Toolkit и улучшена работа с LoRA и дистилляцией.
maxtext-v0.2.4majorbreaking
🕐 2 мес назад · релиз на GitHub ↗

Миграция на Flax NNX, поддержка DeepSeek-V4, Qwen3-VL, LoRA/QLoRA, контекстная параллельность и FP4-квантование.

  • Breaking: Выполнен переход на Flax NNX: конфигурации pure_nnx и enable_nnx включены по умолчанию.
  • Added: Добавлена полная поддержка моделей DeepSeek-V4 и Qwen3-VL, включая конфигурации и конвертацию чекпоинтов.
  • Added: Реализована нативная поддержка LoRA и QLoRA для Gemma4, Gemma3, Qwen3 и Llama3.
  • Added: Расширены возможности контекстной параллельности: добавлены стратегии Ulysses/USP, Ring Attention и YaRN RoPE.
  • Added: Внедрено FP4 (E2M1) квантование, экспериментальное квантование внимания и оптимизации производительности.
maxtext-v0.2.3major
🕐 4 мес назад · релиз на GitHub ↗

MaxText v0.2.3: поддержка Qwen3.5/Omni, DPO/RL, новая система оценки и оптимизация MoE.

  • Added: Добавлена поддержка предобучения Qwen3.5 (35B/397B) и мультимодального SFT для Qwen3-Omni.
  • Added: Внедрена система оценки на базе vLLM и полная поддержка пайплайнов DPO/ORPO и RL-рецептов.
  • Added: Оптимизация MoE-слоев через ragged gather reduce для улучшения производительности и использования памяти.
  • Fixed: Исправлено искажение лог-вероятностей в RL из-за ошибки в маске внимания и проблемы с градиентами JAX/NNX.
  • Deprecated: Удалена устаревшая реализация DPO и функция сбора стека вызовов.