Быстрая и экономически эффективная настройка LLaMA 2 с помощью AWS Trainium

Быстрая и экономически эффективная настройка LLaMA 2 с использованием AWS Trainium

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

Недавно Meta выпустила Llama 2 для исследователей и коммерческих компаний, добавив её в список других LLM, включая MosaicML MPT и Falcon. В этом посте мы рассмотрим, как настроить Llama 2 на AWS Trainium, специализированном ускорителе для тренировки LLM, чтобы сократить время и стоимость обучения. Мы рассмотрим предоставленные скрипты для настройки, используемые конфигурации и полученные результаты пропускной способности.

О модели Llama 2

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

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

Для совместимости Llama 2 с длиной последовательности 2 000 и 4 000 мы реализовали скрипт с использованием NeMo Megatron для Trainium, который поддерживает параллелизм данных (DP), параллелизм тензоров (TP) и параллелизм конвейера (PP). Для конкретики мы используем общий скрипт GPT Neuron Megatron-LM с новой реализацией некоторых функций, таких как развязывание вектора вложения слов, вращающееся вложение, RMSNorm и активация Swiglu, для поддержки скрипта обучения Llama 2.

• Создайте приложение для суммирования HCLS-документов с помощью Falcon, используя Amazon SageMaker JumpStart
• Прогнозирование временных рядов с помощью рекуррентных нейронных сетей
• Объявляем о новых инструментах, чтобы помочь каждому бизнесу принять Генеративный ИИ

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

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

wget https://data.together.xyz/redpajama-data-1T/v1.0.0/book/book.jsonlpython 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:

tensorboard --logdir ./

Для полного примера кода и скриптов, о которых мы упомянули, обратитесь к учебнику Llama 7B tutorial и NeMo коду в Neuron SDK для более подробных шагов.

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

Мы провели настройку 7B модели на наборе данных OSCAR (Open Super-large Crawled ALMAnaCH coRpus) и QNLI (Question-answering NLI) в среде Neuron 2.12 (PyTorch). Для каждой длины последовательности 2 000 и 4 000 мы оптимизировали некоторые конфигурации, такие как batchsize и gradient_accumulation, для обеспечения эффективности обучения. В качестве стратегии доводки мы выбрали полную доводку всех параметров (около 500 шагов), которая может быть расширена до предварительного обучения с дольшей длительностью и большими наборами данных (например, 1Т RedPajama). Параллельность последовательности также может быть включена для успешной доводки моделей с большей длиной последовательности 4 000 NeMo Megatron. В таблице ниже показаны конфигурация и результаты пропускной способности эксперимента по доводке Llama 7B. Пропускная способность практически линейно масштабируется при увеличении количества экземпляров до 4.

Distributed Library	Набор данных	Длина последовательности	Количество экземпляров	Асинхронная параллельность	Параллельность данных	Параллельность канала	Глобальный размер пакета	Пропускная способность (посл./с)
Neuron NeMo Megatron	OSCAR	4096	1	8	4	1	256	3.7
.	.	4096	2	8	4	1	256	7.4
.	.	4096	4	8	4	1	256	14.6
.	QNLI	4096	4	8	4	1	256	14.1

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

Заключение

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