Ao tentar executar um pipeline de treinamento de LoRA com lora-scripts, é comum encontrar o erro CUDA out of memory, mesmo em GPUs com capacidade razoável. Isso geralmente ocorre porque as configurações padrão não estão ajustadas para a memória disponível da sua placa de vídeo. Este guia técnico explica como otimizar o uso de memória GPU em três etapas práticas.

Anatomia do Consumo de Memória GPU

Durante o treinamento, a memória GPU é dividida entre quatro componentes principais:

• Pesos do modelo base: A carga inicial, que pode variar de 2GB a 10GB, dependendo da arquitetura.
• Estados do otimizador: Dados como momentos e variâncias, consumindo até 2-3x o tamanho dos pesos do modelo.
• Ativações intermediárias e gradientes: Resultados temporários durante a propagação forward e backward, fortemente influenciados pelo tamanho do lote e pela resolução.
• Dados do lote atual: As imagens ou textos sendo processados naquele momento.

Quando a soma desses componentes excede a memória disponível, ocorre um erro de falta de memória. O objetivo é reconfigurar o treinamento para que o total se mantenha dentro dos limites da GPU.

Passo 1: Redução Imediata via Tamanho do Lote e Resolução

A alteração mais impactante e rápida é ajustar o tamanho do lote (tamanho_lote) e, se aplicável, a resolução das imagens. Edite o arquivo de configuração (por exemplo, configuracao_treino.yaml):

# Seção de treino
treino:
  tamanho_lote: 2  # Valor padrão pode ser 4 ou 8; reduza para 2 ou 1
  passos_acumulacao_gradiente: 1  # Mantenha como 1 inicialmente

# Seção de dados (para modelos de difusão)
dados:
  resolucao: 512  # Se a memória for limitada, tente 384 (requer reprocessamento das imagens)

Diretrizes práticas:

• Para GPUs com 8GB de VRAM (ex.: RTX 3070), comece com tamanho_lote: 2. Para 12GB (ex.: RTX 3060), tente tamanho_lote: 4.
• Se precisar de um lote efetivo maior com baixo consumo, use passos_acumulacao_gradiente: 4 junto com tamanho_lote: 1. Isso simula um lote de 4, mas ocupa memória como lote de 1, ao custo de tempo de treinamento.
• Reduzir resolucao libera memória significativamente, mas pode exigir reprocessar os dados e afetar a aprendizagem de detalhes finos.

Passo 2: Otimização do Modelo e Configuração de LoRA

Se o consumo ainda for alto, ajuste os parâmetros do modelo e do treino:

# Seção do modelo
modelo:
  # Configuração LoRA
  classificacao_lora: 8  # Rank do LoRA; reduza para 4 para economizar memória
  alfa_lora: 32  # Geralmente 2-4 vezes a classificação_lora

# Seção de treino
treino:
  precisao_mista: "fp16"  # Essencial para economia de memória e aceleração
  # Em GPUs modernas (RTX 30/40), "bf16" pode ser testado

Ações recomendadas:

• Ative a precisão mista (precisao_mista: "fp16") para reduzir o consumo de memória em até 50% sem perda significativa de qualidade.
• Diminua a classificacao_lora (de 8 para 4) para reduzir o número de parâmetros treináveis.
• Para modelos de linguagem grandes (LLMs), use versões quantizadas (ex.: GGUF q4_0) para reduzir drasticamente o tamanho do modelo base.

Passo 3: Técnicas Avançadas e Otimizações de Sistema

Para cenários com memória extremamente limitada (ex.: 4GB), considere:

1. Otimizadores eficientes em memória: Substitua AdamW por alternativas como Adafactor ou 8-bit Adam, que consomem menos memória para estados.
2. Checkpoint de gradiente: Ative checkpoint_gradiente: true na configuração. Isso recalcula ativações durante a retropropagação, reduzindo o uso de memória em ~30%, com aumento de tempo de treino (~20%).
3. Ajustes de sistema: No Linux, limpe o cache do sistema com sudo sh -c 'sync; echo 3 > /proc/sys/vm/drop_caches'. No Windows, aumente o arquivo de memória virtual (1.5-2x da RAM física) em um SSD.

Exemplo de Configuração para GPUs de 6GB/8GB

Para uma RTX 2060 (6GB) ou similar, treinando um LoRA de Stable Diffusion:

# configuracao_baixa_vram.yaml
modelo:
  modelo_base: "./modelos/v1-5-pruned.safetensors"
  classificacao_lora: 4  # Rank baixo para economia
  alfa_lora: 8

dados:
  diretorio_dados: "./dados/meu_estilo"
  resolucao: 512
  tamanho_lote: 1
  passos_acumulacao_gradiente: 4  # Lote efetivo de 4

treino:
  epocas: 10
  taxa_aprendizado: 1e-4
  precisao_mista: "fp16"
  checkpoint_gradiente: true
  # use_otimizador_8bit: true  # Descomente se suportado

sistema:
  semente: 42

Execute com:

python treinar.py --config configuracao_baixa_vram.yaml

Use nvidia-smi para monitorar o consumo de memória em tempo real.

Roteiro de Otimização

Siga esta ordem ao enfrentar erros de memória:

1. Reduza imediatmaente o tamanho_lote.
2. Confirme que precisao_mista: "fp16" está ativada.
3. Ajuste a classificacao_lora e use modelos quantizados.
4. Ative checkpoint_gradiente para cenários críticos.
5. Considere reduzir resolucao ou trocar o otimizador.

Para LoRA, treinar com tamanho_lote: 1 e mais épocas frequentemente produz resultados satisfatórios, priorizando a viabilidade do treinamento sobre grandes lotes.