maxtext
v1.0.0
MaxText — это высокопроизводительный , масштабируемый LLM с открытым исходным кодом, написанный на чистом Python/Jax и предназначенный для TPU и графических процессоров Google Cloud для обучения и вывода . MaxText достигает высоких значений MFU и масштабируется от одного хоста до очень больших кластеров, оставаясь при этом простым и «не требующим оптимизации» благодаря возможностям Jax и компилятору XLA.
MaxText стремится стать отправной точкой для амбициозных проектов LLM как в области исследований, так и производства. Мы рекомендуем пользователям начать с экспериментов с MaxText «из коробки», а затем разветвить и модифицировать MaxText в соответствии со своими потребностями.
Мы использовали MaxText, чтобы продемонстрировать высокопроизводительное, хорошо сходящееся обучение в int8 и масштабировать обучение до ~51 тыс. чипов.
Основные поддерживаемые функции:
При первом запуске MaxText мы предоставляем конкретные инструкции.
MaxText поддерживает обучение и вывод различных открытых моделей. Чтобы узнать больше, следуйте руководствам пользователя в папке «Начало работы».
Некоторые дополнительные полезные руководства:
Помимо руководств по началу работы, всегда есть и другие возможности MaxText, которые постоянно добавляются! Полный набор сквозных тестов находится в end_to_end. Мы запускаем их с ночной частотой. Они могут быть хорошим источником для понимания MaxText. Альтернативно вы можете увидеть непрерывные модульные тесты, которые выполняются почти непрерывно.
Более подробную информацию о воспроизведении этих результатов можно найти в MaxText/configs/README.md.
| Количество параметров | Тип ускорителя | ТФЛОП/чип/сек | Модель использования провалов (MFU) |
|---|---|---|---|
| 32Б | в5п-128 | 3.28e+02 | 71,47% |
| 64Б | в5п-128 | 3.23e+02 | 70,31% |
| 128Б | в5п-256 | 3.15e+02 | 68,68% |
| 128Б | в5п-512 | 3.15e+02 | 68,53% |
| 256Б | v5p-1024 | 3.16e+02 | 68,82% |
| 512Б | v5p-1024 | 2.94e+02 | 63,99% |
| 1024Б | v5p-2048 | 2.49e+02 | 64,05% |
| 1024Б | v5p-4096 | 2.97e+02 | 64,80% |
| 1160Б | v5p-7680 | 2.95e+02 | 64,27% |
| 1160Б | v5p-12288 | 3.04e+02 | 66,23% |
Для моделей 16B, 32B, 64B и 128B. Полные конфигурации запуска см. в MaxText/configs/v5e/ как 16b.sh , 32b.sh , 64b.sh , 128b.sh
| Аппаратное обеспечение | 16Б Тфлопс/сек/чип | 16Б МФУ | 32 млрд Тфлопс/сек/чип | 32Б МФУ | 64Б Тфлопс/сек/чип | 64Б МФУ | 128 терафлопс/сек/чип | 128Б МФУ |
|---|---|---|---|---|---|---|---|---|
| 1x v5e-256 | 120 | 61,10% | 132 | 66,86% | 118 | 59,90% | 110 | 56,06% |
| 2x v5e-256 | 117 | 59,37% | 128 | 64,81% | 112 | 56,66% | 110 | 55,82% |
| 4x v5e-256 | 117 | 59,14% | 126 | 64,10% | 110 | 55,85% | 108 | 54,93% |
| 8x v5e-256 | 115 | 58,27% | 125 | 63,67% | 108 | 54,96% | 104 | 52,93% |
| 16x v5e-256 | 111 | 56,56% | 123 | 62,26% | 105 | 53,29% | 100 | 50,86% |
| 32x v5e-256 | 108 | 54,65% | 119 | 60,40% | 99 | 50,18% | 91 | 46,25% |
MaxText во многом вдохновлен MinGPT/NanoGPT, элегантными автономными реализациями GPT, написанными на PyTorch и ориентированными на графические процессоры Nvidia. MaxText более сложен, поддерживает больше стандартных моделей и масштабируется до десятков тысяч микросхем. В конечном итоге MaxText имеет MFU, более чем в три раза превышающий 17%, о которых сообщалось в последний раз для этой кодовой базы, он легко масштабируется и реализует кеш-ключ-значение для эффективного авторегрессионного декодирования.
MaxText больше похож на Nvidia/Megatron-LM, очень хорошо настроенную реализацию LLM, ориентированную на графические процессоры Nvidia. Обе реализации достигают сопоставимых MFU. Разница в базах кода подчеркивает разные стратегии программирования. MaxText — это чистый Python, в котором для достижения высокой производительности в значительной степени используется компилятор XLA. Megatron-LM, напротив, представляет собой смесь Python и CUDA, в которой для достижения высокой производительности используются хорошо оптимизированные ядра CUDA.
MaxText также сравним с Pax. Как и Pax, MaxText обеспечивает высокопроизводительную и масштабируемую реализацию LLM в Jax. Pax фокусируется на включении мощных параметров конфигурации, позволяя разработчикам изменять модель путем редактирования параметров конфигурации. Напротив, MaxText представляет собой простую и конкретную реализацию различных LLM, которая побуждает пользователей расширять возможности путем разветвления и прямого редактирования исходного кода.
При запуске задания «Одна программа, несколько данных» (SPMD) на ускорителях весь процесс может зависнуть, если возникнет какая-либо ошибка или какая-либо виртуальная машина по какой-либо причине зависнет или выйдет из строя. В этом сценарии запись трассировки стека поможет выявить и устранить проблемы с заданиями, выполняемыми на виртуальных машинах TPU.
Следующие конфигурации помогут отладить ошибку или ситуацию, когда программа где-то зависла или зависла, путем сбора трассировок стека. Измените значения параметров соответствующим образом в MaxText/configs/base.yml :
collect_stack_trace: True , чтобы включить сбор трассировки стека при сбоях или когда программа зависает. Этот параметр будет периодически выгружать трассировки программы, чтобы помочь в отладке. Чтобы отключить это, установите collect_stack_trace: False .stack_trace_to_cloud: False , чтобы отображать трассировки стека на консоли. stack_trace_to_cloud: True создаст временный файл в /tmp/debugging в TPU для хранения трассировок стека. На виртуальных машинах TPU работает агент, который периодически загружает трассировки из временного каталога в облачное журналирование в проекте gcp. Вы можете просмотреть трассировки в обозревателе журналов в Cloud Logging, используя следующий запрос: logName="projects/<project_name>/logs/tpu.googleapis.com%2Fruntime_monitor"
jsonPayload.verb="stacktraceanalyzer"
stack_trace_interval_seconds обозначает продолжительность в секундах между каждым событием сбора трассировки стека. Установка stack_trace_interval_seconds: 600 будет собирать трассировку стека каждые 600 секунд (10 минут).Вот соответствующий пакет PyPI: https://pypi.org/project/cloud-tpu-diagnostics.
Чтобы заранее скомпилировать тренировочный прогон, мы предоставляем инструмент train_compile.py . Этот инструмент позволяет скомпилировать основной train_step в train.py для целевого оборудования (например, большого количества устройств v5e) без использования всего кластера.
Вы можете использовать только ЦП или одну виртуальную машину из другого семейства для предварительной компиляции кластера TPU. Эта компиляция помогает достичь двух основных целей:
Он будет отмечать любую информацию о нехватке памяти (OOM), например, когда per_device_batch_size установлен слишком высоко, с идентичной трассировкой стека OOM, как если бы он был скомпилирован на целевом оборудовании.
Предварительную компиляцию можно сохранить, а затем загрузить для быстрого запуска и перезапуска на целевом оборудовании.
Инструмент train_compile.py тесно связан с train.py и использует тот же файл конфигурации configs/base.yml . Хотя вам не обязательно запускать TPU, вам необходимо установить jax[tpu] в дополнение к другим зависимостям, поэтому мы рекомендуем запустить setup.sh для их установки, если вы еще этого не сделали.
После установки перечисленных выше зависимостей вы готовы к предварительной компиляции:
# Run the below on a single machine, e.g. a CPU
python3 MaxText/train_compile.py MaxText/configs/base.yml compile_topology=v5e-256 compile_topology_num_slices=2
global_parameter_scale=16 per_device_batch_size=4
Это позволит скомпилировать модель MaxText с 16B параметрами на двух модулях v5e.
Вот пример, который сохраняет, а затем загружает скомпилированный train_step , начиная с сохранения:
Шаг 1. Запустите AOT и сохраните скомпилированную функцию.
# Run the below on a single machine, e.g. a CPU
export LIBTPU_INIT_ARGS="--xla_enable_async_all_gather=true"
python3 MaxText/train_compile.py MaxText/configs/base.yml compile_topology=v5e-256
compile_topology_num_slices=2
compiled_trainstep_file=my_compiled_train.pickle global_parameter_scale=16
per_device_batch_size=4 steps=10000 learning_rate=1e-3
Шаг 2. Запустите train.py и загрузите скомпилированную функцию.
Чтобы загрузить скомпилированный train_step, вам просто нужно передать compiled_trainstep_file=my_compiled_train.pickle в train.py :
# Run the below on each host of the target hardware, e.g. each host on 2 slices of v5e-256
export LIBTPU_INIT_ARGS="--xla_enable_async_all_gather=true"
python3 MaxText/train.py MaxText/configs/base.yml run_name=example_load_compile
compiled_trainstep_file=my_compiled_train.pickle
global_parameter_scale=16 per_device_batch_size=4 steps=10000 learning_rate=1e-3
base_output_directory=gs://my-output-bucket dataset_path=gs://my-dataset-bucket
На этапе сохранения примера 2 выше мы включили экспорт флага компилятора LIBTPU_INIT_ARGS и learning_rate , поскольку они влияют на скомпилированный объект my_compiled_train.pickle. Размеры модели (например, global_parameter_scale , max_sequence_length и per_device_batch ) фиксированы при первоначальной компиляции с помощью compile_train.py . Вы увидите ошибку размера, если попытаетесь запустить сохраненный скомпилированный объект с размерами, отличными от тех, с которыми вы скомпилировали. Однако следует отметить, что график скорости обучения также фиксируется при запуске compile_train , который определяется как steps , так и learning_rate . Параметры оптимизатора, такие как adam_b1 передаются компилятору только как сформированные объекты — таким образом, их реальные значения определяются при запуске train.py , а не во время компиляции. Если вы передадите разные формы (например, per_device_batch ), вы получите четкое сообщение об ошибке, сообщающее, что скомпилированная подпись имеет ожидаемую форму, отличную от той, которая была введена. Если вы попытаетесь запустить на оборудовании, отличном от целей компиляции, запрошенных через compile_topology , вы получите сообщение об ошибке, сообщающее, что не удалось сопоставить устройства из скомпилированного устройства с вашими реальными устройствами. Использование других флагов XLA или LIBTPU, отличных от того, что было скомпилировано, вероятно, будет работать в среде, в которой вы скомпилировали, без ошибок. Однако в этом случае нет гарантированного поведения; вам следует работать в той же среде, в которой вы компилировали.
Предварительная компиляция также поддерживается для графических процессоров с некоторыми отличиями от TPU:
Графический процессор не поддерживает компиляцию между оборудованием: для запуска AoT-компиляции по-прежнему требуется хост графического процессора, но один хост графического процессора может скомпилировать программу для более крупного кластера того же оборудования.
Для графических процессоров A3 Cloud максимальный размер «среза» — это один хост, а параметр compile_topology_num_slices представляет количество компьютеров A3, для которых требуется предварительная компиляция.
В этом примере показаны флаги, используемые для многохостовой компиляции графического процессора, ориентированной на кластер из 4 хостов A3:
Шаг 1. Запустите AOT и сохраните скомпилированную функцию.
# Run the below on a single A3 machine
export XLA_FLAGS="--xla_gpu_enable_async_collectives=true"
python3 MaxText/train_compile.py MaxText/configs/base.yml compile_topology=a3
compile_topology_num_slices=4
compiled_trainstep_file=my_compiled_train.pickle global_parameter_scale=16
attention=dot_product per_device_batch_size=4 steps=10000 learning_rate=1e-3
Шаг 2. Запустите train.py и загрузите скомпилированную функцию.
Чтобы загрузить скомпилированный train_step, вам просто нужно передать compiled_trainstep_file=my_compiled_train.pickle в train.py :
# Run the below on each of the 4 target A3 hosts.
export XLA_FLAGS="--xla_gpu_enable_async_collectives=true"
python3 MaxText/train.py MaxText/configs/base.yml run_name=example_load_compile
compiled_trainstep_file=my_compiled_train.pickle
attention=dot_product global_parameter_scale=16 per_device_batch_size=4 steps=10000 learning_rate=1e-3
base_output_directory=gs://my-output-bucket dataset_path=gs://my-dataset-bucket
Как и в случае с TPU, обратите внимание, что среда компиляции должна соответствовать среде выполнения, в данном случае путем установки одного и того же XLA_FLAGS .
MaxText поддерживает автоматическую загрузку журналов, собранных в каталоге, в экземпляр Tensorboard в Vertex AI. Следуйте руководству пользователя, чтобы узнать больше.
