Накладні витрати GPU: чому наша модель MNIST навчається швидше на CPU

Графічні процесори можуть прискорити навчання нейронної мережі, обробляючи багато обчислень паралельно. Однак маленькі моделі можуть не забезпечити достатньо роботи, щоб компенсувати витрати на диспетчерські операції, переміщення даних і синхронізацію результатів.

Ми порівняли наш Модель MNIST на графічному процесорі Quadro RTX 5000 і процесорі Intel Core i9-10885H. Під час одного запуску Keras ЦП був швидшим: 4,4 с проти 6,5 с за епоху. Окрема розгортка зафіксувала лише 0,14 с для циклу навчання JAX на ЦП, але його робоче навантаження за часом відрізнялося від Keras. Ми розглянемо як джерела накладних витрат, так і те, що ці вимірювання можуть нам сказати.

Почнемо з того, що відбувається під час навчання.

Розбивка одного тренувального кроку

Крок навчання запускає прохід вперед, обчислює градієнт втрат, поширює його назад і оновлює ваги. Цей уривок обчислює комбінований градієнт softmax–крос-ентропії безпосередньо, без обчислення скалярних втрат:

# Forward pass
z1 = xb @ w1 + b1                                # matrix multiply + bias (hidden layer)
a1 = np.maximum(0, z1)                            # ReLU activation
z2 = a1 @ w2 + b2                                # matrix multiply + bias (output layer)
exp_z = np.exp(z2 - z2.max(axis=1, keepdims=True))
probs = exp_z / exp_z.sum(axis=1, keepdims=True)  # softmax → probabilities

# Loss gradient
dz2 = probs.copy()
dz2[np.arange(bs), yb] -= 1                       # how far off were the predictions
dz2 /= bs

# Backward pass — compute gradients
dw2 = a1.T @ dz2                                  # gradient for W2
da1 = dz2 @ w2.T                                  # gradient flowing backward
dz1 = da1 * (z1 > 0)                              # ReLU gradient
dw1 = xb.T @ dz1                                  # gradient for W1

# Update weights
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
# ... and so on for each operation

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

Ось результати для одного кроку навчання (batch_size=32, усереднене за 1000 кроків, час у мікросекундах мкс — мільйонні частки секунди, менше — краще):

ЦП працює швидше в цьому тесті на операцію. NumPy використовує оптимізовані процедури BLAS для множення матриць, тоді як вимірювання GPU TensorFlow також включають диспетчеризацію та синхронізацію. Інструкції SIMD дозволяють ЦП обробляти кілька значень на інструкцію; fused multiply-add поєднує множення та додавання. Це порівняння вимірює реалізацію та метод синхронізації разом, а не ізолює арифметичну швидкість двох пристроїв.

Давайте подивимося на xb @ w1 + b1 і a1 @ w2 + b2 — матриця множиться з кожного шару — і чому графічні процесори розроблені, щоб зробити їх швидкими. Ось передача вперед від нашої реалізації NumPy:

class HiddenLayer:
    def forward(self, x):
        self.z = self.W @ x + self.b        # matrix multiply + bias
        self.out = np.maximum(0, self.z)     # ReLU activation
        return self.out

class OutputLayer:
    def forward(self, x):
        self.z = self.W @ x + self.b        # matrix multiply + bias
        exp = np.exp(self.z - np.max(self.z))
        self.probs = exp / np.sum(exp)       # softmax → probabilities
        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 images, 784 pixels each   784 weights per neuron     32 × 128 values

Пакетування повторно використовує однакові ваги для 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 мкс. Це передбачає максимальну пропускну здатність і виключає доступ до пам’яті та відправлення. Це не виміряний час ядра, і ділення кількості операцій на пік FLOPS дає час, а не відсоток використання GPU.

Куди насправді йде час

Етап навчання графічного процесора включає в себе кілька витрат, окрім арифметичних. Їх відносна важливість залежить від навантаження та режиму виконання:

1. Відправка ядра. Центральний процесор і середовище виконання повинні передати роботу графічному процесору. Для невеликих операцій відправка може займати значну частину загального часу. Операція фреймворку може використовувати декілька ядер, у той час як компілятор може об’єднати кілька операцій в одне ядро; кількість виразів Python не визначає кількість запусків.

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

3. Python і фреймворк працюють. Відправка, розподіл пам’яті, метрики та обробка даних додають витрат. Під час активного виконання Python бере участь в окремих операціях. Виконання та компіляція графів можуть усунути велику частину цієї повторюваної роботи. Компіляція XLA відбувається, коли скомпільована функція спочатку потрібна або повторюється, а не для кожної операції на кожному кроці.

Профілювання допомагає розрізнити ці витрати, за умови, що ми зберігаємо тривалість API хосту, виконання графічного процесора та загальний час, що минув.

Профілювання накладних витрат CUDA

Вимірювання за операціями показують, що накладні витрати мають значення для цього робочого навантаження. Щоб перевірити його, нам також потрібен профіль повного конвеєра навчання.

Щоб з’ясувати це, ми використали профайлер nsys (Nsight Systems) NVIDIA, який перехоплює кожен виклик 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 на хості, а не тривалість ядра GPU. Виклики можуть перекривати роботу графічного процесора або діяльність інших потоків ЦП. Їх суму не можна відняти від 6,5 с, щоб визначити час, що залишився, як накладні витрати Python. Для такої поломки потрібна шкала часу з чітко визначеним діапазоном вимірювань. Nsight Systems розрізняє API, чергу та час ядра.

Записи cuGraphLaunch показують, що цей запуск використовував CUDA Graphs для подання записаної роботи. Повторне відтворення графіка може зменшити витрати на диспетчеризацію, але підрахунок викликів не встановлює один запуск на навчальний пакет і не доводить, що обчислення GPU є незначними.

У попередньому порівнянні Keras епоха ЦП займала 4,4 с проти 6,5 с на GPU. Виконання ЦП дозволяє уникнути викликів CUDA та передачі на дискретний графічний процесор, хоча воно все ще має витрати на структуру, планування та доступ до пам’яті.

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

У цьому тесті вимірювання GPU приблизно в 2,5 рази швидше. Обидві цифри включають витрати на програмне забезпечення та доступ до пам’яті; також не є чистим арифметичним вимірюванням. Вони також використовують інший метод визначення часу, ніж у першій таблиці, тому очевидне розворот не є протиріччям.

Супровідний сценарій скомпільованого кроку виконує різну роботу на двох пристроях: шлях GPU оновлює ваги, а шлях CPU лише обчислює градієнти. Таким чином, ці кроки не можуть підтримувати пряме порівняння CPU–GPU. Справедливе порівняння має виконати те саме обчислення та дочекатися завершення обох пристроїв.

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

Використання графічного процесора в nvidia-smi вимірює частку інтервалу вибірки, протягом якого виконується принаймні одне ядро. Він не вимірює частку зайнятих ядер або досягнутий пік FLOPS. Тому невеликі ядра, що працюють безперервно, можуть давати високі показники використання без використання повної пропускної здатності GPU.

Пришвидшення одного тренування

Ми випробували п’ять способів змінити навантаження на навчання: більші партії, Keras на ЦП, цикл NumPy, компіляція JAX і прямий порт CuPy. Вони впливають на різні частини вартості виконання.

1. Більший розмір партії

Більші партії зменшують кількість оновлень на епоху. З 48 000 навчальними зображеннями розмір партії 32 дає 1500 кроків. Розмір партії 4096 дає 12 кроків, якщо включено останню часткову партію. Натомість цикли NumPy та JAX у нашій розгортці відкидають цей залишок, обробляючи 11 повних пакетів або 45 056 зображень.

Більший пакет також дає більше роботи кожному множенню матриці, що може допомогти використовувати паралелізм GPU. Кількість активних SM залежить від вибраного ядра і має бути виміряна; його неможливо прочитати безпосередньо з розміру партії.

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

2. Keras на процесорі — пропустіть накладні витрати CUDA

Вимкнення GPU за допомогою tf.config.set_visible_devices([], 'GPU') перед ініціалізацією пристрою запускає TensorFlow на CPU. Це видаляє роботу CUDA зі шляху виконання, але це не передбачає збереження витраченого часу, що дорівнює сумі тривалостей CUDA API.

3. Чистий NumPy — також пропускайте накладні витрати на структуру

Навчальний цикл NumPy уникає середовища виконання TensorFlow і використовує BLAS для множення матриць. Він все ще оплачує виклики Python, розподіл масиву та доступ до пам’яті.

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

jit від [JAX] компілює функцію, щоб повторні виклики уникали диспетчеризації кожної операції через Python. Компіляція може об’єднати операції та зменшити накладні витрати, але одна скомпільована функція не обов’язково є одним ядром GPU. Час компіляції слід відокремити від часу стійкого стану.

5. CuPy — що, якщо ми просто перемістимо NumPy на GPU?

Ми також перенесли цикл NumPy для кожного зображення на CuPy. Невеликі векторні операції вимагали повторних відправок GPU. Ця версія зайняла 443 секунди протягом п’яти епох, що в шість разів перевищує час, повідомлений для відповідного циклу NumPy. Цей результат застосовується до цієї реалізації; пакетна реалізація CuPy була б іншим порівнянням.

Все разом

У наведеній нижче таблиці записується окрема розгортка за розмірами партій. Робочі навантаження не ідентичні: виклики Keras включають перевірку та метрику, тоді як NumPy та JAX зациклюються лише на час навчання та видаляють неповні пакети. Це час написання сценаріїв, а не контрольоване ранжування швидкості фреймворку.

Результат процесора розміром партії 32 тут становить 11,31 с, а не попередні 4,4 с. Ці окремі вимірювання не слід об’єднувати в одне твердження про прискорення. Контрольоване порівняння використовуватиме ті самі дані, оновлення, перевірку, розминку та часові межі.

GPU Keras наближається до 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):
    # Same model as before: 784→128→10, random weights
    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):
    # Forward pass + cross-entropy loss — same math as our NumPy version
    w1, b1, w2, b2 = params
    h = jnp.maximum(0, x @ w1 + b1)       # hidden layer + ReLU
    logits = h @ w2 + b2                    # output layer
    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):
    # One training step: compute gradients, update weights
    grads = grad(loss_fn)(params, x, y)     # JAX auto-differentiates loss_fn
    return tuple(p - lr * g for p, g in zip(params, grads))

# 5 learning rates, 5 sets of weights, updated in one vectorized call
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 і навчальні дані доступні. Запустіть його як сценарій із основним захисником і враховуйте використання пам’яті кожним робочим процесом і потоки ЦП під час вибору кількості процесів:

import multiprocessing as mp

def train_one_config(args):
    # Train one model with one learning rate — runs in its own process
    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. Мала модель може працювати швидше на ЦП. Наше початкове порівняння Keras зафіксувало 4,4 с на епоху на ЦП і 6,5 с на ГП; окрема розгортка дала різний час.
2. Пікова кількість TFLOPS не є вимірюванням використання. Лише арифметична нижня межа не може сказати нам фактичну тривалість ядра чи кількість активних ядер.
3. Виміряйте повне робоче навантаження. Синхронізація, перевірка, видалені пакети та компіляція впливають на значення часу.
4. Порівняйте еквівалентне навчання. Більші пакети та скомпільовані цикли можуть скоротити час, що минув, але порівняйте якість перевірки, перш ніж заявляти про прискорення.
5. Порівняйте свої пристрої. Лише кількість параметрів не визначає, чи буде швидшим ЦП чи ГП.

Word2vec пропонує ще один приклад того, чому шаблон доступу має значення. При негативній вибірці оновлення стосується вибраних рядків матриць вбудовування, а не множення великих щільних матриць. Відносна продуктивність реалізацій CPU та GPU залежить від того, як ці пошуки та оновлення групуються; це не універсальна перевага ЦП.