В командах ML-инженеров часто пользуются метрикой «GPU Utilization» (Загруженность процессора), чтобы понять, насколько активно задействуется в работе процессор. Чтобы узнать эту информацию, обычно достаточно выполнить команду nvidia-smi в строке терминала. Во многих интегрированных наблюдательных инструментах загруженность процессора также отслеживается как основная характеристика производительности. Но иногда, как ни удивительно, эта метрика даёт не слишком точное представление о производительности GPU. На самом деле, GPU можно загрузить на 100%, выполняя лишь операции чтения и записи (в памяти), но при этом 0 вычислений. Эта статья – не о том, как мы это выяснили, а о том, что нам удалось узнать по ходу дела.

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

• Насытили GPU, изменив заданные по умолчанию значения загрузчика данных (а именно, num_workers, batch_size, pin_memory, коэффициент предвыборки, т.д.)

• По максимуму использовали тензорное ядро, задействовав смешанную точность (fp16, bf16)

• Использовали оптимизатор с функцией слияния от apex/deepspeed (напр., FusedAdam, FusedAdamW, т.д.)

• Использовали инстансы/сети, специально рассчитанные на обучение моделей (H100SXM, A100SXM). Также по возможности использовали более новые инстансы H100 > A100 > V100

При помощи этих простых изменений мы добились 100%-й загруженности GPU и значительной потребляемой мощности, и это отлично! Чтобы проверить, а можно ли добиться большего, мы вычислили фактические показатели использования (MFU) на учебных рабочих нагрузках.

Кратко напомню: MFU означает загруженность модели в пересчёте на FLOPS (количество операций над числами с плавающей точкой в секунду). Это одна из лучших метрик, по которым можно судить о производительности GPU, впервые предложена в статье о PaLM от Google. Это «отношение наблюдаемой пропускной способности (токенов в секунду) к теоретической максимальной пропускной способности системы при пиковом значении FLOP». Проще говоря, эта метрика означает, сколько операций над числами с плавающей точкой успевает провести компьютер при предложенной вами рабочей нагрузке в сравнении с максимальными возможностями вашего GPU. Единственный реальный недостаток MFU в том, что порой эту метрику гораздо сложнее вычислить, чем, например, загруженность GPU, поскольку MFU зависит как от заданных параметров, так и от используемых фреймворков.

К сожалению, при обучении модели достигалась величина лишь ~20% MFU. Для справки: в настоящее время при обучении большинства LLM удаётся выйти на уровень примерно 35% - 45% MFU. Поэтому мы задумались: как так получается, что мы используем всего 20% от теоретического максимума той вычислительной мощности, что заложена в нашем GPU, но при этом сам графический процессор у нас загружен на 100%?

Чтобы ответить на этот вопрос, давайте выясним, что именно отслеживается при измерении загруженности GPU.

В самом деле, а что такое загруженность GPU?

Загруженность GPU достаточно зыбко определяется в документации Nvidia как «актуальная степень вовлечения в работу как для вычислительных ресурсов GPU, так и интерфейса для работы с памятью». Это просто перл.

Удивительно, но нашлось и более качественное определение — в документации по NVML от Datadog. Здесь этот параметр характеризуется как «Процент времени за период, охваченный последней выборкой, в течение которого одно или несколько ядер GPU были заняты работой». Чтобы понять, почему же такое определение — превратное, давайте кратко освежим в памяти, как именно работают GPU.

В GPU есть ядра и менеджеры многопроцессорности. В GPU от Nvidia такие менеджеры многопроцессорности называются «потоковыми мультипроцессорами» (SM), а аппаратном обеспечении от AMD они именуются «вычислительными блоками» (CU). Ниже показан процессор GH100 GPU, в котором 144 таких блока.

Такие потоковые мультипроцессоры подобны прорабам над группами рабочих, в данном случае — ядер. Когда вы запускаете ядро CUDA, работа выполняется именно на ядрах CUDA, это делают один или несколько SM. Как показано ниже, даже у одного SM на чипе GH100 много CUDA-ядер.

Таким образом, метрика «загруженность GPU» позволяет судить лишь о том, выполняет ли ядро работу в конкретный момент времени. Она не показывает, использует ли процессор все доступные ядра, а также распараллеливает ли рабочую нагрузку, чтобы задействовать максимум возможностей GPU. Может случиться и так, что вы будете наблюдать 100%-ю загруженность GPU, в то время как фактически будете только читать данные из памяти и записывать в неё информацию, выполняя при этом 0 FLOPS.

Теперь поясним: эта метрика может вводить в заблуждение лишь тех, у кого нет базового образования в области системного программирования (например, многих ML-инженеров). Как упоминается здесь, определение загруженности GPU в таком виде действительно применимо в рамках методологии «USE».

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

Более глубокий анализ

Изыскивая резервы производительности, на следующем этапе нам определённо стоит отпрофилировать цикл обучения модели. Рассмотрим, как устроен такой цикл в профилировщике Pytorch, чтобы лучше представлять ситуацию.

Как показано ниже, для ядра Softmax регистрируется высокая загруженность GPU, но при этом низкая эффективность SM. Нам это показалось серьёзным тревожным звоночком, поскольку печально известно, каким узким местом при обучении LLM порой оказывается наивная реализация softmax. В такой ситуации применяется множество операций слияния ядер, например FlashAttention, призванных помочь с ограничениями softmax при работе с памятью. Зная это, можем сделать следующий вывод: возможно, статистика по производительности SM указывает на общую неэффективность выполнения модели.

Но что именно показывает параметр эффективности SM?

Эффективность SM (также иногда называется активностью SM) – это одна из метрик Nvidia GPU, характеризующая, какова была эффективность (в процентах) у всех SM, которые были активны в заданный интервал времени. Как было указано выше, SM можно считать «прорабами» над группами ядер CUDA. Например, в GPU Nvidia H100 насчитывается 132 SM, каждый из которых управляет 128 ядрами процессора — всего получаем 16 896 ядер. Измеряя эффективность SM, можно определить, используют ли ядра CUDA наши потоковые мультипроцессоры. Если у нас есть ядро CUDA, которое непрерывно работает в течение 10 секунд, но использует всего 1 SM, то на карте H100 в таком случае будет зарегистрирована 100%-я загруженность, но эффективность SM в то же время составит 1 / 132 = 0,7%.

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

Выполняем оптимизации

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

Говоря о слиянии, мы имеем в виду, что будем пользоваться не нативным определением набора слоёв, применяемым в PyTorch, а вместо этого возьмём ядро GPU, написанное на CUDA или Triton. В таком случае все слои комбинируются в одном ядре. Ускорение достигается за счёт того, что каждое ядро тратит меньше времени на чтение и запись в память GPU, чем приходилось бы тратить при выполнении математических операций в определённых слоях (напр. Softmax). Пример такого ядра, полученного слиянием — Flash Attention.

Написали ли мы эти ядра самостоятельно? Конечно же, нет. Для большинства из них уже существуют библиотечные реализации. Например, слои для Flash Attention реализованы в nn.Modules, поэтому вам не придётся пользоваться ядрами и самостоятельно писать torch.autograd.function с нуля. Кроме того, зачастую эти реализации уже оптимизированы с аппаратной точки зрения, так что они не только работают быстрее, но и расходуют меньше памяти.

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

В результате нам удалось добиться четырёхкратного ускорения и показателя 38% MFU в случае с данным конкретным клиентом, тогда как исходный показатель MFU составлял 20%. Большинство сделанных нами оптимизаций относится именно к тем ядрам, в которых производилось слияние. Также мы добились многих оптимизаций, подыскивая нужный уровень параллелизма для данной модели, учитывая как размер самой модели, так и имевшуюся в распоряжении клиента полосу передачи данных Infiniband шириной 3,2 Тб/с.

Заключение

Большинству команд, работающих с ИИ, настоятельно рекомендуем наряду с загруженностью GPU также отслеживать эффективность потоковых мультипроцессоров в их кластере GPU. Так получается гораздо более репрезентативная картина, по которой можно судить, сколько ещё производительности можно выжать из GPU. В свою очередь, показатель загруженности GPU позволяет судить о том, насколько много машина работает вхолостую. Конечно, при этом было бы также неплохо вычислять MFU, но вряд ли у вас получится отслеживать эту метрику всё время, слой за слоем. Кстати, в Nvidia DCGM (менеджере GPU для ЦОД) информация об активности SM предоставляется по умолчанию.

Существует и множество более детализированных метрик, например, степень занятости SM (в профилировщике Pytorch она называется «Achieved Occupancy»). По ней можно судить, сколько работы выполняет каждый SM. Правда, понять эти метрики не так просто, гораздо удобнее попытаться вывести эффективность SM на максимум. Если вы хотите подробнее изучить эту тему, рекомендуем вам почитать о ней в блоге по Pytorch Profiler, документации по DCGM, руководстве по профилированию ядра от Nsight и в документации по Nsight.

Спасибо, что дочитали. Удачи вам, выжимайте производительность из ваших GPU до последней капли!