YaFSDP — инструмент обучения LLM, который сокращает использование графического процессора на 20% — уже доступен

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

Недавно Яндекс представил новое решение: YaFSDP — инструмент с открытым исходным кодом, который обещает произвести революцию в обучении LLM, значительно сократив потребление ресурсов графического процессора и время обучения. В сценарии предварительного обучения, включающем модель с 70 миллиардами параметров, использование YaFSDP может сэкономить ресурсы примерно 150 графических процессоров. Это означает потенциальную ежемесячную экономию примерно от 0,5 до 1,5 миллиона долларов США, в зависимости от поставщика виртуальных графических процессоров или платформы.

Яндекс выложил ЯФСДП в открытый доступ на GitHub .

Проблема распределенного обучения LLM

Обучение LLM на нескольких графических процессорах включает в себя сложные операции, которые приводят к неэффективности и высокому потреблению памяти. Одной из основных проблем является необходимость отправлять и получать огромные объемы данных между графическими процессорами. Например, в типичной операции all_reduce необходимо передать вдвое больше данных градиента, чем сетевых параметров. В случае модели Llama 70B это означает передачу 280 ГБ данных за итерацию.

Более того, веса, градиенты и состояния оптимизатора дублируются на всех графических процессорах, что приводит к огромной нагрузке на память. Модель Llama 70B и оптимизатор Adam требуют более 1 ТБ памяти, что намного превышает типичный объем памяти в 80 ГБ для большинства графических процессоров. Эта избыточность сильно замедляет процесс обучения и часто делает непрактичным размещение моделей даже среднего размера в памяти графического процессора.

Представляем ЯФСДП

YaFSDP от Яндекса предлагает высокоэффективное решение этих проблем. Сосредоточив внимание на оптимизации потребления памяти и устранении узких мест в общении, YaFSDP повышает эффективность обучения LLM. Он работает путем сегментирования слоев вместо отдельных параметров, обеспечивая эффективную связь и избегая избыточных операций. Кроме того, YaFSDP предварительно выделяет буферы для всех необходимых данных, гарантируя, что распределитель Torch не приведет к снижению эффективности.

YaFSDP работает, используя два буфера для промежуточных весов и градиентов: нечетные слои используют один буфер, а четные слои — другой.

Веса из разных слоев хранятся в одной памяти. Если слои имеют одинаковую структуру, они всегда будут идентичны. Крайне важно убедиться, что когда вам нужен слой X, буфер содержит веса для слоя X. Все параметры будут храниться в соответствующем фрагменте памяти внутри буфера.

Потребление памяти

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

• Веса, градиенты и состояния оптимизатора : они зависят от количества процессов, и их потребление памяти имеет тенденцию приближаться к нулю по мере увеличения количества процессов. Распределяя эти компоненты по графическим процессорам, YaFSDP сводит к минимуму дублирование и, таким образом, снижает использование памяти.

• Буферы потребляют постоянный объем памяти и сохраняют промежуточные значения во время вычислений.

• Активации зависят от размера модели и количества токенов, обрабатываемых одним графическим процессором.

Контрольная точка активации

Контрольная точка активации — это метод, который сохраняет только необходимые активации во время прямого прохода и пересчитывает их во время обратного прохода. Это значительно уменьшает объем памяти, поскольку сохраняются только необходимые данные. Например, при обучении модели Llama 2 70B с размером пакета в 8192 токена объем хранилища активации можно уменьшить с более чем 110 ГБ до всего 5 ГБ.

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

Оптимизация связи

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

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

Прямой проход на третьем слое не начинается до тех пор, пока не завершится операция all_gather (условие 1). Аналогично, операция all_gather на третьем слое не начнется до тех пор, пока не завершится прямой проход на первом слое, который использует тот же буфер (условие 2). Поскольку в этой схеме нет циклов, взаимоблокировка невозможна.

Результаты экспериментов и повышение производительности

Внедрение YaFSDP продемонстрировало значительное повышение эффективности обучения. В сценарии предварительного обучения с моделью, имеющей 70 миллиардов параметров, YaFSDP удалось сэкономить ресурсы примерно 150 графических процессоров. Это приводит к значительной ежемесячной экономии средств: от 0,5 до 1,5 миллиона долларов США, в зависимости от поставщика или платформы виртуального графического процессора.

YaFSDP сокращает время обучения до 26% по сравнению с существующими методами, такими как FSDP, и оптимизирует использование памяти, позволяя более эффективно обучать более крупные модели.

Яндекс выложил ЯФСДП в открытый доступ на GitHub . Инженеры ML могут использовать этот инструмент для повышения эффективности процессов обучения LLM. Открыв исходный код YaFSDP, Яндекс стремится способствовать инновациям и сотрудничеству в сообществе искусственного интеллекта, позволяя разработчикам обучать модели быстрее и с меньшими затратами.

YaFSDP представляет собой значительный прогресс в обучении LLM. Решение критических проблем потребления памяти и неэффективности связи позволяет быстрее и эффективнее обучать большие языковые модели.