Использование памяти GPU

Обучение модели с 4 млрд параметров в смешанной точности при размере батча 16 требует примерно 85 ГБ памяти GPU. Понимание того, что занимает память GPU, и вычислительных операций, выполняемых во время обучения, помогает определить, где можно сократить использование памяти.

┌─────────────────────────── TENSORS ──────────────────────────────────┐
│                                                                      │
│  MODEL WEIGHTS       ████████████████████████  6 bytes/param         │
│  (mixed precision)   ├── fp16 copy (2B) — forward/backward pass      │
│                      └── fp32 copy (4B) — stable weight updates      │
│                                                                      │
│  OPTIMIZER STATES    ████████████████████████████████  8 bytes/param │
│  (Adam)              ├── fp32 momentum  (4B)                         │
│                      └── fp32 variance  (4B)                         │
│                      ↳ quantized Adam (bitsandbytes) → 2 bytes/param │
│                                                                      │
│  GRADIENTS           ████████████████  4 bytes/param  (fp32)         │
│                      └── computed in backward pass, per parameter    │
│                                                                      │
│  ACTIVATIONS         ████ varies — batch × seq_len × depth × hidden  │
│  (forward cache)     └── cached for backward; can OOM even if model  │
│                          fits (long seqs / large batches)            │
│                          bf16/fp16 if using mixed precision          │
│                                                                      │
│  TEMPORARY TENSORS   ▓ short-lived (softmax, matmul scratch)         │
│                      └── peak spikes can cause OOM                   │
│                                                                      │
│  OTHER OVERHEAD      ▒ beam search caches, large embedding tables    │
└──────────────────────────────────────────────────────────────────────┘

Память GPU хранит два типа данных: сохранённые тензоры и операции, которые их обрабатывают.

Тензоры

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

• Веса модели хранятся на GPU. При обучении с смешанной точностью требуются две копии весов — одна в fp16 для прямого/обратного прохода и одна в fp32 в качестве «основной копии» для стабильного обновления весов. Это эквивалентно 6 байтам на параметр.

• Состояния оптимизатора, такие как Adam, хранят два дополнительных тензора на каждый параметр — импульс и дисперсию, оба в fp32. Это дополнительные 8 байт на параметр.

  Вы можете использовать другой оптимизатор, например квантованную версию Adam из bitsandbytes, чтобы сжать его до 2 байт на параметр.

• Тензоры градиентов вычисляются для каждого параметра на этапе обратного распространения. Они хранятся в формате fp32, то есть занимают 4 байта на параметр.

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

• Временные тензоры создаются операциями, такими как softmax и умножение матриц, и освобождаются после каждой операции. Если пиковая нагрузка отдельной операции очень высока, это может вызвать кратковременный всплеск потребления памяти, из-за которого вы исчерпаете доступную память.

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

Ops

Существует три основных типа операций обучения.

• Умножения матриц (matmul) являются основным типом операций и охватывают линейные слои, проекции QKV, проекции выходных значений внимания и слои FFN. Основными потребителями памяти являются матрицы оценок внимания, размер которых растёт пропорционально квадрату длины последовательности. Именно поэтому длинные последовательности обходятся так дорого, и для того, чтобы избежать размещения полной матрицы в памяти, существуют различные бэкенды внимания.

• Операции сведения, такие как softmax и layer norm, читают полный тензор, вычисляют статистику вдоль измерения, а затем снова читают тензор, чтобы применить её. Это требует многократного обращения к памяти для каждой операции.

• Операции поэлементного применения, такие как функции активации и dropout, применяют функцию к каждому элементу независимо, а их потребление памяти пропорционально размеру тензора.

Дальнейшие шаги

• См. руководство Gradient accumulation, чтобы узнать, как имитировать обучение с большим эффективным размером батча, не исчерпывая память GPU.
• См. руководство Gradient checkpointing, чтобы узнать, как снизить потребление памяти, сохраняя только часть промежуточных активаций.
• См. руководство Обучение с смешанной точностью, чтобы узнать, как использовать типы данных с пониженной точностью для сокращения потребления памяти и ускорения обучения.
• См. руководство Kernels, чтобы узнать, как ускорить обучение с помощью пользовательских слиянных ядер.