Материал подготовлен автоматически по первоисточникам: ссылки на них — в конце статьи.
Повторная компиляция может свести на нет экономию памяти, которую обещает новый порядок операций нейросети. В препринте, который не прошёл рецензирование и приводит числа из замеров самих авторов, заново сгенерированная программа не достигла расчётного пика в 22 из 54 конфигураций на CPU. Вместо передачи расписания обратно компилятору предлагается переставлять команды уже готовой программы и отдельно доказывать, что лучший порядок найден.
Почему повторная компиляция меняет результат
Планировщик выбирает порядок операций по графу модели, но после этого компилятор ещё группирует ядра, переиспользует освободившиеся буферы и расставляет освобождение памяти. Новый порядок влияет на эти решения. В итоге исполняемая программа отличается от той, для которой рассчитывали пик памяти.
Предложенный подход сначала компилирует модель обычным способом, а затем профилирует каждую команду сгенерированной функции. Профилировщик записывает размеры хранилищ, моменты их создания и освобождения, псевдонимы буферов, а также чтение и запись памяти.
Из этих событий строится новый граф зависимостей. После поиска подходящего порядка система переставляет исходные команды, но не меняет ядра, их аргументы и уже выбранное переиспользование буферов. Выделение памяти переносится непосредственно перед первой использующей её командой, а освобождение — в самый ранний безопасный момент.
Такой порядок точно воспроизводит расчётный пик, если выделение памяти не зависит от предыдущей истории работы распределителя. Условие выполнялось на CPU и соответствует статическим аренам памяти во встраиваемых средах. Распределитель CUDA хранит освобождённые блоки в кеше и округляет запросы, поэтому на GPU расчётный и измеренный пики могут разойтись: там нужна контрольная проверка готовой программы.
Метод проверяли на серверном выводе, небольших моделях класса микроконтроллеров, скомпилированном обучении и отдельном запуске на GPU. Эксперименты с небольшими моделями охватывали граф операторов и планировщик статической арены, а не прошивку устройства; полученное расписание также относится к конкретной форме входных данных.
Сертификат показывает, можно ли сэкономить ещё
Найденное расписание задаёт верхнюю границу: столько памяти достаточно для исполнения. Чтобы доказать оптимальность, нужна нижняя граница — объём, меньше которого не сможет использовать ни один допустимый порядок.
Для неё работа вводит две оценки. Первая рассматривает исполнение от начала графа и ищет неизбежно дорогой переход из уже выполненной части. Вторая проходит ту же задачу с конца. Обе сводятся к последовательности минимальных разрезов графа; если лучшая нижняя граница совпадает с пиком найденного расписания, сертификат подтверждает, что дальнейшая перестановка операций ничего не даст.
Сертификаты подтвердили оптимальность для 94% отложенных графов сгенерированных функций и для всех 16 исследованных моделей класса микроконтроллеров, причём проверка занимала меньше секунды. Это объясняет нулевой результат оптимизации на большинстве графов вывода: исходный порядок уже нельзя улучшить. Исключением стали сети с ветвлениями — у DeepLabV3 перестановка снизила пик на 10%, а в скомпилированном обучении максимальная экономия достигла 7,6%.
Когда менять планы разработки
Работа не предлагает заменять планировщик компилятора новым поисковым алгоритмом. Она меняет точку применения результата: оптимизировать стоит команды сгенерированной программы, а не абстрактный граф с последующей повторной компиляцией. В производственном контуре после перестановки всё равно нужно повторно измерить пик и сравнить выходы с исходной программой.
Для CPU и статических арен подход позволяет превратить оценку памяти в проверяемое свойство сборки. Команда может зафиксировать форму входа, получить расписание, проверить сертификат и прекратить поиск, если нижняя граница совпала с измеренным пиком. На GPU сертификат остаётся полезен для порядка операций, но кеширующий распределитель не позволяет считать расчётный объём гарантированным без запуска на устройстве.
Если оптимальный порядок всё равно не укладывается в бюджет, следующий рычаг — удалить промежуточное значение и вычислить его повторно перед поздним использованием. В общем случае выбрать такие значения вычислительно трудно даже для узких графов, а приведённый точный алгоритм для ограниченной ширины авторы рассматривают как теоретический результат, а не готовый инструмент.
На исследованных небольших моделях повторное вычисление целых тензоров не помогло: пик задавала одна операция вместе с её входами и выходом. На ReDimNet структура оказалась другой — ранний результат долго хранился ради позднего повторного использования. Повторное вычисление этого результата снизило пик на 5,2–6,1% при дополнительных вычислениях менее 0,4% и сохранило побитовое совпадение выходов.
Практический порядок действий получается коротким: сначала переставить команды без изменения ядер, затем доказать оптимальность порядка и только после этого добавлять повторные вычисления. Если пик создаёт одна крупная операция, потребуется делить тензоры и исполнять слой фрагментами — перестановка целых операций эту границу не обойдёт.
Источники
Похоже на вашу задачу?
Расскажите, что собираете. За полчаса разложим на этапы и назовём сроки — это бесплатно и ни к чему не обязывает.
