Итак, вы хотите улучшить эффективность работы своей модели глубокого обучения. Как подойти к такой задаче? Народ в таких случаях часто набрасывается на «сборную солянку» из всяких хитрых приёмов, которые, вроде бы, кому‑то когда‑то помогли, или хватает что‑то, встреченное в каком‑нибудь твите, вроде «Используйте операции, изменяющие исходные данные! Задайте значение None для градиентов! Устанавливайте PyTorch 1.10.0, но ни в коем случае не 1.10.1!».

Понятно — почему люди часто прибегают к таким вот спонтанным действиям в подобных ситуациях. Ведь «эффективность работы» современных систем, их «производительность» (в особенности — систем глубокого обучения) часто кажутся нам понятиями, которые ближе к алхимии, чем к науке. Тем не менее — рассуждения о производительности, в основе которых лежат базовые принципы работы компьютерных систем, способны устранить надобность в широком круге «магических» приёмов и в результате значительно облегчить путь к решению проблемы.

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

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

1. Вычисления: время, потраченное GPU на выполнение реальных операций с плавающей точкой («флопсы»).

2. Память: время, потраченное на передачу тензоров в пределах GPU.

3. Затраты на вспомогательные операции (оверхед): всё остальное.

Так же, как и при обучении моделей, знание о том, в каком режиме работает система, позволяет сузить круг возможных оптимизаций, которые могут привести к ощутимым улучшениям. Например, всё время тратится на перемещение данных в памяти (то есть — система пребывает в режиме, в котором эффективность её работы ограничена памятью). В таком случае увеличение «флопсов» GPU ничего не даст. А, с другой стороны, если всё время тратится на выполнение больших «тяжёлых» матричных вычислений (то есть — эффективность работы системы упирается в скорость вычислений), тогда переписывание логики модели на C++ ради уменьшения оверхеда пользы не принесёт.

Поэтому — если нужно, чтобы GPU постоянно делал бы «брррр», стоит поговорить о трёх областях, в которых тратится время системы — о вычислениях, о пропускной способности памяти, и об оверхеде.

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

Вычисления

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

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

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

Ниже показана таблица № 2 из той публикации.

Вот — таблица №3 из публикации.

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

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

Ситуация, когда сложность полного использования вычислительных ресурсов всё время растёт, с одной стороны, гарантирует постоянную занятость ML‑инженерам. А с другой стороны — делает крайне важным понимание того, что именно является узким местом некоей системы.

Скажу ещё пару слов о «флопсах». Все современные ускорители машинного обучения обладают специализированными аппаратными подсистемами для умножения матриц. Например — это тензорные ядра (Tensor Core) NVIDIA.

Поэтому, если умножением матриц вы не занимаетесь, то на NVIDIA A100 вам удастся добиться лишь 19,5 терафлопсов, а не упомянутых в документации 312. Обратите внимание на то, что такая ситуация не является чем‑то уникальным именно для GPU. На самом деле, TPU даже менее универсальны, чем GPU

Тот факт, что GPU показывают гораздо меньшую производительность на задачах, где не используется умножение матриц, поначалу может показаться проблемой. А как же другие операторы, вроде нормализации слоёв и функций активации? Но дело тут в том, что процент таких операторов, в плане необходимых для них флопсов, крайне мал. Например — посмотрим на таблицу, взятую отсюда, где подсчитывается процент флопсов, требующихся в BERT для выполнения операторов различных типов. Здесь «Tensor contraction» — это умножение матриц.

Можно видеть, что операции, не относящиеся к умножению матриц, требуют лишь 0,2% имеющихся флопсов. Поэтому неважно то, что GPU обрабатывает другие операции в 15 раз медленнее, чем операции умножения матриц.

Но в данном случае на нормализацию и поэлементные операции уходит, соответственно, в 250 и в 700 раз меньше флопсов, чем на умножение матриц.

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

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

Пропускная способность памяти

Затраты времени, связанные с пропускной способностью памяти — это, в сущности, «стоимость» перемещения данных из одного места в другое. Например — передача данных из CPU в GPU, из одного сетевого узла в другой, или даже из глобальной памяти CUDA в разделяемую память CUDA. И, в частности, последний из упомянутых случаев нам особенно интересен, о нём мы и будем здесь говорить. Именно его обычно имеют в виду, когда говорят о временных затратах, связанных с «пропускной способностью», или с «пропускной способностью памяти».

Другие два случая (их обычно называют, соответственно, затратами, связанными с «передачей данных» и с «передачей данных по сети»), тоже, безусловно, важны. Но если я заведу тут разговор о производительности распределённых систем, то никогда эту статью не допишу.

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

Хотя реальная работа выполняется на заводе, он плохо подходит на роль большого хранилища. Во многом это так из‑за того, что, так как на заводе мы занимаемся реальной работой, все системы хранения, которые в нём есть, оптимизированы в расчёте на то, чтобы ими можно было бы пользоваться быстро (SRAM). Завод не рассчитан на то, чтобы в нём было бы много таких хранилищ.

Где же тогда хранить результаты работы завода и необходимые ему материалы? Обычно — на некоем складе. Возможно, этот склад расположен там, где земельные участки стоят недорого, и там, где есть много места (DRAM). С этого склада мы можем отправлять на заводы материалы. В него же будет поступать продукция, производимая заводами (пропускная способность памяти).

Стоимость перемещения чего‑либо в вычислительный блок и из него и называют стоимостью «пропускной способности памяти». Отмечу, что DRAM GPU — это то, что показывается в nvidia‑smi. И именно об этой памяти идёт речь во всеми любимых сообщениях об ошибках «CUDA Out of Memory».

Тут стоить отметить, что каждый раз, когда мы выполняем что‑либо на ядре GPU, нам нужно переместить данные из DRAM GPU и обратно (то есть — воспользоваться складом).

Теперь представьте себе, что произойдёт, когда выполняется унарная операция, вроде torch.cos. Надо отправить данные со склада на завод, провести какие‑то простые вычисления над каждым элементом данных, а затем отправить то, что получилось, обратно. Отправка чего‑либо куда‑либо — это довольно дорогое удовольствие. Как результат — почти всё время, потраченное на это дело, ушло на отправку и получение данных, а не на сами вычисления.

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

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

Да вы посмотрите — довольно-таки бестолковая получилась схема. Почему мы снова и снова пересылаем одни и те же данные между глобальной памятью и вычислительным модулем? Лучше было бы хранить промежуточные результаты вычислений в самом вычислительном модуле (на заводе), а когда всё будет посчитано — отправлять их в память!

Это — так называемый «фьюзинг операторов» (operator fusion) — самая важная оптимизация в компиляторах, применяемых в глубоком обучении. Если описать это простыми словами, то, вместо того, чтобы писать данные в глобальную память, делая это только для того, чтобы снова их прочитать, мы опускаем излишние операции по работе с памятью, выполняя несколько этапов вычислений за один заход.

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

x1 = x.cos() # Прочитать из x в глобальной памяти, записать в x1
x2 = x1.cos() # Прочитать из x1 в глобальной памяти, записать в x2

А если применить фьюзинг операторов — понадобится лишь 2 операции чтения/записи глобальной памяти! Получается, что этот приём ускорит работу в 2 раза.

x2 = x.cos().cos() # Прочитать из x в глобальной памяти, записать в x2

Гораздо лучше.

Правда, тут есть несколько нюансов, из‑за которых реализация фьюзинга несколько усложняется. Во‑первых — GPU, когда он выполняет текущую операцию, нужно знать о том, что должно произойти дальше. Поэтому такую оптимизацию нельзя произвести в стандартном режиме работы PyTorch (Eager Mode). В этом режиме PyTorch выполняет операторы по одному за раз. Во‑вторых — для этого нужно сгенерировать CUDA‑код, что сразу сваливает на нас целую кучу проблем.

Не все фьюжны операторов так же просты, как в случае с поэлементными операторами. Поэлементные операторы можно «сфьюзить», превратив их в редукцию (reduction) или в умножение матриц. Даже само умножение матриц можно рассматривать в качестве фьюзинга умножения с автоматическим расширением размерности и размеров массивов (broadcasting multiply) и редукции.

Если вас интересует написание собственных ядер CUDA — весьма вероятно то, что именно в этом вы и увидите главный плюс фьюзинга. Если имеются 2 любых оператора PyTorch — это уже открывает возможности для их фьюзинга, благодаря чему экономится полоса пропускания при работе с глобальной памятью, устраняются затраты времени, необходимые на обмен данными между этими операторами. Кроме того, множество существующих компиляторов часто способны применять «простые» фьюзы. Например — NVFuser и XLA. Но автоматике далеко до той изобретательности, которой обладает человек. Поэтому, если вы хотите попробовать себя в написании собственных ядер CUDA — отличной отправной точкой для ваших экспериментов станет язык и компилятор Triton.

И наконец — фьюзинг операторов даёт удивительные побочные эффекты. Вот один из них: фьюзинговая конструкция x.cos().cos() выполняется почти за то же самое время, что и обычный вызов x.cos(). Именно поэтому различные функции активации обладают практически одинаковой «стоимостью», несмотря на то, что gelu, естественно, включает в себя гораздо больше операций, чем relu.

Этот факт ведёт к некоторым интересным последствиям для рематериализации/оптимизации работы с памятью при хранении результатов работы функций активации (rematerialization/activation checkpointing). Смысл тут в том, что выполнение дополнительных пересчётов может вести к уменьшению потребности в пропускной способности памяти, а значит — к сокращению времени работы кода. Получается, что можно снизить и потребность системы в памяти, и длительность выполнения программ посредством рематериализации. Именно этот подход мы использовали при создании приятной min‑cut‑оптимизации в AOTAutograd. Подробности об этом смотрите здесь (я, может быть, ещё напишу об этом!).

Рассуждения о «стоимости» пропускной способности памяти

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

Для простых операторов вполне нормально непосредственно рассуждать о необходимой им пропускной способности памяти. Например, видеоускоритель Nvidia A100 обладает пропускной способностью глобальной памяти в 1,5 терабайта в секунду и способен выполнять вычисления со скоростью 19,5 терафлопсов в секунду. Поэтому, если используются 32-битные (то есть — 4-байтные) числа с плавающей запятой, можно загрузить 400 миллиардов чисел за то же самое время, за которое GPU способен выполнить 20 триллионов операций. Более того, для выполнения простого унарного оператора (вроде умножения тензора на 2), нам, на самом деле, надо записать тензор обратно в глобальную память.

Поэтому… до тех пор, пока в унарном операторе не производится что‑то около сотни операций, основное время будет тратиться не на вычисления, а на работу с памятью.

Воспользовавшись фьюжн‑компилятором, вроде NVFuser, довольно легко измерить соответствующие показатели самостоятельно! Код можно найти на Colab.

Возьмём PyTorch‑функцию:

def f(x: Tensor[N]):
    for _ in range(repeat):
        x = x * 2
    return x

Если измерить производительность этой функции с помощью фьюжн‑компилятора — можно получить разные значения флопсов и пропускной способности памяти, достижимые для различных значений repeat. Увеличение repeat — это лёгкий способ увеличения объёма вычислений без увеличения количества обращений к памяти. Такой подход ещё известен как повышение интенсивности вычислений (compute intensity).

Если говорить конкретнее, то представим, что мы оценили производительность этого кода и нашли количество итераций, выполняемых в секунду. Затем, в виде функции от N (размера тензора), мы выполняем 2*N операций доступа к памяти и производим вычисления, оцениваемые как N*repeat FLOP. Получается, что пропускная способность памяти, которая при этом достигается, будет вычисляться по формуле bytes_per_elem * 2 * N * itrs_per_second (байты_на_элемент * 2 * N * итерации_в_секунду), а вычислительные ресурсы, «флопсы», которые на это тратятся, будут вычисляться по формуле N * repeat * itrs_per_second (N * repeat * итерации_в_секунду).

Теперь построим графики, отражающие время выполнения программы (Runtime), «флопсы» (FLOPS) и пропускную способность памяти (Memory Bandwidth). Обратите внимание на то, что при построении этих графиков используется двойная логарифмическая шкала.

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

В результате мы начинаем с достижения жалкого показателя в 0,2 терафлопса. После каждого удвоения интенсивности вычислений мы видим линейный рост этого показателя, который продолжается до тех пор, пока не доходит до пикового значения в 9,75 терафлопсов. (Возможно — это не та цифра, которую можно увидеть в спецификациях на Nvidia A100, где упомянуты 19,5 терафлопсов. Причина этого в том, что GPU, кроме прочего, имеют ещё более специализированное аппаратное обеспечение для фьюзинговых инструкций умножения и сложения (FMA). Поэтому при выполнении полностью универсальных вычислений A100 доходит лишь до 9,75 терафлопсов.) После того, как мы приблизимся к пиковому значению терафлопсов, наши вычисления можно будет счесть задачей, производительность которой зависит от скорости вычислений.

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

В данном случае чётко видно — когда мы ограничены скоростью вычислений, а когда — пропускной способностью памяти. При repeat < 32 мы полноценно нагружаем полосу пропускания памяти в то время как вычислительные мощности оказываются недоиспользованными. А как только мы оказываемся в ситуации, когда repeat > 64, получается, что вычислительные мощности нагружаются полноценно (то есть — мы доходим до показателей FLOPS, близких к пиковым), а показатели использования пропускной способности памяти начинают падать.

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

Один из распространённых подходов к оценке того, насколько система ограничена скоростью вычислений, заключается в том, чтобы измерить достигнутый ей показатель FLOPS и посчитать то, сколько процентов он составляет от пикового показателя FLOPS. Например, если некая система достигает 80% от пикового показателя — тогда можно сказать, что она зависит от скорости вычислений, по меньшей мере, на 80%, что очень прилично! Остальное время её работы, вероятно, тратится на выполнение операций, зависящих от пропускной способности памяти. (Существует множество подходов к подсчёту «флопсов». Но сейчас их очень просто и приятно считать в PyTorch. Тут вы можете найти обсуждение этого вопроса.)

Однако, на производительность систем глубокого обучения воздействует не только пропускная способность памяти. Есть ещё один фактор, который может помешать GPU делать «брррр».

Оверхед

Затраты на вспомогательные операции, или оверхед, это время, которое уходит на выполнение действий, которые не относятся к передаче тензоров и к вычислениям. Как классифицировать время, проведённое в интерпретаторе Python? Это — оверхед. А время, проведённое во фреймворке PyTorch? Это — тоже оверхед. Время, уходящее на запуск ядер CUDA (но не на их выполнение) — это, снова, оверхед.

Главная причина того, что оверхед — это настолько опасная проблема, заключается в том, что современные GPU — это реально быстрые устройства. Видеоускоритель Nvidia A100 может выполнять 312 триллионов операций с плавающей точкой в секунду (312 терафлопсов). В сравнении с этими цифрами Python — это очень меееееееееееееедленная система. Локальные бенчмарки показывают, что Python может выполнять 32 миллиона операций сложения в секунду.

Всё это значит, что за то время, за которое Python «сделает» 1 флопс, A100 мог бы выдать 9,75 миллиона флопсов.

Сгустим краски: интерпретатор Python — это не единственный источник оверхеда. Фреймворки вроде PyTorch тоже предусматривают использование множества уровней абстракции, через которые нужно пробиться для того чтобы попасть в вычислительное ядро. Если провести такой же эксперимент с PyTorch — получится, что этот фреймворк может выполнить лишь 280 тысяч операций в секунду. Конечно, маленькие тензоры — это не те сущности, для обработки которых создан PyTorch. Но… если вы пользуетесь маленькими тензорами (в научных вычислениях, например), то вы сможете обнаружить, что PyTorch, в сравнении с С++ — это очень медленно.

Например — взгляните на следующий Flame‑график, иллюстрирующий то, как в PyTorch производится единственная операция сложения. Там выделен крошечный участок графика. Именно он и олицетворяет реальные вычисления. Всё остальное — чистый оверхед.

Увиденное вас, возможно, шокирует и заставит задаться вопросом о том, почему кто‑то вообще использует PyTorch. Но не забывайте, что современные модели глубокого обучения часто основаны на выполнении огромных операций. Более того — фреймворки вроде PyTorch выполняются асинхронно. То есть — пока PyTorch выполняет ядро CUDA, он может заниматься и другими задачами, ставя в очередь дополнительные ядра CUDA. Поэтому, до тех пор, пока PyTorch может «бежать впереди» ядер CUDA — большая часть оверхеда фреймворка оказывается полностью скрытой!

Как понять, что система находится в таком режиме? Ну, так как оверхед обычно не растёт с ростом масштабов задачи (в то время как потребление вычислительных ресурсов и памяти растут), легче всего это можно выяснить, просто увеличив размер данных. Если это не приведёт к пропорциональному увеличению времени выполнения программы — это значит, что эффективность работы системы ограничена оверхедом. Например, если при удвоении размера пакета время выполнения программы вырастет лишь на 10% — это с высокой долей вероятности означает, что производительность системы упирается в оверхед. (Это — далеко не единственная причина, по которой увеличение размера пакета может не привести к пропорциональному увеличению времени выполнения программы. В определённых режимах такой ход, кроме прочего, ведёт к росту интенсивности вычислений. Например, в MLP обычно выполняют умножение матриц вида [B, D] x [D, D]. Если B меньше, чем D (скажем, размер пакета — 1, в то время как размер скрытого измерения — 128), тогда может оказаться так, что, незначительно увеличивая общую нагрузку на память, мы удваиваем объём необходимых вычислений. Я не смог додуматься до того, как объяснить эту особенность простыми словами.)

Ещё один способ оценки оверхеда заключается в использовании профилировщика PyTorch. Ниже розовые линии указывают на то, как ядра CPU соотносятся с ядрами GPU.

Ещё отмечу, что запись «GPU-Util» (не «Volatile GPU-Util») в nvidia-smi показывает то, какой процент нижнего ряда реально выполняет ядро GPU. Поэтому это — ещё один хороший способ на глаз оценить масштабы оверхеда.

Главная причина существования этого оверхеда кроется во всей той гибкости, которой обладают фреймворки вроде PyTorch. В сущности — речь идёт о том, что фреймворку нужно потратить много времени на то, чтобы «понять, что ему делать».

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

1. Python должен найти то, что __add__ диспетчеризует в a.

2. PyTorch должен определить множество атрибутов тензора (таких, как dtype, device, а также — нужно ли применять autograd) для принятия решения о том, какое именно ядро вызывать.

3. PyTorch нужно вызвать подходящее ядро.

По сути — этот оверхед произрастает из гибкости, позволяющей делать на каждом шаге что-то новое. Если такая гибкость не нужна — один из способов с ней разобраться заключается в том, чтобы отследить соответствующие фрагменты кода с помощью jit.trace, FX или jax.jit. Кроме того, сделать это можно даже и на более низком уровне, воспользовавшись чем-то вроде CUDA Graphs.

К сожалению, за это приходится платить потерей гибкости. Есть один подход, который мне очень нравится, позволяющий, так сказать, взять лучшее из двух миров. Заключается он в том, чтобы написать что-то большее, что-то уровня «настоящего» JIT, подойдя к этому через анализ того, что происходит на уровне виртуальной машины. Подробнее об этой идее можно почитать в материале о TorchDynamo.

Итоги

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

Часто я вижу, как исследователи и другие люди, заинтересованные в ускорении своего PyTorch-кода, слепо пробуют разные подходы, не понимая того, в каком режиме работают их системы.

Конечно, пожалуй, если пользователям некоего фреймворка нужно обо всём этом думать — это означает недоработку со стороны самого этого фреймворка. Компилятор PyTorch или API профилировщика никогда не относились к технологиями, которыми… невероятно легко пользоваться, хотя этому вопросу сейчас уделяется пристальное внимание.

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