Вычислительный канат современного ИИ
Большие языковые модели (LLMs) являются столпами современного искусственного интеллекта, демонстрируя выдающиеся возможности, которые преобразуют отрасли и научные открытия. Их мастерство в генерации человекоподобного текста, обеспечении работы сложных диалоговых агентов и даже помощи в сложных исследовательских задачах сделало их незаменимыми инструментами. В основе этих мощных моделей лежит архитектура transformer, характеризующаяся чередующимися слоями. Входные данные, разбитые на токены, проходят через последовательность механизмов внимания, которые взвешивают важность различных токенов, а затем через сети прямого распространения (FFNs), которые обрабатывают полученную информацию. Эта слоистая, последовательная обработка является фундаментальной для того, как transformers обучаются и генерируют выходные данные.
Однако именно эта архитектура, будучи эффективной, представляет собой растущую проблему по мере увеличения размера и сложности моделей. Последовательный характер означает, что каждый слой обычно должен ждать завершения вычислений предыдущего слоя, прежде чем начать свои. Эта пошаговая обработка создает неотъемлемое узкоеместо, особенно на этапе инференса – стадии, когда обученная модель фактически используется для генерации прогнозов или текста. По мере того как модели, подобные тем, что лежат в основе продвинутых ИИ-ассистентов, включают сотни миллиардов или даже триллионы параметров, вычислительные ресурсы и время, необходимые для инференса, резко возрастают. Этот растущий спрос приводит к значительной задержке (latency), снижению пропускной способности (throughput) и увеличению операционных расходов, что препятствует широкому развертыванию и применению самых мощных LLMs в реальном времени. Следовательно, повышение эффективности инференса стало первостепенной задачей в сообществе исследователей ИИ, стимулируя поиск инновационных стратегий, которые могут оптимизировать вычисления без ущерба для выдающейся производительности, предлагаемой этими моделями. Основная проблема заключается в смягчении ограничений, налагаемых последовательным выполнением, особенно в распределенных средах, где вычисления охватывают несколько GPU, добавляя коммуникационные издержки ко времени обработки.
Навигация по ландшафту оптимизации: существующие инструменты и их пределы
В постоянных усилиях по созданию более экономичных и быстрых LLMs исследователи разработали набор методов оптимизации. Каждый из них предлагает путь к эффективности, но часто сопряжен со своими компромиссами, что мешает какому-либо одному методу стать универсальным решением. Понимание этих компромиссов имеет решающее значение для осознания необходимости новых подходов, таких как FFN Fusion.
Один из выдающихся методов – квантование (quantization). Он включает в себя снижение числовой точности, используемой для представления весов и активаций модели. Вместо стандартных 32-битных чисел с плавающей запятой модели могут использовать 16-битные, 8-битные или даже представления с еще меньшим количеством бит. Это напрямую сокращает объем памяти, занимаемый моделью, и может значительно ускорить вычисления, поскольку операции с числами меньшей точности обычно выполняются быстрее и требуют меньше энергии. Однако квантование не лишено риска. Снижение точности может привести к потере информации, потенциально ухудшая точность модели. Этот риск становится более выраженным при очень низких битовых ширинах, требуя тщательной реализации и иногда дообучения для смягчения падения точности. Задача состоит в том, чтобы найти золотую середину, которая максимизирует прирост эффективности, удерживая деградацию производительности в допустимых пределах.
Другая распространенная стратегия – прунинг (pruning). Этот метод основан на принципе, что многие параметры в большой нейронной сети могут быть избыточными или вносить минимальный вклад в конечный результат. Алгоритмы прунинга идентифицируют и удаляют эти менее важные соединения или нейроны, что приводит к созданию меньшей, более разреженной модели. Как и квантование, прунинг сокращает требования к памяти и вычислительную нагрузку. Однако точное определение того, какие параметры “безопасно” удалять, является сложной задачей. Агрессивный прунинг может непреднамеренно удалить критически важные компоненты, что приведет к существенной потере точности. Часто необходимо дообучение модели после прунинга для восстановления производительности, что усложняет рабочий процесс. Тщательная калибровка необходима для обеспечения эффективности обрезанной модели.
Более архитектурно отличный подход – модель Mixture-of-Experts (MoE). Вместо обработки каждого входа всей сетью, модели MoE состоят из нескольких “экспертных” подсетей (обычно FFNs). Для каждого входного токена механизм маршрутизации (gating mechanism) динамически выбирает небольшое подмножество этих экспертов для выполнения вычислений. Это условное вычисление означает, что для любого данного входа активируется только часть общих параметров модели, что приводит к значительной экономии вычислений, особенно во время обучения и инференса очень больших моделей. Модели MoE могут масштабироваться до триллионов параметров, сохраняя при этом разумные вычислительные затраты. Однако их эффективность сильно зависит от рабочей нагрузки. Они превосходно справляются с очень большими размерами батчей, где селективная схема активации приводит к хорошему использованию оборудования. При меньших или средних размерах батчей модели MoE могут страдать от недостаточного использования вычислительных ресурсов, поскольку параллельное оборудование может не быть постоянно загружено разреженно активируемыми экспертами. Кроме того, реализация и балансировка нагрузки моделей MoE могут быть сложнее, чем развертывание стандартных “плотных” архитектур.
Хотя квантование, прунинг и модели MoE представляют собой ценные достижения в оптимизации LLM, их неотъемлемые ограничения подчеркивают необходимость альтернативных или дополнительных стратегий. Поиск продолжается методов, которые могут обеспечить широкое повышение эффективности в различных сценариях, в идеале с меньшими компромиссами в отношении точности или сложности реализации, особенно для плотных архитектур моделей, которые остаются популярными из-за их относительной простоты в обучении и развертывании.
FFN Fusion: Переосмысление параллелизма в Transformers
На фоне этого ландшафта методов оптимизации исследователи из NVIDIA представили убедительный новый подход под названием FFN Fusion. Этот метод напрямую противостоит последовательному узкому месту, присущему архитектуре transformer, не путем изменения параметров или выборочной активации частей, а путем фундаментального переосмысления того, как последовательности вычислений могут быть распараллелены. Инновация проистекает из ключевого наблюдения о поведении слоев FFN в глубоких моделях transformer.
Используя диагностический инструмент под названием Puzzle, исследователи проанализировали внутреннюю работу больших моделей. Когда они экспериментально удаляли слои внимания, они заметили, что модели часто сохраняли удивительно длинные последовательности последовательных слоев FFN. Что еще более важно, анализ показал, что вычисления, выполняемые этими смежными FFNs, часто демонстрировали минимальную взаимозависимость. По сути, выход одного FFN в последовательности часто не сильно изменял направление или основную информацию, необходимую для немедленно следующего FFN. Это предполагало, что эти FFNs, традиционно выполняемые один за другим, могут обладать потенциалом для одновременного, параллельного выполнения без значительного нарушения общей функции модели.
Это понимание легло в основу FFN Fusion. Основная идея элегантно проста, но мощна: идентифицировать последовательности последовательных слоев FFN с низкой вычислительной зависимостью и объединить их в один, более широкий слой FFN, который выполняет эквивалентные вычисления параллельно. Вместо цепочки типа Вход -> FFN1 -> FFN2 -> FFN3 -> Выход, объединенная структура становится Вход -> Объединенный_FFN (Эквивалент FFN1+FFN2+FFN3 параллельно) -> Выход. Эта архитектурная трансформация эффективно сокращает последовательную глубину сети, заменяя несколько шагов одним, более широким вычислительным шагом. Нацеливаясь на эти последовательности FFN с низкой зависимостью, FFN Fusion стремится уменьшить задержку и вычислительные затраты, сохраняя при этом репрезентативную мощность и точность модели. Разработка Ultra-253B-Base из Llama-3.1-405B-Instruct послужила яркой демонстрацией потенциала этой техники.
Архитектурная алхимия: Как работает FFN Fusion
Магия FFN Fusion заключается в умном манипулировании базовой математической структурой сетей прямого распространения. Речь идет не просто о запуске существующих слоев бок о бок; это включает создание нового, унифицированного слоя, который воспроизводит коллективное поведение исходной последовательности, но делает это одновременно.
Рассмотрим последовательность из k последовательных слоев FFN. В стандартном transformer вход x проходит через FFN1, его выход становится входом для FFN2, и так далее, до FFNk. Каждый шаг явно зависит от завершения предыдущего. FFN Fusion разрывает эту цепь зависимостей. Математически, FFN обычно включает два линейных преобразования с нелинейной функцией активации (например, GeLU или SwiGLU) между ними: FFN(x) = W_out * Activation(W_in * x). FFN Fusion использует тот факт, что линейные преобразования часто можно объединить.
Процесс слияния работает путем конкатенации весов отдельных слоев FFN. В частности, матрицы входных весов (W_in) последовательных FFNs объединяются (например, блочно-диагонально) в одну, большую матрицу входных весов для объединенного слоя. Аналогично, матрицы выходных весов (W_out) конкатенируются для формирования одной, более широкой матрицы выходных весов. Функция активации применяется поэлементно внутри этой большей структуры. Эта конструкция гарантирует, что объединенный FFN работает с исходным входом x одновременно по параллельным путям, соответствующим исходным FFNs. Выходы из этих параллельных путей затем неявно агрегируются структурой конкатенированных выходных весов.
Теоретическое обоснование подтверждает, что эта объединенная структура может сохранять ту же репрезентативную способность, что и исходная последовательность FFNs, при условии, что зависимости между исходными слоями действительно были низкими. Ключ в том, чтобы определить, какие последовательности подходят для слияния. Чтобы сделать это систематически, исследователи NVIDIA использовали технику анализа зависимостей. Они измеряли косинусное расстояние (cosine distance) между выходными скрытыми состояниями последовательных слоев FFN для репрезентативного набора входных токенов. Малое косинусное расстояние указывает на то, что выходной вектор одного FFN указывает в очень похожем направлении, что и выходной вектор следующего FFN в последовательности. Это сходство предполагает низкую функциональную зависимость – второй FFN не сильно изменяет представление информации, установленное первым. Последовательности FFNs, демонстрирующие стабильно низкие косинусные расстояния между слоями, были определены как основные кандидаты на слияние, поскольку их объединение с меньшей вероятностью нарушит выученные представления модели и общую производительность. Этот подход, основанный на данных, позволяет целенаправленно применять FFN Fusion к тем частям модели, где он будет наиболее эффективным и наименее разрушительным.
От Бегемота к Спринтеру: Трансформация Ultra-253B-Base
Практическая мощь FFN Fusion была ярко продемонстрирована на примере ее применения к одной из крупнейших общеизвестных моделей того времени, Llama-3.1-405B-Instruct. Эта модель, насчитывающая 405 миллиардов параметров, представляла собой значительную вычислительную задачу для инференса. Исследователи предприняли процесс архитектурного усовершенствования, сочетая FFN Fusion со стратегическим прунингом, чтобы создать новую, более эффективную модель под названием Ultra-253B-Base.
Процесс трансформации включал несколько шагов:
- Анализ: Используя свои инструменты анализа зависимостей (измерение косинусных расстояний), исследователи определили последовательности последовательных слоев FFN в архитектуре Llama-405B, которые демонстрировали низкую межуровневую зависимость.
- Слияние (Fusion): Эти идентифицированные последовательности FFN затем были объединены в единые, более широкие слои FFN, как описано ранее (конкатенация весов). Это напрямую сократило количество последовательных шагов в сети.
- Прунинг (Pruning): Одновременно или впоследствии параметры, считающиеся менее критичными (потенциально идентифицированные с помощью стандартных методов прунинга или на основе информации от процесса слияния), были удалены из модели.
Этот комбинированный подход привел к созданию Ultra-253B-Base, модели с 253 миллиардами параметров. Это представляет собой существенное сокращение – более чем на 37% меньше параметров, чем у исходной модели 405B. Архитектурные изменения, достигнутые за счет слияния, были ключом к обеспечению такого значительного сокращения размера при сохранении производительности. Целью была не просто меньшая модель, а фундаментально более быстрая и вычислительно экономичная, благодаря увеличенному параллелизму, разблокированному FFN Fusion. Это исследование послужило важным доказательством концепции, показав, что крупномасштабные модели могут быть существенно реструктурированы для повышения эффективности.
Измерение выигрыша: Производительность, скорость и экономия ресурсов
Настоящая проверка любого метода оптимизации заключается в его измеримом влиянии. Для Ultra-253B-Base результаты, полученные от применения FFN Fusion и прунинга к базовой модели Llama-405B, были убедительными, демонстрируя значительные улучшения по нескольким параметрам без существенных компромиссов в возможностях.
Скорость и стоимость инференса: Наиболее поразительные выигрыши наблюдались в эффективности инференса. По сравнению с исходной моделью с 405B параметрами, Ultra-253B-Base достигла:
- Улучшение задержки инференса в 1.71 раза. Это означает, что модель могла генерировать ответы значительно быстрее, что критически важно для приложений реального времени.
- Снижение вычислительных затрат на токен в 35 раз при измерении с размером батча 32. Это резкое снижение количества вычислительных операций (FLOPs) на токен напрямую приводит к снижению энергопотребления и уменьшению требований к оборудованию для обслуживания модели.
Бенчмарки производительности модели: Критически важно, что эти улучшения эффективности не были достигнуты за счет интеллекта или возможностей модели. Ultra-253B-Base была тщательно оценена на наборе стандартных бенчмарков для LLM, достигнув результатов, которые были весьма конкурентоспособны, а в некоторых случаях превосходили исходную, гораздо большую модель:
- MMLU (Massive Multitask Language Understanding): 85.17%
- MMLU-Pro (Более сложная версия): 72.25%
- Arena Hard (Оценка предпочтений человека на сложных запросах): 84.92%
- HumanEval (Способность к генерации кода): 86.58%
- MT-Bench (Качество многоходового диалога): 9.19
Эти показатели свидетельствуют о том, что объединенная и обрезанная модель сохранила очень высокий уровень понимания, рассуждения, способности к кодированию и качества диалога, сравнимый со своим 405B-параметровым предшественником, несмотря на наличие всего 253 миллиардов параметров.
Эффективность использования памяти: Помимо скорости вычислений и затрат, FFN Fusion также способствовала экономии памяти. Архитектурные изменения, потенциально в сочетании с другими оптимизациями, ставшими возможными благодаря слиянию, привели к 2-кратному сокращению размера кэша ключ-значение (KV cache), необходимого во время инференса. KV cache хранит промежуточные активации (ключи и значения внимания) и может потреблять значительный объем памяти GPU, особенно для длинных входных последовательностей. Уменьшение этого требования вдвое делает возможным запуск модели на менее требовательном к памяти оборудовании или обработку более длинных контекстов в рамках тех же ограничений памяти.
Эти количественные результаты подчеркивают эффективность FFN Fusion. Она позволила создать модель, которая была не только меньше, но и фундаментально более эффективной с точки зрения скорости, вычислительных операций и использования памяти, при этом сохраняя первоклассную производительность на сложных бенчмарках.
Сохранение знаний: Ключевая роль обучения и дообучения
Архитектурное изменение массивной, предварительно обученной языковой модели, такой как Llama-405B, с помощью таких методов, как FFN Fusion и прунинг, неизбежно нарушает тонкий баланс ее выученных параметров. Хотя математическая эквивалентность направлена на сохранение функции локально, глобальное поведение сети может измениться. Чтобы гарантировать, что полученная модель Ultra-253B-Base не только стала более эффективной, но и сохранила свой высокий уровень производительности, был необходим тщательно организованный процесс обучения после модификации.
Этот процесс включал две основные фазы:
Дистилляция знаний (Knowledge Distillation): Первым шагом была передача знаний из исходной, большей модели (или подходящей модели-учителя) обратно в модифицированную архитектуру. Это было достигнуто путем дистилляции, когда модель Ultra-253B-Base обучалась имитировать выходы или внутренние представления модели-учителя. На этом этапе использовался значительный набор данных, в частности 54 миллиарда токенов, обработанных с контекстным окном 8k. Дистилляция помогает объединенной и обрезанной модели восстановить нюансы и возможности, которые могли быть слегка нарушены во время архитектурных изменений.
Поэтапное дообучение (Staged Fine-Tuning): После дистилляции модель прошла серию этапов дообучения, специально разработанных для адаптации к обработке постепенно увеличивающихся длин контекста. Это критически важно для современных LLMs, от которых часто ожидается обработка и генерация текста на основе обширного ввода. Дообучение проходило поэтапно:
- Дообучение с контекстным окном 16k.
- Дальнейшее дообучение с контекстным окном 32k.
- Заключительный этап дообучения с контекстным окном 128k.
Этот поэтапный подход позволяет модели постепенно адаптировать свои параметры, включая вновь сформированные объединенные слои FFN и оптимизированные механизмы KV cache, для эффективного управления зависимостями и потоком информации на очень длинных последовательностях. Каждый этап основывается на предыдущем, обеспечивая стабильность и надежную производительность при различных размерах контекста.
Этот тщательный режим обучения, сочетающий крупномасштабную дистилляцию с поэтапным дообучением на длинном контексте, сыграл важную роль в преодолении разрыва между архитектурной эффективностью и высококачественной производительностью. Он гарантировал, что преимущества в скорости, стоимости и памяти, обеспечиваемые FFN Fusion, не поставили под угрозу точность и возможности модели на требовательных бенчмарках.
Широкие горизонты: Обобщаемость и будущие направления
Успешная трансформация Llama-405B в Ultra-253B-Base предоставляет убедительные доказательства потенциала FFN Fusion, но ее истинная ценность заключается в более широкой применимости и идеях, которые она предлагает для будущего дизайна LLM. Исследование показало, что это был не просто единичный трюк, применимый только к огромным моделям.
Валидация на разных масштабах: Исследователи NVIDIA явно протестировали методологию FFN Fusion на моделях разного размера. Они успешно применили эту технику к моделям с 70B параметрами, достигнув аналогичного прироста эффективности по сравнению с их исходными аналогами. Они также сообщили о валидации на масштабе 49B, что еще больше подкрепляет идею о том, что независимость FFN и потенциал для слияния не являются исключительными характеристиками самых больших моделей, а могут быть более общим свойством архитектуры transformer, потенциально становясь более выраженным на больших масштабах, где естественным образом возникают более глубокие последовательности FFN. Это предполагает, что FFN Fusion может стать стандартным инструментом в арсенале оптимизации LLM, применимым к моделям различных размеров.
FFN против слияния полных блоков: Исследование также пролило свет на специфическую роль слоев FFN по сравнению со слоями внимания в блоке transformer. В то время как последовательные слои FFN часто демонстрировали низкую зависимость, что делало их идеальными для слияния, попытки распараллелить целые блоки transformer (включая как слои внимания, так и FFN) оказались более сложными. Анализ показал более сильные взаимозависимости, связанные с механизмами внимания. Одновременное слияние целых блоков приводило к более значительной деградации производительности, предполагая, что слои внимания играют более критическую, последовательно зависимую роль в интеграции информации между токенами. Это открытие помогает очертить границы эффективного распараллеливания – последовательности FFN являются плодородной почвой, в то время как механизмы внимания могут требовать других стратегий оптимизации.
Последствия для архитектуры LLM: FFN Fusion предлагает больше, чем просто метод оптимизации постфактум; она дает ценные идеи для проектирования будущих LLM. Открытие того, что последовательности FFN часто можно рассматривать как распараллеливаемые единицы, бросает вызов строго последовательному предположению, часто лежащему в основе дизайна transformer. Это может вдохновить на создание новых архитектур, которые изначально более дружелюбны к параллелизму. Будущие модели могут быть спроектированы со структурами FFN, явно предназначенными для слияния или параллельного выполнения, что потенциально приведет к совместному проектированию аппаратного и программного обеспечения, где архитектуры GPU будут дополнительно оптимизированы для использования этого типа параллелизма. Систематический метод использования косинусного расстояния для количественной оценки межуровневой зависимости также предоставляет ценный аналитический инструмент для понимания и перепроектирования структур нейронных сетей. Демонстрируя, что значительный прирост эффективности возможен за счет продуманного архитектурного редизайна, сфокусированного на распараллеливании существующих компонентов, FFN Fusion прокладывает путь к разработке LLM, которые являются одновременно мощными и более вычислительно устойчивыми. Это подчеркивает путь ксмягчению эскалации требований к ресурсам со стороны передового ИИ.