Reducción del uso de memoria

Los flujos de trabajo de entrenamiento a menudo se pueden optimizar para reducir el consumo de memoria, y TRL ofrece varias funciones integradas para ayudar a lograrlo.

A continuación, describimos estas técnicas y recomendamos experimentar con diferentes combinaciones para determinar qué configuración funciona mejor para su configuración específica.

Cada método incluye ejemplos para los entrenadores compatibles. Si no está seguro de si una técnica es compatible con su entrenador, consulte la documentación correspondiente del entrenador.

Para conocer estrategias adicionales, como gradient checkpointing, que es compatible con todos los trainers, consulte la guía de rendimiento de transformers.

Truncamiento

Las longitudes de las secuencias en el conjunto de datos pueden variar ampliamente. Cuando los datos se agrupan en lotes, las secuencias se rellenan para igualar la más larga del lote, lo que puede provocar un alto uso de memoria, incluso si la mayoría de las secuencias son relativamente cortas.

Para reducir el uso de memoria, es importante truncar las secuencias a una longitud razonable. Si bien los entrenadores de TRL truncan las secuencias de forma predeterminada, es posible que desee ajustar la longitud de truncación predeterminada para adaptarla mejor a su caso de uso específico.

from trl import DPOConfig

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

from trl import SFTConfig

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

Cómo elegir el valor de max_length?

Si max_length es demasiado pequeño, una parte significativa de tus tokens se descartará y no contribuirá al entrenamiento. Si es demasiado grande, el uso de memoria puede aumentar drásticamente, lo que podría provocar errores de memoria insuficiente (OOM). Sin packing o padding-free, un max_length grande también puede resultar en un entrenamiento ineficiente, ya que muchos tokens serán de relleno.

Para ayudarte a elegir un valor apropiado, proporcionamos una utilidad para visualizar la distribución de la longitud de la secuencia en tu conjunto de datos.

Empaquetado

Esta técnica solo está disponible para el entrenamiento de SFT y configuraciones que utilizan FlashAttention (o sus variantes).

Truncamiento tiene varias desventajas:

1. Pérdida de información: Es posible que se descarten tokens importantes al final de las secuencias.
2. Elegir la longitud de truncamiento: Si es demasiado corta, se pierde información; si es demasiado larga, se reduce la eficiencia.

Packing mitiga estos problemas al agrupar varias secuencias en la misma fila de entrenamiento, llenando cada fila hasta max_length.

TRL implementa el empaquetado mediante el empaquetado de contenedores Best-Fit Decreasing (BFD), que agrupa las secuencias de manera eficiente y minimiza el relleno. Cuando una secuencia supera max_length, distintas estrategias determinan cómo se manejan los tokens que exceden el límite.

TRL admite tres estrategias:

• "bfd" (predeterminado): utiliza Best-Fit Decreasing packing. Si una secuencia supera max_length, los tokens excedentes se descartan.

• "bfd_split": Utiliza Best-Fit Decreasing packing, pero las secuencias largas se dividen en fragmentos ≤ max_length antes del empaquetado. Esto preserva todos los tokens y sigue el enfoque propuesto en Fewer Truncations Improve Language Modeling.

• "wrapped": Todos los tokens se concatenan en un flujo y se dividen en bloques de longitud fija. Esto minimiza el relleno, pero puede mezclar ejemplos no relacionados. Esta estrategia corresponde al preprocesamiento concatenate-then-split descrito en la literatura (p. ej., Fewer Truncations Improve Language Modeling). Tiene el inconveniente de romper la continuidad de la secuencia para una gran parte del conjunto de datos, lo que perjudica el rendimiento, como se discutió en el Qwen3-Coder-Next Technical Report.

Si todas las secuencias son más cortas que max_length, bfd y bfd_split se comportan de manera idéntica, ya que no se requiere truncamiento ni división.

from trl import SFTConfig

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

PEFT para el ajuste fino eficiente de parámetros

Los métodos de ajuste fino eficiente en parámetros (PEFT), como LoRA, son algunas de las técnicas más eficaces para reducir el uso de memoria durante el entrenamiento. En lugar de entrenar todos los parámetros del modelo, los métodos PEFT entrenan solo un pequeño número de parámetros de adaptador, lo que reduce significativamente los requisitos de memoria y permite el ajuste fino de modelos más grandes en hardware limitado.

Para obtener detalles completos sobre el uso de PEFT con TRL, incluidos diversos métodos de adaptadores, opciones de cuantización y configuraciones avanzadas, consulte PEFT Integration.

Para usar PEFT para reducir el uso de memoria:

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,
)

PEFT puede combinarse con otras técnicas de reducción de memoria, como la cuantización (de 4 bits o 8 bits), para lograr un mayor ahorro de memoria. Consulte PEFT Integration para ver ejemplos de cuantización.

Sin relleno

El loteo sin relleno es un enfoque alternativo para reducir el uso de memoria. En este método, primero se muestrea un lote y luego se aplana en una única secuencia, evitando el relleno. A diferencia del empaquetado, que puede dar lugar a secuencias incompletas al combinar partes de diferentes muestras, el loteo sin relleno garantiza que todas las secuencias permanezcan 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"})

Descarga de activaciones

El descargado de activaciones (activation offloading) es una técnica de eficiencia de memoria que reduce el uso de VRAM de la GPU moviendo temporalmente los tensores de activación a la RAM del CPU durante la pasada hacia adelante (forward pass) y devolviéndolos solo cuando se necesitan para la pasada hacia atrás (backward pass). Esto reduce significativamente el uso de memoria pico a costa de un ligero aumento en el tiempo de entrenamiento.

Para habilitar el descargado de activaciones en tu configuración de entrenamiento SFT:

from trl import SFTConfig

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

Por debajo, el descargado de activaciones implementa saved_tensors_hooks de PyTorch para interceptar las activaciones durante la pasada hacia adelante. Gestiona de forma inteligente qué tensores descargar en función de su tamaño y contexto, evitando descargar los tensores de salida que serían ineficientes. Para optimizar el rendimiento, puede, mediante una bandera (que es verdadera por defecto), usar flujos de CUDA para superponer la computación con las transferencias CPU-GPU.

Relleno de secuencias hasta un múltiplo

Esta técnica se admite actualmente para los entrenadores SFT y Reward.

Cuando está habilitada, esta opción garantiza que todas las secuencias se rellenen hasta un múltiplo del valor especificado.
Esto puede mejorar la eficiencia computacional en algunos hardware al alinear las longitudes de las secuencias con límites de memoria favorables.

from trl import SFTConfig

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

from trl import RewardConfig

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

Asignador de memoria con caché de PyTorch

En ejecuciones largas, especialmente en RL en línea (GRPO, RLOO, Online DPO), la memoria de la GPU puede fragmentarse. Establecer expandable_segments:True permite que el asignador de caché de PyTorch amplíe los segmentos existentes, lo que reduce la brecha entre la memoria asignada y la reservada:

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

Con el modo de suspensión de vLLM (vllm_enable_sleep_mode=True), vLLM solo lee PYTORCH_CUDA_ALLOC_CONF para evitar este ajuste, por lo que no establezca únicamente PYTORCH_ALLOC_CONF.

Desactivar la recopilación del modelo para la generación en métodos en línea

Al usar DeepSpeed ZeRO-3, los pesos del modelo se fragmentan en varias GPU. Los métodos en línea implican generar completados desde el modelo como parte del proceso de entrenamiento. Durante esta etapa, los pesos del modelo se recopilan temporalmente en una sola GPU para la generación. En el caso de modelos muy grandes, esta recopilación puede provocar errores de OOM, tal como se describe en este problema: #2250.

Si te encuentras con este problema, puedes desactivar la recopilación de los pesos del modelo para la generación estableciendo el siguiente 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)

Este ajuste impide que se recojan los pesos del modelo, evitando errores de OOM, pero puede resultar en velocidades de generación más lentas.

Modo de suspensión de vLLM

Al usar vLLM como backend de generación para los métodos de entrenamiento en línea, se puede habilitar el sleep mode para descargar los parámetros y la caché de vLLM a la RAM de la CPU durante el paso de optimización y recargarlos en la VRAM de la GPU cuando se necesiten para la sincronización de pesos y la generación.

from trl import GRPOConfig

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

from trl import RLOOConfig

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

Descargar los pesos y la caché de vLLM ayuda a mantener bajo el uso de memoria de la GPU, lo cual puede ser especialmente beneficioso al entrenar modelos grandes o al usar recursos de GPU limitados. Sin embargo, despertar el motor de vLLM del modo de suspensión introduce cierta latencia de transferencia entre el host y el dispositivo, lo que puede afectar ligeramente la velocidad de entrenamiento.

Punto de control de gradientes

El *gradient checkpointing* intercambia cómputo por memoria al no almacenar todas las activaciones intermedias durante la pasada hacia adelante, recalcándolas en su lugar durante la pasada hacia atrás.

from trl import SFTConfig

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

El checkpointing de gradientes está habilitado de forma predeterminada en todos los trainers para optimizar el uso de memoria. Puedes desactivarlo estableciendo gradient_checkpointing=False si es necesario.

Puntos de control de activación selectivos

Con SFTTrainer, puede guardar la salida de la atención durante la pasada hacia adelante en lugar de recalcularla en la pasada hacia atrás. Esto recupera la mayor parte de la ralentización del checkpointing en contextos largos, a cambio de un tensor adicional del tamaño del estado oculto por capa. Fuerza el checkpointing no reentrante.

from trl import SFTConfig

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

Para conocer más técnicas de optimización de memoria, consulte la Guía de rendimiento de Transformers.