Container Compute Service:training-nv-pytorch 26.02

Principales fonctionnalités

• Mise à niveau de DeepSpeed vers la version 0.18.5 et de Transformers vers la version 4.57.6.

• Mise à niveau de vLLM vers la version 0.15.0 et de flashinfer-python vers la version 0.6.1.

Corrections de bugs

Aucune.

Composants principaux

Images publiques

Images VPC

Remarques d'utilisation

Vérifiez les exigences suivantes avant d'exécuter l'image training-nv-pytorch :

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

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

• Ajustez la variable d'environnement intégrée NCCL_SOCKET_IFNAME de cette image en fonction de votre scénario :

  • Si un seul pod demande 1, 2, 4 ou 8 GPU pour une tâche d'entraînement ou d'inférence, vous devez définir NCCL_SOCKET_IFNAME=eth0. Il s'agit de la configuration par défaut de cette image.

  • Si un seul pod demande les 16 GPU d'un nœud pour une tâche d'entraînement ou d'inférence, vous pouvez utiliser le réseau haute performance HPN. Dans ce cas, vous devez définir NCCL_SOCKET_IFNAME=hpn0.

Étape 1 : Tirer l'image

Remplacez [tag] par le tag correspondant à la version du pilote NVIDIA et à l'architecture CPU de vos nœuds, comme indiqué dans la section Contenu.

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

Étape 2 : Démarrer le conteneur et exécuter une tâche d'entraînement

L'image training-nv-pytorch inclut un outil intégré d'entraînement de modèles nommé ljperf. L'exemple suivant utilise cet outil pour démarrer un conteneur et exécuter une tâche d'entraînement.

# Start the container and connect to it
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 3 : Activer la compilation et l'optimisation de la mémoire GPU

Activez les optimisations suivantes de l'image training-nv-pytorch :

• Activer l'optimisation de la compilation : Utilisez l'API Transformers Trainer.

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

• Activer l'optimisation de la mémoire GPU pour le recalcul : Définissez la variable d'environnement suivante.

  export CHECKPOINT_OPTIMIZATION=true

Optimisation de la compilation PyTorch

torch.compile(), introduit dans PyTorch 2.0, offre souvent des gains significatifs pour les charges de travail à petite échelle sur un seul GPU. Toutefois, l'entraînement des LLM repose sur l'optimisation de la mémoire GPU et sur des frameworks distribués tels que FSDP ou DeepSpeed ; par conséquent, torch.compile() peut offrir des avantages limités, voire dégrader les performances.

• Contrôlez la granularité de la communication dans le framework DeepSpeed. Cela aide le compilateur à capturer un graphe de calcul plus complet et à appliquer des optimisations de compilation plus larges.

• Utilisez une version optimisée de PyTorch :

  • Le frontend du compilateur PyTorch est amélioré pour garantir que la compilation réussisse même si une rupture de graphe se produit dans le graphe de calcul.

  • La correspondance de motifs et la prise en charge des formes dynamiques sont renforcées pour améliorer les performances après compilation.

Grâce à ces optimisations, l'entraînement des LLM de 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 de mémoire GPU, construit à partir de données de performances à grande échelle — incluant différents 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 tests de référence — recommande le nombre optimal de couches de recalcul d'activation. Cette approche est intégrée à PyTorch, vous permettant d'obtenir les 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.