CUDA Graphs pour PyTorch

Les CUDA Graphs capturent une séquence d'opérations GPU une seule fois et la rejouent avec une surcharge CPU minimale. Cette compétence guide l'application des CUDA Graphs aux charges de travail d'entraînement et d'inférence PyTorch en utilisant les APIs PyTorch natives, Transformer Engine et Megatron-LM.

Quand l'utiliser

Recourez à cette compétence quand vous rencontrez :

• Déclencheurs : l'utilisateur souhaite optimiser avec CUDA Graphs, réduire la surcharge de lancement de kernel, ou accélérer les boucles d'entraînement/inférence
• Symptômes : faible utilisation GPU (<80%), nombreux petits lancements de kernel (<50 µs chacun), entraînement limité par le CPU, latence élevée de lancement de kernel visible dans les profils Nsight Systems
• Mots-clés : « CUDA graph », « torch.cuda.graph », « make_graphed_callables », « reduce-overhead », « graph capture », « graph replay », « kernel launch overhead », « CudaGraphManager », « FullCudaGraphWrapper », « full-iteration graph », « stream capture »

N'utilisez PAS cette compétence pour :

• Le tuning de performance PyTorch général non lié à la surcharge de lancement de kernel
• Le développement de kernels CUDA ou code CUDA C++ personnalisé
• L'élimination de synchronisation host-device uniquement (utilisez la compétence perf-torch-sync-free à la place)
• Le profilage Nsight Systems (utilisez la compétence perf-nsight-systems)
• La compilation de graphes TensorFlow/JAX (APIs complètement différentes)

Pré-requis

Guide de sélection d'API

Choisissez l'API en fonction de votre framework et de vos besoins de performance.

Diagramme de décision :

1. Utilisez-vous Megatron-LM avec FP8/PP ?
   • Oui, voulez performance maximale avec charge de travail statique --> Flux 6 (FullCudaGraphWrapper)
   • Oui, voulez graphing automatique par couche --> Flux 5 (CudaGraphManager)
   • Oui, voulez contrôle manuel sur ce qui est graphé --> Flux 4 (TE make_graphed_callables)
2. Utilisez-vous Transformer Engine sans Megatron ?
   • Oui, besoin de FP8 ou PP --> Flux 4 (TE make_graphed_callables)
3. PyTorch générique ?
   • Voulez zéro effort, OK avec graphes fragmentés --> Flux 2 (torch.compile)
   • Voulez support d'autograd, boucle d'entraînement --> Flux 3 (PyTorch make_graphed_callables)
   • Voulez contrôle manuel complet --> Flux 7 (torch.cuda.graph)

Stratégie : Commencez par l'API de plus haut niveau disponible pour votre framework. Passez à des APIs de bas niveau uniquement si vous avez besoin de plus de contrôle, rencontrez des limitations, ou n'atteignez pas l'amélioration de performance attendue.

Flux de travail

Flux 1 : Profiler et décider si les graphes aident

Objectif : déterminer si les CUDA Graphs bénéficieront à votre charge de travail avant d'investir des efforts.

1. Profilez avec Nsight Systems :

   nsys profile --cuda-graph-trace=graph python train.py

2. Vérifiez l'utilisation GPU -- si déjà >95%, les graphes n'aideront pas beaucoup.
3. Recherchez des écarts entre les lancements de kernel (surcharge CPU) et de nombreux petits kernels (<50 µs chacun). Ce sont les cibles pour le graphing.
4. Annotez les régions d'intérêt pour corréler le temps d'inactivité du GPU avec le code :

   with torch.cuda.nvtx.range("forward"):
       output = model(input)

5. Estimez le bénéfice : comptez les kernels par itération. Les charges de travail avec des centaines de petits kernels et <80% d'utilisation GPU sont de bons candidats.

Résultat attendu : régions de goulot d'étranglement identifiées avec faible occupation GPU entre les kernels. Passez au flux de travail approprié du Guide de sélection d'API.

Flux 2 : torch.compile(mode="reduce-overhead")

Objectif : capture automatique de CUDA Graph sans effort manuel.

Quand l'utiliser : expérience rapide, limites de graphe inconnues, déjà en utilisant torch.compile.

Étapes :

1. Décorrez l'étape d'entraînement avec @torch.compile(mode="reduce-overhead") :

   @torch.compile(mode="reduce-overhead")
   def train_step(model, x, target, criterion):
       output = model(x)
       loss = criterion(output, target)
       loss.backward()
       return loss

2. Exécutez la boucle d'entraînement normalement -- les graphes sont capturés automatiquement.
3. Profilez avec Nsight Systems pour voir les graphes capturés :

   nsys profile --cuda-graph-trace=graph python train.py

4. Si vous voyez trop de petits graphes (fragmentation de graphe), vérifiez les ruptures de graphe : .item(), print(), contrôle de flux dépendant des données. Corrigez-les ou passez à Flux 3+.

Compromis :

• Zéro effort, mais peut créer de petits graphes fragmentés.
• Contrôle limité sur ce qui est graphé.
• La fragmentation de graphe limite les gains de performance par rapport aux approches manuelles.

Flux 3 : torch.cuda.make_graphed_callables()

Objectif : entraînement avec support d'autograd. Graphes forward/backward séparés.

Quand l'utiliser : entraînement avec boucles personnalisées, non-FP8, besoin d'autograd.

Étapes :

1. Préparez des entrées d'exemple correspondant à la forme du batch d'entraînement :

   sample_input = torch.randn(batch_size, seq_len, hidden_size, device="cuda")

2. Créez le modèle graphé :

   graphed_model = torch.cuda.make_graphed_callables(
       model, (sample_input,), num_warmup_iters=3
   )

3. Utilisez graphed_model comme remplacement clé en main dans la boucle d'entraînement :

   for data, target in dataloader:
       optimizer.zero_grad()
       output = graphed_model(data)
       loss = criterion(output, target)
       loss.backward()
       optimizer.step()

4. Si vous utilisez AMP, définissez cache_enabled=False :

   for data, target in dataloader:
       optimizer.zero_grad()
       with torch.amp.autocast("cuda", cache_enabled=False):
           output = graphed_model(data)
           loss = criterion(output, target)
       loss.backward()
       optimizer.step()

5. Si vous utilisez DDP, construisez DDP sur un flux latéral et utilisez 11 itérations de warmup :

   os.environ["TORCH_NCCL_ASYNC_ERROR_HANDLING"] = "0"
   s = torch.cuda.Stream()
   with torch.cuda.stream(s):
       model = DistributedDataParallel(model)
   torch.cuda.current_stream().wait_stream(s)
   
   graphed_model = torch.cuda.make_graphed_callables(
       model, (sample_input,), num_warmup_iters=11
   )

Limitations :

• Pas de backward double (gradients d'ordre supérieur).
• Pas de hooks de module durant la capture.
• La structure du module est figée après le graphing (pas d'ajout/suppression de paramètres).
• La signature de l'argument doit correspondre exactement à sample_args.

Flux 4 : TE make_graphed_callables

Objectif : graphing par callable avec support FP8 et parallélisme de pipeline.

Quand l'utiliser : entraînement FP8, PP avec planification manuelle, modèles non-Megatron nécessitant FP8, ou tout modèle PyTorch nécessitant CUDA Graphs conscients de FP8.

Étapes :

1. Importez et configurez :

   from transformer_engine.pytorch.graph import make_graphed_callables
   from transformer_engine.pytorch.fp8 import fp8_autocast

2. Préparez les entrées d'exemple (une par callable par microbatch par chunk) :

   sample_args = tuple(
       (torch.randn(batch_size, seq_len, hidden_size, device="cuda"),)
       for _ in range(num_callables * num_microbatches)
   )

3. Définissez le calendrier de pipeline si vous utilisez PP (IDs de chunk indexés à partir de 1, positif=fwd, négatif=bwd) :

   # Exemple : 2 chunks, 3 microbatches
   layer_order = [1, 2, 1, 2, 1, 2, -2, -1, -2, -1, -2, -1]

4. Enveloppez les couches dans CUDA Graphs :

   graphed_layers = make_graphed_callables(
       tuple(layers),
       sample_args=sample_args,
       fp8_enabled=True,
       fp8_recipe=fp8_recipe,
       fp8_weight_caching=True,
       _order=layer_order,  # None pour pas de PP
   )

5. Boucle d'entraînement -- enveloppez avec fp8_autocast durant la rejeu :

   with fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
       for layer in graphed_layers[start:end]:
           x = layer(x, is_first_microbatch=(mb_idx == 0))
   # FP8 scaling auto-mise à jour à la sortie de fp8_autocast
   optimizer.step()

Points clés :

• Capture AOT : graphes capturés avant la boucle d'entraînement quand vous appelez make_graphed_callables().
• L'ordre de rejeu doit correspondre à _order : la boucle d'entraînement doit exécuter les graphes dans le même ordre entrelacé que celui spécifié durant la capture.
• fp8_autocast requis durant la rejeu : sans lui, l'état FP8 n'est pas correctement configuré.
• Mise en cache des poids : fp8_weight_caching=True met en cache la quantification des poids FP8 entre microbatches ; passez l'argument kwarg is_first_microbatch pour contrôler quand les poids sont requantifiés.

Pour les détails complets de l'API, consultez references/api-te-megatron.md.

Flux 5 : MCore CudaGraphManager (Par-couche)

Objectif : graphing automatique par couche pour l'entraînement Megatron-LM.

Quand l'utiliser : entraînement Megatron-LM, en particulier avec PP > 1. Choix par défaut pour les utilisateurs de Megatron.

Étapes :

1. Activez via les drapeaux CLI (pas de changements de code nécessaires) :

   python pretrain_gpt.py \
       --enable-cuda-graph \
       --cuda-graph-num-warmup-steps 3

2. Ou activez via la config Python :

   config = TransformerConfig(
       enable_cuda_graph=True,
       cuda_graph_num_warmup_steps=3,
   )

3. La boucle d'entraînement inchangée -- les graphes sont capturés automatiquement après les itérations de warmup.

Points clés :

• Couches Megatron uniquement : fonctionne avec TransformerLayer et MambaLayer.
• Capture JIT : enregistre l'ordre d'exécution durant le warmup, capture les graphes après le warmup, puis rejoue sur les itérations suivantes.
• Gestion FP8 automatique : utilise fp8_autocast(..., _graph=True) pour sauter la réduction amax par couche ; la réduction se fait une seule fois après tous les graphes backward.
• Support PP automatique : gère automatiquement l'entrelacement de microbatch.
• Économies de mémoire : définissez cuda_graph_share_io_buffers=True pour partager les buffers I/O entre couches (nécessite pas d'opérations entre couches).
• Stratégie de pool mémoire : par défaut utilise des pools séparés par microbatch pour la réutilisation de graphes. Définissez cuda_graph_use_single_mempool=True pour un pool partagé (nombre de graphes plus élevé mais peut réduire la fragmentation).

Flux 6 : MCore FullCudaGraphWrapper (Itération complète)

Objectif : performance maximale. Capture forward+backward pour tous les microbatches en un seul graphe.

Quand l'utiliser : priorité de performance maximale, charges de travail statiques, entraînement Megatron-LM.

Étapes :

1. Activez via les drapeaux CLI :

   python pretrain_gpt.py \
       --enable-cuda-graph \
       --cuda-graph-scope full_iteration \
       --cuda-graph-warmup-steps 1 \
       --te-rng-tracker \
       --no-check-for-nan-in-loss-and-grad

2. Assurez-vous que tout le code forward+backward est capturable (pas de .item(), pas de vérification de NaN, pas de contrôle de flux dynamique).
3. L'optimiseur reste en mode eager par défaut (en dehors du graphe). Peut être inclus dans le graphe pour une performance maximale.

Points clés :

• Seulement 2 graphes au total : un pour l'entraînement, un pour la validation.
• --te-rng-tracker requis : la RNG standard utilise des scalaires CPU qui ne peuvent pas être capturés ; la RNG TE utilise des tenseurs device compatibles avec les graphes.
• --no-check-for-nan-in-loss-and-grad obligatoire : la vérification de NaN utilise .item() qui nécessite une synchronisation CPU-GPU, interdite durant la capture.
• StaticBufferLoader : pré-alloue les buffers d'entrée pour tous les microbatches durant le warmup.
• Optimiseur dans/en dehors du graphe : dans = performance maximale (tous les kernels d'optimiseur capturés). En dehors = plus flexible (peut changer l'optimiseur/LR sans recapture).
• Capture JIT : graphe capturé durant l'entraînement à l'itération warmup_steps + 1.

Flux 7 : torch.cuda.graph() (Manuel)

Objectif : contrôle complet sur la capture et la rejeu. Pipelines personnalisés, capture d'itération complète sans Megatron.

Quand l'utiliser : besoin de contrôle fin, capture d'itération complète non-Megatron, pipelines personnalisés.

Motif d'inférence :

1. Pré-allouez les tenseurs d'entrée/sortie statiques :

   static_input = torch.randn(batch_size, *shape, device="cuda")

2. Warmup sur un flux latéral (3 itérations, 11 pour DDP) :

   s = torch.cuda.Stream()
   with torch.cuda.stream(s):
       for _ in range(3):
           _ = model(static_input)
   torch.cuda.current_stream().wait_stream(s)

3. Capturez le graphe :

   g = torch.cuda.CUDAGraph()
   with torch.cuda.graph(g):
       static_output = model(static_input)

4. Boucle de rejeu -- mettez à jour les entrées via .copy_(), clonez les sorties :

   for data in loader:
       static_input.copy_(data)
       g.replay()
       result = static_output.clone()

Motif d'entraînement complet (fwd+bwd+optimiseur en un seul graphe) :

model = MyModel().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = torch.nn.CrossEntropyLoss()

static_input = torch.randn(batch_size, *shape, device="cuda")
static_target = torch.randint(0, num_classes, (batch_size,), device="cuda")

# Warmup
s = torch.cuda.Stream()
with torch.cuda.stream(s):
    for _ in range(3):
        optimizer.zero_grad()
        with torch.amp.autocast("cuda", cache_enabled=False):
            out = model(static_input)
            loss = criterion(out, static_target)
        loss.backward()
torch.cuda.current_stream().wait_stream(s)

# Capture
g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
    optimizer.zero_grad()
    with torch.amp.autocast("cuda", cache_enabled=False):
        static_output = model(static_input)
        static_loss = criterion(static_output, static_target)
    static_loss.backward()

# Boucle de rejeu
for data, target in loader:
    static_input.copy_(data)
    static_target.copy_(target)
    g.replay()
    optimizer.step()

Configuration DDP :

os.environ["TORCH_NCCL_ASYNC_ERROR_HANDLING"] = "0"

s = torch.cuda.Stream()
with torch.cuda.stream(s):
    model = DistributedDataParallel(model)

# 11 itérations de warmup pour DDP
with torch.cuda.stream(s):
    for _ in range(11):
        out = model(static_input)
        out.sum().backward()
torch.cuda.current_stream().wait_stream(s)

# Capture sur le même flux latéral
with torch.cuda.graph(g):
    static_output = model(static_input)

Partage de pool mémoire pour plusieurs graphes :

g1 = torch.cuda.CUDAGraph()
with torch.cuda.graph(g1):
    out1 = model_a(static_in_a)

# Le deuxième graphe partage le pool mémoire du premier graphe
g2 = torch.cuda.CUDAGraph()
with torch.cuda.graph(g2, pool=g1.pool()):
    out2 = model_b(static_in_b)

Enregistrement RNG personnalisé :

gen = torch.cuda.default_generators[0]
g = torch.cuda.CUDAGraph()
g.register_generator_state(gen)
with torch.cuda.graph(g):
    out = model(static_input)  # État RNG correctement capturé

Navigation entre les flux de travail

• torch.compile donne une accélération insuffisante --> escaladez à make_graphed_callables (Flux 3) pour des graphes plus grands et moins nombreux.
• make_graphed_callables ne peut pas gérer FP8/PP --> TE make_graphed_callables (Flux 4).
• Besoin de Megatron par-couche automatique --> CudaGraphManager (Flux 5).
• Voulez performance maximale --> FullCudaGraphWrapper (Flux 6) ou capture manuelle d'itération complète (Flux 7).
• Quelque chose trop difficile à grapher --> capture partielle (graphez ce que vous pouvez, laissez le reste en mode eager).
• L'utilisateur veut la meilleure perf absolue --> passez directement à Flux 6 (Megatron) ou Flux 7 (manuel).
• Commencez petit, expandez progressivement : commencez avec un module/couche. Vérifiez la correctitude. Puis expandez à plus de couches, la passe forward complète, ajoutez backward, et finalement l'itération complète avec l'optimiseur.

Rendre le code compatible avec les graphes

Ces principes s'appliquent à tous les flux de travail. Le code dans la région capturée doit satisfaire trois contraintes.

Principe 1 : GPU uniquement

Seules les opérations GPU sont capturées. Le code côté CPU (logique Python, I/O, logging) s'exécute durant la capture mais est éliminé durant la rejeu.

Violations :

• I/O fichier : data = torch.load("file.pt") ne rechargera pas lors de la rejeu
• Prétraitement CPU : tokens = tokenizer.encode(text) ne re-tokénisera pas
• Logging : print(f"Step {i}") n'imprimera pas durant la rejeu
• RNG CPU : random.randint(0, 10) ne régénérera pas
• Comptabilité CPU : buffer.append(tensor) ne remplira pas durant la rejeu

Fix : déplacez toutes les opérations côté CPU en dehors de la région graphée.

Principe 2 : Sans synchronisation

Pas de synchronisation CPU-GPU dans le graphe. Le CPU met le travail en queue continuellement sans attendre les résultats du GPU.

Violations :

• .item() pour obtenir les valeurs scalaires
• .cpu() pour déplacer les tenseurs pour l'inspection
• torch.cuda.synchronize() ou stream.synchronize()
• print(tensor) (synchronise implicitement)

Fix : invoquez la compétence perf-torch-sync-free pour la détection et l'élimination systématiques des points de synchronisation. Utilisez torch.cuda.set_sync_debug_mode("warn") pour trouver les syncs cachées.

Principe 3 : Statique

Toutes les opérations, le contrôle de flux, les adresses mémoire et les formes doivent être fixes dans toutes les rejeus.

Violations et fixes :

Pour les motifs détaillés, consultez references/patterns-dynamic.md.

Checklist de compatibilité

Vérifiez chaque élément avant de tenter la capture :

• [ ] Pas de .item(), .cpu(), .numpy(), print(tensor) dans le graphe
• [ ] Pas de torch.cuda.synchronize() ou stream.synchronize()
• [ ] Pas de if tensor_value: -- utilisez torch.where() à la place
• [ ] Toutes les entrées pré-allouées, mises à jour via .copy_()
• [ ] Toutes les formes fixes (utilisez padding ou bucketing pour les tailles variables)
• [ ] Scalaires Python --> tenseurs GPU avec .fill_()
• [ ] Tenseurs de sortie clonés avant la prochaine rejeu
• [ ] cache_enabled=False avec torch.amp.autocast
• [ ] Générateurs RNG personnalisés enregistrés avec graph.register_generator_state()
• [ ] Utilisez graphsafe_get_state() / graphsafe_set_state() pour la RNG
• [ ] Warmup complété (3 standard, 11 pour DDP)
• [ ] DDP : TORCH_NCCL_ASYNC_ERROR_HANDLING=0, construisez sur le flux latéral
• [ ] DDP : NCCL >= 2.9.6 pour capture de graphe complète
• [ ] Les libraries/extensions utilisent torch.cuda.current_stream(), pas le flux par défaut
• [ ] Pas d'allocation de mémoire pinned durant la capture (déclenche une requête d'événement cachée)
• [ ] activation_checkpointing : preserve_rng_state=False
• [ ] Les tenseurs globaux utilisés dans le graphe restent vivants (ne sont pas supprimés/réassignés)
• [ ] Pas de fonctions torch.compile dans la capture manuelle sans warmup préalable
• [ ] Le clipping de gradients utilise clip_grad_norm_ sans synchronisation (PyTorch >= 1.13)

Pour la checklist complète avec références, consultez references/patterns-compatibility.md.

Formats de sortie

Indicateurs de succès :

• g.replay() se complète sans erreur
• Les sorties correspondent au mode eager dans la tolérance (torch.allclose)
• Le profil Nsight Systems montre un seul lancement de graphe remplaçant de nombreux kernels
• L'utilisation GPU augmente, la latence d'entraînement/inférence diminue

Métriques clés :

Gestion des erreurs

Pour le dépannage détaillé, consultez references/troubleshooting.md.

Trouver plus d'informations

Utilisez cette hiérarchie de recherche en 3 tiers -- commencez au Tier 1 et escaladez uniquement si nécessaire.

Tier 1 : Ce fichier (SKILL.md)

Vous le lisez maintenant. Les flux de travail, la checklist de compatibilité et le tableau d'erreurs ci-dessus couvrent les tâches les plus courantes. Cherchez d'abord dans ce fichier avant d'aller plus profond.

Tier 2 : Répertoire references/

Le répertoire references/ à côté de ce fichier contient du matériel de référence distillé -- détails d'API, motifs et pages de dépannage.

Comment chercher :

1. Grep votre mot-clé dans references/ -- les en-têtes sont conçus pour être grep-friendly.
2. Lisez seulement le fichier que grep vous pointe. Ne lisez pas chaque fichier.

Références disponibles :

• references/api-pytorch.md -- APIs PyTorch CUDA Graph (torch.cuda.graph, make_graphed_callables, torch.compile reduce-overhead)
• references/api-te-megatron.md -- implémentations TE make_graphed_callables, CudaGraphManager, FullCudaGraphWrapper
• references/patterns-compatibility.md -- Principes GPU-only, sync-free et static avec checklist complète
• references/patterns-dynamic.md -- Contrôle de flux dynamique, tenseurs, scalaires, formes : contournements et motifs
• references/troubleshooting.md -- Échecs de capture, erreurs numériques, problèmes mémoire, problèmes de performance

Tier 3 : Documentation originale

Si les Tiers 1-2 ne répondent pas à la question, consultez les sources originales :

• Guide NVIDIA : https://docs.nvidia.com/dl-cuda-graph/latest/index.html
• Docs PyTorch : https://docs.pytorch.org/docs/stable/notes/cuda.html (section CUDA Graphs)
• Docs TE : https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/index.html
• Docs Megatron Core : https://docs.nvidia.com/megatron-core/developer-guide/latest/index.html

Retournez au Tier 2 après et considérez si la réponse devrait être distillée dans le répertoire references pour la prochaine fois.