07 - Разгрузка оперативной памяти

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

Вступление

При обучении (или тестировании) многослойной нейронной сети в любой момент времени требуется только один слой. Этот факт может быть использован для замены слоев между оперативной и видеопамятью по мере необходимости. Наивная реализация могла бы

загрузить первый слой модели в видеопамять
выполнить этот слой
выгрузить этот слой обратно в оперативную память

Затем повторять этот процесс до тех пор, пока не будут выполнены все слои. Хотя этот метод и будет работать, он значительно повлияет на производительность. Более продвинутый метод может использовать асинхронные возможности современных графических процессоров для переноса слоев модели между оперативной и видеопамятью во время выполнения других слоев. В Pytorch это можно сделать, используя функцию Stream в сочетании с non_blocking=True тензорными копиями.

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

Выгрузка слоя

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

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

Слои в оперативной памяти хранятся в закрепленной памяти.
Передача выполняется в отдельном потоке CUDA, и флаг non_blocking=True передается каждому вызову функции копирования.

После передачи каждого слоя в поток передачи записывается событие CUDA. Это событие используется для того, чтобы поток по умолчанию мог дождаться завершения передачи перед выполнением слоев.

Схема разгрузки

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

Это лучше всего объяснить на примере. Давайте представим модель с 9 слоями и долей выгрузки в 33%, то есть 3 слоя выгружаются в оперативную память. На следующей диаграмме показано расположение слоев после каждого выполнения слоя. Имейте в виду, что это только идеализированный случай. В реальном мире слои постоянно перемещаются, и разница во времени может означать, что перенос слоя занимает больше времени, чем выполнение самого слоя.

Каждый слой представлен столбцом, каждая строка представляет определенный момент времени.

-> направление выполнения (прямой или обратный проход)
■ выполненный блок
X блок в видеопамяти
_ блок в оперативной памяти

->  ■ X X X X X _ _ _  (начинается первый шаг)
->  _ ■ X X X X X _ _
->  _ _ ■ X X X X X _
->  _ _ _ ■ X X X X X
->  X _ _ _ ■ X X X X
->  X X _ _ _ ■ X X X
->  X X X _ _ _ ■ X X
->  X X X X _ _ _ ■ X
->  X X X X X _ _ _ ■
->  ■ X X X X X _ _ _  (начинается второй шаг)
->  _ ■ X X X X X _ _
->  _ _ ■ X X X X X _
->  _ _ _ ■ X X X X X
->  X _ _ _ ■ X X X X
->  X X _ _ _ ■ X X X
->  X X X _ _ _ ■ X X
->  X X X X _ _ _ ■ X
->  X X X X X _ _ _ ■

Схема выполнения во время выборки Схема выполнения во время обучения

-> направление выполнения (прямой или обратный проход)
■ выполненный блок
X блок в видеопамяти
_ блок в оперативной памяти

->  ■ X X X X X _ _ _  (начинается прямой проход)
->  _ ■ X X X X X _ _
->  _ _ ■ X X X X X _
->  _ _ _ ■ X X X X X
->  _ _ _ X ■ X X X X
->  _ _ _ X X ■ X X X
->  _ _ _ X X X ■ X X
->  _ _ _ X X X X ■ X
->  _ _ _ X X X X X ■
<-  _ _ _ X X X X X ■  (начинается обратный проход)
<-  _ _ X X X X X ■ _
<-  _ X X X X X ■ _ _
<-  X X X X X ■ _ _ _
<-  X X X X ■ X _ _ _
<-  X X X ■ X X _ _ _
<-  X X ■ X X X _ _ _
<-  X ■ X X X X _ _ _
<-  ■ X X X X X _ _ _

Шаги оптимизатора

Хотя такая схема разгрузки может значительно сократить использование видеопамяти, у нее есть довольно большой недостаток. Для выполнения шагов оптимизации обычно требуется, чтобы все параметры модели были размещены на одном устройстве (например, на графическом процессоре). Чтобы устранить эту проблему, можно использовать метод, называемый "Плавный обратный проход". Шаги оптимизатора выполняются сразу после вычисления градиента. Таким образом, каждый параметр обновляется только тогда, когда он находится в VRAM. Это ограничение применяется только для полной точной настройки. В случае обучения LoRa все параметры, которые можно обучить, сохраняются в VRAM.

Активация разгрузка

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

Таким образом, при прямом проходе активации удаляются, когда в них нет необходимости. При обратном проходе они загружаются обратно в VRAM только при необходимости.

Управление памятью

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

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

Распределение уровней

Уровни в модели обычно имеют четко определенный порядок выполнения. Это означает, что мы можем заранее знать шаблоны распределения и определить оптимальную стратегию распределения. OneTrainer использует кольцевой буфер для распределения. Тензоры могут быть распределены в обоих направлениях в зависимости от направления выполнения модели. Опять же, это лучше всего объяснить на примере. На следующей диаграмме показан шаблон распределения в видеопамяти для 12-слойной модели с выгрузкой, установленной на 50%, что означает, что в любой момент выгружается 6 слоев.

Каждый слот в буфере представлен столбцом, каждая строка представляет определенный момент времени.

Схема распределения во время обучения

-> направление выполнения (прямой или обратный проход)
Номер - индекс блока
* выполненный блок

->  *00* 01  02  03  04  05  (начинается прямой проход)
->   06 *01* 02  03  04  05
->   06  07 *02* 03  04  05
->   06  07  08 *03* 04  05
->   06  07  08  09 *04* 05
->   06  07  08  09  10 *05*
->  *06* 07  08  09  10  11
->   06 *07* 08  09  10  11
->   06  07 *08* 09  10  11
->   06  07  08 *09* 10  11
->   06  07  08  09 *10* 11
->   06  07  08  09  10 *11*
<-   06  07  08  09  10 *11* (начинается обратный проход)
<-   06  07  08  09 *10* 05
<-   06  07  08 *09* 04  05
<-   06  07 *08* 03  04  05
<-   06 *07* 02  03  04  05
<-  *06* 01  02  03  04  05
<-   00  01  02  03  04 *05*
<-   00  01  02  03 *04* 05
<-   00  01  02 *03* 04  05
<-   00  01 *02* 03  04  05
<-   00 *01* 02  03  04  05
<-  *00* 01  02  03  04  05

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

Распределение ресурсов для активации

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