Методы эффективного использования памяти для решения проблемы «ошибки памяти CUDA…» во время обучения

Статья вдохновлена ​​курсом Эффективные системы глубокого обучения, преподаваемым в Школе анализа данных Яндекса.

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

И так это начинается…

Вы, наверное, задавались вопросом, почему, когда вы пытаетесь реализовать документ, в котором используется какая-то большая модель (также известная как gpt-2-xl) с ›500 миллионами параметров, вы даже не можете уместить ее на своем GPU или используйте тот же размер пакета, что и в статье, во время обучения. Тогда, возможно, вы сдались и начали использовать более легкую версию модели или обучали ее на меньшем размере партии, что не позволяло получить сравнимые с бумагой результаты.

Но,есть несколько приемов, которые помогут вам справиться с описанной проблемой.

Давайте обсудим некоторые подходы и посмотрим, как их использовать для тонкой настройки модели GPT-2-XL с 1,5 миллиардами параметров в конце статьи.

Суть проблемы

Давайте разберемся в сути проблемы нехватки памяти GPU, необходимой для загрузки модели на GPU.

Предположим, у вас есть модель с параметрами 1e9 FP32 (32 бита с плавающей запятой). Вы хотите обучить эту модель на своем прекрасном графическом процессоре, используя, например, оптимизатор Adam.

Итак, давайте посчитаем. Думаю, вы будете шокированы.

Представьте, что у вас есть NVIDIA GeForce RTX 3060 с 12 ГБ памяти. Во-первых, параметры 1e9 FP32 занимают около 4 ГБ памяти вашего графического процессора. Также такое же количество памяти будет зарезервировано для градиентов. Итак, у нас уже 8 ГБ всего зарезервировано, еще не начав обучение и не загрузив оптимизатор, потому что он тоже не свободен по памяти. Оптимизатор Адама должен хранить первый и второй моменты для каждого параметра, т. е. 8 ГБ дополнительной памяти.

В конце концов, у нас должно быть около 16 ГБ свободной памяти графического процессора только для того, чтобы правильно загрузить модель в графический процессор, который в нашем случае имеет только 12 ГБ свободной памяти. Выглядит ужасно, да?

Но есть некоторые подходы, которые мы можем использовать, чтобы попытаться решить проблему. Вот некоторые из них ниже:

  • накопление градиента/микропакетирование;
  • контрольные точки градиента;
  • модельно-параллельное обучение;
  • конвейерная обработка;
  • тензорный параллелизм;
  • тренировка смешанной точности;
  • разгрузка памяти;
  • оптимизатор 8-битного квантования.

Сегодня мы узнаем о них.

Пойдем!

Градиентная контрольная точка

Обзор

Что делать, если модель больше, чем GPU, то есть мы не можем уместить размер партии 1? Решение есть — градиентный чекпойнт. Давайте посмотрим на это понятие.

Для простой нейронной сети с прямой связью с n слоями граф вычислений для получения градиентов выглядит следующим образом:

Активации слоев нейронной сети соответствуют узлам, отмеченным f. Во время прямого прохода все эти узлы оцениваются по порядку. Градиент потерь по отношению к активациям и параметрам этих слоев указан узлами, отмеченными b. При обратном проходе все эти узлы оцениваются в обратном порядке. Результаты, полученные для узлов f, необходимы для вычисления узлов b, поэтому все узлы f сохраняются в памяти после прямого прохода. Только когда обратное распространение продвинулось достаточно далеко, чтобы вычислить все зависимости узла f, его можно стереть из памяти. Это означает, что объем памяти, необходимый для простого обратного распространения, растет линейно с количеством слоев нейронной сети n.

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

Градиентная контрольная точка

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

Эта стратегия оптимальна с точки зрения памяти. Однако обратите внимание, что количество оценок узлов теперь масштабируется с n², тогда как ранее оно масштабировалось как n: каждый из n узлов пересчитывается на порядка n раз. Медленность вычислений делает этот метод непрактичным для использования в глубоком обучении.

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

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

Пример:

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

Модельно-параллельное обучение, конвейерная обработка, тензорный параллелизм, разгрузка памяти

Это очень большая и сложная тема, мы обсудим ее в следующих постах.

Накопление градиента/микропакетирование

Обзор

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

Что такое накопление градиента?

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

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

Накопление градиента помогает имитировать больший размер партии. Представьте, что вы хотите использовать 64 изображения в одном мини-пакете, но «Ошибка памяти CUDA…», как только вы превысите размер 8. В этом случае вы можете использовать пакеты из 8 изображений и обновлять веса один раз после 64 / 8 = 8. партии, обрабатываемые моделью. Если вы накопите градиенты из каждой из этих 8 партий, результаты будут (почти) одинаковыми, и вы сможете проводить тренировки! Йоа!

Пример:

Стандартный цикл обучения без накопления градиента обычно выглядит так:

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

Красиво, да? Градиенты вычисляются, когда мы вызываем loss.backward(), и накапливаются PyTorch, пока мы не вызовем optimizer.zero_grad().

Важный

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

Обучение смешанной точности

Обзор

Обучение смешанной точности означает преобразование некоторых или всех параметров, которые являются FP32 числами, в меньшие форматы, такие как FP16, TF16 (tensor float), BF16 (bfloat).

Ключевые преимущества

Ключевые преимущества обучения смешанной точности:

  • Уменьшено использование памяти;
  • Более высокая производительность (из-за более высокой арифметической интенсивности или меньшего объема связи);
  • Может использовать специализированное оборудование для еще более быстрых вычислений;

Но сейчас нас интересует только первое преимущество — уменьшение использования памяти. Давайте посмотрим, как это сделать с моделями PyTorch.

Пример:

В результате после выполнения .half() модель становится в 2 раза меньше.

Различные форматы, в которые модель может быть преобразована (например, BF16, TF16), и масштабирование потерь мы обсудим в следующих постах.

Но вы должны помнить, что есть некоторые операции, которые нельзя сделать в FP16, т.е. Softmax. В PyTorch есть torch.autocast, который помогает обрабатывать такие ситуации.

Оптимизатор 8-битного квантования

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

Переход от 32-разрядных оптимизаторов к 8-разрядным оптимизаторам сокращает диапазон возможных значений с 2³² до 2⁸= 256. Это оказывает огромное влияние на объем памяти, который должен быть зарезервирован оптимизатором.

Research представляет новый 8-разрядный оптимизатор Adam, который поддерживает 32-разрядную производительность при незначительном потреблении памяти. Вот что говорят авторы в своей статье:

Наши 8-битные оптимизаторы состоят из трех компонентов: (1) поблочное квантование, которое изолирует выбросы и более равномерно распределяет ошибку по всем битам; (2) динамическое квантование, которое квантует как малые, так и большие значения с высокой точностью; и (3) стабильный уровень встраивания для повышения стабильности во время оптимизации моделей с встраиванием слов.

С этими компонентами выполнение обновления оптимизатора с 8-битными состояниями выполняется просто. Мы деквантуем 8-битные состояния оптимизатора до 32-битных, выполняем обновление, а затем квантуем состояния обратно до 8-битных для хранения. Мы делаем это 8-битное преобразование в 32-битное поэлементно в регистрах, что означает отсутствие медленных копий в память графического процессора или дополнительной временной памяти для выполнения квантования и деквантования. Для графических процессоров это делает 8-битные оптимизаторы быстрее, чем обычные 32-битные оптимизаторы…

Давайте посмотрим на вдохновляющие результаты использования 8-битного Адама:

Как мы видим, использование квантованного Адама экономит около 8,5 ГБ памяти графического процессора. Выглядит фантастически!

Теперь, когда мы поняли полезность его использования, давайте посмотрим, как использовать его из python.

Пакет Bitsandbytes от Facebook — это легкая оболочка для пользовательских функций CUDA, в частности 8-битных оптимизаторов и функций квантования. Это позволяет нам использовать 8-битного Адама.

Пример:

Как вы можете видеть выше, использование квантизированного оптимизатора довольно просто, но результат огромен.

Сочетание всех вышеперечисленных подходов для тонкой настройки GPT-2-XL на графическом процессоре

В конце концов, поскольку мы изучили все вышеперечисленные методы, давайте применим их для решения реальной проблемы. Нам предстоит доработать модель GPT-2-XL с › 1,5 миллиардами параметров. Очевидно, что его нельзя загрузить на графическом процессоре NVIDIA GeForce RTX 3060 с 12 ГБ памяти.

Давайте перечислим все методы, которые мы можем использовать:

  • Градиентный чекпойнт;
  • Обучение смешанной точности (я делаю трюк: использую два образца одной и той же модели. Первый — .half-ed и загружается на GPU, назовем его gpu_model. Второй — просто на CPU, назовем его cpu_model. Оцениваем модель GPU , затем загрузите градиенты из gpu_modelв cpu_model, затем выполните optimizer.step(), загрузите обновленные параметры в gpu_model);
  • Накопление градиента с batch_size=64, minibatch_size=4. Не забудьте масштабировать потери на accumulation_steps;
  • 8-битный оптимизатор Адама.

Давайте использовать их все. Посмотрите на код:

‹код скоро будет доступен›

В результате использование всех вышеперечисленных методов позволило нам настроить 16-гигабайтную модель GPT-2-XL на нашем графическом процессоре. Я думаю, это потрясающе!

Заключение

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

Мы обсудим другие концепции в следующих постах.

Спасибо за прочтение статьи!

Использованная литература:

  • 8-БИТНЫЕ ОПТИМАЙЗЕРЫ ЧЕРЕЗ БЛОЧНОЕ КВАНТИЗИРОВАНИЕ, бумага;
  • биты и байты, git;
  • градиент-чекпойнт, git;
  • обучение смешанной точности, git;