Большие языковые модели (или LLM) стали темой ежедневных разговоров. Об их быстром внедрении свидетельствует количество времени, необходимое для достижения 100 миллионов пользователей, которое увеличилось с «4,5 лет для Facebook» до рекордно низкого уровня в «2 месяца для ChatGPT». Генеративный предварительно обученный преобразователь (GPT) использует причинно-следственные авторегрессионные обновления для прогнозирования. Эти модельные архитектуры продемонстрировали потрясающую производительность при выполнении различных задач, таких как распознавание речи, генерация текста и ответы на вопросы. Несколько последних моделей, таких как Неокс, Сокол, Лама используйте архитектуру GPT в качестве основы. Обучение LLM требует колоссального количества вычислительного времени, которое стоит миллионы долларов. В этом посте мы суммируем процедуру обучения GPT. Неокс на AWS Trainium, специальном ускорителе машинного обучения (ML), оптимизированном для обучения глубокому обучению. Мы расскажем, как мы с минимальными затратами (3,2 млн токенов/$) обучили такие модели с помощью AWS Trainium без потери качества моделей.

Обзор решения

Модели GPT NeoX и Pythia

GPT NeoX и Пифия — это модели причинного языка с открытым исходным кодом от Eleuther-AI, содержащие примерно 20 миллиардов параметров в NeoX и 6,9 миллиардов в Pythia. Обе модели декодера имеют схожий архитектурный дизайн с Chat GPT3. Однако у них также есть несколько дополнений, которые также широко используются в последних моделях, таких как Llama. В частности, они имеют вращательно-позиционную заделку (ВЕРЕВКА) с частичным поворотом по размерам головы. Оригинальные модели (NeoX и Pythia 6.9B) обучены на общедоступной Набор данных свай с дедупликацией и использованием серверной части Megatron и Deepspeed.

Мы демонстрируем предварительное обучение и тонкую настройку этих моделей на экземплярах Trn1 на базе AWS Trainium, используя Нейрон НеМо библиотека. Чтобы обеспечить подтверждение концепции и быстрое воспроизведение, мы будем использовать меньший подмножество набора данных Википедии, токенизированное с помощью токенизатора кодирования пар байтов GPT2 (BPE).

Прохождение

Загрузите предварительно токенизированный набор данных Википедии, как показано:

export DATA_DIR=~/examples_datasets/gpt2

mkdir -p ${DATA_DIR} && cd ${DATA_DIR}

wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-vocab.json
wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-merges.txt
aws s3 cp s3://neuron-s3/training_datasets/gpt/wikipedia/my-gpt2_text_document.bin . --no-sign-request
aws s3 cp s3://neuron-s3/training_datasets/gpt/wikipedia/my-gpt2_text_document.idx . --no-sign-request
aws s3 cp s3://neuron-s3/training_datasets/gpt/wikipedia/license.txt . --no-sign-request

И NeoX 20B, и Pythia 6.9B используют ROPE с частичным вращением, например, вращая 25% размеров головы и оставляя остальную часть неподвижной. Чтобы эффективно реализовать частичное вращение в ускорителе AWS Trainium, вместо объединения вращающихся и невращающихся размеров мы добавляем нулевые частоты для невращающихся размеров, а затем поворачиваем полный набор размеров головы. Этот простой трюк помог нам повысить пропускную способность (обрабатываемых последовательностей в секунду) в AWS Trainium.

Этапы обучения

Для проведения обучения мы используем управляемый SLURM многоузловой кластер Amazon Elastic Compute Cloud (Amazon EC2) Trn1, каждый узел которого содержит экземпляр trn1.32xl. Каждый trn1.32xl имеет 16 ускорителей, по два рабочих на каждый ускоритель. После загрузки последней Нейрон НеМо пакет, используйте предоставленный неокс и пифия сценарии предварительного обучения и точной настройки с оптимизированными гиперпараметрами и выполните следующие действия для обучения с четырьмя узлами.

1. Компиляция: Предварительная компиляция модели с тремя итерациями обучения для создания и сохранения графиков:

   sbatch --nodes 4 compile.slurm ./neoX_20B_slurm.sh

2. Запуск: выполнить обучение, загрузив кэшированные графики с первых шагов.

   sbatch --nodes 4 run.slurm ./neoX_20B_slurm.sh

3. Мониторинг результатов

   tensorboard --logdir=nemo_experiments/megatron_neox

Те же шаги необходимо выполнить для запуска модели Pythia 6.9B с заменой neox_20B_slurm.sh к pythia_6.9B_slurm.sh.

Предварительное обучение и эксперименты по тонкой настройке

Мы демонстрируем предварительное обучение моделей GPT-NeoX и Pythia на AWS Trainium с использованием Нейрон НеМо библиотеку на 10 тысяч итераций, а также покажем тонкую настройку этих моделей за 1 тысячу шагов. Для предварительного обучения мы используем токенизатор GPT2 BPE внутри NeMo и следуем тому же принципу. конфигурация как использовано в исходной модели. Точная настройка AWS Trainium требует изменения нескольких параметров (например, коэффициент деления размера словаря), которые предоставляются в сценариях тонкой настройки для учета различий Megatron и NeMo, а также изменений графического процессора и AWS Trainium. Пропускная способность распределенного обучения с различным количеством узлов показана в Таблице 1.

Таблица 1. Сравнение средней пропускной способности моделей GPT NeoX и Pythia для обучения до 500 шагов с изменением количества узлов. Цена trn1,32xl основана на действующей почасовой ставке за 3 года резервирования.

Далее мы также оцениваем траекторию потерь при обучении модели на AWS Trainium и сравниваем ее с соответствующим запуском на кластере P4d (ядра графического процессора Nvidia A100). Наряду с потерями обучения мы также сравниваем полезный показатель, такой как норма градиента, которая представляет собой 2-норму градиентов модели, вычисляемых на каждой итерации обучения для мониторинга прогресса обучения. Результаты обучения показаны на рисунках 1, 2, а тонкая настройка NeoX 20B – на рисунке 3.

Рисунок 1. Потери при обучении усреднены по всем работникам (слева) и норме градиента (справа) при обучении на каждом этапе. NeoX 20B обучается на 4 узлах с небольшим набором вики-данных на графическом процессоре и Trainium с одинаковыми гиперпараметрами обучения (глобальный размер пакета = 256). Графический процессор использует BF16 и смешанную точность по умолчанию, тогда как AWS Trainium использует полный BF16 со стохастическим округлением. Траектории норм потерь и градиента совпадают для графического процессора и AWS Trainium.

Фигура 2. Потери при обучении усреднены по всем работникам (слева) и норме градиента (справа) при обучении на каждом этапе. Подобно GPT NeoX на рисунке 1, Pythia 6.9B обучается на 4 узлах с небольшим набором вики-данных на графическом процессоре и Trainium с теми же гиперпараметрами обучения (глобальный размер пакета = 256). Траектории нормы потерь и градиента совпадают для графического процессора и Trainium.

Рисунок-3. Точная настройка модели GPT NeoX 20B на графическом процессоре и AWS Trainium с усредненными потерями при обучении для всех рабочих процессов (слева) и нормой градиента (справа). Небольшой набор данных Wiki используется для демонстрации точной настройки. Траектории норм потерь и градиента совпадают для графического процессора и AWS Trainium.

В этом посте мы показали экономичное обучение LLM на оборудовании глубокого обучения AWS. Мы обучили модели GPT NeoX 20B и Pythia 6.9B на AWS Trn1 с помощью библиотеки Neuron NeMo. Нормализованная по стоимости пропускная способность для 20 миллиардов моделей с AWS Trainium составляет примерно 3,2 миллиона токенов на каждый потраченный доллар. Наряду с экономичным обучением на AWS Trainium мы получаем аналогичную точность модели, о чем свидетельствуют потери шагов обучения и траектория нормы градиента. Мы также доработали доступные контрольные точки для модели NeoX 20B на AWS Trainium. Дополнительную информацию о распределенном обучении с помощью NeMo Megatron на AWS Trainium см. Справочник по AWS Neuron для NeMo Megatron. Хороший ресурс для начала тонкой настройки модели Llama можно найти здесь, Тонкая настройка Llama2. Чтобы начать работу с управляемым AWS Trainium на Amazon SageMaker, см. раздел Обучение моделей машинного обучения с помощью AWS Trainium и Amazon SageMaker.

Об авторах

Гаурав Гупта в настоящее время является научным сотрудником в лабораториях искусственного интеллекта Amazon Web Services (AWS). Доктор Гупта получил докторскую степень в Университете Южной Калифорнии в Витерби. Его исследовательские интересы охватывают область последовательного моделирования данных, изучения уравнений в частных производных, теории информации для машинного обучения, дробных динамических моделей и сложных сетей. В настоящее время он работает над прикладными и математическими проблемами поведения при обучении LLM, моделями видения с PDE, теоретико-информационными мультимодальными моделями. Доктор Гупта имеет публикации в ведущих журналах/конференциях, таких как Neurips, ICLR, ICML, Nature, IEEE Control Society, киберфизическое общество ACM.

Бен Снайдер — учёный-прикладник в AWS Deep Learning. Его исследовательские интересы включают фундаментальные модели, обучение с подкреплением и асинхронную оптимизацию. Вне работы он любит кататься на велосипеде и ходить в походы за городом.

Амит (справа) Мамидала — старший разработчик приложений машинного обучения в AWS Annapurna Labs. Доктор Мамидала защитил докторскую диссертацию в Университете штата Огайо в области высокопроизводительных вычислений и связи. Во время своей работы в исследовательском отделе IBM д-р Мамидала внес свой вклад в создание класса компьютеров BlueGene, которые часто возглавляли рейтинг Top500 самых мощных и энергоэффективных суперкомпьютеров. Проект был награжден Национальной медалью технологий и инноваций 2009 года. После непродолжительной работы инженером по искусственному интеллекту в финансовом хедж-фонде доктор Мамидала присоединился к лабораториям Аннапурны, сосредоточившись на обучении моделям большого языка.

Цзюнь (Люк) Хуан — ведущий научный сотрудник AWS AI Labs. Доктор Хуан занимается искусственным интеллектом и наукой о данных. Он опубликовал более 180 рецензируемых статей на ведущих конференциях и журналах. В 2009 году он был удостоен премии NSF за развитие ранней карьеры. До прихода в AWS он работал в исследовательской компании Baidu в качестве выдающегося ученого и руководителя лаборатории больших данных Baidu. Он основал StylingAI Inc., стартап в области искусственного интеллекта, и работал генеральным директором и главным научным сотрудником в 2019–2021 годах. Прежде чем присоединиться к промышленности, он был профессором Чарльза Э. и Мэри Джейн Спар на факультете EECS в Университете Канзаса.

Шрути Копаркар — старший менеджер по маркетингу продуктов в AWS. Она помогает клиентам исследовать, оценивать и внедрять инфраструктуру ускоренных вычислений Amazon EC2 для нужд машинного обучения.