Ajouter une hiérarchie fédérée à l'entraînement shardé

Le réseau devient le goulot

Le sharded data parallelism répartit paramètres et données entre GPU, mais ses synchronisations deviennent coûteuses quand le cluster mélange plusieurs niveaux d'interconnexion. Le problème n'est plus seulement de calculer les gradients : il faut limiter les échanges qui traversent les groupes de machines.

Fédérer des groupes de shards

Gianluca Mittone et Marco Aldinucci proposent FL+FSDP et FL+HSDP. Les GPU sont organisés en groupes de fédération ; chaque groupe effectue son entraînement shardé, puis les groupes échangent des agrégations inspirées de FedAvg. Cette hiérarchie réduit le trafic inter-groupe et borne la croissance du batch global par la taille des groupes.

Un test sur Llama

Sur un pré-entraînement Llama3.1 8B exécuté sur 512 GPU A100, les auteurs rapportent jusqu'à 8,04 fois plus de traitement de données et une perplexité d'évaluation 4,48 fois plus faible que les méthodes comparées, à hyperparamètres identiques. L'analyse de communication et les expériences visent à montrer que le gain vient de la topologie d'échanges, pas d'un simple changement de modèle.

Limites et statut de la source

Les résultats dépendent du réseau, du découpage des groupes, de la fréquence d'agrégation et de la charge Llama étudiée. Une hiérarchie qui réduit le trafic peut aussi modifier la dynamique de convergence ; les auteurs prévoient d'étendre l'approche à d'autres axes de parallélisme et d'automatiser le réglage selon la topologie. Le document est un preprint ArXiv.

Sources

Gianluca Mittone et Marco Aldinucci, « Accelerating Sharded Data Parallelism at Scale with Federated Learning », ArXiv 2609.20359v1 : résumé et preprint, PDF.