Серия MegEngine Architecture: Анализ статической памяти

Автор: Чжоу Жуйлян | Архитектор MegEngine, Megvii Technology

Прежде чем официально перейти к этой статье, давайте сначала интуитивно почувствуем улучшение использования памяти, вызванное оптимизацией статической памяти, с помощью набора простых сравнений. Для модели vgg16/resnet50/mobilenetv2 1batch/128batch используйте инструменты tflite, чтобы получить пиковую память tflite и сравнить ее с пиковой памятью до и после оптимизации памяти megengine:

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

1. Уменьшить использование памяти — зачем и как

why

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

how

Управление памятью является важной темой исследований для систем глубокого обучения. Решений этой проблемы много, типовые решения следующие:

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

2. Управление статической памятью в MegEngine

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

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

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

А. Получить необработанную информацию о потоке данных

MegEngine скомпилирует граф, чтобы получить топологически отсортированную последовательность выполнения операторов, которая будет выполняться в порядке последовательности выполнения операторов во время выполнения. Благодаря этой последовательности операторов и информации об узлах ввода/вывода этих операторов легко наблюдать, как данные передаются и используются при работе всей модели. Это информация о потоке данных полного графа последовательной программы.

б. Установите минимальный блок управления памятью Memory Chunk

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

(1) Пройдите последовательность выполнения оператора (opr) и сделайте вывод о зависимости между входными (входными)/выходными (выходными) узлами оператора в соответствии с методом, определенным в операторе Выходной узел создает фрагменты памяти:

(Схема отношения между адресами только для чтения и доступными для записи)

Как показано на рисунке выше, input0/output0 - это зависимость чтения, а input1/output1 - зависимость записи. Отношение отображения между их адресами показано в правой части рисунка. Каждый прямоугольник соответствует концепции фрагмента памяти. . Зависимости и правила создания между фрагментами памяти следующие:

• Чтение зависимостей, требующих от выходного узла повторного использования памяти входного узла и гарантии того, что содержимое не будет изменено. В этом случае выходной узел не создает новый блок памяти, а напрямую повторно использует блок памяти входного узла. Как показано на рисунке, два узла, input0 и output0, совместно используют один блок памяти, а именно chunk0.
• При записи зависимостей выходной узел повторно использует всю или часть памяти входного узла. В этом случае выходной узел создаст новый блок памяти (chunk2) и запишет адрес смещения блока памяти (chunk2) после выполнения операции записи в исходный блок памяти (chunk1) и размер занимаемого пространства. .
• Это не имеет ничего общего с предыдущим вводом.Создайте отдельный блок памяти для записи информации, такой как размер занимаемого пространства.

(2) Получите жизненный цикл каждого фрагмента памяти:

• Сначала запишите начальный жизненный цикл (begin) и конечный жизненный цикл (end) всех фрагментов памяти как 0.
• Последовательность выполнения оператора обхода по порядку
  • Получить выходной узел.Если выходной узел подсчитывается впервые, начальный жизненный цикл чанка памяти узла записывается как id порядкового номера выполнения оператора, что означает, что данные чанка памяти генерируются после выполнения id-го оператора;
  • Получите входной узел, сравните время завершения (конец) фрагмента памяти узла с id+1 и обновите конец до большего значения из двух, то есть end=max(end, id+1). id+1 означает, что он не может быть освобожден до тех пор, пока не будет выполнена id-я операция.

Следует отметить, что в приведенных выше шагах описания «первой статистики» и «max(end, id+1)» связаны с зависимостями чтения, что приводит к существованию нескольких узлов ввода/вывода, совместно использующих фрагмент памяти, а чанк памяти будет подсчитан много раз, а топологические свойства последовательности выполнения операторов и приведенные выше методы используются для обеспечения корректности статистики жизненного цикла.

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

C. Алгоритм оптимизации памяти оптимизирует фрагмент памяти

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

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

(1) Имитационная реализация пространства приложения блока памяти

Во-первых, введем в алгоритм метод управления памятью и логику free()/alloc()/alloc_overwrite().

Для имитации процесса планирования памяти поддерживаются две очереди информации о памяти:

• Очередь информации о всей памяти (далее именуемая очередью памяти): записывает всю память, на которую подана заявка, в порядке обращения, включая два состояния: используется (отмечено серым цветом) и простаивает (отмечено белым цветом).
• Очередь информации о свободной памяти (далее именуемая очередью свободной памяти): память в свободном состоянии записывается в порядке размера пространства от меньшего к большему.

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

(связанный список управления памятью)

Логика free(): освободить указанное адресное пространство и сохранить очередь памяти и очередь свободной памяти. В основном есть два случая, как показано на следующем рисунке:

• Случай 1: нет другого свободного места до и после свободного места после свободного места, и нет необходимости объединять свободное пространство.

• Случай 2: есть другие свободные места до и после свободного места после свободного места, и необходимо объединить свободное пространство.

)

Логика alloc(): подайте заявку на размер пространства и поддерживайте очередь памяти и очередь свободной памяти. Есть в основном три случая, как показано на следующем рисунке:

• Случай 1: если очередь свободной памяти пуста, запросите размер памяти напрямую и запишите его в конец очереди памяти;

• Случай 2: Очередь свободной памяти не пуста, и есть пространство, превышающее или равное размеру, запросите минимальное пространство, большее или равное размеру из очереди свободной памяти. Перехватите необходимое пространство и поддерживайте очередь памяти и очередь свободной памяти.

• Случай 3: Очередь свободной памяти не пуста, но размер используемого пространства меньше размера, тогда вынимается самое большое свободное место в очереди свободной памяти (то есть память, записанная в конце очереди) и увеличен до размеров.

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

Логика alloc_overwrite(): Зависимость от записи упоминается в пункте b(1) выше.Можно видеть, что зависимый от записи узел может использовать только часть исходного пространства зависимого от записи узла объекта (или все его, на основе в записи, когда установлена ​​зависимость записи) Можно вывести адрес смещения и информацию о размере). Следовательно, после операции записи может быть свободное место до и после, освободить эту часть пространства и сохранить очередь памяти и очередь свободной памяти. После успеха запишите кусок памяти переднего пространства.

В качестве примера предположим, что существует зависимость записи, как показано в b(1): chunk2 записывает chunk1, смещение = 6M, размер = 10M, а размер chunk1 равен 20M. Из этого можно сделать вывод, что после завершения операции записи 20-мегабайтное пространство исходного фрагмента 1 разбивается на первое 6-мегабайтное свободное пространство, среднее 10-мегабайтное пространство, закрытое для записи, и последнее 4-мегабайтное свободное пространство. Затем операция освобождения свободного места до и после показана на следующем рисунке:

Далее, в соответствии с жизненным циклом блока памяти, полученного в b(2), создайте структуру данных time=>(alloc, free) (время — это последовательность opr, alloc — массив указателей блока памяти alloc на этом этапе). момент времени, а free — массив указателей на свободный фрагмент памяти в данный момент времени), и в соответствии с характером зависимости записи, полученной в b(1), моделируем весь процесс управления памятью:

(Алгоритм имитирует схемы распределения и свободного потока)

Как показано на рисунке выше, весь процесс можно разделить на три этапа, которые также соответствуют трем упомянутым выше методам:

• Вызовите метод free(), чтобы освободить фрагмент памяти, который необходимо освободить в данный момент времени и не зависит от записи.
• Вызовите метод alloc() для применения к блоку памяти, который необходимо применить в данный момент времени и не является поведением, зависящим от записи.
• После выполнения двух вышеуказанных шагов вызовите alloc_overwrite() для фрагмента памяти с зависимым от записи поведением для обработки.

(2) Получите общий размер пула памяти и назначьте адрес смещения каждому фрагменту памяти.

Для того, чтобы облегчить понимание того, как записывается адрес смещения блока памяти. Для специального обзора вышеупомянутого содержимого все фрагменты памяти можно разделить на две категории: фрагменты памяти (также называемые memory_chunk_root), которые применяются с помощью alloc() без поведения зависимости от записи, и те, которые применяются с помощью alloc_overwrite(). Блоки памяти зависимого поведения. Блок памяти может создать набор связанных списков на основе зависимостей записи, как показано на рисунке ниже.В качестве примера возьмем один из связанных списков (записи в память_chunk_n зависят от memroy_chunk_n-1,..., записи в memory_chunk_1 зависят от memory_chunk_root , который составляет связанный список зависимостей записи и анализирует все Зависимость записи фрагмента памяти может получить несколько связанных списков зависимостей записи). Поскольку информация о смещении и размере записывается в зависимости записи, каждый memory_chunk в связанном списке может вывести offset_root с memory_chunk_root. Затем для блока памяти во всем связанном списке необходимо только настроить начальный адрес (root_addr_begin) memory_chunk_root, чтобы завершить настройку адреса блока памяти во всем связанном списке.

(перезаписать связанный список и отношение памяти)

Далее мы опишем, как вывести адреса смещения всех фрагментов памяти. Размер переднего пространства получается в соответствии с фрагментом памяти переднего пространства, записанным в предыдущем моделировании. Отрегулируйте выделенное пространство фрагмента памяти в соответствии с размером переднего пространства, чтобы получить адрес смещения фрагмента памяти. В фактической реализации есть рекурсивные вызовы, а использование карты уменьшает дублирование запросов и вычислений. Упомянутое здесь предварительное пространство представляет собой блок памяти предварительного пространства записи моделирования и сумму блоков памяти всех предварительных пространств, полученных в результате итерации блока памяти.Интуитивно понятно, что при выделении пространства адресное пространство начинается с 0, начинается пространство, которое уже занято перед ним.

Как показано на рисунке ниже, вертикальное направление представляет собой логическое адресное пространство, красная область слева представляет собой адресное пространство, занятое передним пространством, а memory_chunk_n справа — это memory_chunk, который необходимо настроить на этот раз. Начальное значение root_addr_begin равно 0, и при настройке возможны две ситуации:

• Случай 1: Занятый размер переднего пространства

• Случай 2: Занятый размер переднего пространства > root_addr_begin + offset_root, что означает, что фактическое использование пространства memory_chunk_n и переднего пространства будет конфликтовать, поэтому root_addr_begin необходимо настроить на root_addr_begin + offset_root = размер занимаемого переднего пространства чтобы избежать конфликтов и избежать конфликтов Создать пустую трату пространства.

(Схема распределения логических адресов)

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

• Общий размер пула памяти: максимальное значение, достигнутое root_addr_begin, плюс его собственный размер среди всех memory_chunk_roots.
• Адрес смещения фрагмента памяти: root_addr_begin + offset_root

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

3. Резюме

Наконец, краткое изложение преимуществ оптимизации памяти MegEngine с использованием «анализа потока данных в последовательных графах программы для повторного использования памяти»:

• Использование памяти модели может быть уменьшено.
• Это полностью прозрачно для модели и не повлияет на точность модели.
• Это чисто ориентированная на память оптимизация, которая не приведет к потерям в вычислениях/коммуникациях.
• Оптимизация памяти выполняется до запуска модели, и затраты времени на управление памятью практически нулевые.
• Пиковая видеопамять статической памяти может быть получена заранее, в определенной степени, чтобы избежать ситуации, когда программа выполняется наполовину из-за нехватки видеопамяти, что приводит к сбою обучения и пустой трате времени.
• Унифицированное приложение памяти снижает затраты на управление памятью на устройствах хранения.