Container Compute Service:training-nv-pytorch 26.05

Fonctionnalités principales

• Mise à niveau des composants liés à Torch vers la version 2.11.0.

• Mise à niveau des composants d'entraînement intégrés : megatron-core passe en version 0.17.0 et transformer_engine en version 2.14.

• Mise à niveau du composant d'inférence vllm vers la version 0.20.2.

Corrections de bogues

Aucune correction pour cette version.

Ressources

Prérequis du pilote

• La version 26.05 prend en charge CUDA 13.0.2, ce qui nécessite un pilote NVIDIA version 580 ou ultérieure. Pour consulter la liste complète des pilotes compatibles, reportez-vous au guide CUDA Application Compatibility. Pour plus d'informations, voir CUDA Compatibility and Upgrading.

Optimisation de la compilation PyTorch

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

• Contrôle de la granularité des communications dans le framework DeepSpeed. Cela permet au compilateur de capturer un graphe de calcul plus complet et d'appliquer des optimisations de compilation plus larges.

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

  • Le frontend du compilateur PyTorch a été 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 ont été renforcées afin d'améliorer les performances post-compilation.

Grâce à ces optimisations, l'entraînement de LLM à 8 milliards de paramètres obtient 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 s'appuie sur des données de performance à grande échelle. Celles-ci incluent divers modèles, clusters et paramètres d'entraînement, ainsi que des métriques système comme l'utilisation de la mémoire GPU collectée lors des benchmarks. Ce modèle recommande le nombre optimal de couches de recalcul des activations. Intégrée directement dans PyTorch, cette approche vous permet de bénéficier des gains de performance liés à l'optimisation mémoire avec un effort minimal. Cette fonctionnalité est désormais prise en charge par le framework DeepSpeed.

Évaluation des performances E2E

À l'aide de CNP, un outil cloud natif de benchmarking des performances IA, nous avons réalisé des comparaisons de performances de bout en bout entre notre image et l'image de base standard, en utilisant des modèles open source courants. Nous avons également mené des études d'ablation afin de mesurer la contribution de chaque composant optimisé à la performance globale d'entraînement.

Comparaison Image vs image de base & évaluation itérative

Analyse de la contribution des composants GPU principaux aux performances E2E

Les tests suivants ont été effectués sur un cluster GPU multi-nœuds pour comparer les performances d'entraînement E2E. Les configurations comparées sont les suivantes :

1. Base : L'image NGC PyTorch.

2. Image ACS AI : Base+ACCL : L'image de base intégrant la bibliothèque de communication ACCL.

3. Image ACS AI : AC2+ACCL : L'image de référence utilisant AC2 BaseOS, sans aucune optimisation activée.

4. Image ACS AI : AC2+ACCL+CompilerOpt : L'image de référence basée sur AC2 BaseOS avec uniquement l'optimisation torch.compile activée.

5. Image ACS AI : AC2+ACCL+CompilerOpt+CkptOpt : L'image de référence sous AC2 BaseOS combinant les optimisations torch.compile et checkpointing sélectif du gradient.

1. Sélectionner une image

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

2. Activer le compilateur et le checkpointing via l'API

• Activer l'optimisation de compilation

  Utilisez l'API Trainer de transformers :

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

• Activer l'optimisation mémoire par checkpointing du gradient

  export CHECKPOINT_OPTIMIZATION=true

3. Démarrer le conteneur

L'image intègre un outil d'entraînement de modèles nommé ljperf. L'exemple suivant montre comment démarrer un conteneur et exécuter une tâche d'entraînement.

# Start and enter the container
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. Recommandations d'utilisation

• Cette image contient des modifications apportées à des bibliothèques telles que PyTorch et DeepSpeed. Ne les réinstallez pas.

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

• La variable d'environnement intégrée NCCL_SOCKET_IFNAME doit être ajustée dynamiquement selon le scénario 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. Il s'agit de la configuration par défaut de cette image.

  • Lorsqu'un Pod unique sollicite les 16 GPU d'une machine pour l'entraînement ou l'inférence, configurez NCCL_SOCKET_IFNAME=hpn0 afin d'utiliser le réseau haute performance (HPN).