Container Compute Service:training-nv-pytorch 26.03

Points forts

• Mise à jour de torch vers la version 2.10.

• Mise à jour de vllm vers la version 0.17.0.

• Mise à jour de megatron-core vers la version 0.16.0.

• Mise à jour de deepspeed vers la version 0.18.8.

• Mise à jour de transformer_engine vers la version 2.12.

Corrections de bugs

Aucune.

Composants principaux

Images publiques

Images VPC

Remplacez l'URI de ressource d'image de conteneur IA egslingjun-registry.cn-wulanchabu.cr.aliyuncs.com/egslingjun/{image:tag} spécifié dans le fichier YAML de la console ACS par acs-registry-vpc.{region-id}.cr.aliyuncs.com/egslingjun/{image:tag} afin de télécharger rapidement les images de conteneur IA PG1 via le VPC.

• Où {region-id} correspond à la région disponible dans laquelle votre service ACS est activé, par exemple cn-beijing et cn-wulanchabu.

• {image:tag} représente le nom et le tag de l'image.

Étape 1 : Télécharger l'image

Exécutez la commande suivante dans le shell de l'hôte :

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

Remplacez [tag] par le tag correspondant à votre version de NVIDIA Driver et à l'architecture de votre processeur, soit 26.03-cu130-serverless soit 26.03-cu128-serverless. Pour connaître la version du pilote et les architectures prises en charge par chaque tag, consultez la section Contenu de l'image.

Étape 2 : Démarrer le conteneur

Exécutez la commande suivante dans le shell de l'hôte. Cette commande démarre le conteneur et ouvre un shell de conteneur, dans lequel vous effectuerez les étapes restantes.

# 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]

Étape 3 : Activer les optimisations

Activez les deux optimisations avant le début de l'entraînement. Définissez l'optimisation de compilation dans votre script d'entraînement et l'optimisation de mémoire par recomputation dans le shell du conteneur ouvert à l'étape 2.

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

export CHECKPOINT_OPTIMIZATION=true

Étape 4 : Exécuter la démonstration d'entraînement

L'image inclut ljperf, un outil intégré d'entraînement de modèles. Exécutez la démonstration dans le shell du conteneur.

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

Recommandations d'utilisation

• Cette image contient des bibliothèques modifiées, 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.

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ôle de 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.

• Utilisation d'une build PyTorch optimisée :

  • Le frontend du compilateur PyTorch a été 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 ont été 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 la recomputation

Un modèle prédictif de la surcharge de mémoire GPU, construit à partir de données de performance à 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 recomputation d'activations. 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.