Pourquoi désactiver Gradient Checkpointing provoque un OOM dans AI Toolkit

Si vous n'êtes pas absolument sûr de la raison pour laquelle Gradient Checkpointing est désactivé, la règle la plus sûre est simple :

Réactivez-le.

Pour les grands modèles d'images d'AI Toolkit, désactiver gradient_checkpointing est l'un des moyens les plus rapides de transformer un entraînement stable en un crash OOM.

Ce guide explique :

ce que fait Gradient Checkpointing
pourquoi le désactiver peut faire exploser la VRAM
quelles familles de modèles sont les plus sensibles
quand il est raisonnable de tester GC désactivé
quoi changer à la place si votre vrai objectif est la vitesse

Réponse rapide

Gradient Checkpointing est principalement un compromis mémoire vs vitesse.

Activé = pic de VRAM plus bas, généralement plus sûr
Désactivé = pic de VRAM plus élevé, parfois plus rapide, beaucoup plus facile de tomber en OOM

Pour la plupart des utilisateurs, surtout sur les grands modèles d'images, ce n'est pas un réglage de qualité. C'est un réglage de stabilité.

Checklist rapide

✅ Cliquez sur Show Advanced
✅ Sous train:, vérifiez que gradient_checkpointing: true
✅ Relancez avec le même Batch Size et les mêmes résolutions d'abord
✅ Si l'OOM persiste, réduisez le Batch Size et supprimez le bucket de résolution le plus élevé
✅ Si votre vrai problème est l'échantillonnage de prévisualisation, réduisez Sample Width / Height / Sample Steps plutôt que de désactiver GC

1) Ce que fait réellement Gradient Checkpointing

Pas besoin de maîtriser la théorie PyTorch pour l'utiliser correctement.

Une façon pratique de voir les choses :

avec GC activé, AI Toolkit conserve moins de données d'activation intermédiaires en mémoire et en recalcule une partie à la demande
avec GC désactivé, davantage de ces données restent en VRAM, ce qui permet un entraînement plus rapide si la mémoire est disponible

Ça semble anodin jusqu'à ce que vous combiniez :

haute résolution
plusieurs buckets de résolution
batch size important
échantillonnage coûteux
architectures lourdes comme Qwen Edit, Z-Image, FLUX / modèles de type Flex

La marge de mémoire supplémentaire disparaît alors très vite.

2) Pourquoi la désactivation provoque un OOM « soudain »

Les utilisateurs pensent souvent :

« Je n'ai changé qu'un seul paramètre. »

Mais ce paramètre détermine la quantité de mémoire d'activation conservée pendant l'entraînement.

Une configuration qui était « lourde mais viable » avec GC activé devient « limite ou impossible » avec GC désactivé.

C'est pourquoi l'échec semble si brutal :

même dataset
même modèle
mêmes prompts de prévisualisation
même batch size
seul GC a changé
OOM au step 2, step 3, ou lors de la première prévisualisation

3) Familles de modèles à haut risque avec GC désactivé

D'après les patterns OOM récurrents observés dans AI Toolkit, ces combinaisons doivent être considérées comme dangereuses par défaut :

Famille de modèle	Haut risque avec GC désactivé	Premier essai le plus sûr
Z-Image	Buckets 1024–1536 avec batch important	GC activé, commencer conservateur, puis monter en charge
Qwen-Edit	Workflows 1024 avec batch important / multiples conditions lourdes	GC activé, Batch Size 1, réduire la charge de prévisualisation
FLUX-dev / grands modèles type Flex	batches importants ou entraînement multi-bucket haute résolution	GC activé, Batch Size 1–4 selon la marge disponible

Un modèle mental utile :

Z-Image devient vite dangereux avec haute résolution + batch important + GC désactivé
Qwen Edit devient vite dangereux avec 1024 + conditionnement lourd + GC désactivé
FLUX / type Flex deviennent vite dangereux avec batch important + GC désactivé

4) Où modifier ce paramètre dans RunComfy AI Toolkit

Dans l'interface actuelle de RunComfy, vous pouvez voir directement :

Batch Size
Gradient Accumulation
Steps
Unload TE
Cache Text Embeddings

Cependant, gradient_checkpointing se vérifie le plus facilement via la configuration avancée.

Étape par étape

Ouvrez votre job
Cliquez sur Show Advanced
Trouvez le bloc train:
Définissez :

train:
  gradient_checkpointing: true

Sauvegardez le job
Relancez sans rien changer d'autre

Si la même configuration fonctionne maintenant, vous avez confirmé que GC était le problème.

5) Que faire à la place si votre vrai objectif est la vitesse

Beaucoup d'utilisateurs désactivent GC pour accélérer l'entraînement.

C'est compréhensible — mais c'est généralement la mauvaise première piste d'optimisation.

Essayez ceci avant de désactiver GC :

A. Rendre l'échantillonnage de prévisualisation moins coûteux

Dans le panneau Sample :

réduisez Width / Height
réduisez Sample Steps
augmentez Sample Every
activez temporairement Disable Sampling pour vérifier la stabilité de l'entraînement

Cela donne souvent plus de gain de vitesse réel qu'une configuration risquée avec GC désactivé.

B. Garder le batch petit, monter avec l'accumulation si nécessaire

Dans le panneau Training :

gardez Batch Size conservateur
augmentez Gradient Accumulation uniquement si vous voulez un batch effectif légèrement plus grand sans pic de VRAM important

C. Supprimer d'abord le bucket le plus élevé

Dans Datasets :

conservez 512 / 768
ne réintroduisez 1024 / 1536 qu'après un entraînement stable

D. Utiliser le mode basse mémoire prévu par le modèle

Pour certaines architectures, la bonne approche n'est pas « GC désactivé », mais :

low_vram: true
quantification spécifique au modèle
optimisation du text-encoder spécifique au modèle

Suivez le guide du modèle, pas une supposition générique.

6) Quand est-il vraiment acceptable de tester GC désactivé ?

Traitez GC désactivé comme une expérience avancée, pas comme un réglage par défaut.

Conditions raisonnables :

vous avez déjà un entraînement stable
le Batch Size est encore conservateur
vous n'êtes pas déjà au maximum du plus grand bucket
les prévisualisations sont légères ou désactivées
vous avez suffisamment de marge VRAM
vous ne changez qu'une seule variable, pas plusieurs

Un bon protocole de test :

stabilisez le job avec GC activé
gardez le reste de la configuration identique
désactivez GC
surveillez :

OOM aux premiers steps
OOM pendant l'échantillonnage
pics intermittents de bucket

Si l'un de ces problèmes survient, réactivez GC et arrêtez-vous là.

7) Ce que GC désactivé n'est pas

Gradient Checkpointing n'est pas :

un booster magique de qualité
un interrupteur pour améliorer la ressemblance
le bon premier réflexe face à des échantillons décevants
une bonne première expérience quand vous êtes déjà en OOM

Si vos échantillons sont faibles, examinez :

la qualité du dataset
les captions
le rank
les steps
la parité de prévisualisation

Ne partez pas du principe que GC désactivé est la solution.

8) FAQ

Activer GC nuit-il à la qualité des résultats ?

En pratique, ce n'est généralement pas le compromis que les utilisateurs remarquent.

Le vrai compromis est surtout mémoire vs vitesse.

J'ai activé GC et j'ai toujours un OOM. Que faire ?

Vos prochains leviers sont :

Batch Size
bucket de résolution le plus élevé
coût de l'échantillonnage de prévisualisation
pour la vidéo, Num Frames

Faut-il un jour lancer un premier entraînement avec GC désactivé ?

Pour la plupart des utilisateurs : non.

Prouvez d'abord la stabilité, ensuite expérimentez.

Résumé en une ligne

Si vous rencontrez des OOM avec AI Toolkit et que gradient_checkpointing est désactivé, corrigez ça en premier.

GC activé est le réglage sûr par défaut. GC désactivé est une expérience de vitesse avancée.