[Перевод] Масштабирование LLM: от одного чипа до ЦОДа. Глава 1. Теоретические основы

Предисловие

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

Прежде чем начнете читать

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

• как обучаются нейросети, в частности механизм прямого и обратного распространения ошибки

• базовое понимание того, как работают языковые модели в целом и механизм внимания в частности

• ну и там по мелочи, например знать, как перемножаются вектора и матрицы :)

На что тратится время при вычислениях?

Допустим взяли мы какую-то нейронку, подготовили для нее данные, запустили тренировку и вот спустя N часов получили обученную модель. А собственно почему именно N часов? Чем именно "железо" занималось все это время? А занималось оно в основном двумя вещами:

Вычисления. Работа нейросети это в основном последовательность матричных умножений, при этом каждая матрица (обычно) состоит из чисел с плавающей точкой (floating-point). Умножение матриц это операции умножения и сложения этих чисел, то есть последовательность операций над числами с плавающей точкой (Floating-point Operations - FLOPs). Поэтому чтобы посчитать, сколько времени в целом займут вычисления, нужно просто прикинуть, сколько в целом FLOPs для этого всего потребуется, а затем поделить на мощность нашего ускорителя (предполагаем, что все это мы делаем на какой-нибудь видеокарте (GPU) или специально заточенном под это дело ускорителе (TPU)). Получим такую формулу:

Где T_math это собственно искомое время, Computation FLOPs это общий объем вычислений, а Accelerator FLOPs/с - производительность нашего ускорителя, т.е. сколько в теории операций с плавающей точкой он сможет выполнить за секунду. Например NVIDIA H100 имеет производительность около 9.89 * 10^14 bfloat16 FLOPs/с (bfloat16 или bf16 - это такой формат для чисел с плавающей точкой, имеет разрядность 16 бит и часто используется при работе с нейросетями). Предположим, что все вычисления мы тоже проводим в формате bf16 (в общем случае это скорее всего будет не так, но об этом поговорим в другой раз) и общий объем вычислений равен 10^12 FLOPs. Тогда получим, что чисто на вычисления мы потратим (в идеале):

Коммуникация. Как известно, большинство современных компьютеров и GPU/TPU в целом основаны на архитектуре фон Неймана. То есть состоят из процессора, памяти и шины между ними. Поэтому чтобы что-то вычислить на процессоре, необходимо сначала это что-то доставить на процессор из памяти через шину данных. И на это тоже нужно время. А как его посчитать?Прикинуть для начала, сколько в сумме нужно перекинуть данных туда-сюда, а затем поделить на пропускную способность шины данных. Для NVIDIA H100 это примерно 3.35 Тб/с или 3.35 * 10^12 байт/с.

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

Поэтому для времени, необходимого для передачи данных внутри / между чипами, получаем такую формулу

Где T_comms это искомое время, Communication Bytes это общий объем передаваемых данных, а Network/Memory Bandwidth Bytes/s - пропускная способность шины внутри ускорителя или между ускорителями. Например для передачи 1 ГБ между чипом и памятью NVIDIA H100 потребуется (опять таки в идеальных условиях):

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

Арифметическая интенсивность

А собственно как должно соотноситься время, затрачиваемое на вычисления (T_math) ко времени, затрачиваемому на передачу данных (T_comms)? То есть должно ли одно из них быть больше другого или они должны быть равны? Когда мы тренируем нейросеть, мы хотим, чтобы тренировка закончилась как можно быстрее, то есть идея в том, чтобы выполнить весь необходимый объем вычислений за как можно меньшее время. Следовательно, чем большее время наше оборудование будет заниматься вычислениями вместо простоя, тем лучше.

Поэтому когда T_math​ > T_comms​, мы можем сказать, что тратим на вычисления больше времени, чем на передачу данных, и это хорошо, поскольку означает, что наше оборудование бОльшую часть времени занято нужными нам вычислениями, а не перегонкой данных туда-сюда. Такое состояние называется вычислительно-ограниченным или compute-bound. Когда же мы наблюдаем, что T_comms > T_bound, значит оборудование находится в коммуникационно-ограниченном состоянии или communication-bound. Это означает, что как минимум какую-то часть времени ускоритель простаивает в ожидании получения или отправки данных.

А можно как-то заранее сказать, будут ли данные операции на данном ускорителе compute-bound или communication-bound? Да, можно, и для этого нам придется ввести такое понятие как "арифметическая интенсивность".

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

То есть арифметическая интенсивность (АрИ) измеряет, сколько в среднем операций приходится на один байт переданных данных.

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

А как должны соотноситься АрИ алгоритма и АрИ ускорителя? Поскольку мы хотим, чтобы время, затрачиваемое на вычисления (T_math), было больше времени, затрачиваемого на передачу данных (T_comms) то, исходя из данных выше определений T_math и T_comms, в примере с Nvidia H100 получаем:

То есть АрИ алгоритма должна быть не просто высокой - она должна быть выше АрИ ускорителя. Давайте рассмотрим в качестве примера произведение двух векторов и посчитаем АрИ для этой операции. Пусть оба вектора имеют размер N и представлены в формате bf16:

x • y: bf16[N], bf16[N] → bf16[1]

Для выполнения данной операции нам нужно:

• загрузить оба вектора из памяти (то есть два вектора по 2N байт, т.к. формат bf16, а значит 16 бит или 2 байта на элемент)

• провести N умножений и N - 1 сложений

• затем записать обратно в память 2 байта.

То есть АрИ векторного умножения стремится к 1/2 или, другими словами, векторное умножение производит 0.5 FLOPs на байт данных. Следовательно загружать ускоритель подобными операциями такая себе идея, потому что большую часть времени он будет простаивать.

Roofline график

Можно отобразить соотношение между вычислениями и данными с использованием т.н. roofline графика. Он показывает, какую производительность (FLOPS/c) можно выжать из алгоритма (ось y) с учетом его АрИ (ось x). Собственно вот пример такого графика:

Как видим из графика, производительность Алгоритма 1 (на графике Algo 1) как бы "упирается в крыши", сформированные пропускными способностями BW1 и BW2. Говорят, что он ограничен пропускной способностью шины (Bandwidth Bound). Причем более низкая BW1 ограничивает Алгоритм 1 сильнее, чем более высокая BW2. То есть чем ниже пропускная способность шины ускорителя, тем сложнее любому запущенному на нем алгоритму выжать из него все соки. Но это и не удивительно: чем ниже пропускная способность, тем выше АрИ ускорителя, а как мы помним из предыдущего раздела, АрИ алгоритма должен быть больше АрИ ускорителя.

А теперь взглянем на Алгоритм 2 (на графике Algo 2). У него настолько высокая АрИ, что ему плевать на низкую пропускную способность шины, он в любом случае упирается в вычислительные возможности ускорителя. В таком случае говорят, что Алгоритм 2 вычислительно ограничен (Compute Bound). В случае ускорителя Nvidia H100 это означает, что АрИ Алгоритма 2 больше 295. То есть на один байт данных он проводит более 295 арифметических операций, и это позволяет ему загрузить ускоритель по полной.

АрИ умножения матриц

Основная операция, которой занимается ускоритель при работе с нейросетями - умножение матриц (или Matrix Multiplication - MM). Давайте же посчитаем АрИ этой операции.

Пусть у нас есть матрица X размера [B, D], а также матрица Y размера [D, F]. Мы хотим их перемножить и получить на выходе матрицу Z размера [B, F]. Веса матриц хранятся в формате bf16, так что опять имеем по два байта на вес:

X • Y: bf16[B, D], bf16[D, F] → Z: bf16[B, F]

Итак, сколько операций потребуется? Нужно провести векторное умножение каждой строки X и каждого столбца матрицы Y. Оба этих вектора имеют размер D, а значит для их перемножения нужно провести около 2D операций. И такие перемножения нужно провести BF раз. Получаем, что для перемножения двух матриц размеров [B, D] и [D, F] нужно 2BDF FLOPs.

А сколько потербуется передать данных? Нужно загрузить две матрицы размером 2BD и 2DF байт, а затем записать в память матрицу размером 2BF байт. В сумме получаем 2BD + 2DF + 2BF байт. Таким образом АрИ матричного умножения получается:

Если предположим, что B это размер батча (в контексте трансформера это количество токенов, подающихся на вход модели за один проход), то формулу можно упростить, предположив, что B сильно меньше чем D и F. Такое предположение выглядит разумным, если учесть, что при работе с трансформером B обычно меньше 1024 токенов (именно токенов, а не последовательностей, и на один ускоритель, а не в целом), тогда как D и F > 8000. Тогда слагаемые BD и BF из знаменателя можно выкинуть и получим в итоге:

То есть для максимально эффективной работы алгоритма перемножения матриц на Nvidia H100 размер батча должен быть больше 295, очень удобное правило!

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

bfloat16[B, D] * int8[D, F] -> bfloat16[B, F]

В таком случае АрИ нашего алгоритма бы возросла, так как данных нам нужно передавать меньше, а количество вычислений осталось бы таким же. В итоге требуемый размер батча упал бы вдвое, до примерно 150.

АрИ для нескольких ускорителей

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

Поэтому рассмотрим случай, когда две матрицы X ~ bf16[B, D] и Y ~ bf16[D, F] поровну распределены между двумя ускорителями Nvidia H100. Тогда можно перемножить по половинке матрицы на каждом GPU:

GPU0: A = X[:, :D // 2] @ Y[:D // 2, :]

GPU1: B = X[:, D // 2:] @ Y[D // 2:, :]

А затем скопировать результаты обоих умножений на третий GPU, где их и сложить.

Допустим все ускорители соединены шиной NVLink с односторонней пропускной способностью 450 Гб/с или 4.5 * 10^11 байт/с. Каждый ускоритель имеет производительность 9.89 * 10^14 FLOPs/с. Какими в таком случае у нас будут T_math и T_comms?

T_math будет в два раза меньше, так как ускорителей теперь два, а не один.

А что насчет T_comms? Нам нужно передать две матрицы размером 2BF от двух GPU на третий. У каждой H100 только одно соединение NVLink, поэтому скорость передачи по сети останется такой же.

Таким образом, такой алгоритм становится Compute Bound, если АрИ(MM x 2) > АрИ(GPU x 2) или:

Отсюда получаем, что D > 8792. Заметим, что как только мы перешли от вычислений внутри ускорителя к распределенным вычислениям, выяснилось, что наша вычислительная эффективность теперь зависит от D, а не от B! И связано это в первую очередь с тем, что D теперь участвует только в вычислениях внутри ускорителя, но не участвует в передаче данных между ними. То есть чем большего размера будет D, тем больше вычислений будет сделано, а объем передаваемых данных при этом не изменится => вырастет АрИ нашего алгоритма, что нам собственно и надо.

И это всего лишь один из примеров. Матричные операции можно параллелить по-разному, и соответствующие ограничения при этом тоже будут разные. Но об этом мы поговорим в следующих главах. До встречи!

Бонус. Код для построения Roofline графика.

А вот немного кода, который позволит вам построить собственный Roofline график у себя дома без СМС и регистрации. Тут приведен случай матричного умножения внутри ускорителя, но ничто не мешает переделать функцию roofline под нужный вам случай.

import matplotlib.pyplot as plt import numpy as np gpu_flops = 9.89e14 gpu_bw = 3.35e12 bs = np.arange(1, 512) def roofline(B, D, F): total_flops = 2*B*D*F flops_time = total_flops / gpu_flops comms_time = (2*B*D + 2*D*F + 2*B*F) / gpu_bw total_time = np.maximum(flops_time, comms_time) return total_flops / total_time roofline_big = roofline(bs, 16384, 16384) roofline_med = roofline(bs, 2048, 2048) roofline_small = roofline(bs, 1024, 1024) plt.figure(figsize=(16, 8)) plt.plot(bs, roofline_big, label='F=D=16384') plt.plot(bs, roofline_med, label='F=D=2048') plt.plot(bs, roofline_small, label='F=D=1024') plt.legend() plt.xlabel('batch size') plt.ylabel('peak bfloat16 FLOPs/s on Nvidia H100') plt.grid()

На рисунке ниже приведен результат работы кода. По оси X отмечены разные размеры батча (по мере возрастания), по оси Y - какую долю общего потраченного времени мы занимались непосредственно вычислениями (исходим из предположения, что вычислять и передавать данные можно всегда одновременно). Обратите внимание, что алгоритм перемножения матриц с большими размерами F и D становится Compute Bound при размере батча около 300. Если увеличивать размеры матриц и дальше, то это значение будет стремиться к теоретически посчитанному нами значению 295 для Nvidia H100. Но чем меньше размеры матриц, тем больший размер батча потребуется для полной загрузки GPU.