Résumé technique : Adam à Décroissance Cosine-Gated (CGAD) pour DiLoCo Découplé

1. Énoncé du problème

Dans les systèmes DiLoCo Découplé (DDiLoCo) asynchrones, les workers effectuent plusieurs étapes d'optimisation locales avant d'échanger des pseudo-gradients avec un serveur de paramètres global. Un défi critique dans cette architecture est la vieillissement (staleness) : les pseudo-gradients appliqués à la ronde externe $t$ peuvent avoir été calculés $\tau$ rondes plus tôt en raison de la latence réseau, de workers traînards ou des budgets d'asynchronie.

L'optimiseur externe standard pour DiLoCo est la descente de gradient stochastique (SGD) avec momentum de Nesterov. Cette méthode applique une étape de mise à jour uniforme à chaque pseudo-gradient entrant sans conditionner sur son âge ($\tau$). À mesure que $\tau$ augmente, le tampon de momentum accumule une vue obsolète du paysage de perte. Par conséquent, l'optimiseur effectue des pas qui dépassent des directions qui ne descendent plus la perte actuelle, entraînant une instabilité ou une divergence, en particulier à mesure que l'échelle du modèle augmente. Les analyses théoriques existantes du momentum-SGD asynchrone suggèrent que les taux de convergence se dégradent avec un facteur proportionnel à $\tau_{max}^2$ (le carré du délai maximal).

2. Méthodologie : Adam à Décroissance Cosine-Gated (CGAD)

Les auteurs proposent l'Adam à Décroissance Cosine-Gated (CGAD), un remplacement direct pour l'optimiseur externe de Nesterov qui conditionne explicitement les mises à jour sur l'âge du pseudo-gradient.

Mécanisme central

CGAD modifie la mise à jour Adam standard en multipliant chaque pseudo-gradient entrant $g_f$ pour le fragment $f$ par un facteur $\sigma(\tau)$ avant qu'il n'entre dans les tampons de premier et de second moment d'Adam :
$$\sigma(\tau) = \gamma(\tau) e^{-\alpha \tau}$$

Où :

• Décroissance exponentielle ($e^{-\alpha \tau}$) : Modélise la dégradation de l'information des gradients obsolètes. À mesure que $\tau$ croît, la contribution du gradient aux tampons de momentum est réduite de manière exponentielle.
• Porte Cosine ($\gamma(\tau)$) : Une fonction de coupure lisse définie comme suit :
  $$\gamma(\tau) = \frac{1}{2} \left(1 + \cos\left(\frac{\pi \tau}{\tau_{cut}}\right)\right) \cdot \mathbb{I}[\tau \le \tau_{cut}]$$
  Cette porte est égale à 1 pour les gradients frais, décroît lissment vers 0 à un seuil de coupure choisi $\tau_{cut}$, et annule strictement les contributions pour $\tau > \tau_{cut}$.

Choix de conception clés

1. Porte pré-tampon : Le facteur d'échelle $\sigma(\tau)$ est appliqué au gradient avant qu'il ne mette à jour les tampons de premier moment ($m$) et de second moment ($v$). Cela empêche les gradients obsolètes d'augmenter le second moment ($v$), ce qui atténuerait autrement les mises à jour fraîches ultérieures.
2. Variante par fragment (PA-CGAD) : Pour les planificateurs de synchronisation partielle (par exemple, Streaming DiLoCo), la méthode suit l'âge des fragments individuels ($a_f$) et porte les mises à jour en fonction de $\max(\tau, a_f)$, permettant un modèle d'obsolescence unifié pour les délais réseau et les délais de planification.
3. Hyperparamètres : La méthode introduit deux hyperparamètres, $\alpha$ (taux de décroissance) et $\tau_{cut}$ (seuil de coupure), avec des valeurs par défaut ($\alpha=0,2, \tau_{cut}=32$) qui se transfèrent à travers les échelles de modèles.

3. Contributions théoriques

L'article fournit une borne de convergence non asymptotique pour une mise à jour adaptative à porte idéalisée sur des objectifs non convexes lisses.

• Indépendance biais-délai : Le terme de biais-délai de la borne de convergence dépend de l'hyperparamètre de porte $\alpha$ (spécifiquement $O(1/e^\alpha)$) plutôt que du délai maximal réalisé $\tau_{max}$.
• Contraste avec l'analyse standard : Les analyses standard du momentum-SGD asynchrone produisent un terme de complexité évoluant avec $\tau_{max}^2$. En remplaçant cela par un terme dépendant de $\alpha$, le mécanisme de porte découple théoriquement la pénalité de convergence de la distribution du pire délai, à condition que $\tau_{cut} \ge 1/\alpha$.
• Portée : Les auteurs notent que ceci est une motivation pour le mécanisme de porte plutôt qu'un théorème de convergence complet pour l'Adam empirique, car il suppose des pas normalisés bornés plutôt que de les dériver des statistiques empiriques.

4. Résultats expérimentaux

Des expériences ont été menées sur l'entraînement préalable de modèles de langage de type Llama à trois échelles : 25 M, 1 B et 7 B paramètres, en utilisant des délais contrôlés jusqu'à $\tau=16$.

Stabilité et convergence

• Échelle 1 B : À $\tau=16$, la recette Nesterov standard a divergé (perte $\approx 319$), tandis que CGAD a atteint une perte de 7,44, comparable à sa performance à $\tau=0$ (7,48). La méthode Nesterov a perdu toute convergence à $\tau \ge 8$.
• Échelle 25 M : CGAD et d'autres bases sensibles à l'obsolescence (Adam-Decay, SDM, Nesterov retardé) sont restés stables, tandis que Nesterov a divergé à $\tau \ge 4$. À cette échelle, les différences entre les méthodes sensibles à l'obsolescence étaient marginales.

Assurance d'échelle et robustesse

La découverte la plus significative émerge à l'échelle 7 B dans des conditions contraintes en mémoire (poids bf16, AdamW 8 bits) :

• Variance graine-à-graine : La base la plus proche, Adam-Decay (CGAD sans la coupure cosinus), a présenté une augmentation de 27 fois de l'écart-type graine-à-graine ($\sigma$) à mesure que le modèle passait de 25 M à 7 B (de 0,07 à 1,89 à $\tau=8$). En revanche, le $\sigma$ de CGAD est resté stable (environ 0,38).
• Risque unique : En utilisant « moyenne + $\sigma$ » comme proxy pour le risque de déploiement unique, le risque d'Adam-Decay à 7 B / $\tau=8$ (13,19) a dépassé la perte de niveau aléatoire (10,83) de 2,4 nats, impliquant une forte probabilité d'échec dans un seul run de production. Le risque de CGAD (11,28) est resté bien en dessous de ce seuil.
• Mécanisme d'échec : Les auteurs attribuent l'échec d'Adam-Decay à 7 B à l'accumulation du bruit d'arrondi au carré dans le tampon de second moment ($v$) provenant de gradients obsolètes qui ne sont jamais entièrement annulés. La coupure cosinus dans CGAD empêche cela en annulant strictement les contributions au-delà de $\tau_{cut}$.

Sensibilité au réglage

• Nesterov vs CGAD : Même lorsque Nesterov a été réajusté (en supprimant complètement le momentum), il n'a pas pu égaler les performances d'une configuration CGAD ajustée à des délais élevés. Nesterov ajusté nécessitait des hyperparamètres différents pour différents régimes de délai, tandis que CGAD maintenait la robustesse sur toute la plage.
• Sensibilité aux hyperparamètres : Le $\alpha$ optimal change légèrement avec l'échelle (0,2 pour 25 M, 0,4 pour 1 B à $\tau=8$), mais le $\alpha$ par défaut de 0,2 reste robuste sur tout le spectre de délais.

5. Importance et affirmations

L'article affirme que CGAD est le choix de mise à l'échelle le plus robuste pour l'entraînement DiLoCo asynchrone, abordant spécifiquement la fragilité du momentum de Nesterov sous de grands délais contrôlés.

• Assurance d'échelle : La contribution principale est la coupure cosinus, qui agit comme une « assurance d'échelle ». Alors que la décroissance exponentielle seule (Adam-Decay) suffit pour les petits modèles ou les paramètres de pleine précision, la coupure stricte est critique pour les modèles de pointe (7 B+) fonctionnant avec des états d'optimiseur compressés (8 bits/bf16). Elle empêche l'accumulation de variance obsolète qui conduit à une variance catastrophique graine-à-graine et à un échec unique.
• Compatibilité directe : CGAD ne nécessite qu'une modification d'une ligne aux implémentations Adam existantes et ajoute deux hyperparamètres avec des valeurs par défaut qui se transfèrent à travers les échelles.
• Impact pratique : Pour l'entraînement asynchrone de pointe, la méthode offre la différence entre un seul run d'entraînement qui est livré avec succès et un qui échoue (se situe au-dessus de la perte aléatoire).

Les auteurs restent modestes concernant la portée théorique, reconnaissant que la borne de convergence s'applique à une mise à jour idéalisée et que la supériorité empirique est démontrée spécifiquement au sein de la famille DiLoCo et des contextes d'entraînement préalable de type Llama. Ils notent que le nombre d'étapes internes ($H=8$) utilisé dans les expériences est inférieur aux paramètres de production DiLoCo ($H=100-500$), bien que des études de sensibilité suggèrent que le classement persiste à des $H$ plus élevés.