Repolode
Предложить
google

google/paxml

Pax — фреймворк на базе JAX для настройки и запуска крупномасштабных экспериментов в машинном обучении, главным образом с языковыми моделями на TPU.

Открыть на GitHub
Звёзды
560
Форки
73
Открытые issues
30
Последний коммит
12 дней назад

Что это?

Pax (Paxml) — фреймворк для настройки и запуска экспериментов в машинном обучении на базе JAX. В его основе — полностью настраиваемые эксперименты и параллельное выполнение; фреймворк предназначен для обучения крупных моделей, например языковых моделей на базе Transformer только с декодером. Модели, задачи и алгоритмы обучения задаются через гиперпараметры на основе dataclass. Они реализованы поверх библиотеки Praxis и поддерживают конфигурацию с помощью Fiddle.

Проект рассчитан на работу с Google Cloud TPU VM, в том числе с конфигурациями v4 на одном хосте, срезами TPU Pod и многосрезовыми конфигурациями. В комплект входят примеры конфигураций экспериментов для обучения на датасете C4 с моделями на 1B, 16B и большим числом параметров, а также бенчмарк утилизации вычислительных ресурсов модели (MFU) для запусков с масштабированием по слабому закону на TPU v4 Pod. Данные можно загружать через SeqIO, Lingvo или собственные подклассы BaseInput; в проекте также есть рекомендации по шардингу на нескольких хостах и оценке моделей.

В репозитории указано, что NVIDIA поддерживает отдельную обновлённую версию Pax с поддержкой GPU и ссылается на свой репозиторий.

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

  • 01Настраиваемые эксперименты и гиперпараметры на базе dataclass и Fiddle.
  • 02Поддержка параллельных вычислений с pjit/SPMD и pmap, а также многосрезовых конфигураций TPU.
  • 03Примеры конфигураций для обучения на C4 моделей с 1B и 16B параметров, а также моделей в стиле GPT-3 XL.
  • 04Несколько способов загрузки данных, включая SeqIO, адаптеры Lingvo и пользовательские источники данных.
  • 05Бенчмарки утилизации вычислительных ресурсов модели на TPU v4 Pod.

Сценарии использования

  • Обучение крупных языковых моделей только с декодером на Google Cloud TPU v4 Pod.
  • Запуск воспроизводимых экспериментов на базе JAX с настраиваемыми моделями и параметрами параллельных вычислений.
  • Сравнение эффективности обучения (MFU) при разных масштабах моделей и числе ускорителей.

Быстрый старт

  1. 1

    Создайте Cloud TPU VM (сначала задайте переменные окружения)

    gcloud compute tpus tpu-vm create $TPU_NAME --zone=$ZONE --version=$VERSION --project=$PROJECT --accelerator-type=$ACCELERATOR
  2. 2

    Подключитесь к TPU VM по SSH

    gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE
  3. 3

    Установите Pax из PyPI вместе с JAX для TPU

    python3 -m pip install paxml jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
  4. 4

    Запустите пример тестовой модели

    python3 .local/lib/python3.8/site-packages/paxml/main.py --exp=tasks.lm.params.lm_cloud.LmCloudSpmd2BLimitSteps --job_log_dir=gs://<your-bucket>

Что важно знать

  • Решение предназначено для оборудования Cloud TPU; инструкции по настройке и примеры в README рассчитаны на TPU VM.
  • Для настройки требуются ресурсы Google Cloud, например проект GCP и хранилище bucket.

Темы

#jax#large-language-models#tpu#distributed-training#parallelism#machine-learning#c4#gpt#llm#model-flops

Частые вопросы

Что такое paxml?
Pax — фреймворк на базе JAX для настройки и запуска крупномасштабных экспериментов в машинном обучении, главным образом с языковыми моделями на TPU. Pax (Paxml) — фреймворк для настройки и запуска экспериментов в машинном обучении на базе JAX. В его основе — полностью настраиваемые эксперименты и параллельное выполнение; фреймворк предназначен для обучения крупных моделей, например языковых моделей на базе Transformer только с декодером. Модели, задачи и алгоритмы обучения задаются через гиперпараметры на основе dataclass. Они реализованы поверх библиотеки Praxis и поддерживают конфигурацию с помощью Fiddle.
paxml можно использовать бесплатно? Какая у него лицензия?
paxml распространяется с открытым исходным кодом под лицензией Apache-2.0. Перед использованием в коммерческих продуктах проверьте условия лицензии.
paxml всё ещё поддерживается?
Активно поддерживается. Последний коммит в paxml — 26 окт. 2026 г..
Для чего можно использовать paxml?
Обучение крупных языковых моделей только с декодером на Google Cloud TPU v4 Pod. Запуск воспроизводимых экспериментов на базе JAX с настраиваемыми моделями и параметрами параллельных вычислений. Сравнение эффективности обучения (MFU) при разных масштабах моделей и числе ускорителей.
Как установить paxml?
Создайте Cloud TPU VM (сначала задайте переменные окружения): gcloud compute tpus tpu-vm create $TPU_NAME --zone=$ZONE --version=$VERSION --project=$PROJECT --accelerator-type=$ACCELERATOR Подключитесь к TPU VM по SSH: gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE Установите Pax из PyPI вместе с JAX для TPU: python3 -m pip install paxml jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html Запустите пример тестовой модели: python3 .local/lib/python3.8/site-packages/paxml/main.py --exp=tasks.lm.params.lm_cloud.LmCloudSpmd2BLimitSteps --job_log_dir=gs://<your-bucket>
Где скачать paxml?
Исходный код и инструкции по установке — на GitHub: https://github.com/google/paxml

Описание создано ИИ на основе README проекта. Статистика GitHub обновляется ежедневно.

Похожие проекты