Reduzindo o uso de memória

Os fluxos de trabalho de treinamento podem ser frequentemente otimizados para reduzir o consumo de memória, e o TRL oferece vários recursos integrados para ajudar a alcançar isso.

Abaixo, descrevemos essas técnicas e recomendamos experimentar diferentes combinações para descobrir qual configuração funciona melhor para o seu ambiente específico.

Cada método inclui exemplos para os treinadores suportados. Se você não tiver certeza de que uma técnica é compatível com o seu treinador, consulte a documentação correspondente do treinador.

Para estratégias adicionais, como gradient checkpointing, que é suportado por todos os trainers, consulte o guia de desempenho do transformers.

Truncamento

Os comprimentos das sequências no conjunto de dados podem variar amplamente. Quando os dados são agrupados em lotes, as sequências são preenchidas para corresponder à mais longa do lote, o que pode causar um alto uso de memória, mesmo que a maioria das sequências seja relativamente curta.

Para reduzir o uso de memória, é importante truncar as sequências para um comprimento razoável. Embora os treinadores do TRL truncuem sequências por padrão, você pode querer ajustar o comprimento de truncamento padrão para se alinhar melhor ao seu caso de uso específico.

from trl import DPOConfig

training_args = DPOConfig(..., max_length=...)

from trl import SFTConfig

training_args = SFTConfig(..., max_length=...)

Como escolher o valor de max_length?

Se max_length for muito pequeno, uma parte significativa dos seus tokens será descartada e não contribuirá para o treinamento. Se for muito grande, o uso de memória pode aumentar drasticamente, o que pode levar a erros de memória insuficiente (OOM). Sem packing ou padding-free, um max_length grande também pode resultar em treinamento ineficiente, pois muitos tokens serão de padding.

Para ajudá-lo a escolher um valor apropriado, fornecemos uma utilidade para visualizar a distribuição do comprimento da sequência no seu conjunto de dados.

Pacote

Essa técnica está disponível apenas para treinamento SFT e configurações que usam FlashAttention (ou suas variantes).

Truncamento tem várias desvantagens:

1. Perda de informações: Tokens importantes no final das sequências podem ser descartados.
2. Escolhendo o comprimento de truncamento: muito curto perde dados; muito longo reduz a eficiência.

O Packing mitiga esses problemas ao agrupar várias sequências na mesma linha de treinamento, preenchendo cada linha até max_length.

O TRL implementa o packing usando empacotamento de caixas Best-Fit Decreasing (BFD), que agrupa sequências de forma eficiente enquanto minimiza o padding. Quando uma sequência excede max_length, diferentes estratégias determinam como os tokens excedentes são tratados.

O TRL suporta três estratégias:

• "bfd" (padrão): usa Best-Fit Decreasing packing. Se uma sequência exceder max_length, os tokens excedentes são descartados.

• "bfd_split": Usa Best-Fit Decreasing packing, mas sequências longas são divididas em blocos ≤ max_length antes do empacotamento. Isso preserva todos os tokens e segue a abordagem proposta em Fewer Truncations Improve Language Modeling.

• "wrapped": Todos os tokens são concatenados em um fluxo e divididos em blocos de comprimento fixo. Isso minimiza o preenchimento, mas pode misturar exemplos não relacionados. Essa estratégia corresponde ao pré-processamento concatenate-then-split descrito na literatura (por exemplo, Fewer Truncations Improve Language Modeling). A desvantagem é que ela quebra a continuidade da sequência para uma grande fração do conjunto de dados, o que prejudica o desempenho, conforme discutido no Qwen3-Coder-Next Technical Report.

Se todas as sequências forem mais curtas que max_length, bfd e bfd_split se comportam de forma idêntica, já que não é necessário truncamento ou divisão.

from trl import SFTConfig

training_args = SFTConfig(
    ...,
    packing=True,
    packing_strategy="bfd",
    max_length=512,
)

PEFT para ajuste fino eficiente de parâmetros

Métodos de ajuste fino eficiente em parâmetros (PEFT), como LoRA, estão entre as técnicas mais eficazes para reduzir o uso de memória durante o treinamento. Em vez de treinar todos os parâmetros do modelo, os métodos PEFT treinam apenas um pequeno número de parâmetros de adaptador, reduzindo significativamente os requisitos de memória e permitindo o ajuste fino de modelos maiores em hardware limitado.

Para detalhes completos sobre o uso do PEFT com o TRL, incluindo vários métodos de adaptador, opções de quantização e configurações avançadas, consulte Integração do PEFT.

Para usar o PEFT para reduzir o uso de memória:

from datasets import load_dataset
from peft import LoraConfig
from trl import SFTTrainer

dataset = load_dataset("trl-lib/Capybara", split="train")

peft_config = LoraConfig()

trainer = SFTTrainer(
    model="Qwen/Qwen2.5-0.5B",
    train_dataset=dataset,
    peft_config=peft_config,
)

O PEFT pode ser combinado com outras técnicas de redução de memória, como quantização (4 bits ou 8 bits), para obter ainda mais economia de memória. Consulte Integração do PEFT para exemplos de quantização.

Sem padding

O loteamento sem preenchimento (padding-free batching) é uma abordagem alternativa para reduzir o uso de memória. Neste método, um lote é primeiro amostrado e, em seguida, achatado em uma única sequência, evitando o preenchimento. Diferentemente do empacotamento (packing), que pode resultar em sequências incompletas ao combinar partes de diferentes amostras, o loteamento sem preenchimento garante que todas as sequências permaneçam completas e intactas.

from trl import DPOConfig

training_args = DPOConfig(..., padding_free=True, model_init_kwargs={"attn_implementation": "kernels-community/flash-attn2"})

from trl import SFTConfig

training_args = SFTConfig(..., padding_free=True, model_init_kwargs={"attn_implementation": "kernels-community/flash-attn2"})

Offloading de ativação

O offloading de ativações é uma técnica de eficiência de memória que reduz o uso de VRAM da GPU movendo temporariamente os tensores de ativação para a RAM do CPU durante a passagem direta e trazendo-os de volta apenas quando necessário para a passagem reversa. Isso reduz significativamente o uso de memória de pico, com um leve aumento no tempo de treinamento.

Para habilitar o offloading de ativações na sua configuração de treinamento SFT:

from trl import SFTConfig

training_args = SFTConfig(..., activation_offloading=True)

Por baixo dos panos, o offloading de ativações implementa o saved_tensors_hooks do PyTorch para interceptar as ativações durante a passagem de frente. Ele gerencia de forma inteligente quais tensores devem ser offloadados com base no tamanho e no contexto, evitando o offloading de tensores de saída que seriam ineficientes. Para otimização de desempenho, ele pode, por meio de uma flag (que é verdadeira por padrão), usar streams CUDA para sobrepor o cálculo com as transferências CPU-GPU.

Preenchimento de sequências para um múltiplo

Esta técnica é atualmente suportada para os treinadores SFT e Reward.

Quando habilitada, esta opção garante que todas as sequências sejam preenchidas até um múltiplo do valor especificado.
Isso pode melhorar a eficiência computacional em alguns hardwares, alinhando os comprimentos das sequências a fronteiras de memória favoráveis.

from trl import SFTConfig

training_args = SFTConfig(..., pad_to_multiple_of=2048)

from trl import RewardConfig

training_args = RewardConfig(..., pad_to_multiple_of=2048)

Alocador de cache do PyTorch

Em execuções longas, especialmente em RL online (GRPO, RLOO, Online DPO), a memória da GPU pode se fragmentar. Definir expandable_segments:True permite que o alocador de cache do PyTorch expanda segmentos existentes, o que reduz a diferença entre a memória alocada e a reservada:

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export PYTORCH_ALLOC_CONF=expandable_segments:True  # canonical name since PyTorch 2.10

Com o modo de suspensão do vLLM (vllm_enable_sleep_mode=True), o vLLM apenas lê PYTORCH_CUDA_ALLOC_CONF para contornar essa configuração, portanto, não defina apenas PYTORCH_ALLOC_CONF.

Desativando a coleta de modelos para geração em métodos online

Ao usar o DeepSpeed ZeRO-3, os pesos do modelo são divididos em várias GPUs. Os métodos online envolvem a geração de conclusões pelo modelo como parte do processo de treinamento. Durante essa etapa, os pesos do modelo são temporariamente reunidos em uma única GPU para a geração. Para modelos muito grandes, essa reunião pode levar a erros de OOM, conforme descrito nesta issue: #2250.

Se você encontrar esse problema, pode desabilitar a coleta dos pesos do modelo para geração definindo o seguinte parâmetro:

from trl import GRPOConfig

training_args = GRPOConfig(..., ds3_gather_for_generation=False)

from trl.experimental.online_dpo import OnlineDPOConfig

training_args = OnlineDPOConfig(..., ds3_gather_for_generation=False)

from trl import RLOOConfig

training_args = RLOOConfig(..., ds3_gather_for_generation=False)

Esse ajuste impede que os pesos do modelo sejam reunidos, evitando erros de OOM, mas pode resultar em velocidades de geração mais lentas.

Modo de suspensão do vLLM

Ao usar vLLM como backend de geração para métodos de treinamento online, você pode habilitar o sleep mode para descarregar os parâmetros e o cache do vLLM para a RAM do CPU durante a etapa de otimização e recarregá-los para a VRAM do GPU quando necessário para a sincronização de pesos e a geração.

from trl import GRPOConfig

training_args = GRPOConfig(..., vllm_enable_sleep_mode=True)

from trl import RLOOConfig

training_args = RLOOConfig(..., vllm_enable_sleep_mode=True)

O offloading dos pesos e do cache do vLLM ajuda a manter o uso de memória da GPU baixo, o que pode ser particularmente benéfico ao treinar modelos grandes ou ao usar recursos de GPU limitados. No entanto, acordar o mecanismo do vLLM do modo de espera introduz alguma latência de transferência entre host e dispositivo, o que pode impactar ligeiramente a velocidade de treinamento.

Checkpointing de gradiente

O gradient checkpointing troca computação por memória, pois não armazena todas as ativações intermediárias durante a passagem direta, recalculando-as durante a passagem reversa.

from trl import SFTConfig

training_args = SFTConfig(..., gradient_checkpointing=True)

O checkpointing de gradiente é ativado por padrão em todos os trainers para otimizar o uso de memória. Você pode desativá-lo definindo gradient_checkpointing=False, se necessário.

Checkpointing seletivo de ativação

Com o SFTTrainer, você pode salvar a saída de atenção durante a passagem direta, em vez de recalculá-la na passagem reversa. Isso recupera a maior parte da lentidão do checkpointing em contextos longos, com um tensor extra do tamanho do estado oculto por camada. Isso força o checkpointing não reentrante.

from trl import SFTConfig

training_args = SFTConfig(
    ...,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"selective": True},
)

Para mais técnicas de otimização de memória, consulte o Transformers Performance Guide.