Как стать автором
Обновить

Mamba 2 + Transformer = Nemotron H

Уровень сложностиПростой
Время на прочтение12 мин
Количество просмотров458

Содержание

  1. TL;DR

  2. Введение

  3. Дизайн гибридной архитектуры

  4. Методология обучения

  5. Сжатие модели с помощью MiniPuzzle

  6. Повышение производительности и эффективности

  7. Применение и универсальность

  8. Заключение


🔥Подпишитесь, что бы не пропустить новые материалы🔥

📌 Telegram @TheWeeklyBrief — краткие обзоры и подкасты 📰🎧

📌 GitHub Pages — углублённые статьи, код и формулы 📝💻


1. TL;DR

Основная тема

Внедрение и характеристика Nemotron-H, семейства больших языковых моделей (LLM) от NVIDIA, использующих гибридную архитектуру, сочетающую слои Transformer и Mamba для повышения эффективности и точности, особенно при обработке длинных последовательностей.

Ключевые идеи и факты

Гибридная архитектура

  • Nemotron-H сочетает сильные стороны Transformer и Mamba;

  • Слои Mamba, основанные на моделях пространства состояний (SSM), имеют постоянную вычислительную и объемную сложность на токен, что делает их очень эффективными для длинных последовательностей;

  • Слои самовнимания Transformer сохраняются в стратегически расположенных местах для захвата глобальных отношений;

"Семейство моделей Nemotron-H решает это ограничение, представляя гибридную архитектуру, которая сочетает в себе сильные стороны Transformer с эффективностью слоев Mamba."

"Разработанные NVIDIA модели Nemotron-H стратегически заменяют большую часть слоев самовнимания в Transformer слоями Mamba..."

"В отличие от самовнимания, вычислительная и объемная сложность которого масштабируются квадратично с длиной последовательности, слои Mamba предлагают постоянную вычислительную и объемную сложность на токен..."

  • Примерно 8% слоев в моделях Nemotron-H-8B/56B составляют слои самовнимания, равномерно распределенные по всей модели.

Повышенная эффективность логического вывода

  • Гибридная архитектура приводит к значительному увеличению пропускной способности при логическом выводе, особенно для длинных последовательностей (65 536 токенов).

"Nemotron-H-56B предлагает в 2,4 раза более высокую пропускную способность, чем Llama-3.1-70B, при более высоких уровнях точности." (Рисунок 1)

"Nemotron-H-56B достигает до 3 раз более высокой пропускной способности при выводе, чем Qwen-2.5-72B и Llama-3.1-70B"

"Nemotron-H-8B обеспечивает в 1,8 раза более высокую пропускную способность, чем Qwen-2.5-7B при аналогичных уровнях точности"

Высокая точность

  • Несмотря на изменения в архитектуре, модели Nemotron-H демонстрируют конкурентоспособную или превосходящую точность по сравнению с чистыми моделями Transformer аналогичного размера в широком спектре тестов.

"Nemotron-H-56B превосходит Llama-3.1-70B в 16 из 17 оцененных задач"

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

Методология обучения

Данные

  • Обучение проводилось на разнообразной смеси данных, включая веб-сканирование (59% для 56B), код (20%) и академический контент (8,8%), для развития широкого спектра возможностей;

  • Объем данных для модели 56B составил около 20 триллионов токенов.

Обучение FP8

  • Использование 8-битной арифметики с плавающей запятой (FP8) для обучения значительно снижает требования к памяти и вычислительные затраты, сохраняя при этом качество модели;

  • Метод включает текущее масштабирование, сохранение точности BF16 для определенных слоев и постепенную сходимость с обучением BF16.

"Значительным нововведением в разработке Nemotron-H является использование 8-битной арифметики с плавающей запятой (FP8) для обучения, что снижает требования к памяти и вычислительные затраты при сохранении качества модели"

"Результаты показывают, что обучение FP8 может соответствовать или превосходить производительность обучения BF16 по различным тестам" (Рисунок 5)

Сжатие модели с помощью MiniPuzzle

  • Разработана новая структура сжатия, сочетающая обрезку, поиск нейронной архитектуры и дистилляцию знаний для дальнейшего повышения эффективности развертывания;

  • MiniPuzzle включает оценку важности слоев, поиск архитектуры-кандидата и дистилляцию;

  • Nemotron-H-56B был успешно сжат до Nemotron-H-47B, что привело к сокращению параметров на 16%, сохранению сопоставимой точности и увеличению пропускной способности при логическом выводе на 20%.

Применение и универсальность

Nemotron-H разработан как универсальные базовые модели с потенциалом для различных применений:

  • Vision-Language: Базовые модели были расширены для создания моделей vision-language (VLM), демонстрирующих самые современные результаты в соответствующих тестах (VQAv2, GQA, VizWiz);

  • Генерация кода: Модели демонстрируют сильные возможности в задачах, связанных с кодом, что связано со значительным объемом данных кода в обучающем наборе;

  • Обработка длинного контекста: Гибридная архитектура особенно хорошо подходит для эффективной обработки длинных контекстов;

  • Адаптация данных для способностей: Распределение данных обучения может быть скорректировано для развития специфических возможностей, например, компетенций в STEM-областях, без изменения архитектуры.

Вывод

Семейство Nemotron-H представляет собой значительный шаг вперед в разработке LLM, успешно решая проблемы эффективности традиционных Transformer-ов при сохранении высокой точности. Гибридная архитектура, инновационная методология обучения FP8 и структура сжатия MiniPuzzle делают Nemotron-H эффективным и мощным решением для обработки длинных контекстов и различных приложений. Ожидается, что доступность этих моделей в популярных фреймворках будет способствовать их более широкому внедрению и дальнейшим исследованиям гибридных архитектур.


2. Введение

Большие языковые модели (LLM) продемонстрировали замечательные возможности в различных задачах, но их вычислительные потребности во время логического вывода остаются серьезной проблемой, особенно для обработки длинных последовательностей. Семейство моделей Nemotron-H решает это ограничение, представляя гибридную архитектуру, которая сочетает в себе сильные стороны Transformer с эффективностью слоев Mamba.

Figure_01
Рисунок 1. Сравнение пропускной способности и точности.

Рисунок 1: Сравнение моделей Nemotron-H с другими современными LLM с точки зрения пропускной способности (токенов/с/GPU) и точности на эталонном тесте MMLU. Nemotron-H-56B предлагает в 2,4 раза более высокую пропускную способность, чем Llama-3.1-70B, при более высоких уровнях точности.

Разработанные NVIDIA модели Nemotron-H стратегически заменяют большую часть слоев самовнимания в Transformer слоями Mamba, которые основаны на моделях пространства состояний (SSM). В отличие от самовнимания, вычислительная и объемная сложность которого масштабируются квадратично с длиной последовательности, слои Mamba предлагают постоянную вычислительную и объемную сложность на токен, что делает их особенно эффективными для генерации длинных последовательностей.

Ключевое новшество Nemotron-H заключается в тщательном балансировании этих двух архитектурных парадигм для поддержания или улучшения точности при значительном увеличении скорости логического вывода. Этот подход отвечает критической потребности в сообществе LLM в моделях, которые могут эффективно обрабатывать длинные контексты без ущерба для производительности.

3. Дизайн гибридной архитектуры

Архитектура Nemotron-H сочетает в себе слои Mamba-2 с традиционными компонентами Transformer для создания сбалансированной гибридной конструкции. Структура модели стратегически располагает слои самовнимания, чтобы использовать их сильные стороны в захвате глобальных отношений, в то время как слои Mamba используются для эффективной обработки последовательностей.

Figure_02
Рисунок 2. Архитектуры моделей Nemotron-H-8B/56B.

Рисунок 2. Архитектуры моделей Nemotron-H-8B/56B. Примерно 8% от общего числа слоев в модели составляют слои самовнимания (self-attention); эти слои равномерно распределены по всей модели. Остальная часть модели состоит из чередующихся слоев Mamba-2 и FFN.

Как показано на рисунке 2, как Nemotron-H-8B, так и Nemotron-H-56B следуют схожей схеме: серия начальных пар слоев Mamba-2 и FFN, за которой следует средняя секция, которая включает в себя один слой внимания среди нескольких пар Mamba-2 и FFN, и завершается дополнительными слоями Mamba-2 и FFN. Ключевое различие заключается в количестве повторений — в то время как Nemotron-H-8B имеет 4 повторения в своей средней секции, Nemotron-H-56B имеет 10.

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

В архитектуре используются несколько ключевых компонентов:

  1. Слои Mamba-2: обеспечивают эффективное моделирование последовательностей с постоянной вычислительной сложностью на токен;

  2. Слои самовнимания: стратегически расположены для захвата глобальных отношений;

  3. Прямые нейронные сети (FFN): обрабатывают выходы слоев Mamba и внимания.

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

4. Методология обучения

Модели Nemotron-H обучались с использованием комбинации инновационных подходов для обеспечения высокой производительности и эффективности:

4.1 Курирование и подготовка данных

Обучающие данные состояли из разнообразной смеси источников, тщательно сбалансированных для развития различных возможностей:

Figure_03
Рисунок 3. Распределение источников данных предварительного обучения для моделей Nemotron-H

Рисунок 3. Распределение источников данных предварительного обучения для моделей Nemotron-H, показывающее баланс между веб-сканированием, кодом, академическими и другими типами данных.

Модель 56B была обучена примерно на 20 триллионах токенов, при этом данные веб-сканирования составляли наибольшую часть (59%), за ними следовали код (20%) и академический контент (8,8%). Смесь данных была разработана для обеспечения всестороннего охвата общих знаний при одновременном развитии сильных возможностей в специализированных областях, таких как кодирование и математика.

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

4.2 Рецепт обучения FP8

Значительным нововведением в разработке Nemotron-H является использование 8-битной арифметики с плавающей запятой (FP8) для обучения, что снижает требования к памяти и вычислительные затраты при сохранении качества модели:

Относительная разница в потерях при обучении между FP8 и BF16
Рисунок 4. Относительная разница в потерях при обучении между FP8 и BF16

Рисунок 4. Относительная разница в потерях при обучении между FP8 и BF16, показывающая сходимость по мере прогресса обучения.

Рецепт обучения FP8 включает в себя:

  • Текущее масштабирование для каждого тензора для повышения стабильности;

  • Сохранение первых и последних четырех GEMM модели в точности BF16;

  • Постепенную сходимость с обучением BF16 с течением времени.

Результаты показывают, что обучение FP8 может соответствовать или превосходить производительность обучения BF16 по различным тестам:

Сравнение обучения FP8 и BF16
Рисунок 5. Сравнение обучения FP8 и BF16

Рисунок 5. Сравнение обучения FP8 и BF16 по различным тестам, показывающее сопоставимую или лучшую производительность при обучении FP8.

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

5. Сжатие модели с помощью MiniPuzzle

Чтобы еще больше повысить эффективность развертывания, исследователи разработали MiniPuzzle, новую структуру сжатия, которая сочетает в себе обрезку, поиск нейронной архитектуры и дистилляцию знаний:

Рабочий процесс структуры сжатия MiniPuzzle
Рисунок 6. Рабочий процесс структуры сжатия MiniPuzzle

Рисунок 6. Рабочий процесс структуры сжатия MiniPuzzle, показывающий переход от предварительно обученной модели к сжатой модели посредством оценки важности, поиска нейронной архитектуры и дистилляции.

Подход MiniPuzzle состоит из нескольких этапов:

1. Оценка важности: анализ вклада каждого слоя в производительность модели

def importance_estimation(model: Any, dataset: Any) -> List[float]:
    """
    Description:
    ---------------
        Вычисляет оценки важности для каждого слоя модели на основе
        влияния слоя на функцию потерь при его временном отключении.
    """

    # Вычисление оценок важности для каждого слоя
    scores: List[float] = []
    for layer in model.layers:
        # Зануление выходов слоя и измерение влияния на потери
        layer_score = measure_impact_on_loss(model, layer, dataset)
        scores.append(layer_score)

    return scores

2. Анализ важности слоев: понимание того, какие слои вносят наибольший вклад в производительность модели

Оценки важности для каждого слоя
Рисунок 7. Оценки важности для каждого слоя

Рисунок 7. Оценки важности для каждого слоя, показывающие различные вклады различных типов слоев модели.

3. Условный поиск нейронной архитектуры: изучение кандидатов на сжатую архитектуру

Шаблоны выбора слоев для различных кандидатов
Рисунок 8. Шаблоны выбора слоев для различных кандидатов

Рисунок 8. Шаблоны выбора слоев для различных кандидатов на архитектуру, показывающие, какие слои сохраняются в каждой потенциальной сжатой модели.

4. Компромисс между памятью и производительностью: оценка моделей на основе использования памяти и точности

Поиск компромисса между памятью и производительностью
Рисунок 9. Поиск компромисса между памятью и производительностью

Рисунок 9. Компромисс между предполагаемой нагрузкой на память и производительностью в тестах для архитектур-кандидатов.

5. Дистилляция знаний: обучение сжатой модели, чтобы соответствовать или превосходить возможности исходной

Благодаря этому процессу Nemotron-H-56B был успешно сжат до Nemotron-H-47B, что позволило сократить количество параметров на 16%, сохранив сопоставимую точность и повысив пропускную способность при выводе на 20%.

6. Повышение производительности и эффективности

Модели Nemotron-H демонстрируют значительное повышение производительности и эффективности по сравнению с сопоставимыми моделями на основе Transformer:

Пропускная способность при выводе

Гибридная архитектура обеспечивает значительно более быстрый вывод, особенно для длинных последовательностей:

  • Nemotron-H-56B достигает до 3 раз более высокой пропускной способности при выводе, чем Qwen-2.5-72B и Llama-3.1-70B;

  • Nemotron-H-8B обеспечивает в 1,8 раза более высокую пропускную способность, чем Qwen-2.5-7B при аналогичных уровнях точности.

Сравнение Nemotron-H-8B с моделями аналогичного размера
Рисунок 10. Сравнение Nemotron-H-8B с моделями аналогичного размера

Рисунок 10. Сравнение Nemotron-H-8B с моделями аналогичного размера с точки зрения пропускной способности и точности.

Эти улучшения эффективности особенно заметны при обработке длинных последовательностей (65 536 токенов в представленных примерах), что подчеркивает преимущество постоянной вычислительной сложности слоев Mamba на токен.

Точность по результатам тестов

Несмотря на архитектурные изменения, модели Nemotron-H сохраняют высокую производительность в широком диапазоне тестов:

  • Nemotron-H-56B превосходит Llama-3.1-70B в 16 из 17 оцененных задач;

  • Модели демонстрируют особенно высокую производительность в задачах математического мышления.

Сравнение Nemotron-H и других моделей в тесте MMLU
Рисунок 11. Сравнение Nemotron-H и других моделей в тесте MMLU

Рисунок 11. Сравнение Nemotron-H и других моделей в тесте MMLU, показывающее конкурентоспособную производительность.

Модели были оценены с помощью комплексного набора тестов, включая MMLU, GSM8K, MATH, HumanEval и различные задачи на рассуждение, неизменно демонстрируя конкурентоспособную или превосходящую производительность по сравнению с моделями Transformer аналогичного размера.

7. Применения и универсальность

Модели Nemotron-H были разработаны как универсальные базовые модели, которые можно адаптировать для различных приложений:

Возможности Vision-Language

Базовые модели были расширены для создания моделей vision-language (VLM) в соответствии с архитектурой NVLM-D. Эти VLM продемонстрировали самую современную производительность в тестах, таких как VQAv2, GQA и VizWiz, что показывает адаптируемость гибридной архитектуры к мультимодальным задачам.

Генерация кода

Модели демонстрируют особенно высокую производительность в задачах, связанных с кодом. Включение значительной доли данных кода (20%) в обучающую смесь способствует их способности понимать и генерировать высококачественный код на нескольких языках программирования.

Обработка длинного контекста

Одним из наиболее значительных преимуществ гибридной архитектуры является ее способность эффективно обрабатывать длинные контексты. Модель Nemotron-H-8B была специально настроена для возможностей длинного контекста, демонстрируя высокую производительность в тесте RULER и других задачах оценки длинного контекста.

Распределение данных для различных способностей

Исследователи тщательно настроили распределение данных для различных этапов обучения, чтобы развить определенные способности:

Распределение данных обучения
Рисунок 12. Распределение данных обучения

Рисунок 12. Распределение данных обучения, оптимизированное для STEM-компетенций, с повышенным акцентом на математический и кодовый контент.

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

8. Заключение

Семейство моделей Nemotron-H представляет собой значительный прогресс в разработке БЯМ, успешно сочетая сильные стороны архитектур Transformer и Mamba для создания моделей, которые являются одновременно точными и эффективными. Ключевые достижения включают:

  1. Гибридную архитектуру, которая стратегически интегрирует слои Mamba с механизмом самовнимания для баланса производительности и эффективности.

  2. Рецепт обучения FP8, который снижает вычислительные и требования к памяти, сохраняя при этом качество модели.

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

  4. Демонстрация значительного ускорения вывода (до 3 раз) по сравнению с Transformer-моделями аналогичного размера.

  5. Конкурентоспособная или превосходящая точность в широком диапазоне бенчмарк-задач.

Успех моделей Nemotron-H указывает на то, что гибридные архитектуры представляют собой перспективное направление для будущего развития БЯМ, особенно по мере того, как приложения все больше требуют эффективной обработки длинных контекстов. Решая вычислительные узкие места традиционных Transformer-ов, сохраняя при этом их сильные стороны, Nemotron-H предлагает практическое решение для развертывания мощных языковых моделей в средах с ограниченными ресурсами.

Планируемый выпуск этих моделей с поддержкой в популярных фреймворках, таких как Hugging Face, NeMo и Megatron-LM, позволит более широкому AI-сообществу воспользоваться этими достижениями и в дальнейшем изучить потенциал гибридных архитектур для языкового моделирования.

Теги:
Хабы:
0
Комментарии1

Публикации

Ближайшие события