Container Compute Service:training-nv-pytorch 25.06

Nouvelles fonctionnalités

• Cette version met à niveau PyTorch et ses composants associés vers la v2.7.1.8, Triton vers la v3.3.0 et vLLM vers la v0.9.1. Elle ajoute également la prise en charge de la nouvelle architecture Blackwell.

Corrections de bugs

• Mise à niveau de PyTorch vers la v2.7.1.8 pour corriger la dégradation des performances d'optimisation de la VRAM observée dans les images précédentes.

25.06

• egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com/egslingjun/training-nv-pytorch:25.06-serverless

acs-registry-vpc.{region-id}.cr.aliyuncs.com/egslingjun/{image:tag}

Optimisation de la compilation PyTorch

Introduite dans PyTorch 2.0, la fonction torch.compile() offre souvent des gains de performance significatifs pour les charges de travail mono-GPU à petite échelle. Cependant, comme l'entraînement de LLM repose 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ôle de la granularité de communication dans le framework DeepSpeed. Cela permet au compilateur de capturer un graphe de calcul plus complet et d'appliquer des optimisations plus étendues.

• Utilisation d'une version optimisée de PyTorch :

  • Amélioration du frontend du compilateur PyTorch pour garantir la réussite de la compilation même en cas de rupture de graphe.

  • Renforcement de la correspondance de motifs et de la prise en charge des formes dynamiques afin d'améliorer les performances après 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 de 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. Le framework DeepSpeed prend désormais en charge cette fonctionnalité.

ACCL

ACCL est la bibliothèque de communication haute performance d'Alibaba Cloud conçue pour Lingjun. ACCL-N en est la version dédiée aux GPU. Personnalisée à partir de NVIDIA NCCL, cette bibliothèque de communication haute performance est entièrement compatible avec NCCL. Elle corrige des problèmes présents dans la version upstream de NCCL et inclut des améliorations de performance et de stabilité.

Évaluation des performances de bout en bout

L'outil d'analyse des performances IA de la plateforme Cloud Native Platform (CNP) a permis de réaliser des comparaisons complètes des performances de bout en bout. Cette image a été comparée à une image de base standard exécutant des modèles open source courants et des configurations de framework standards. Des études d'ablation ont également évalué la contribution de chaque composant d'optimisation aux performances globales de l'entraînement du modèle.

Contribution des composants principaux du GPU aux performances E2E

Les tests suivants ont été exécutés sur un cluster GPU multi-nœuds pour comparer les performances d'entraînement de bout en bout selon les configurations suivantes :

1. Base : Image NGC PyTorch

2. Base+ACCL : Image de base avec ACCL.

3. Image IA ACS : AC2+ACCL : AC2 BaseOS avec ACCL et aucune autre optimisation.

4. Image IA ACS : AC2+ACCL+CompilerOpt : AC2 BaseOS avec ACCL et optimisation torch.compile.

5. Image IA ACS : AC2+ACCL+CompilerOpt+CkptOpt : AC2 BaseOS avec ACCL, torch.compile et optimisations sélectives de checkpointing de gradient.

1. Récupérer l'image

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

2. Activer le compilateur et le checkpointing de gradient sélectif

• Activer l'optimisation de compilation

  Utilisez l'API Transformers Trainer :

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

• Activer le checkpointing de gradient sélectif

  export CHECKPOINT_OPTIMIZATION=true

3. Démarrer le conteneur

Cette section utilise l'outil d'entraînement de modèle intégré ljperf pour décrire les étapes de démarrage d'un conteneur et d'exécution d'une tâche d'entraînement.

# Start the container and enter its command line
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 

4. Remarques d'utilisation

• L'image contient des versions modifiées de bibliothèques telles que PyTorch et DeepSpeed. Ne les réinstallez pas.

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

• Ajustez dynamiquement la variable d'environnement intégrée NCCL_SOCKET_IFNAME de cette image selon le cas d'utilisation :

  • Si un Pod unique demande seulement 1, 2, 4 ou 8 GPU pour des tâches d'entraînement ou d'inférence, définissez NCCL_SOCKET_IFNAME=eth0 (configuration par défaut dans cette image).

  • Si un Pod unique demande l'ensemble des 16 GPU d'une machine entière pour des tâches d'entraînement ou d'inférence (permettant l'utilisation du réseau haute performance HPN), définissez NCCL_SOCKET_IFNAME=hpn0.