Большие языковые модели (LLM) захватили воображение и внимание разработчиков, ученых, технологов, предпринимателей и руководителей в нескольких отраслях. Эти модели можно использовать для ответов на вопросы, обобщения, перевода и многого другого в таких приложениях, как диалоговые агенты для поддержки клиентов, создание контента для маркетинга и помощники по программированию.

Недавно Мета выпустила Лама 2 как для исследователей, так и для коммерческих организаций, пополнив список других программ LLM, включая MosaicML. МПТ и Сокол. В этом посте мы расскажем, как настроить Llama 2 на AWS Trainium, специально созданном ускорителе для обучения LLM, чтобы сократить время и затраты на обучение. Мы рассматриваем сценарии тонкой настройки, предоставляемые AWS Neuron SDK (с использованием NeMo Megatron-LM), различные использованные нами конфигурации и наблюдаемые нами результаты пропускной способности.

О модели Лама 2

Аналогично предыдущему Лама 1 Модель и другие модели, такие как GPT, Llama 2 использует архитектуру Transformer, использующую только декодер. Он бывает трех размеров: 7 миллиардов, 13 миллиардов и 70 миллиардов параметров. По сравнению с Llama 1, Llama 2 удваивает длину контекста с 2000 до 4000 и использует внимание сгруппированных запросов (только для 70B). Предварительно обученные модели Llama 2 обучены на 2 триллионах токенов, а ее точно настроенные модели обучены на более чем 1 миллионе аннотаций, сделанных человеком.

Распределенное обучение Llama 2

Чтобы разместить Llama 2 с длиной последовательности 2000 и 4000, мы реализовали сценарий, используя НеМо Мегатрон для Trainium, который поддерживает параллелизм данных (DP), тензорный параллелизм (TP) и конвейерный параллелизм (PP). Говоря конкретнее, с новой реализацией некоторых функций, таких как встраивание развязок слов, вращающееся встраивание, RMSNorm и активация Swiglu, мы используем общий сценарий GPT-нейрон Мегатрон-ЛМ для поддержки сценария обучения Llama 2.

Наша процедура обучения высокого уровня следующая: для нашей среды обучения мы используем многоэкземплярный кластер, управляемый системой SLURM для распределенного обучения и планирования в рамках NeMo.

Сначала загрузите модель Llama 2 и наборы обучающих данных и предварительно обработайте их с помощью токенизатора Llama 2. Например, чтобы использовать набор данных RedPajama, используйте следующую команду:

wget https://data.together.xyz/redpajama-data-1T/v1.0.0/book/book.jsonl

python nemo/scripts/nlp_language_modeling/preprocess_data_for_megatron.py

Подробные инструкции по загрузке моделей и аргументам сценария предварительной обработки см. Загрузите набор данных и токенизатор LlamaV2.

Далее скомпилируйте модель:

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

После компиляции модели запустите задание обучения с помощью следующего скрипта, который уже оптимизирован с использованием наилучшей конфигурации и гиперпараметров для Llama 2 (включен в пример кода):

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

Наконец, мы следим за TensorBoard, чтобы отслеживать прогресс обучения:

Полный пример кода и сценариев, которые мы упомянули, можно найти в Llama 7B. руководство и Код НеМо в Neuron SDK, чтобы выполнить более подробные шаги.

Эксперименты по точной настройке

Мы доработали модель 7B на наборах данных OSCAR (открытый сверхбольшой сканируемый корпус ALMAnaCH) и QNLI (NLI с ответами на вопросы) в среде Neuron 2.12 (PyTorch). Для каждой длины последовательности 2000 и 4000 мы оптимизировали некоторые конфигурации, такие как batchsize и gradient_accumulation, для эффективности обучения. В качестве стратегии тонкой настройки мы приняли полную тонкую настройку всех параметров (около 500 шагов), которую можно расширить до предварительного обучения с более длинными шагами и большими наборами данных (например, 1T RedPajama). Также можно включить параллелизм последовательностей, чтобы позволить NeMo Megatron успешно настраивать модели с большей длиной последовательности — 4000. В следующей таблице показаны результаты эксперимента по точной настройке Llama 7B по конфигурации и пропускной способности. Пропускная способность увеличивается почти линейно при увеличении числа экземпляров до 4.

Последний шаг — проверка точности базовой модели. Мы реализовали эталонный сценарий для экспериментов с графическим процессором и подтвердили совпадение кривых обучения для графического процессора и Trainium, как показано на следующем рисунке. На рисунке показаны кривые потерь в зависимости от количества шагов обучения в наборе данных QNLI. Смешанная точность была принята для графического процессора (синий) и bf16 со стохастическим округлением по умолчанию для Trainium (оранжевый).

Заключение

В этом посте мы показали, что Trainium обеспечивает высокую производительность и экономичную тонкую настройку Llama 2. Дополнительные ресурсы по использованию Trainium для распределенного предварительного обучения и точной настройки ваших генеративных моделей ИИ с помощью NeMo Megatron см. Справочник по AWS Neuron для NeMo Megatron.

Об авторах

Хао Чжоу — научный сотрудник Amazon SageMaker. До этого он работал над разработкой методов машинного обучения для обнаружения мошенничества для Amazon Fraud Detector. Он увлечен применением методов машинного обучения, оптимизации и генеративного искусственного интеллекта для решения различных реальных проблем. Он получил степень доктора электротехники в Северо-Западном университете.

Картик Гопалсвами — учёный-прикладник в AWS. До появления AWS он работал научным сотрудником в Uber и Walmart Labs, уделяя особое внимание смешанной целочисленной оптимизации. В Uber он сосредоточился на оптимизации сети общественного транспорта с помощью SaaS-продуктов по требованию и совместных поездок. В Walmart Labs он работал над оптимизацией цен и упаковки. Картик имеет докторскую степень в области промышленной и системной инженерии, а также степень по исследованиям операций в Университете штата Северная Каролина. Его исследования сосредоточены на моделях и методологиях, сочетающих исследование операций и машинное обучение.

Синь Хуан — старший научный сотрудник Amazon SageMaker JumpStart и встроенных алгоритмов Amazon SageMaker. Он занимается разработкой масштабируемых алгоритмов машинного обучения. Его исследовательские интересы лежат в области обработки естественного языка, объяснимого глубокого обучения табличных данных и надежного анализа непараметрической пространственно-временной кластеризации. Он опубликовал множество статей на конференциях ACL, ICDM, KDD и Королевском статистическом обществе: серия A.

Парк Ёнсук — старший научный сотрудник в AWS Annapurna Labs, занимающийся разработкой и обучением базовых моделей на ускорителях искусственного интеллекта. До этого д-р Парк работал над исследованиями и разработками Amazon Forecast в AWS AI Labs в качестве ведущего научного сотрудника. Его исследования касаются взаимодействия машинного обучения, фундаментальных моделей, оптимизации и обучения с подкреплением. Он опубликовал более 20 рецензируемых статей на ведущих площадках, включая ICLR, ICML, AISTATS и KDD, а также организовал семинары и представил учебные пособия в области временных рядов и обучения LLM. Прежде чем присоединиться к AWS, он получил степень доктора электротехники в Стэнфордском университете.

Йида Ван — ведущий научный сотрудник группы AWS AI компании Amazon. Его исследовательские интересы связаны с системами, высокопроизводительными вычислениями и анализом больших данных. В настоящее время он работает над системами глубокого обучения, уделяя особое внимание компиляции и оптимизации моделей глубокого обучения для эффективного обучения и вывода, особенно крупномасштабных базовых моделей. Миссия состоит в том, чтобы соединить модели высокого уровня из различных платформ и аппаратных платформ низкого уровня, включая центральные процессоры, графические процессоры и ускорители искусственного интеллекта, чтобы разные модели могли работать с высокой производительностью на разных устройствах.

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

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