Entraînement au-delà de 1 M de jetons

Un agent est capable de travailler sur une tâche pendant des heures : lire des fichiers, exécuter des commandes et faire avancer chaque tour de la session. Cette session accumulée atteint des centaines de milliers de jetons, ce qui explique pourquoi les modèles de pointe annoncent désormais des fenêtres de contexte d'un million de jetons.

Rendre un modèle compétent sur un contexte aussi long implique de l’entraîner sur des séquences de cette taille, et c’est là que la difficulté se situe : une séquence d’un million de jetons ne tient pas sur un GPU, et elle ne tient pas non plus sur huit, sans aide.

L'exemple de ce guide entraîne le modèle sur exactement une telle séquence par étape, sur un nœud unique de 8 GPU.

L'exécuter

Il vous faut un nœud équipé de 8 H100 (ou mieux) et transformers>=5.17.0 : l’exemple utilise l’option offload du gradient checkpointing, ajoutée dans cette version.

accelerate launch \
    --config_file examples/sft_qwen3_8b_1m_context/context_parallel_8gpu.yaml \
    examples/sft_qwen3_8b_1m_context/sft_qwen3_8b_1m_context.py

Le script effectue un fine-tuning de Qwen3-8B sur des livres issus de PG-19, concaténés bout à bout jusqu’à ce que chacun atteigne environ un million de jetons. Après une dizaine de minutes de chargement et de tokenisation, la première étape est atteinte :

{'loss': 4.311, 'grad_norm': 29.25, 'num_tokens': 1049000.0, 'epoch': 0.25}
  8%|▊ | 1/12 [06:20<1:09:41, 380.10s/it]

Trois chiffres méritent d'être lus :

• num_tokens vaut 1,049e6. Il s’agit d’une seule séquence, et non d’un lot de séquences courtes. Chaque jeton porte son attention sur tous les jetons précédents, sur l’ensemble du million.
• La perte démarre à 4,3, ce qui est une perte de départ normale. Une exécution légèrement mal configurée pour le long contexte démarre plutôt autour de 10, et nous reviendrons sur les raisons.
• 380 s par étape. Une étape correspond à une séquence, soit un peu plus de six minutes par séquence.

De manière naïve, cette séquence nécessite 288 Go par GPU. Le reste de ce guide raconte comment ce chiffre est ramené à 56.

Ce qui vient de se passer

Tout ce qui suit s’appuie sur une configuration que nous connaissons déjà : un GPU, quelques milliers de jetons par séquence. Les problèmes commencent à apparaître à mesure que la séquence s’allonge, et ils surviennent dans un ordre fixe.

La perte est la première chose à surveiller

Prenez Qwen3-4B sur une seule carte de 80 Go et augmentez la séquence. La première chose à manquer de mémoire n'est pas le modèle. C'est la perte. Vérifiez le profil mémoire.

Le pic se situe entre les passes avant et arrière, c'est-à-dire là où la perte est calculée. Pour comprendre pourquoi le pic se trouve à cet endroit, il faut examiner ce que font les dernières étapes du décodage et comment la perte est construite à partir de celles-ci.

La dernière couche du décodeur produit un état caché de forme$L \times H$, où$L$ est la longueur de la séquence et$H$ est la taille cachée. Cet état caché est ensuite multiplié par les poids de la tête de modélisation du langage pour produire une matrice de logits de forme$L \times V$, où$V$ est la taille du vocabulaire. Cette matrice est énorme, car la séquence est longue et le vocabulaire est vaste (plus de cent mille entrées). La perte est ensuite calculée à l'aide d'une fonction d'entropie croisée à partir de cette matrice et des étiquettes cibles, et sa matérialisation correspond précisément au pic.

Peut-on calculer la perte sans matérialiser toute la matrice de logits ? Oui, nous pouvons procéder par morceaux.

La matrice de logits massive : découpez-la

Plutôt que de calculer les logits pour l’ensemble de la séquence d’un coup, nous prenons quelques centaines de lignes à la fois, calculons la perte pour celles-ci, puis additionnons les morceaux. La matrice complète des logits n’est jamais conservée en mémoire en une seule fois. Les lignes sont la dimension qui se découpe proprement : la perte est une somme sur les lignes, et le softmax à l’intérieur s’exécute le long du vocabulaire d’une ligne, de sorte qu’un bloc de lignes contient tout ce dont sa propre perte a besoin.

Une grande multiplication devient plusieurs plus petites, et l'économie de mémoire est suffisamment importante pour que la perte ne soit plus jamais le pic. TRL utilise 256 lignes par bloc, ce qui maintient chaque multiplication assez grande pour exploiter efficacement le GPU.

Le nouveau profil ressemble à ceci.

Le pic a disparu.

Comment activer cela dans TRL ? Il n’y a rien à faire : la perte par morceaux est le comportement par défaut. Si vous souhaitez opter pour la perte simple, définissez loss_type="nll" dans SFTConfig.

Jusqu’où cela nous permet-il de pousser la séquence ? En termes de mémoire, ce seul changement nous fait passer de 32k jetons à 160k.

Le modèle n'a jamais vu la position 100 000 : rééchelonnez les positions avec YaRN

Nous pouvons désormais stocker 160k jetons en mémoire. Mais le modèle peut-il réellement les lire ?

Qwen3-4B a été entraîné sur des séquences de 40 960 jetons. La configuration l'indique :

>>> from transformers import AutoConfig
>>> AutoConfig.from_pretrained("Qwen/Qwen3-4B").max_position_embeddings
40960

Toute position au-delà est une position que le modèle n'a jamais rencontrée. Rien ne plante, mais l'exécution s'entraîne simplement sur des jetons que le modèle ne peut pas situer. Voici la perte sur chaque jeton d'une séquence de 160k jetons :

Le modèle ne s'effondre pas à 40 960. Il continue un certain temps et ne commence à perdre le fil qu'au-delà de 70 k environ. À partir de là, il grimpe régulièrement, pour finir au-dessus de l'entropie unigramme du même texte (6,45) : pire qu'un modèle qui ignore totalement le contexte et prédit à partir des seules fréquences de jetons.

Chaque jeton est associé à sa position, et le modèle a appris ce que signifient les positions comprises entre 0 et 40 960. Si vous lui demandez quelque chose concernant la position 100 000, il n’a aucune base pour répondre.

La solution consiste à ne jamais lui transmettre une position qu’elle ne connaît pas. Nous rééchelonnons : la position 100 000 devient la position 25 000, et chaque position de la séquence revient dans la plage que le modèle connaît. Les jetons ne changent pas, seule la position attribuée à chacun change. C’est ce que fait YaRN, et le facteur est simplement la distance supplémentaire que nous visons : 40 960 entraînés, 163 840 souhaités, donc 4.

Comment activer cela dans TRL ? Passez la configuration RoPE lors du chargement du modèle :

training_args = SFTConfig(
    ...,
    model_init_kwargs={
        "rope_parameters": {
            "rope_type": "yarn",
            "rope_theta": 1_000_000,  # has to match what the model ships with
            "factor": 4.0,
            "original_max_position_embeddings": 40960,
        },
    },
)

La même mesure, avec l'exécution redimensionnée en haut :

L'exécution avec rééchelonnement reste stable jusqu'à 160k. Sur les 20k derniers jetons, la perte moyenne est de 2,8, contre 7,3 sans rééchelonnement.

Les activations sont ce qu'il reste : déchargez-les

Avec la perte chunkée et les positions re-échelonnées, un seul GPU nous permet d'atteindre 160k tokens. Puis il sature à nouveau. Qu'est-ce qui occupe la carte cette fois-ci ?

Le modèle et son optimiseur représentent 22 Go fixes, quelle que soit la longueur de la séquence. Au-delà, chaque bande rouge correspond aux activations enregistrées d’une couche, soit 0,47 Go chacune. Elles s’accumulent au fur et à mesure que la passe avant les écrit, restent en mémoire, puis sont retirées une par une au fur et à mesure que la passe arrière les consomme. La partie grise en haut représente tout ce qui est recalculé et jeté en quelques millisecondes.

TRL les réduit déjà pour vous, et cette fonctionnalité est activée par défaut. Le gradient checkpointing ne conserve que ce qui entre dans chaque couche et jette tout ce que la couche calcule en interne : les scores d'attention, l'intermédiaire large du MLP, les normalisations. Lorsque le passage arrière en a besoin, il exécute la couche une seconde fois pour les récupérer. Ce qui subsiste est un tenseur enregistré par couche, de taille sequence x hidden chacun, et à de longues longueurs de séquence, c'est toujours la plus grande chose sur la carte.

L'observation utile est que ces tenseurs ne sont pas du tout nécessaires pendant la passe avant. Ils sont écrits, laissés de côté, puis lus bien plus tard. Ils n'ont donc pas besoin de rester sur le GPU en attendant : ils peuvent être déplacés en mémoire CPU et ramenés lorsque la passe arrière les atteint.

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

Même étape, même longueur : au lieu de 38 bandes empilées, seules 4 sont présentes à la fois, et le pic passe de 59,9 Go à 48,8 Go. Ce qui permet d’augmenter la longueur de séquence :

En dessous de 48k, les deux sont identiques, car à cette longueur il y a à peine quoi que ce soit à déplacer. Au-delà, l’exécution avec offload reste plus basse, et là où la séquence s’arrêtait à 160k, elle atteint désormais 256k sur la même carte.

Celle-ci n’est pas gratuite. Chaque activation sauvegardée passe par l’hôte puis revient, et ce trafic doit s’insérer entre les calculs qui l’entourent.

Un seul GPU ne suffit pas : répartissez la séquence sur plusieurs

Perte par morceaux, activations déchargées : une carte entraîne désormais sur 256k jetons. Comment atteindre un million ? Pour voir ce qui fait obstacle, examinez l’intérieur d’une couche unique pendant la passe arrière (pics gris dans le profil ci-dessus).

Le gradient checkpointing signifie que le forward ne conserve presque rien : chaque couche stocke son entrée et jette le reste. La couche est recalculée plus tard, lorsque le backward l’atteint, et c’est à ce moment que tout existe simultanément. Dans le MLP, une ligne du modèle,

down_proj(act_fn(gate_proj(x)) * up_proj(x))

quatre tenseurs sont construits à la largeur intermédiaire complète, et l’implémentation standard conserve les quatre jusqu’à ce que la rétropropagation les ait dépassés. Pour un million de jetons, sur un modèle dont le MLP élargit chaque jeton de 2560 à 9728 :

gate_proj(x)                             19 GB = 1,048,576 x 9728 x 2 bytes
act_fn(gate_proj(x))                     19 GB
up_proj(x)                               19 GB
act_fn(gate_proj(x)) * up_proj(x)        19 GB
                                  total  76 GB

76 Go, à l’intérieur d’une seule couche, sur une carte qui en possède 80 et qui a déjà alloué 30 au modèle et à son optimiseur. À ce stade, la réponse ne consiste pas à stocker la séquence de manière plus astucieuse. Il s’agit de donner moins de données à chaque GPU.

Attribuez à chaque GPU une portion de la séquence. Avec quatre d’entre eux, chaque GPU détient 262 144 jetons au lieu de 1 048 576. Les quatre tenseurs se réduisent en conséquence : 19 Go au lieu de 76, ce qui, avec les 30 Go du modèle, tient sur la carte.

L’objection évidente est l’attention. Chaque jeton doit porter attention sur tous les jetons précédents, et aucun GPU ne détient les jetons antérieurs. Les GPU échangent donc ce qui leur manque, et le résultat est exactement ce qu’un GPU unique et immense aurait calculé.

Il existe deux façons d’exécuter cet échange : le parallélisme de contexte (CP) et le parallélisme de séquence Ulysses (SP).

Sous FSDP2, utilisez CP. Chaque GPU conserve sa propre tranche, et les autres tranches lui parviennent à tour de rôle, de sorte que chaque token finit par voir tous les tokens précédents.

parallelism_config:
  parallelism_config_cp_size: 4

Sous DeepSpeed, utilisez SP. Au lieu de déplacer des tranches, il réorganise le lot juste avant l'attention afin que chaque GPU détienne tous les jetons mais seulement un quart des têtes d'attention, puis il réorganise le lot dans l'ordre initial par la suite.

parallelism_config:
  parallelism_config_sp_size: 4

Les deux réglages ne se croisent pas aujourd'hui : cp_size nécessite FSDP2 et sp_size ne fonctionne que sous DeepSpeed. On choisit le backend et la méthode s'ensuit. Le reste de cette section suit le chemin FSDP2, qui est celui utilisé par l'exemple en haut de ce guide : ses GPU ne sont pas des travailleurs indépendants sur des lots distincts, ils forment un seul groupe partageant une séquence unique.

Deux éléments changent une fois activé :

1. Les séquences doivent être remplies jusqu'à un multiple de cp_size * 2, donc pad_to_multiple_of=8 (dans SFTConfig) pour quatre GPU.
2. L’exigence de causal-SDPA exclut le packing, qui repose sur un masque en diagonale par blocs pour empêcher les documents de se lire mutuellement, et TRL lève une erreur si vous demandez les deux. Elle exclut également les modèles dont les couches utilisent une attention à fenêtre glissante ou par morceaux, que Accelerate refuse : OpenAI GPT-OSS, Gemma 3 et 4, Qwen3.5 et suivants. Qwen3 et Qwen3-MoE utilisent une attention complète sur l’ensemble de la séquence, ce qui explique pourquoi ce sont les modèles présentés ici.

Faire circuler les tranches coûte moins que vous ne le pensez. Avec 131k jetons sur un nœud, une étape prend 34,6 s sur deux GPU, 17,8 s sur quatre et 9,5 s sur huit. Chaque doublement du groupe divise par deux le temps d’étape, et le passage de deux à huit récupère 3,7x sur un potentiel de 4x.

Et c'est le dernier des quatre leviers :

Une seule carte s’arrête juste après 256k. Quatre d’entre elles couvrent le million complet !, en atteignant 63,6 Go avec de la marge. La séquence avec laquelle ce guide s’ouvre tient.

Définissez PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True. Pour un million de jetons, l’exécution ci-dessus nécessite 63,6 Go sur une carte de 80 Go et échoue toujours sans ce paramètre : les tenseurs demandés sont volumineux et contigus, et l’allocateur ne parvient pas à leur trouver de place parmi les blocs qu’il détient déjà.

Lecture complémentaire

• Parallélisme de contexte, le guide Accelerate pour cp_size et le maillage d’appareils qu’il construit.
• Ulysses et ring attention, sur la manière dont les deux échanges diffèrent et ce que chacun coûte en communication.
• YaRN, le recalage des positions utilisé dans la section RoPE.