JAX Toolbox предоставляет общедоступный CI, образы Docker для популярных библиотек JAX и оптимизированные примеры JAX для упрощения и улучшения вашего опыта разработки JAX на графических процессорах NVIDIA. Он поддерживает библиотеки JAX, такие как MaxText, Paxml и Pallas.
Мы поддерживаем и тестируем следующие платформы JAX и модели архитектуры. Более подробную информацию о каждой модели и доступных контейнерах можно найти в соответствующих файлах README.
| Рамки | Модели | Варианты использования | Контейнер |
|---|---|---|---|
| Макстекст | GPT, LLaMA, Джемма, Мистраль, Микстрал | предварительная подготовка | ghcr.io/nvidia/jax:maxtext |
| пакмл | GPT, LLaMA, МО | предобучение, тонкая настройка, LoRA | ghcr.io/nvidia/jax:pax |
| т5х | Т5, ВИТ | предварительная подготовка, тонкая настройка | ghcr.io/nvidia/jax:t5x |
| т5х | Изображение | предварительная подготовка | ghcr.io/nvidia/t5x:imagen-2023-10-02.v3 |
| большое видение | ПалиДжемма | доработка, оценка | ghcr.io/nvidia/jax:gemma |
| Левантер | GPT, LLaMA, MPT, Рюкзаки | предварительная подготовка, тонкая настройка | ghcr.io/nvidia/jax:levanter |
| Компоненты | Контейнер | Строить | Тест |
|---|---|---|---|
 | ghcr.io/nvidia/jax:base |   | [нет тестов] |
 | ghcr.io/nvidia/jax:jax |   |        |
 | ghcr.io/nvidia/jax:levanter |   |   |
 | ghcr.io/nvidia/jax:equinox |   | [тесты отключены] |
 | ghcr.io/nvidia/jax:triton |  |   |
 | ghcr.io/nvidia/jax:upstream-t5x |   |  |
 | ghcr.io/nvidia/jax:t5x |   |  |
 | ghcr.io/nvidia/jax:upstream-pax |   |  |
 | ghcr.io/nvidia/jax:pax |   |  |
 | ghcr.io/nvidia/jax:maxtext |   |  |
 | ghcr.io/nvidia/jax:gemma |  |   |
Во всех случаях ghcr.io/nvidia/jax:XXX указывает на последнюю ночную сборку контейнера для XXX . Для получения стабильной ссылки используйте ghcr.io/nvidia/jax:XXX-YYYY-MM-DD .
Помимо общедоступных CI-тестов, мы также проводим внутренние CI-тесты на H100 SXM 80 ГБ и A100 SXM 80 ГБ.
В образ JAX встроены следующие флаги и переменные среды для настройки производительности XLA и NCCL:
| Флаги XLA | Ценить | Объяснение |
|---|---|---|
--xla_gpu_enable_latency_hiding_scheduler | true | позволяет XLA перемещать коммуникационные коллективы для увеличения перекрытия с вычислительными ядрами |
--xla_gpu_enable_triton_gemm | false | используйте cuBLAS вместо ядер Trition GeMM |
| Переменная среды | Ценить | Объяснение |
|---|---|---|
CUDA_DEVICE_MAX_CONNECTIONS | 1 | используйте единую очередь для работы графического процессора, чтобы снизить задержку потоковых операций; Хорошо, поскольку XLA уже заказывает запуски |
NCCL_NVLS_ENABLE | 0 | Отключает NVLink SHARP (1). В будущих выпусках эта функция будет снова включена. |
Пользователи могут установить различные другие флаги XLA для повышения производительности. Подробное объяснение этих флагов можно найти в документации по производительности графического процессора. Флаги XLA можно настроить для каждого рабочего процесса. Например, каждый скрипт в contrib/gpu/scripts_gpu устанавливает свои собственные флаги XLA.
Список ранее использовавшихся флагов XLA, которые больше не нужны, также можно найти на странице производительности графического процессора.
| Впервые за ночь с новым базовым контейнером | Базовый контейнер |
|---|---|
| 06.11.2024 | nvidia/cuda:12.6.2-devel-ubuntu22.04 |
| 2024-09-25 | nvidia/cuda:12.6.1-devel-ubuntu22.04 |
| 2024-07-24 | nvidia/cuda:12.5.0-devel-ubuntu22.04 |
См. эту страницу для получения дополнительной информации о том, как профилировать программы JAX на графическом процессоре.
Решение:
docker run -it --shm-size=1g ...
Объяснение: bus error могла возникнуть из-за ограничения размера /dev/shm . Вы можете решить эту проблему, увеличив размер общей памяти, используя опцию --shm-size при запуске контейнера.
Описание проблемы:
slurmstepd: error: pyxis: [INFO] Authentication succeeded slurmstepd: error: pyxis: [INFO] Fetching image manifest list slurmstepd: error: pyxis: [INFO] Fetching image manifest slurmstepd: error: pyxis: [ERROR] URL https://ghcr.io/v2/nvidia/jax/manifests/returned error code: 404 Not Found
Решение: обновите enroot или примените однофайловый патч, как указано в примечаниях к выпуску enroot v3.4.0.
Объяснение: Docker традиционно использовал Docker Schema V2.2 для списков манифестов с несколькими арками, но с 20 октября перешел на использование формата Open Container Initiative (OCI). Enroot добавил поддержку формата OCI в версии 3.4.0.
АВС
Добавить интеграцию EFA
Пример кода SageMaker
GCP
Начало работы с многоузловыми приложениями JAX с графическими процессорами NVIDIA в Google Kubernetes Engine.
Лазурный
Ускорение приложений ИИ с помощью платформы JAX на виртуальных машинах Azure NDm A100 v4.
ЗКИ
Выполнение рабочей нагрузки глубокого обучения с помощью JAX в многоузловых кластерах с несколькими графическими процессорами на OCI.
ДЖАКС | Контейнер NVIDIA NGC
Интеграция нулевой конфигурации Slurm и OpenMPI
Добавление пользовательских операций графического процессора
Сортировка регрессий
Равноденствие для JAX: основа экосистемы науки и машинного обучения
Масштабирование Grok с помощью JAX и H100
JAX Supercharged на графических процессорах: высокопроизводительные курсы LLM с JAX и OpenXLA
Что нового в JAX | GTC весна 2024 г.
Что нового в JAX | ГТК весна 2023 г.
