Обзор технического отчёта о моделях Llama 3

Введение

Недавно Meta выпустила технический отчёт, в котором исследователи описали тренировку новых моделей Llama 3 и 3.1, а также оценили их качество.

В нашей статье мы не затронем описание тренировки моделей на большом кластере (об этом можно дополнительно прочитать в отчёте), однако подробно рассмотрим сбор данных, трейнинг и оценку качества моделей. Давайте разбираться со всем по порядку 😊

Есть две модификации Llama 3 — это Llama 3 и Llama 3.1. Формат первой характеризуется наличием 8 и 70 миллиардов параметров. Формат второй — уже 405 миллиардами параметров. В Llama 3.1 появились также новые приятные фичи, например, длинный контекст и вызов тулзов.

Все модели предобучались на 15 триллионах токенов. Максимальная длина контекстного окна равна 128 тысячам токенов.

По качеству на тестах самая большая модель выдаёт результаты, сопоставимые с результатами GPT-4:

Обучение LLM

Обучение современных LLM включает три этапа:

1. Pre-training (пре-трейнинг) — обучение на неструктурированных данных;
2. Instruction tuning (пост-трейнинг) — работа с инструкциями;
3. Aligment (пост-трейнинг) — работа с безопасностью.

Пре-трейнинг

Сначала большой многоязычный корпус разбивается на токены. Модель обучается предсказывать каждый следующий токен. На этапе предобучения она изучает структуру языка и получает огромное количество знаний о мире из текстов, которые «‎читает»‎.

Для эффективности процесса нужны данные большого размера: модель с 405 миллиардами параметров обучается на 15.6 триллионах токенов, используя контекстное окно в 8 тысяч токенов. Затем идёт следующий этап предобучения, где контекстное окно расширяется уже до 128 тысяч токенов (но об этом мы поговорим чуть позже).

Пост-трейнинг

На этапе пре-трейнинга модель уже довольно хорошо понимает язык и мир, но ведёт себя не совсем так, как нам бы хотелось. Поэтому во время пост-трейнинга её несколько раз файнтюнят с помощью фидбека людей. Это делают через supervise finetuning (SFT) на инструкциях (почитать можно здесь и здесь) и Direct Preference Optimization (DPO). Также модель обучается использовать инструменты, или tools (подробнее разберём это в следующей части). Данный этап отвечает за улучшение и других областей — написание кода и рассуждение (reasoning). А ещё здесь принимаются меры по обеспечению безопасности модели.

Обучение Llama мультимодальности

В итоге мы имеем модель, которая знает 8 языков, умеет писать высококачественный код и производить довольно сложные рассуждения. А ещё работает с картинками, видео и речью. Общая структура выглядит таким образом:

Обучение мультимодальности включает несколько этапов:

1. Предобучение языковой модели;
2. Предобучение мультимодального энкодера;
3. Обучение адаптера для зрения (картинки, видео);
4. Обучение адаптера для речи;
5. Файнтюнинг всей модели.

Давайте рассмотрим каждый этап подробнее 🙂

Предобучение языковой модели

Эта стадия содержит:

1. Отбор крупного обучающего корпуса;
2. Разработку архитектуры модели и соответствующих законов масштабирования для определения её оптимальных размеров;
3. Выбор методов для эффективного пре-трейнинга (у нас много данных, поэтому огромная модель);
4. Нахождение «‎рецепта»‎ пре-трейнинга.

Данные для пре-трейнинга

Источники данных в статье не указаны, но время публикации их большей части — весь период до окончания 2023 года. Из данных были исключены персональная информация, вредный контент и порнография.

We remove domains that contain large amounts [хоть и непонятно, сколько именно] of personally identifiable information (PII), and domains with known adult content.

Отбор данных происходил следующим образом:

1. Авторы очистили информацию на сайтах, которые, по их мнению, содержат личные данные, вредный контент или контент для взрослых [«вредность»‎ определялась по внутренним стандартам Meta].
2. Написали специальный парсер для парсинга HTML. Он тестировался людьми и сравнивался с парсерами других компаний. Особый акцент был на страницах с математикой и кодом. Для изображений сохранялся alt. Маркдауновские теги убирались: markdown вредит моделям, обученным на интернет-данных.
3. Сохранили только последнюю версию документа на уровне URL; удалили дубликаты документов с помощью MinHash; через ccNet убрали дубликаты строчек, которые встречаются более 6 раз в каждом блоке из 30 миллионов документов. Так мы избавляемся от навигационных меню и других шаблонов с веб-сайтов, но иногда упускаем важную информацию. Тем не менее фильтрация заметно улучшает качество модели, если судить по эмпирическим оценкам.
4. Удалили все строки с частыми повторами (логгеры, сообщения об ошибке) через duplicated n-gram coverage ratio.
5. Убрали статьи с «плохими‎» словами [пример].
6. Применили дивергенцию Кульбака-Лейблера по распределению токенов для фильтрации документов, содержащих большое количество выбросов в сравнении с распределением обучающего корпуса (то есть документов с непонятным распределением нестандартных токенов).
7. Применили разные классификаторы для выбора высококачественных токенов: быстрые классификаторы как fasttext, обученные для распознавания текста (Touvron et al., 2023a), а также более ресурсоёмкие классификаторы на основе Roberta, обученные на предсказаниях Llama 2.
8. Создали набор очищенных веб-документов для обучения классификатора на основе Llama 2 и описали требования к качеству. В итоге чат-модель Llama 2 смогла определять соответствие документов этим требованиям. Для повышения эффективности использовалась DistilRoberta, чтобы генерировать оценки качества каждого документа. После различных комбинаций нашлось экспериментально лучшее сочетание.

Данные для генерации кода и логического вывода, или reasoning

Авторы создали специальные пайплайны для извлечения веб-страниц, связанных с программированием и математикой [аналогично DeepSeek-AI et al.(2024)]. В датасете они подобрали (натюнили) промты для Llama 2, разметили с их помощью нескольких веб-страниц и обучили классификатор на DistilRoberta. Затем распарсили полученные HTML’ки (но не сказано, каким образом).

Данные для языков помимо английского

Процесс схож с пайплайном для английского языка:

1. Удаление «‎плохой»‎ (личной, взрослой, вредной) информации.
2. Определение языка документа (всего их 176) через классификатор на основе fasttext (предположительно простая линейная модель на эмбеддингах).
3. Удаление дубликатов на уровне URL, документов и строк.
4. Удаление низкокачественных документов с помощью учёта особенностей конкретного языка и дополнительных моделей.
5. Проверка документов через классификатор, основанный на многоязычной Llama 2 [to ensure that high-quality content is prioritized].
6. Контроль количества токенов, приходящих из других языков, для баланса между качеством работы на английском и остальных языках.

Cмешение разных типов данных, или Data mix

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

Обучать огромную модель несколько раз на основе разных сплитов трудозатратно. Поэтому авторы решили обучить относительно маленькие модели, чтобы оценить, как при данном сплите по категориям будет вести себя большая модель. После этого предсказывалось её качество — предполагалось обучение с таким же сплитом, уточнялось соотношение классов, повторялся процесс обучения маленьких моделей.

Предсказание поведения большой модели на основе поведения маленьких не так очевидно, однако оно даёт возможность экстраполировать, как поведут себя модели при увеличении количества параметров. Подробнее о scaling laws можно прочитать в статье.

В итоге соотношение данных следующее: 50% токенов относится к общим вещам (general knowledge), 25% токенов — к математике и рассуждениям, 17% токенов — к коду, 8% токенов — к другим (человеческим) языкам.

Отжиг данных, или Data annealying

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

Подобно Li et al. (2024b), они провели отжиг с использованием смеси данных, на которой валидировали влияние новых высококачественных источников на результат. Поскольку процесс начинаем с одного и того же чек-пойнта, мы и без полного переобучения модели понимаем, как собрать финальный Data mix для достижения наилучших метрик.

Данные часто используемых бенчмарков не включались в датасет для отжига. Это позволило оценить возможности модели Llama 3 решать задачу при наличии примеров решения схожих задач (Few-Short learning, в контексте LLM подаются на вход в виде промта) и обобщении на основе данных из других доменов.

Как и в OpenAI (2023a), эффективность отжига оценивалась на тренировочных наборах GSM8k (Cobbe et al., 2021) и MATH (Hendrycks et al., 2021b). Авторы выяснили, что отжиг улучшает производительность предварительно обученной модели Llama 3 8B на валидационных наборах GSM8k и MATH на 24,0% и 6,4% соответственно. Однако улучшения для модели 405B незначительны, следовательно, большая флагманская модель обладает сильными способностями к обучению и рассуждению в контексте и не требует специфических тренировочных примеров из данной области для достижения высокой производительности.

Отжиг для оценки качества данных

Подобно Blakeney et al. (2024), отжиг позволяет судить о ценности небольших домен-специфических наборов данных. Она измерялась с помощью многих экспериментов, где на старте с 50% всех эпох авторы учили модель, снижая lr до 0 в процессе обучения на 40B токенов. Здесь 30% веса назначались новому набору данных и 70% — основному набору данных.

Использование отжига для оценки новых источников данных более эффективно, чем проведение экспериментов по законам масштабирования для каждого небольшого набора данных.

Предобучение мультимодального энкодера

Есть два энкодера: один для изображений и видео, второй для речи. Первый обучается на парах «‎картинка — текст». Второй — на предсказании маскированных токенов в дискретном представлении речи (это self supervised: сначала речь переводится в токены, затем энкодер учится предсказывать маскированные токены). Подробнее об этом мы поговорим в следующей части.

Обучение адаптера для зрения (картинки, видео)

Адаптер состоит из нескольких слоёв cross-attention, подающих представления из энкодера в языковую модель. При обучении на парах «‎картинка — текст» веса адаптера и энкодера разморожены, а веса языковой модели заморожены.

Видео-адаптер тренируется поверх адаптера для изображений на парах «‎видео — текст». Это позволяет модели собирать информацию на основе разных кадров (отдельные картинки она распознавать уже умеет, значит, мы даём ей последовательность картинок, модель учится понимать её смысл).

Обучение адаптера для речи

Он обучается последним. Его задача — перевод представления речи из энкодера в языковую модель. При обучении параметры энкодера и адаптера разморожены, а параметры языковой модели заморожены. Также есть отдельная надстройка, которая умеет переводить текст в речь (поговорим также об этом в следующей части 😉 ).

Архитектура и размеры модели

У llama 3 стандартная архитектура трансформера [Attention is all you need, см. рисунок 3), которая несильно отличается от архитектуры Llama и Llama 2.

Основные отличия от Llama и Llama 2:

1) Grouped-query Attention (GQA) с 8-ю головами. В основе лежит идея усреднения части голов Key и Value, следовательно, увеличение скорости инференса и уменьшение размера key-value cache на декодинге [для подробного изучения можно посмотреть нашу лекцию по эффективным механизмам внимания].

2) Матрица внимания. Она строится так, чтобы самовнимание между различными документами в пределах одной последовательности занулялось (они, скорее всего, не связаны между собой). Это несильно влияет на пре-трейнинг, однако заметно помогает при обучении на очень длинных последовательностях.

3) Словарь модели. Он содержит 128 тысяч токенов (речь идёт о языковой модели без учёта токенов картинок и речи). Словарь токенов сочетает 100 тысяч токенов от токенизатора tiktoken с 28 тысячами дополнительных токенов для наилучшей поддержки языков помимо английского. По сравнению с токенизатором Llama 2 новый токенизатор улучшает коэффициент сжатия на образце данных для английского языка с 3.17 до 3.94 символов на токен. Это позволяет модели «‎читать» больше текста при тех же вычислительных затратах на обучение. А ещё добавление 28 тысяч токенов из языков, отличных от английского, улучшает как коэффициенты сжатия, так и производительность на downstream задачах без влияния на токенизацию английского языка.

4) Базовая частота RoPE. Она увеличивается до 500,000. Это позволяет лучше поддерживать длинные контексты. Как показали Xiong et al. (2023), данное значение хорошо работает для контекста длиною до 32,768.

RoPE (Rotary Position Embedding) — метод кодирования позиций, используемый в моделях трансформеров для обработки последовательностей данных.

В отличие от традиционных методов позиционного кодирования, RoPE использует матрицы поворота для кодирования информации о позиции каждого токена в последовательности. Это позволяет модели лучше захватывать относительные позиции токенов, что особенно полезно для работы с длинными последовательностями [подробнее тут].

В итоге Llama 3 на 405 миллиардах параметров содержит 126 слоёв, размерность представления токенов равна 16384,128 голам внимания и 8 голам Key/Value.

Чтобы определить оптимальный размер модели, авторы исследовали зависимость качества от размера (scaling laws) для небольших моделек, а затем экстраполировали (предсказывали качество большой модели, обученной на том же наборе данных).

Для параллелизации использовалась техника, которую они назвали 4D-параллелизацией:

1. Параллелизм на уровне тензоров (Tensor parallelism) разделяет отдельные тензоры весов на несколько частей, распределяя их по разным устройствам.
2. Параллелизм на уровне пайплайна (Pipeline parallelism) вертикально разделяет модель на стадии по слоям так, что разные устройства могут параллельно обрабатывать различные стадии полного конвейера модели.
3. Параллелизм на уровне контекста (Context parallelism) разделяет входной контекст на сегменты, что снижает нагрузку на память при очень длинных входных последовательностях.
4. Полностью шардированный параллелизм данных (FSDP; Rajbhandari и др., 2020; Ren и др., 2021; Zhao и др., 2023b) шардирует модель, оптимизатор и градиенты, реализуя при этом параллелизм данных. Он, в свою очередь, параллельно обрабатывает данные на нескольких GPU и синхронизируется после каждого шага обучения. В их реализации FSDP для Llama 3 шардируются состояния оптимизатора и градиенты, но для шардов модели не выполняется повторное шардирование после прямого прохода, чтобы избежать дополнительной коммуникации all-gather во время обратного прохода.

Технические детали обучения модели

Есть три основных стадии обучения (которое они называют пре-трейнингом):

1. Обычный пре-трейнинг;
2. Пре-трейнинг с длинным контекстом;
3. Отжиг.

Обычный пре-трейнинг

Llama 3 на 405B обучается при использовании AdamW с максимальной скоростью обучения 8 × 10⁻⁵, линейного разогрева (постепенного увеличения learning rate) в течение 8000 шагов и косинусного расписания скорости обучения, уменьшающегося до 8 × 10⁻⁷ за 1 200 000 шагов.

На ранних этапах применяется меньший размер батча для улучшения стабильности обучения. Затем его размер увеличивается для повышения эффективности. Начальный размер батча — 4M токенов, а последовательность равна 4096. После обучения на 252M токенов и 2.87 триллионах токенов размер батча удваивается до 8M и 16M соответственно.

Такой режим обучения очень стабилен: есть немногочисленные пики потерь, а также не нужно корректировать расхождения модели при обучении.

Чтобы повысить производительность модели на конкретных downstream задачах, авторы меняли соотнощение данных в сравнении с пре-трейном. В частности увеличивали процент данных на языках, отличных от английского, чтобы улучшить результаты.

Кроме того, они увеличивали долю математических данных, чтобы улучшить математическое мышление модели.

Также на поздних стадиях обучения, чтобы расширить порог знаний модели, авторы добавляли более свежие веб-данные и уменьшали долю данных, идентифицированных как данные низкого качества.

Пре-трейнинг с длинным контекстом

На последних этапах пре-трейнинга размер последовательностей специально увеличивался, чтобы прийти к контекстному окну в 128K токенов. Почему размер контекста увеличивается только в конце? Потому что вычислительная сложность self-attention растёт квадратично с увеличением размера контекста.

Контекст увеличивался не сразу, а в 6 этапов, начиная с 8K и заканчивая 128K. Успешность адаптации модели к новому размеру окна измерялась за счёт её производительности на оценках с коротким контекстом и успешного решения задачи «needle in a haystack» для новой длины контекста.

«‎Needle in a haystack»‎ («‎иголка в стоге сена») — задача поиска специально введеной информации в промте. Например, редких токенов.

Этот этап предварительного обучения на длинных контекстах выполнялся с использованием приблизительно 800 миллиардов тренировочных токенов.

Отжиг

Во время предварительного обучения на последних 40 миллионах токенов скорость линейно уменьшалась до 0, при этом поддерживалась длина контекста в 128K токенов.

При фазе уменьшения скорости обучения корректировался состав данных — увеличивалась доля источников данных очень высокого качества.

Финальный этап — усреднение чек-пойнтов, полученных во время отжига (подробнее об этом можно прочитать в статье B. T. Polyak. New stochastic approximation type procedures. Automation and Remote Control, 7(7), 1991).

Заключение

Итак, сегодня мы познакомились с общей идей отчёта, рассмотрели подготовку данных и пре-трейнинг моделей семейства Llama 3 и 3.1, трейнинг адаптеров для видео и речи.

Оставайтесь с нами — уже в следующей части мы обсудим файнтюнинг модели и его результаты 🙂