Planifier une exécution PyTorch reproductible

Prouver une étape complète avant de passer à l'échelle

Figez les révisions du modèle et du jeu de données, les versions des dépendances, la configuration d'entraînement, les graines aléatoires et le condensé du conteneur. Exécutez le chargement du modèle, une étape d'entraînement complète, l'évaluation, la sauvegarde du point de contrôle et la restauration du point de contrôle sur un accélérateur.

Utilisez la même longueur de séquence et les mêmes réglages de micro-lot que ceux prévus pour l'exécution plus longue. Un petit pilote qui omet l'état de l'optimiseur ou utilise des séquences plus courtes peut masquer le problème de mémoire que vous essayez de mesurer.

Mesurer une étape d'optimisation représentative

Dans votre programme d'entraînement existant, entourez une étape représentative d'une mesure de mémoire de pointe. L'exemple suppose que le modèle, le lot et l'optimiseur ont déjà été initialisés sur le GPU. Incluez l'étape de l'optimiseur car certains états sont créés de manière paresseuse.

torch.cuda.reset_peak_memory_stats()
optimizer.zero_grad(set_to_none=True)
loss = model(**batch).loss
loss.backward()
optimizer.step()
torch.cuda.synchronize()
peak_gib = torch.cuda.max_memory_allocated() / 2**30
print(f"Peak tensor memory: {peak_gib:.2f} GiB")

Ceci indique l'allocation maximale de tenseurs, pas l'empreinte complète sur l'appareil. Inspectez également la mémoire réservée, la surcharge d'exécution et le rapport de mémoire de l'appareil du système.

Choisissez DDP ou le sharding pour la bonne raison

DistributedDataParallel conserve une réplique du modèle sur chaque processus et synchronise les gradients. Il est utile lorsque l'état d'entraînement complet tient sur chaque GPU et que vous souhaitez traiter davantage de données en parallèle ; il ne fusionne pas la VRAM en un seul pool contigu.

Fully Sharded Data Parallel peut distribuer les paramètres, les gradients et l'état de l'optimiseur. Utilisez-le lorsque la pression mémoire justifie le partitionnement, et tenez compte de la communication supplémentaire et de la complexité des points de contrôle distribués.

La précision mixte et le point de contrôle d'activation sont des choix distincts. Validez le comportement numérique pour une précision inférieure ; mesurez le coût de calcul du recalcul des activations lorsque le point de contrôle économise de la mémoire.

Lancer un pilote sur un seul nœud

Le script d'entraînement doit initialiser l'exécution distribuée et assigner un processus à chaque GPU à l'aide des informations de rang local. Faites correspondre le nombre de processus aux GPU visibles. Ce qui suit est un exemple de lancement pour un script compatible distribué, pas une implémentation d'entraînement complète.

torchrun --standalone --nproc-per-node=4 \
  train.py --config configs/pilot.yaml

Exécutez ceci uniquement sur un nœud provisionné. La démo locale AnchorGPU n'exécute pas d'entraînement et n'alloue pas de GPU distants.

Sélectionnez le matériel à partir du résultat du pilote

Utilisez l'A100 comme référence 80 GB CUDA à prix réduit, et testez le H100 si la charge de travail dispose d'un chemin spécifique Hopper utile. Le H200 offre une configuration mémoire CUDA plus grande. Le MI300X offre davantage de capacité par accélérateur dans ce catalogue, à condition que l'application complète soit qualifiée sur ROCm.

Pour AMD, vérifiez les extensions personnalisées, les noyaux, les versions des paquets et l'image PyTorch testée avec ROCm comme un ensemble complet. La compatibilité au niveau du framework ne valide pas toutes les opérations optionnelles.

Testez la restauration, pas seulement la sauvegarde

Enregistrez le modèle, l'optimiseur, l'ordonnanceur, l'étape d'entraînement, la configuration et l'état du scaler de précision mixte lorsque cela est pertinent. Enregistrez un point de contrôle, terminez le processus, puis restaurez et reprenez avant de vous engager dans l'exécution longue.

Utilisez une méthode de point de contrôle adaptée à l'état distribué. Testez que les artefacts enregistrés peuvent être lus indépendamment du processus d'origine et que la sortie d'évaluation reste dans la tolérance attendue.

Budgétisez l'ensemble du pipeline

Incluez les téléchargements, le prétraitement, le chargement des données, la validation, la création de points de contrôle, l'exportation et la copie des sorties hors du nœud. Mesurez le débit avec la mémoire, les blocages de chargement des données, le temps de communication, la durée des points de contrôle et la qualité de la validation.

Un terme de 7 jours est utile pour une passe de compatibilité et de profilage délimitée. Choisissez une réservation de 30 jours seulement après avoir pu estimer le travail utile et la marge opérationnelle. Davantage de GPU n'aident que lorsque le calcul supplémentaire dépasse les limites de communication et de pipeline d'entrée.