Réduire l'utilisation de la mémoire

Les flux de travail d’entraînement peuvent souvent être optimisés afin de réduire la consommation de mémoire, et TRL fournit plusieurs fonctionnalités intégrées pour y parvenir.

Ci-dessous, nous décrivons ces techniques et recommandons d’expérimenter avec différentes combinaisons afin de déterminer quelle configuration fonctionne le mieux pour votre configuration spécifique.

Chaque méthode inclut des exemples pour les formateurs pris en charge. Si vous n'êtes pas sûr qu'une technique soit compatible avec votre formateur, veuillez consulter la documentation du formateur correspondant.

Pour d'autres stratégies, telles que le activation checkpointing, pris en charge par tous les formateurs, consultez le guide d'optimisation de transformers.

Troncature

Les longueurs des séquences dans le jeu de données peuvent varier considérablement. Lorsque les données sont regroupées par lots, les séquences sont remplies (padding) pour correspondre à la plus longue du lot, ce qui peut entraîner une forte consommation de mémoire, même si la plupart des séquences sont relativement courtes.

Pour réduire l’utilisation de la mémoire, il est important de tronquer les séquences à une longueur raisonnable. Bien que les formateurs TRL tronquent les séquences par défaut, vous pouvez souhaiter ajuster la longueur de troncature par défaut pour mieux l’adapter à votre cas d’usage spécifique.

from trl import DPOConfig

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

from trl import SFTConfig

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

Comment choisir la valeur de max_length ?

Si max_length est trop faible, une part importante de vos jetons sera écartée et ne contribuera pas à l’entraînement. S’il est trop élevé, l’utilisation de la mémoire peut augmenter brutalement, ce qui peut entraîner des erreurs de mémoire insuffisante (OOM). Sans regroupement ni absence de bourrage, une valeur élevée de max_length peut également résulter en un entraînement inefficace, car de nombreux jetons seront du bourrage.

Pour vous aider à choisir une valeur appropriée, nous mettons à disposition un utilitaire permettant de visualiser la distribution des longueurs de séquences dans votre jeu de données.

Packing

Cette technique n'est disponible que pour l'entraînement SFT et les configurations utilisant FlashAttention (ou ses variantes).

La troncature présente plusieurs inconvénients :

1. Perte d'informations : Les jetons importants situés à la fin des séquences peuvent être supprimés.
2. Choix de la longueur de troncature : trop court, on perd des données ; trop long, on réduit l'efficacité.

Le packing atténue ces problèmes en regroupant plusieurs séquences dans la même ligne d'entraînement, en remplissant chaque ligne jusqu'à max_length.

TRL implémente le packing à l'aide du bin packing Best-Fit Decreasing (BFD), qui regroupe efficacement les séquences tout en minimisant le padding. Lorsqu'une séquence dépasse max_length, différentes stratégies déterminent la manière dont les jetons en dépassement sont gérés.

TRL prend en charge trois stratégies :

• "bfd" (par défaut) : utilise Best-Fit Decreasing packing. Si une séquence dépasse max_length, les jetons en dépassement sont supprimés.

• "bfd_split" : Utilise le Best-Fit Decreasing packing, mais les séquences longues sont découpées en morceaux ≤ max_length avant le packing. Cela préserve tous les jetons et suit l'approche proposée dans Fewer Truncations Improve Language Modeling.

• "wrapped" : tous les jetons sont concaténés en un flux puis découpés en blocs de longueur fixe. Cette approche minimise le bourrage, mais peut mélanger des exemples sans rapport. Cette stratégie correspond au prétraitement concatenate-then-split décrit dans la littérature (par exemple, Fewer Truncations Improve Language Modeling). Son inconvénient est de briser la continuité des séquences pour une grande partie du jeu de données, ce qui nuit aux performances, comme discuté dans le Qwen3-Coder-Next Technical Report.

Si toutes les séquences sont plus courtes que max_length, bfd et bfd_split se comportent de manière identique, car aucun tronquage ni découpage n'est nécessaire.

from trl import SFTConfig

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

PEFT pour l'ajustement fin efficace en paramètres

Les méthodes de fine-tuning à faible coût paramétrique (PEFT), comme LoRA, comptent parmi les techniques les plus efficaces pour réduire l’utilisation de la mémoire lors de l’entraînement. Au lieu d’entraîner tous les paramètres du modèle, les méthodes PEFT n’entraînent qu’un petit nombre de paramètres d’adaptateurs, réduisant ainsi considérablement les besoins en mémoire et permettant le fine-tuning de modèles plus volumineux sur du matériel limité.

Pour des détails complets sur l'utilisation de PEFT avec TRL, y compris les différentes méthodes d'adaptateurs, les options de quantification et les configurations avancées, consultez Intégration PEFT.

Pour utiliser PEFT afin de réduire l'utilisation de la mémoire :

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 peut être combiné avec d'autres techniques de réduction de la mémoire, telles que la quantification (4 bits ou 8 bits), pour obtenir des économies de mémoire encore plus importantes. Consultez Intégration PEFT pour des exemples de quantification.

Sans bourrage

Le lotissement sans bourrage est une approche alternative pour réduire l'utilisation de la mémoire. Dans cette méthode, un lot est d'abord échantillonné puis aplati en une seule séquence, évitant ainsi le bourrage. Contrairement au regroupement, qui peut entraîner des séquences incomplètes en combinant des parties de différents échantillons, le lotissement sans bourrage garantit que toutes les séquences restent complètes et intactes.

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"})

Déchargement des activations

L'activation offloading est une technique d'efficacité mémoire qui réduit l'utilisation de la VRAM du GPU en déplaçant temporairement les tenseurs d'activation vers la RAM du CPU pendant le passage avant, puis en les ramenant uniquement lorsqu'ils sont nécessaires pour le passage arrière. Cela réduit considérablement l'utilisation mémoire de pointe au prix d'une légère augmentation du temps d'entraînement.

Pour activer le déchargement des activations dans votre configuration d'entraînement SFT :

from trl import SFTConfig

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

Sous le capot, l’offloading des activations met en œuvre saved_tensors_hooks de PyTorch pour intercepter les activations lors du passage avant. Il gère intelligemment quels tenseurs offloader en fonction de leur taille et de leur contexte, en évitant d’offloader les tenseurs de sortie qui seraient inefficaces. Pour optimiser les performances, il peut, via un drapeau (vrai par défaut), utiliser des flux CUDA pour superposer le calcul aux transferts CPU-GPU.

Remplissage des séquences jusqu'à un multiple

Cette technique est actuellement prise en charge pour les entraîneurs SFT et Reward.

Lorsqu'elle est activée, cette option garantit que toutes les séquences sont remplies jusqu'à un multiple de la valeur spécifiée.
Cela peut améliorer l'efficacité computationnelle sur certains matériels en alignant les longueurs des séquences sur des limites de mémoire optimales.

from trl import SFTConfig

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

from trl import RewardConfig

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

PyTorch allocateur de cache

Sur les longues exécutions, en particulier le RL en ligne (GRPO, RLOO, Online DPO), la mémoire GPU peut se fragmenter. Définir expandable_segments:True permet à l’allocateur de cache de PyTorch d’agrandir les segments existants, ce qui réduit l’écart entre la mémoire allouée et la mémoire réservée :

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

Avec le mode veille de vLLM (vllm_enable_sleep_mode=True), vLLM ne lit que PYTORCH_CUDA_ALLOC_CONF pour contourner ce paramètre, il ne faut donc pas définir uniquement PYTORCH_ALLOC_CONF.

Désactiver la collecte du modèle pour la génération dans les méthodes en ligne

Lors de l'utilisation de DeepSpeed ZeRO-3, les poids du modèle sont répartis sur plusieurs GPU. Les méthodes en ligne impliquent la génération de complétions par le modèle dans le cadre du processus d'entraînement. Lors de cette étape, les poids du modèle sont temporairement regroupés sur un seul GPU pour la génération. Pour les modèles très volumineux, ce regroupement peut entraîner des erreurs OOM, comme décrit dans ce problème : #2250.

Si vous rencontrez ce problème, vous pouvez désactiver la collecte des poids du modèle pour la génération en définissant le paramètre suivant :

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)

Cet ajustement empêche la collecte des poids du modèle, évitant ainsi les erreurs OOM, mais il peut entraîner des vitesses de génération plus lentes.

Mode veille de vLLM

Lorsque vous utilisez vLLM comme moteur de génération pour les méthodes d’entraînement en ligne, vous pouvez activer le sleep mode pour décharger les paramètres et le cache de vLLM vers la RAM CPU pendant l’étape d’optimisation, puis les recharger en VRAM GPU lorsque cela est nécessaire pour la synchronisation des poids et la génération.

from trl import GRPOConfig

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

from trl import RLOOConfig

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

Le déchargement des poids et du cache de vLLM contribue à maintenir une faible utilisation de la mémoire GPU, ce qui peut être particulièrement bénéfique lors de l’entraînement de grands modèles ou de l’utilisation de ressources GPU limitées. Cependant, le réveil du moteur vLLM depuis le mode veille introduit une certaine latence de transfert hôte-appareil, ce qui peut légèrement affecter la vitesse d’entraînement.

Activation checkpointing

L’activation checkpointing échange des calculs contre de la mémoire en ne stockant pas toutes les activations intermédiaires lors du passage avant, et en les recalculant à la place lors du passage arrière.

from trl import SFTConfig

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

Le checkpointing d'activation est activé par défaut dans tous les formateurs pour optimiser l'utilisation de la mémoire. Vous pouvez le désactiver en définissant activation_checkpointing=False si nécessaire.

Activation checkpointing sélectif

Avec SFTTrainer, vous pouvez enregistrer la sortie de l'attention lors du passage avant au lieu de la recalculer lors du passage arrière. Cela récupère la majeure partie du ralentissement dû au checkpointing pour un contexte long, au prix d'un tenseur supplémentaire de la taille de l'état caché par couche. Cela impose un checkpointing non réentrant.

from trl import SFTConfig

training_args = SFTConfig(
    ...,
    activation_checkpointing=True,
    activation_checkpointing_kwargs={"selective": True},
)

Pour plus de techniques d'optimisation de la mémoire, consultez le Transformers Performance Guide.