Résumé Technique : SPOTLIGHT

Énoncé du Problème

L'apprentissage par renforcement (RL) post-entraînement pour les Diffusion Transformers (DiTs) est extrêmement coûteux, nécessitant souvent des milliers de GPU haut de gamme. Le flux de travail standard de type GRPO comprend deux phases : une phase de rollout (génération d'échantillons avec des graines aléatoires) et une phase d'entraînement (calcul des gradients et mise à jour des poids).

Les stratégies existentes de réduction des coûts font face à deux goulots d'étranglement principaux lorsqu'elles sont appliquées aux DiTs :

1. Surcharge d'exploration des graines (Seed Exploration) : Pour accélérer la convergence, des méthodes comme DanceGRPO et TreeGRPO explorent plusieurs graines aléatoires pour trouver des échantillons à fort contraste (récompenses élevées et faibles). Cependant, cette phase d'exploration se situe sur le chemin critique, ajoutant du temps de calcul à chaque itération et annulant les économies de coûts potentielles.
2. Inactivité et fragmentation des GPU Spot : Bien que les GPU spot offrent des coûts 69 à 77 % inférieurs, ils sont préemptibles et disposent d'une bande passante inter-nœuds limitée.
   • Inactivité : Contrairement aux Large Language Models (LLMs) où les latences de rollout varient (permettant un pipelining avec l'entraînement), les rollouts de DiT se terminent en un temps presque uniforme. Cela crée une barrière de synchronisation où les GPU spot restent inactifs pendant toute la phase d'entraînement.
   • Fragmentation : Les préemptions fréquentes des GPU spot brisent les groupes de Parallélisme de Séquence (Sequence Parallelism - SP), qui sont essentiels pour exécuter de grands modèles DiT. La reconfiguration de ces groupes nécessite généralement plusieurs minutes pour recharger les modèles et réinitialiser les moteurs, gaspillant ainsi des ressources importantes.

Méthodologie : SPOTLIGHT

Les auteurs présentent SPOTLIGHT, le premier système conçu pour exploiter les GPU spot pour le post-entraînement RL des DiTs en synergie entre l'exploration des graines et les ressources spot. Le système repose sur deux idées clés et trois techniques centrales.

Idées Clés

1. L'exploration tolère les poids obsolètes (Stale Weights) : Les auteurs démontrent que le classement relatif des graines aléatoires (quelles graines produisent des récompenses élevées vs faibles) reste très cohérent même en utilisant les poids du modèle de l'itération précédente plutôt que ceux de l'itération actuelle mise à jour. Cela permet à l'exploration des graines de s'exécuter sur des GPU spot en utilisant des poids obsolètes pendant que le cluster réservé effectue la phase d'entraînement, déchargeant ainsi l'exploration du chemin critique.
2. La reconfiguration SP peut réutiliser l'état : Le coût élevé de la reconfiguration des groupes SP provient de l'initialisation du scheduler CPU et du chargement de modèles distants. En découplant le scheduler et en utilisant la copie de poids intra-nœud (via NVLink) à partir de pairs colocalisés, le temps de reconfiguration SP peut être réduit de plusieurs minutes à moins d'une seconde.

Architecture Système et Techniques

SPOTLIGHT est implémenté sur la plateforme de RL open-source ROLL et emploie trois techniques principales :

1. Planificateur d'Exploration basé sur les Bandits :

   • Le système formule la planification de l'exploration comme un problème de bandit multi-bras.
   • Il sélectionne dynamiquement le nombre de séquences et les étapes de débruitage effectives (en utilisant des stratégies de mise en cache comme Teache) pour maximiser la variance de la récompense dans le budget temporel fourni par la phase d'entraînement.
   • Il utilise une stratégie d'Upper Confidence Bound (UCB) pour équilibrer l'exploration de nouvelles configurations et l'exploitation des réglages à haute variance connus.

2. Parallélisme de Séquence Élastique (Elastic SP) :

   • Scheduler Persistant : Le scheduler CPU est découplé des travailleurs GPU et reste résident sur le nœud lors des changements de SP, éliminant les coûts d'initialisation répétés.
   • Chargement de Poids Intra-Nœud : Lorsqu'un nouveau travailleur GPU rejoint un groupe SP sur le même nœud, il copie les poids d'un pair colocalisé via NVLink plutôt que de les extraire d'un nœud distant, réduisant ainsi drastiquement la latence de reconfiguration.

3. Planificateur de Requêtes Pull-Based Sensible à la Préemption :

   • Les travailleurs de rollout tirent (pull) les requêtes d'une file d'attente centralisée plutôt que de recevoir des tâches assignées statiquement, s'adaptant ainsi aux vitesses hétérogènes des travailleurs.
   • Engagement de l'État (State Commitment) : Les états intermédiaires de rollout (latents et métadonnées) sont enregistrés dans un Tensor Store persistant sur les nœuds réservés. Si un GPU spot est préempté, le système peut récupérer l'état en cours et reprendre la requête sur un autre travailleur, minimisant la perte de progression.

Résultats

Le système a été évalué sur le post-entraînement de Qwen-Image (20B) en utilisant les jeux de données DeepSeek-OCR et Geneval à des résolutions de 512×512 et 1280×1280.

• Réduction des Coûts : SPOTLIGHT réduit le coût total d'entraînement de 1,4× à 6,4× par rapport aux meilleures bases de référence (y compris les systèmes utilisant des GPU spot sans ces optimisations et les systèmes uniquement sur ressources réservées).
• Vitesse de Convergence : Il atteint le même score de validation cible 4× plus vite que les bases de référence. Cela est dû à la réduction du nombre d'itérations nécessaires (grâce à l'exploration efficace des graines) et au chevauchement de l'exploration avec l'entraînement.
• Qualité : SPOTLIGHT obtient des scores de validation supérieurs sur tous les jeux de données par rapport aux bases de référence, grâce à la sélection dynamique de graines à haut contraste.
• Efficacité :
  • La surcharge d'exploration est maintenue basse (2–3 % du temps d'itération) en intégrant l'exploration dans la fenêtre d'entraînement.
  • L'Elastic SP permet au système de suivre les changements de capacité spot en temps quasi réel, là où les approches naïves souffrent de délais de plusieurs minutes lors de la reconfiguration.
  • Le système démontre une extensibilité (scalability) quasi linéaire du débit à mesure que le nombre de GPU spot passe de 8 à 64.

Signification et Revendications

L'article affirme que SPOTLIGHT est le premier système à réussir l'exploitation des GPU spot pour le post-entraînement RL des DiTs. Sa signification réside dans :

1. Briser la Dépendance au Chemin Critique : En prouvant que l'exploration des graines peut tolérer des poids obsolètes, le système élimine le compromis entre une convergence plus rapide (via l'exploration) et la latence d'itération.
2. Résoudre l'Inactivité des GPU Spot : Il résout le problème spécifique d'inactivité des rollouts de DiT en utilisant la capacité d'entraînement disponible pour l'exploration, une synergie négligée par les systèmes précédents axés sur les LLM.
3. Robustesse à la Préemption : Grâce au SP élastique et à la persistance de l'état, il rend l'entraînement des DiT résilient aux préemptions fréquentes inhérentes aux marchés spot, transformant un pool de ressources fragmenté en une infrastructure d'entraînement viable et rentable.

Les auteurs concluent que ces contributions permettent un système de post-entraînement RL pour DiT économiquement efficace qui ne sacrifie pas la qualité d'image, rendant l'entraînement de modèles de diffusion de haute qualité plus accessible.