Внедрение и оптимизация динамической оптимизации графической памяти (DTR) в MegEngine

искусственный интеллект глубокое обучение задняя часть

| Автор: Дэн Чжей | Инженер группы инфраструктуры MegEngine

В области глубокого обучения в последние годы многие исследовательские институты и исследователи улучшили производительность модели, увеличив количество параметров модели, и добились очень значительных результатов, которые снова и снова поражают отрасль. Это объективно делает «расширение размера модели» едва ли не единственным показателем, за который конкурируют различные компании. Количество параметров современных моделей с годами увеличилось в сотни и тысячи раз, но объем памяти на GPU почти не вырос. В результате обучение больших моделей часто зависит от огромного количества карт GPU. В результате многие исследователи с отличными идеями и энтузиазмом в исследованиях испытывают трудности с продолжением исследований в области глубокого обучения просто из-за нехватки средств, а важные научные достижения последних лет почти полностью монополизированы несколькими ведущими исследовательскими институтами. В долгосрочной перспективе эта тенденция может не способствовать развитию и прогрессу науки глубокого обучения. Как разработчик обучающих фреймворков для глубокого обучения, мы не только помогаем пользователям использовать больше карт GPU в учебной задаче (т. е. распределенное обучение), но и применяем различные технические средства для повышения эффективности использования видеопамяти на каждом GPU и снижения затрат. капитала для исследователей. Общие способы повышения эффективности использования памяти:

  • Операторы с непересекающимися жизненными циклами совместно используют видеопамять;

  • Уменьшить использование памяти за счет дополнительных передач данных;

  • Уменьшите использование памяти за счет дополнительных вычислений.

Большинство существующих методов требуют, чтобы расчетный граф был статическим.Поскольку все больше и больше фреймворков поддерживают режим динамического графа, вопрос о том, можно ли в максимальной степени использовать ограниченные ресурсы видеопамяти во время обучения динамическому графу, стал оценочной структурой для глубокого обучения. важный показатель эффективности. Недавно выпущенный MegEnginev1.4версии, вводя DTR[1]Технологии и дальнейшая инженерная оптимизация позволяют сократить использование памяти за счет дополнительных вычислений, так что небольшая память может также обучать большие модели и пользоваться преимуществами обучения, обеспечиваемыми большими размерами пакетов. На 2080Ti максимальный размер пакета ResNet-50, ShuffleNet и других сетей может быть более чем в 3 раза больше исходного. В этой статье основное внимание будет уделено тому, как использовать технологию DTR для оптимизации динамической графической памяти в MegEngine с точки зрения инженерной реализации.

1. Введение

1.1 Расчетный граф

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

Взяв в качестве примера y=wx+b, его процесс прямого вычисления заключается в том, что входные данные x и параметр w сначала перемножаются для получения промежуточного результата p, а затем p и параметр b добавляются для получения окончательного результата y на правильно.

Обратному распространению необходимо найти производную y относительно w и b. Во-первых, найдите, что производная y относительно p равна 1, а производная p относительно w равна x. Используя цепное правило, производная y относительно w можно получить как x.

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

1.2 Оптимизация статической памяти изображений

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

  1. Статическое выделение памяти. Поскольку получается весь граф вычислений, можно проанализировать жизненный цикл каждого тензора и каждого оператора. Для операторов, время жизни которых не пересекается, они могут совместно использовать видеопамять.
  2. Градиентные контрольные точки (расчет на видеопамять). Установите несколько контрольных точек градиента и сначала отпустите оставшиеся промежуточные результаты.Если обнаружится, что прямые результаты не находятся в видеопамяти во время процесса обратного распространения в будущем, найдите ближайшую контрольную точку градиента и восстановите освобожденный тензор.

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

2. Динамическая оптимизация графической памяти и стратегия DTR

2.1 Оптимизация динамической графической памяти

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

  1. Подкачка вычислений для видеопамяти, то есть сублинейная оптимизация видеопамяти динамической графики;
  2. Обменяйте пропускную способность на видеопамять и обменивайтесь контентом между GPU и CPU.

На приведенном выше рисунке показаны три тензора, взятые из сети ResNet-50, которые являются выходными тензорами свертки, BatchNorm и ReLu, сравнивая затраты времени на повторные вычисления и обменивая их на пропускную способность. Можно обнаружить, что затраты времени на обмен обычно на два порядка больше, чем затраты времени на вычисления. Поскольку время, затрачиваемое на обмен данными между CPU и GPU, зависит от скорости PCIe, а при одновременном обучении 8 видеокарт 2080Ti скорость обмена, выделяемая каждой карте, составляет всего около 3 ГБ/с. Следовательно, можно определить, что основным направлением оптимизации в динамическом графе по-прежнему является использование вычислений для обмена видеопамятью.

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

  1. реализовать инфраструктуру: запишите путь вычислений, который генерирует каждый тензор, чтобы платформа поддерживала освобождение и восстановление тензора;
  2. пользовательская политика: Предоставляет интерфейс для выпуска тензора, который явно вызывается пользователем. Платформе не нужно предоставлять какую-либо стратегию. Ему нужно только выполнять каждый шаг в соответствии со стратегией пользователя и пересчитывать на месте, когда тензор должен быть восстановлен;
  3. стратегия поиска кадра: Фреймворк автоматически находит стратегию и выполняет ее без вмешательства пользователя, так что пользователь совершенно не знает об оптимизации памяти.

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

  1. Загружается внешними данными, например: входными данными;
  2. является выходом оператора, например выходом сверточного слоя.

Для вывода оператора мы можем записать Compute Path этого тензора, и структура выглядит следующим образом:

  • У каждого тензора будет производитель, если производитель пустой, значит он загружен внешними данными, иначе это путь вычисления, где:
    • op представляет оператор, сгенерировавший этот тензор;
    • inputs представляет входной тензор, требуемый этим оператором;
    • outputs представляет выходной тензор, сгенерированный этим оператором;
    • вычисление_время представляет фактическое время выполнения этого оператора;
  • Что хранится в пользователях, так это все пути вычислений, которые полагаются на тензор в качестве входных данных;
  • ref_cnt представляет количество тензоров, которые полагаются на этот тензор в качестве входных данных.

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

Сначала определите два тензора a и b в MegEngine и вычислите c=a+b. Каждый серый прямоугольник на рисунке представляет собой видеопамять, предполагая, что в видеопамять можно поместить только 3 тензора. На данный момент места достаточно, чтобы записать c и записать путь вычисления c (соответствует желтому прямоугольнику на рисунке выше). Затем вычислите d=a*b, поскольку в данный момент в видеопамяти нет места для размещения d, вам нужно сначала освободить c из видеопамяти.При освобождении c путь вычисления c все еще зарезервирован на хосте. сторону, но видеопамять, занятую c, можно освободить, и есть свободная позиция для установки d (соответствует первому зеленому прямоугольнику на рисунке). Если пользователь хочет напечатать (c) в это время, фреймворк обнаружит, что в данный момент c нет в видеопамяти, и должен немедленно восстановить его. Перед восстановлением, если обнаружено, что видеопамять заполнена, сначала необходимо освободить d, а затем c восстанавливается в соответствии с путем вычисления c и возвращается пользователю (соответствует серому прямоугольнику на рисунке). Если пользователь продолжает печатать (d), сначала отпустите c и восстановите d (соответствует последнему зеленому прямоугольнику на рисунке).

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

2.2 Стратегия DTR

Чтобы платформа могла автоматически вычислять политики, мы представили DTR в MegEngine v1.4 — метод, описанный в статье «Динамический тензорный реинжиниринг», который представляет собой полностью динамическую эвристическую политику. Суть его в том, что когда объем видеопамяти превышает пороговое значение, он динамически выбирает некоторый тензор, чтобы освободить его до тех пор, пока объем видеопамяти не станет ниже порогового значения. При выборе тензор оценивается по трем аспектам:

  1. Чем ниже стоимость перерасчета, тем лучше;
  2. Чем больше занимаемая память, тем лучше;
  3. Чем дольше он остается в видеопамяти, тем лучше.

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

В связи с этим в статье предлагаются два метода оптимизации времени выполнения:

  1. Малый тензор не рассматривается, когда размер тензора меньше 1% от среднего размера тензора в наборе кандидатов, набор кандидатов не добавляется;
  2. Каждый раз, когда необходимо освободить тензор, случайным образом выбираются тензоры sqrt(N) для обхода (N — это размер выпущенного в данный момент набора кандидатов на тензор)

3. Инженерная реализация в MegEngine

3.1 Ядро динамического графа — тензорный интерпретатор

Прежде чем представить реализацию DTR, давайте сначала представим ядро ​​​​динамического графа MegEngine —Tensor Interpreter(интерпретатор), который переводит код Python в следующие четыре основные операции, которые интерпретируются и выполняются по очереди:

  • Поставил: Загрузить внешние данные со стороны хоста в видеопамять для получения тензора
  • ApplyOp: выполнить оператор, параметрами которого являются op (оператор) и входной тензор, и вернуть выходной тензор.
  • Del: удалить тензор, чтобы освободить занимаемое им место в видеопамяти.
  • GetValue: Чтобы получить значение тензора, вам нужно загрузить данные из видеопамяти на хост-сторону.

3.2 Базовая реализация освобождения и восстановления тензора

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

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

3.3 Выполнение оператора после введения DTR

На приведенном выше рисунке показан псевдокод ядра DTR.Для метода ApplyOp в прошлом необходимо выполнить только желтый код, указывающий, что оператор op выполняется на входе input.

Теперь из-за того, что мы представили технологию DTR, возможно, что эти входные тензоры больше не находятся в видеопамяти. Следовательно, их необходимо пометить перед выполнением, и эти входные тензоры не могут быть освобождены до тех пор, пока оператор не будет выполнен. Затем вызовите AutoEvict(), чтобы контролировать текущую занятость видеопамяти, чтобы не превысить порог.Метод состоит в том, чтобы проверить текущую занятость видеопамяти.Если она превышает порог, непрерывно вызывать алгоритм FindBestTensor(), а затем найти оптимальный тензор в соответствии к функции эвристической оценки и отпустите ее.

После выполнения AutoEvict() текущее использование видеопамяти ниже порогового значения. В это время проверьте, находится ли каждый входной тензор в видеопамяти. Если его нет в видеопамяти, вызовите Regenerate(), чтобы восстановить его, а затем текущий оператор может быть выполнен. Процесс Regenerate(x) — это процесс пересчета x.При пересчете прочитайте историю вычислений x-op и входных данных, а затем рекурсивно вызовите ApplyOp для восстановления x.

3.4 Удалить операцию тензора

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

Например, следующее изображение является подграфом расчетного графа.Вы можете видеть, что тензор 9 МБ проходит через оператор свертки для получения тензора 25 МБ, а затем проходит через оператор Elemwise для получения тензора 25 МБ, а затем проходит через оба Оператор BatchNorm и оператор Elemwise получают тензор на 25 МБ.

Обратите внимание, что поскольку оператор Элемвайза здесь представляет собой все сложение, его ввод (два красных тензора) не будет использоваться при выводе. Следовательно, дифференциатору не нужно сохранять два красных тензора, они фактически высвобождаются сразу после прямого вычисления. Преимущество этого в том, что видеопамять можно сэкономить сразу, но после внедрения технологии DTR, если действительно удалить два красных тензора, зеленые тензоры на рисунке никогда не будут освобождены, потому что их вычислительные источники (красный тензор) был потерян, и после освобождения зеленого тензора его невозможно восстановить. Решение состоит в том, чтобы использовать release вместо удаления в прямом процессе, то есть «fake delete» — сохранить tensorInfo, но освободить соответствующую видеопамять под tensorInfo. Таким образом, для выпуска следующих четырех тензоров по 25 МБ необходимо зарезервировать только тензор 9 МБ, и их можно восстановить в любое время в будущем.

На приведенном выше рисунке показана псевдокодовая реализация удаления тензора в MegEngine.Когда интерпретатор получает команду Del, он вызывает функцию Free() для tensorInfo и решает, следует ли удалить ее true или false в зависимости от того, является ли текущий состояние. расчет вперед. Реализация поддельного удаления очень проста, достаточно отметить удаление и освободить видеопамять, управляемую tensorInfo; реализация реального удаления сложнее, сначала обновить ref_cnt входного тензора, который генерирует тензор, а затем вызвать RemoveDep() для проверки всех зависимостей, которые зависят от тензора в качестве входных данных. Если они не находятся в видеопамяти, вы должны теперь вызвать Regenerate для их восстановления, потому что, как только текущий тензор действительно удален, эти тензоры не могут быть восстановлены.

После выполнения вышеуказанных операций вы можете освободить tensorInfo, соответствующий тензору. После выпуска необходимо рекурсивно проверить входной тензор истории вычислений x.Если какой-либо из этих тензоров имеет ref_cnt=0 и помечен для удаления, может быть выполнено истинное удаление.

3.5 Сравнение времени обучения

На следующем рисунке показано сравнение ситуации обучения реализации DTR в MegEngine и реализации исходной статьи в PyTorch на ResNet-1202. Обратите внимание, что видеокарта, используемая в эксперименте, отличается, поэтому MegEngine немного быстрее по данным. MegEngine лучше справляется с управлением памятью, потому что он по-прежнему может запускать обучение с пакетным размером = 100 на видеокарте 11G. В дополнение к максимальному размеру пакета = 140, опробованному в статье, мы пробовали большие размеры пакетов, и все они были работоспособны.

Ниже приведено сравнение времени обучения для различных оптимизаций видеопамяти на платформе MegEngine.Базовый уровень — это результат работы без какой-либо оптимизации видеопамяти в режиме динамического графика. Первая — это две распространенные модели — ResNet-50 и ShuffleNet.Можно обнаружить, что предельный размер пакета после включения DTR-оптимизации превышает статический график Sublinear и baseline, а затраты времени такие же, как у Sublinear, когда пакет размер такой же.

Вышеупомянутые две модели являются статическими, поэтому мы можем использовать сублинейную оптимизацию памяти статических графов для сравнения, а следующая сеть SPOS особенная, это большая сеть с несколькими путями от входа к выходу, которые можно обновлять. В процессе обучения каждый раунд будет случайным образом обновлять определенный путь, что приводит к возможности выполнения разных предложений в каждом раунде. Для такого типа сети более естественно реализовать ее в виде динамического графа. Поэтому с базовой линией сравниваются только результаты DTR-оптимизации динамического графа. Будь то одна карта или восемь карт, предельный размер пакета динамических графиков составляет 100, а если DTR включен, он может достигать 250 или даже больше.

3.6 Проблемы фрагментации и методы оптимизации

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

Для этой задачи мы предлагаем три возможных оптимизации:

  • Параметры обновляются на месте

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

  • Улучшить функцию оценки

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

f(t)=занимает много времени, чтобы вычислитьαпродолжительность пребыванияβ(Видеопамять + размер свободного сегмента)γδколичество перерасчетовf(t)=\frac{\text{время расчета}^\alpha}{\text{продолжительность пребывания}^\beta(\text{видеопамять + размер свободного сегмента})^\gamma}\delta^{\text {количество перерасчетов}}

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

  • статическая стратегия планирования

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

Например, в ResNet-50 на рисунке ниже, когда размер пакета = 400, пиковое значение динамически выделенной видеопамяти составляет 9595 МБ, а пиковое значение статически выделенной видеопамяти — 8549 МБ, уменьшение примерно на 10%. Как только видеопамять будет выделена статически, проблема фрагментации больше никогда не возникнет.

4. Будущее направление работы

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

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

Будь то сублинейная оптимизация для статических графов или оптимизация DTR для динамических графов, ее можно рассматривать как преобразование Seq2Seq в последовательность выполнения — добавление в последовательность некоторых операторов удаления и повторного вычисления. Отличие состоит в том, что статический граф вычисляет оптимальную последовательность освобождения и пересчета после получения всей текущей последовательности, а динамический граф вставляет операторы удаления и пересчета на месте в процессе интерпретации последовательности выполнения. Существует два способа выполнения последовательности: динамическое графическое выполнение императивной интерпретации во время выполнения и статическое графическое выполнение компиляции вычислительного графа. Профиль будет записывать информацию о времени выполнения, такую ​​как время работы каждого оператора и время, в течение которого каждый тензор остается в видеопамяти во время фактической работы, а затем пользователи могут настроить последовательность вычислений в соответствии с результатами профиля. Таким образом, пользователи могут настраивать стратегии для различных моделей, не зная базовой логики выполнения и не изменяя исходный код фреймворка.

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

[1] Kirisame M, Lyubomirsky S, Haan A, et al. Dynamic tensor rematerialization[J]. arXiv preprint arXiv:2006.09616, 2020.

Категории