Container Compute Service:training-nv-pytorch 25.03

Fonctionnalités principales

• L'image de base est mise à jour pour s'aligner sur NGC 25.02.

• PyTorch et ses composants associés passent en version 2.6.0.7, TransformerEngine (TE) en 2,1 et accelerate en 1.5.2.

• ACCL-N passe en version 2.23.4.12.

• vLLM passe en version 0.8.2.dev0 et Ray en 2.44, avec la prise en charge de flash-infer 0.2.3. Transformers passe en version 4.49.0+ali et flash_attn en 2.7.2.

Optimisation du compilateur PyTorch

Bien que la fonctionnalité torch.compile() introduite dans PyTorch 2.0 offre des gains significatifs pour les charges de travail sur GPU unique, ces avantages diminuent lors de l'entraînement de grands modèles de langage (LLM) en raison des exigences d'optimisation de la mémoire GPU et des frameworks d'entraînement distribué tels que FSDP et DeepSpeed. Cette version relève ces défis en :

• Contrôlant la granularité des communications au sein du framework DeepSpeed, ce qui permet au compilateur d'obtenir un graphe de calcul plus complet et d'appliquer des optimisations plus larges.

• Intégrant une version optimisée de PyTorch qui :

  • Améliore le frontend du compilateur PyTorch afin de garantir une compilation réussie même en présence de ruptures de graphe.

  • Renforce les capacités de correspondance de motifs et de gestion des formes dynamiques pour accroître les performances du code compilé.

Conjointement, ces optimisations augmentent le débit E2E d'environ 20 % pour l'entraînement de LLM de 8B.

Optimisation de la mémoire GPU pour le recalcul

Nous avons développé un modèle prédictif de consommation de mémoire GPU fondé sur de vastes données de performance issues de divers modèles, clusters et configurations d'entraînement. Ce modèle recommande le nombre optimal de couches de recalcul d'activation. Intégrée directement à PyTorch, cette fonctionnalité vous permet de tirer facilement parti des avantages de performance liés à l'optimisation de la mémoire GPU. Elle est actuellement adaptée au framework DeepSpeed.

Bibliothèque de communication ACCL

ACCL est une bibliothèque de communication réseau haute performance développée par Alibaba Cloud pour les produits Lingjun. La variante ACCL-N est conçue spécifiquement pour les charges de travail GPU. Personnalisée à partir de NVIDIA NCCL, ACCL-N est entièrement compatible avec la bibliothèque d'origine tout en intégrant des corrections de bogues, des optimisations de performance et des améliorations de stabilité.

Évaluation des gains de performance E2E

À l'aide de l'outil d'évaluation des performances IA cloud-native CNP, nous avons mené une comparaison complète des performances E2E entre cette image et une image de base standard, en utilisant des modèles open source courants et des configurations de framework standards. Nous avons également réalisé une étude d'ablation afin d'évaluer la contribution spécifique de chaque composant optimisé aux performances globales.

Contribution des composants principaux GPU à la performance E2E

Les tests suivants, basés sur la version 25.03, ont été effectués sur un cluster GPU multi-nœuds afin d'évaluer et de comparer les performances E2E. Les éléments comparés incluent :

1. Base : Image NGC PyTorch

2. Image IA ACS : Base+ACCL : L'image de base utilisant la bibliothèque de communication ACCL.

3. Image IA ACS : AC2+ACCL : Cette image utilisant AC2 BaseOS sans aucune optimisation activée.

4. Image IA ACS : AC2+ACCL+CompilerOpt : Cette image utilisant AC2 BaseOS avec uniquement l'optimisation torch.compile activée.

5. Image IA ACS : AC2+ACCL+CompilerOpt+CkptOpt : Cette image utilisant AC2 BaseOS avec les optimisations torch.compile et checkpointing sélectif du gradient activées.

Étape 1 : Récupérer l'image

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

Étape 2 : Activer le compilateur et les optimisations de mémoire pour le recalcul

• Activer l'optimisation du compilateur

  Lorsque vous utilisez l'API Transformers Trainer, définissez torch_compile=True :

  training_args = TrainingArguments(
      bf16=True,
      gradient_checkpointing=True,
      torch_compile=True
  )

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

  export CHECKPOINT_OPTIMIZATION=true

Étape 3 : Lancer le conteneur

L'image inclut l'outil d'entraînement de modèles ljperf. Suivez les étapes ci-dessous pour lancer le conteneur et exécuter une tâche d'entraînement.

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

# Run a training demo
ljperf --action train --model_name deepspeed/llama3-8b 

Étape 4 : Remarques d'utilisation

• L'image contient des versions modifiées de bibliothèques telles que PyTorch et DeepSpeed. Ne les réinstallez pas, car cela écraserait les optimisations incluses. Ces modifications seront intégrées en amont ultérieurement.

• Dans la configuration DeepSpeed, laissez le paramètre zero_optimization.stage3_prefetch_bucket_size vide ou définissez-le sur auto.