Container Compute Service:training-nv-pytorch 26.04

Fonctionnalités principales

• Mise à niveau de vllm vers la version 0.19.0.

• Mise à niveau de transformer_engine vers la version 2.13.

Corrections de bugs

Aucune.

Ressources

Prérequis relatifs aux pilotes

• La version 26.04 prend en charge CUDA 13.0.2 (pilote >= 580) et CUDA 12.8.0 (pilote >= 575). Consultez les pages Compatibilité des applications CUDA et Compatibilité et mises à niveau CUDA.

Optimisation de la compilation PyTorch

Introduite dans PyTorch 2.0, la fonction torch.compile() offre souvent des gains significatifs pour les charges de travail mono-GPU à petite échelle. Toutefois, l'entraînement de LLM reposant sur l'optimisation de la mémoire GPU et des frameworks distribués tels que FSDP ou DeepSpeed, torch.compile() peut apporter des avantages limités, voire dégrader les performances.

• Contrôlez la granularité des communications dans le framework DeepSpeed. Le compilateur capture ainsi un graphe de calcul plus complet et applique des optimisations plus étendues.

• Utilisez une version optimisée de PyTorch :

  • Le frontend du compilateur PyTorch est amélioré pour garantir la réussite de la compilation même en cas de rupture de graphe lors du calcul.

  • La correspondance de motifs et la prise en charge des formes dynamiques sont renforcées afin d'améliorer les performances post-compilation.

Grâce à ces optimisations, l'entraînement de LLM à 8 milliards de paramètres atteint généralement un gain de débit de bout en bout d'environ 20 %.

Optimisation de la mémoire GPU pour le recalcul

Un modèle prédictif de la surcharge mémoire GPU recommande le nombre optimal de couches de recalcul d'activation. Ce modèle s'appuie sur des données de performance à grande échelle, incluant divers modèles, clusters et paramètres d'entraînement, ainsi que des métriques système telles que l'utilisation de la mémoire GPU collectées lors des benchmarks. Intégrée à PyTorch, cette approche permet de bénéficier des gains de performance liés à l'optimisation de la mémoire GPU avec un effort minimal. Cette fonctionnalité est désormais prise en charge dans le framework DeepSpeed.

Évaluation des performances E2E

À l'aide de CNP (outil de benchmarking des performances IA cloud-native), nous avons comparé les performances E2E d'une image de base standard et de modèles open source principaux. Des études d'ablation mesurent la contribution de chaque composant d'optimisation aux performances d'entraînement.

Comparaison d'images et évaluation des itérations

Contribution des composants GPU essentiels aux performances E2E

Configurations testées sur un cluster GPU multi-nœuds :

1. Base : Image NGC PyTorch

2. Image IA ACS (Base+ACCL) : Image de base avec la bibliothèque de communication ACCL.

3. Image IA ACS (AC2+ACCL) : Image de référence avec AC2 BaseOS, sans aucune optimisation activée.

4. Image IA ACS (AC2+ACCL+CompilerOpt) : Image de référence avec AC2 BaseOS, avec uniquement l'optimisation de compilation PyTorch activée.

5. Image IA ACS (AC2+ACCL+CompilerOpt+CkptOpt) : Image de référence avec AC2 BaseOS, avec les optimisations de compilation PyTorch et de checkpointing sélectif des gradients activées.

Étape 1 : Sélectionner une image

docker pull egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com/egslingjun/training-nv-pytorch:[tag]

Étape 2 : Activer le compilateur et les optimisations de recalcul

• Activer l'optimisation de compilation

  Utilisez l'API Transformers Trainer :

  

• Activer l'optimisation de la mémoire GPU par recalcul

  export CHECKPOINT_OPTIMIZATION=true

Étape 3 : Démarrer le conteneur

L'image inclut ljperf, un outil d'entraînement de modèles. Utilisez-le pour démarrer un conteneur et lancer l'entraînement :

# Start the container and open a shell
docker run --rm -it --ipc=host --net=host --privileged egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com/egslingjun/training-nv-pytorch:[tag]

# Run the training demo
ljperf benchmark --model deepspeed/llama3-8b 

Étape 4 : Remarques d'utilisation

• Ne réinstallez pas PyTorch ni DeepSpeed. Cette image contient des correctifs personnalisés pour ces bibliothèques.

• Dans votre configuration DeepSpeed, laissez zero_optimization.stage3_prefetch_bucket_size vide ou définissez-le sur auto.

• Configurez la variable d'environnement intégrée NCCL_SOCKET_IFNAME selon votre scénario :

  • Si un Pod unique utilise 1, 2, 4 ou 8 GPU pour une tâche d'entraînement ou d'inférence, définissez NCCL_SOCKET_IFNAME=eth0. Il s'agit du paramètre par défaut dans cette image.

  • Si un Pod unique utilise les 16 GPU d'un hôte pour une tâche d'entraînement ou d'inférence, définissez NCCL_SOCKET_IFNAME=hpn0 pour utiliser HPN.