По большому счёту, самая первая и самая важная оптимизация, которую можно применить к любой современной системе машинного обучения, заключается в том, чтобы реализовать в этой системе пакетную обработку данных (batching). Для того чтобы получить результат работы системы (inference, инференс) в пакетном режиме — ей, вместо одного элемента входных данных, отправляют N таких элементов. Чаще всего никаких дополнительных нагрузок на систему это не создаёт. Формирование инференса для каждого из элементов, входящих в пакет размера N, занимает в точности столько же времени, сколько нужно для обработки одного элемента входных данных. Почему это так? На первый взгляд может показаться, что обработка пакета данных не может обойтись без некоторых накладных затрат ресурсов. В конце концов — оборудованию приходится выполнять в N раз больше действий.

Если прибегнуть к простейшей модели работы нейронной сети, то получится, что некоторая дополнительная нагрузка на систему, всё же, создаётся. Для выполнения пакетных вычислений нужно выполнить в N раз больше операций. И, на самом деле, если попробовать это на CPU, то окажется, что так оно и есть (среднее время формирования вывода для ResNet-50, Colab).

А вот если запустить тот же самый пример на современном GPU, то окажется, что это уже не так. Вот как это выглядит на Nvidia T4.

Переход от пакета размера 1 к пакетам размера 2 и 3 не приводит к росту времени, необходимого на обработку данных. Последующее увеличение размера пакета приводит к линейному увеличению времени.

Почему это так? Из‑за параллельной обработки данных. Современные видеоускорители распараллеливают вычисления (и они, на самом деле, медленнее GPU, если смотреть на скорость выполнения одного потока).

Когда мы размышляем о чём‑то наподобие «формирования вывода модели для элемента входных данных», мы обычно рассматриваем модель, как некую единую неделимую сущность. А ведь модели — это системы, которые состоят из множества матриц. Когда формируется вывод модели — эти матрицы загружаются в память. А конкретнее — каждый блок матрицы загружается в память устройства, в частности — в разделяемую память (shared memory). В Nvidia A100 всего 192 Кб такой памяти. Затем этот блок используется при вычислении результата для каждого элемента в пакете. Обратите внимание на то, что разделяемая память — это не то же самое, что «оперативная память GPU» (GPU RAM), то есть — HBM (High Bandwidth Memory, память с высокой пропускной способностью). У Nvidia A100, в зависимости от модели, имеется 40 или 80 Гб HBM, но только 192 Кб разделяемой памяти. Это приводит к появлению узкого места, связанного с пропускной способностью памяти, дающего о себе знать при выполнении математических операций. Дело в том, что они предусматривают постоянное выполнения операций чтения/записи в применении к разделяемой памяти. Можно примерно прикинуть время, необходимое для передачи в память значений, соответствующих весам нейронов, вычислив отношение размера модели к пропускной способности памяти. Время, необходимое для выполнения вычислений, можно приблизительно посчитать, поделив флопсы (FLOPS), необходимые модели, на флопсы, которые может обеспечить GPU.

В системах, основанных на многослойных перцептронах (MLP, multilayer perceptron) флопсы модели можно приблизительно найти по следующей формуле: 2 * количество параметров * количество элементов в пакете (2 * m * n * b для пакета размера b и для матрицы размером m x n). В результате время передачи данных равно времени вычислений при выполнении следующего условия:

Pпропускная способность памяти=2∙B∙Pфлопсы

Заметьте, что здесь можно избавиться от количества параметров P:

Это уравнение можно преобразовать, выразив размер пакета через другие показатели:

Когда размер пакета меньше отношения флопсов к пропускной способности памяти — это значит, что производительность системы ограничена пропускной способностью памяти. Когда размер пакета больше этого отношения — это значит, что производительность ограничена флопсами. Обратите внимание на то, что тут мы анализируем MLP, а не свёрточные сети (convolutional network), вроде ResNet-50. Анализ таких сетей — это уже другая, немного более сложная задача.

Видеоускоритель Nvidia T4 даёт нам 65 TFLOPS при работе с типом данных fp32 и обладает пропускной способностью памяти в 300 Гб/с. Это говорит о том, что наше «волшебное» соотношение будет равняться 216. При запуске на этом GPU MLP‑модели (глубина — 8, ширина — 1024), мы получаем примерно то, на что рассчитывали:

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

Это очень хорошо! Мы можем видеть наличие важного порогового значения при работе с широким разнообразием архитектур. Интересно ещё и то, что скорость расчётов для маленьких сетей особо от количества элементов в пакете не зависит. Об этом говорит примерно одинаковое время, необходимое для обработки пакетов с размерами от 1 до 512. Я, опираясь лишь на собственные представления, могу объяснить это тем, что GPU — это, когда дело доходит до вычислений — невероятно быстрые системы. А вот всё остальное (CPU и прочее) работает, так сказать, медленнее. В начале графиков можно видеть очень много «шума». Хорошего объяснения этому у меня нет (ну, разве что, могу пожать плечами и сослаться на оверхед).

Многие ML‑инженеры часто тратят время не на машинное обучение, а на избавление от оверхеда, от дополнительной нагрузки на систему. Эта нагрузка, в основном, создаётся не тем кодом, который отвечает непосредственно за машинное обучение. В исследованиях, где используется обучение с подкреплением (RL, reinforcement learning), особенно — там, где работают над задачами непрерывного обучения (continual learning), когда единственный агент выполняет длительную последовательность действий, GPU обычно использовать не стоит. К GPU имеет смысл прибегать в двух ситуациях. Первая — это когда имеется очень большая сеть. Вторая — это если проводится глубокая оптимизация всех остальных элементов используемого стека технологий. (Хотите довести старого DeepMind‑инженера до дрожи — спросите его про графовые нейросети — было время, когда RL‑окружения реализовывали в TensorFlow‑графах).

Свёрточные сети

В случае со свёрточными сетями количество весов равно количеству фильтров, умноженному на размер фильтра. Если говорить о torch.nn.Conv2d — это будет kernel_size^2 * out_channels. Поэтому, если имеется изображение размером (224, 224) с шагом перемещения 1 и с размером ядра 3, то мы применяем один и тот же фильтр 224 раза. Это означает, что, в частности — для свёрточного слоя, пакетная обработка данных даёт гораздо меньше преимуществ. Дело в том, что одни и те же веса мы используем по очень много раз. В случае c подвыборочным слоем (pooling layer), как и можно ожидать, наблюдается практически линейная зависимость нагрузки от количества пикселей.

Трансформеры

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

То же самое справедливо и для моделей архитектуры «смесь экспертов» (MoE, Mixture of Experts). Во многих реализациях трансформеров KV-кеш размещается внутри класса внутреннего внимания (отличный пример — MaxText). Единственное различие между MoE и классическим декодером в том, что некоторые из слоёв прямого распространения заменены на MoE-слои. KV-кеш в таких системах ведёт себя точно так же, то же самое касается и формирования вывода модели. Но тут есть одна загвоздка.

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

Надеюсь — вам пригодится то, о чём вы прочли. Если у вас будут вопросы по поводу пакетной обработки данных (или, что ещё важнее, если вы найдёте какую-нибудь неточность) — свяжитесь со мной.