Uso de memória de GPU

Treinar um modelo de 4B parâmetros em precisão mista com um tamanho de lote de 16 requer aproximadamente 85GB de memória de GPU. Compreender o que ocupa a memória da GPU e as operações de computação que ocorrem durante o treinamento ajuda a identificar onde você pode reduzir o uso de memória.

┌─────────────────────────── 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    │
└──────────────────────────────────────────────────────────────────────┘

A memória da GPU armazena duas categorias de elementos: os tensores armazenados e as operações que os processam.

Tensores

O treinamento requer e produz muitos tipos de tensor que precisam ser armazenados.

• Os pesos do modelo são armazenados na GPU. No treinamento em precisão mista, são necessárias duas cópias dos pesos: uma em fp16 para a passagem direta/reversa e outra em fp32 como a “cópia principal” para atualizações estáveis dos pesos. Isso equivale a 6 bytes por parâmetro.

• Estados de otimizador, como o Adam, armazenam dois tensores extras por parâmetro, o momento e a variância, ambos em fp32. Isso representa um adicional de 8 bytes por parâmetro.

  Você pode usar um otimizador diferente, como uma versão quantizada do Adam do bitsandbytes, para comprimi-lo para 2 bytes por parâmetro.

• Os tensores de gradiente são calculados para cada parâmetro na passagem reversa. Isso é mantido em fp32, ou seja, 4 bytes por parâmetro.

• As ativações diretas são calculadas na passagem direta e armazenadas em cache para que a passagem reversa possa calcular os gradientes. O tamanho dessas ativações varia conforme o tamanho do lote, o comprimento da sequência, a profundidade do modelo e o tamanho oculto. É por isso que o tamanho do lote ou o comprimento da sequência podem esgotar a memória da GPU, mesmo que o próprio modelo caiba nela.

• Tensores temporários são criados por operações como softmax e multiplicações de matrizes e liberados após cada operação. Se o pico de uma única operação for muito intenso, ele pode criar um pico temporário que faz com que você fique sem memória.

• Algumas funcionalidades adicionam seu próprio overhead. Por exemplo, o beam search mantém múltiplas saídas, e as tabelas de embeddings para vocabulários grandes podem ser muito grandes.

Ops

Existem três tipos principais de operações de treinamento.

• As multiplicações de matrizes (matmuls) são o principal tipo de operação, abrangendo camadas lineares, projeções QKV, projeções de saída da atenção e camadas FFN. Os principais consumidores de memória são as matrizes de pontuação de atenção, que crescem com o quadrado do comprimento da sequência. É por isso que sequências longas são tão caras e existem diferentes backends de atenção para evitar a materialização da matriz completa na memória.

• Operações de redução, como softmax e layer norm, leem um tensor completo, calculam uma estatística ao longo de uma dimensão e, em seguida, leem o tensor novamente para aplicá-la. Isso requer acessar a memória várias vezes por operação.

• Operações elemento a elemento, como ativações e dropout, aplicam uma função a cada elemento de forma independente e seu uso de memória é proporcional ao tamanho do tensor.

Próximas etapas

• Consulte o guia Acumulação de gradientes para aprender a simular o treinamento com um tamanho de lote efetivo maior sem esgotar a memória da GPU.
• Consulte o guia Gradient checkpointing para aprender a reduzir o uso de memória armazenando apenas algumas ativações intermediárias.
• Consulte o guia Mixed precision training para aprender a usar tipos de dados de menor precisão para reduzir a memória e acelerar o treinamento.
• Consulte o guia Kernels para aprender a acelerar o treinamento com kernels fundidos personalizados.