Накладные расходы GPU: почему наша модель MNIST обучается быстрее на CPU

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

Мы сравнили нашу модель MNIST на GPU Quadro RTX 5000 и процессоре Intel Core i9-10885H. В одном запуске Keras CPU оказался быстрее: 4,4 с против 6,5 с на эпоху. В отдельном эксперименте цикл обучения JAX на CPU занял всего 0,14 с, но измерял другой объём работы, чем Keras. Разберём источники накладных расходов и выводы, которые позволяют сделать эти измерения.

Начнём с того, что происходит во время обучения.

Разбираем один шаг обучения

Шаг обучения включает прямой проход, вычисление градиента потерь, обратное распространение и обновление весов. Этот фрагмент сразу вычисляет объединённый градиент softmax и кросс-энтропии, не вычисляя скалярное значение потерь:

# Прямой проход
z1 = xb @ w1 + b1                                # умножение матриц + смещение (скрытый слой)
a1 = np.maximum(0, z1)                            # активация ReLU
z2 = a1 @ w2 + b2                                # умножение матриц + смещение (выходной слой)
exp_z = np.exp(z2 - z2.max(axis=1, keepdims=True))
probs = exp_z / exp_z.sum(axis=1, keepdims=True)  # softmax → вероятности

# Градиент потерь
dz2 = probs.copy()
dz2[np.arange(bs), yb] -= 1                       # насколько предсказания промахнулись
dz2 /= bs

# Обратный проход — вычисляем градиенты
dw2 = a1.T @ dz2                                  # градиент для W2
da1 = dz2 @ w2.T                                  # градиент, текущий назад
dz1 = da1 * (z1 > 0)                              # градиент ReLU
dw1 = xb.T @ dz1                                  # градиент для W1

# Обновляем веса
w1 -= lr * dw1
b1 -= lr * dz1.sum(axis=0)
w2 -= lr * dw2
b2 -= lr * dz2.sum(axis=0)

Куда уходит время? Чтобы выяснить это, мы обернули каждую операцию в time.perf_counter():

t = time.perf_counter()
z1 = xb @ w1 + b1
timings["fwd: X @ W1 + b1"] += time.perf_counter() - t

t = time.perf_counter()
a1 = np.maximum(0, z1)
timings["fwd: ReLU"] += time.perf_counter() - t
# ... и так далее для каждой операции

Операции GPU асинхронны, поэтому мы добавили tf.test.experimental.sync_devices(), чтобы дождаться завершения. Синхронизация после каждой операции сама добавляет расходы и мешает перекрытию работы. Поэтому эти времена не следует считать стоимостью оптимизированного цикла обучения.

Вот результаты для одного шага обучения (batch_size=32, среднее по 1000 шагам, время в микросекундах μs — миллионных долях секунды, меньше — лучше):

В этом тесте отдельных операций CPU быстрее. NumPy использует оптимизированные процедуры BLAS для умножения матриц, а измерения TensorFlow на GPU включают отправку операций и синхронизацию. SIMD-инструкции позволяют CPU обрабатывать несколько значений одной инструкцией, а FMA объединяет умножение со сложением. Мы сравниваем реализации вместе с методом измерения, а не изолированную скорость арифметики устройств.

Посмотрим на xb @ w1 + b1 и a1 @ w2 + b2 — умножения матриц из каждого слоя — и на то, почему GPU спроектированы делать их быстро. Вот прямой проход из нашей реализации на NumPy:

class HiddenLayer:
    def forward(self, x):
        self.z = self.W @ x + self.b        # умножение матриц + смещение
        self.out = np.maximum(0, self.z)     # активация ReLU
        return self.out

class OutputLayer:
    def forward(self, x):
        self.z = self.W @ x + self.b        # умножение матриц + смещение
        exp = np.exp(self.z - np.max(self.z))
        self.probs = exp / np.sum(exp)       # softmax → вероятности
        return self.probs

Для одного изображения self.W @ x умножает матрицу весов 128×784 на вектор из 784 пикселей. Каждый нейрон вычисляет скалярное произведение: 784 умножения и 783 сложения. После добавления смещения получаем 128 значений до активации, которые ReLU преобразует в скрытые активации:

W (128 × 784) @ x (784,) + b (128,) → z (128,)

Для батча из 32 изображений матрица X имеет форму (32, 784). Транспонируя веса из нашей реализации для отдельных изображений, получаем W1 формы (784, 128). Произведение имеет форму (32, 128); затем прибавляем смещение и применяем ReLU:

      X (32 × 784)                W1 (784 × 128)            result (32 × 128)
 ┌─────────────────┐         ┌──────────────────┐       ┌──────────────────┐
 │ img1:  p1 … p784│         │  n1   n2  … n128 │       │ img1: z1  … z128 │
 │ img2:  p1 … p784│    @    │  w    w   …  w   │   =   │ img2: z1  … z128 │
 │ ...             │         │  ...  ...    ... │       │ ...              │
 │ img32: p1 … p784│         │  w    w   …  w   │       │ img32:z1  … z128 │
 └─────────────────┘         └──────────────────┘       └──────────────────┘

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

Все 4096 выходных значений можно вычислять независимо, но оптимизированное умножение матриц на GPU обычно распределяет блоки матриц между группами потоков. По одному числу выходов нельзя определить число активных ядер или мультипроцессоров SM.

Для приблизительной оценки возьмём 3072 ядра FP32 на частоте 1,8 ГГц и две операции с плавающей точкой на умножение со сложением. Получим теоретический пик около 11 TFLOPS. Реальная производительность зависит от варианта GPU, частоты и задачи.

Умножение матриц требует примерно $2 \times 32 \times 784 \times 128 = 6{,}422{,}528$ операций с плавающей точкой, если считать умножение и сложение отдельно. При 11 TFLOPS нижняя граница времени одной лишь арифметики — около 0,58 μs. Оценка предполагает пиковую производительность и не учитывает доступ к памяти и отправку работы. Это не измеренное время ядра. Деление числа операций на пиковые FLOPS даёт время, а не процент загрузки GPU.

Куда на самом деле уходит время

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

1. Отправка ядер на выполнение. CPU и среда выполнения должны передать работу GPU. Для малых операций это может занимать заметную часть времени. Операция фреймворка может использовать несколько ядер, а компилятор может объединить несколько операций в одно ядро. Число выражений Python не определяет число запусков.

2. Передача данных и синхронизация. Входные данные может потребоваться скопировать из памяти CPU в память GPU. Батч из 32 изображений по 784 пикселя в формате float32 занимает около 100 КБ. Время передачи зависит от соединения, выделения памяти и синхронизации. При обычном обучении веса и градиенты могут оставаться на GPU; возвращать градиенты на CPU после каждого шага не нужно.

3. Работа Python и фреймворка. Отправка операций, выделение памяти, метрики и обработка данных имеют свою стоимость. В eager-режиме Python участвует в отдельных операциях. Выполнение графа и компиляция могут убрать значительную часть этой повторяющейся работы. XLA компилирует функцию при первом обращении или повторной трассировке, а не для каждой операции на каждом шаге.

Профилирование помогает различать эти затраты, если отдельно рассматривать длительность API-вызовов на CPU, выполнение на GPU и общее прошедшее время.

Профилируем накладные расходы CUDA

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

Чтобы выяснить это, мы воспользовались профилировщиком NVIDIA nsys (Nsight Systems), который перехватывает каждый вызов CUDA API:

nsys profile -o keras-gpu-profile python mnist-keras.py
nsys stats --force-export=true keras-gpu-profile.nsys-rep

CUDA — это программный слой NVIDIA между вашим кодом и оборудованием GPU. Когда TensorFlow хочет умножить матрицы, он не обращается к GPU напрямую — он вызывает функции CUDA вроде «выдели память», «скопируй эти данные», «запусти это ядро». Каждый такой вызов идёт через драйвер и несёт собственные накладные расходы. Профилировщик nsys записывает каждый из них, так что мы можем точно увидеть, куда уходит время.

Измеренная эпоха на GPU заняла 6,5 с. Ниже приведена соответствующая сводка CUDA API. При 48 000 обучающих изображений и батче 32 эпоха содержит 1500 обучающих шагов, но сводка показывает 1875 запусков графа. Одни счётчики не позволяют определить, какие вызовы относятся к обучению, валидации или подготовке.

Это длительности вызовов CUDA API на CPU, а не ядер на GPU. Вызовы могут перекрываться с работой GPU и других потоков CPU. Нельзя вычесть их сумму из 6,5 с и назвать остаток расходами Python. Для такого разбиения нужна временная диаграмма с чёткими границами измерения. Nsight Systems различает время API, ожидания в очереди и выполнения ядра.

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

В предыдущем сравнении Keras эпоха на CPU заняла 4,4 с, а на GPU — 6,5 с. CPU избегает CUDA-вызовов и передачи данных на отдельный GPU, но сохраняет расходы фреймворка, планирования и доступа к памяти.

Отдельный микробенчмарк TensorFlow измерял (32, 784) @ (784, 128) без синхронизации после каждой операции. Он показал следующие средние времена:

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

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

Во время обучения nvidia-smi может показывать загрузку GPU 90–100% даже для задачи, которая быстрее выполняется на CPU.

Загрузка GPU в nvidia-smi — доля интервала измерения, в которой выполняется хотя бы одно ядро. Она не показывает долю занятых вычислительных ядер или достигнутую долю пиковых FLOPS. Непрерывное выполнение маленьких ядер может давать высокий показатель без полной загрузки вычислительной мощности GPU.

Как ускорить один запуск обучения

Мы попробовали пять изменений: большие батчи, Keras на CPU, цикл NumPy, компиляцию JAX и прямой перенос на CuPy. Они влияют на разные составляющие времени выполнения.

1. Больший размер батча

Большие батчи уменьшают число обновлений за эпоху. Для 48 000 обучающих изображений батч 32 даёт 1500 шагов. Батч 4096 даёт 12 шагов, если включить последний неполный батч. Циклы NumPy и JAX в нашем эксперименте отбрасывают остаток, обрабатывая 11 полных батчей, или 45 056 изображений.

Большой батч также даёт каждому матричному умножению больше работы, что может помочь использовать параллелизм GPU. Число активных SM зависит от выбранного ядра и требует измерения; из одного размера батча его получить нельзя.

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

2. Keras на CPU — обойти накладные расходы CUDA

Отключение GPU через tf.config.set_visible_devices([], 'GPU') до инициализации устройств запускает TensorFlow на CPU. CUDA исчезает из пути выполнения, но это не означает сокращение общего времени на сумму длительностей CUDA API.

3. Чистый NumPy — обойти ещё и накладные расходы фреймворка

Цикл NumPy обходится без среды выполнения TensorFlow и использует BLAS для умножения матриц. При этом сохраняются затраты на вызовы Python, выделение массивов и доступ к памяти.

4. JIT в JAX — скомпилировать весь шаг

jit в JAX компилирует функцию, чтобы повторные вызовы не отправляли каждую операцию через Python. Компиляция может объединять операции и уменьшать расходы, но одна скомпилированная функция не обязательно является одним ядром GPU. Время компиляции следует отделять от замеров после прогрева.

5. CuPy — а что если просто перенести NumPy на GPU?

Мы также перенесли на CuPy цикл NumPy, обрабатывающий изображения по одному. Маленькие векторные операции потребовали многократной отправки работы GPU. Эта версия заняла 443 секунды за пять эпох, более чем в шесть раз дольше соответствующего цикла NumPy. Результат относится к этой реализации; CuPy с обработкой целых батчей был бы другим сравнением.

Собираем всё вместе

Таблица ниже относится к отдельному эксперименту с разными размерами батча. Объём работы различается: вызовы Keras включают валидацию и метрики, а циклы NumPy и JAX измеряют только обучение и отбрасывают неполные батчи. Это времена конкретных скриптов, а не контролируемый рейтинг скорости фреймворков.

Результат CPU при батче 32 здесь равен 11,31 с, а не прежним 4,4 с. Эти отдельные измерения нельзя объединять в одну оценку ускорения. Контролируемое сравнение требует одинаковых данных, обновлений, валидации, прогрева и границ измерения.

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

Keras на CPU сокращает время с 11,31 с до 1,39 с с ростом батча. Здесь измеряется весь вызов fit(), использованный в эксперименте.

NumPy достигает 0,84 с при батче 2048 и занимает 1,03 с при 4096. Одни времена не объясняют причину замедления.

JAX показывает 0,14 с при батче 4096 для цикла, измеряющего только обучение. Отношение 6,26 с к 0,14 с равно примерно 45, но объединяет разные размеры батча и объёмы работы. Оно не доказывает ускорение в 45 раз для эквивалентного обучения или достижения одинаковой валидационной точности.

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

Используем распараллеливание для перебора гиперпараметров

Подходы выше ускоряют один запуск обучения. Но когда вы перебираете гиперпараметры — пробуете 5 скоростей обучения или 4 архитектуры, — каждый запуск полностью независим. Они не делят между собой ни веса, ни градиенты, ни состояние.

vmap в JAX добавляет измерение моделей в функцию обновления, позволяя одному вызову обновить пять независимых наборов параметров на одних изображениях. Вместе с jit это убирает цикл Python по моделям. Компилятор выбирает способ выполнения операций; он не гарантирует одно ядро или одновременное выполнение всех моделей.

В этом примере у всех моделей одинаковые формы параметров. Код показывает векторизованное обновление и предполагает, что x_batch и y_batch уже подготовлены:

import jax
import jax.numpy as jnp
from jax import vmap, jit, grad, random

def init_params(key):
    # Та же модель, что и раньше: 784→128→10, случайные веса
    k1, k2 = random.split(key)
    w1 = random.normal(k1, (784, 128)) * jnp.sqrt(2.0 / 784)
    b1 = jnp.zeros(128)
    w2 = random.normal(k2, (128, 10)) * jnp.sqrt(2.0 / 128)
    b2 = jnp.zeros(10)
    return (w1, b1, w2, b2)

def loss_fn(params, x, y):
    # Прямой проход + перекрёстная энтропия — та же математика, что в версии на NumPy
    w1, b1, w2, b2 = params
    h = jnp.maximum(0, x @ w1 + b1)       # скрытый слой + ReLU
    logits = h @ w2 + b2                    # выходной слой
    log_probs = jax.nn.log_softmax(logits, axis=-1)
    return -jnp.mean(log_probs[jnp.arange(y.shape[0]), y])

def sgd_step(params, x, y, lr):
    # Один шаг обучения: считаем градиенты, обновляем веса
    grads = grad(loss_fn)(params, x, y)     # JAX автоматически дифференцирует loss_fn
    return tuple(p - lr * g for p, g in zip(params, grads))

# 5 скоростей обучения, 5 наборов весов, обучаются одновременно
lr_array = jnp.array([0.001, 0.01, 0.1, 1.0, 10.0])
keys = random.split(random.PRNGKey(42), len(lr_array))
batched_params = vmap(init_params)(keys)
batched_step = jit(vmap(sgd_step, in_axes=(0, None, None, 0)))
batched_params = batched_step(batched_params, x_batch, y_batch, lr_array)

Другой вариант — обучать каждую конфигурацию в отдельном процессе. Этот набросок предполагает доступность Keras и обучающих данных. Запускайте его как скрипт с условием if __name__ == "__main__", а при выборе числа процессов учитывайте память и потоки CPU каждого процесса:

import multiprocessing as mp

def train_one_config(args):
    # Обучаем одну модель с одной скоростью обучения — работает в своём процессе
    lr, X_train, y_train = args
    model = keras.Sequential([
        keras.layers.Dense(128, activation="relu", input_shape=(784,)),
        keras.layers.Dense(10, activation="softmax"),
    ])
    model.compile(optimizer=keras.optimizers.SGD(learning_rate=lr),
                  loss="sparse_categorical_crossentropy")
    history = model.fit(X_train, y_train, epochs=10, batch_size=32,
                        validation_split=0.2, verbose=0)
    return lr, history.history

if __name__ == "__main__":
    with mp.get_context("spawn").Pool(5) as pool:
        results = pool.map(train_one_config,
                           [(lr, X_train, y_train) for lr in [0.001, 0.01, 0.1, 1.0, 10.0]])

Что стоит запомнить

1. Небольшая модель может работать быстрее на CPU. Первое сравнение Keras дало 4,4 с на эпоху на CPU и 6,5 с на GPU; отдельный эксперимент показал другие времена.
2. Пиковые TFLOPS не измеряют загрузку. Нижняя граница времени арифметики не показывает длительность ядра или число активных вычислительных ядер.
3. Измеряйте всю работу. Синхронизация, валидация, отброшенные батчи и компиляция влияют на смысл результата.
4. Сравнивайте эквивалентное обучение. Большие батчи и скомпилированные циклы могут сократить время, но перед заявлением об ускорении сравните качество на валидации.
5. Проверяйте доступные устройства. Одно число параметров не определяет, будет ли CPU быстрее GPU.

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