Résumé Technique : Parallélisme de Sous-réseau de Données (SDP)

Problématique

La montée en charge rapide des réseaux de neurones profonds a créé des demandes de mémoire sans précédent sur les accélérateurs, nécessitant souvent des stratégies d'entraînement distribué. Bien que le Parallélisme de Données (DDP) soit largement utilisé, il réplique le modèle complet sur chaque GPU, ce qui entraîne un surcoût de mémoire élevé et une synchronisation coûteuse des gradients. Inversement, le Parallélisme de Modèle (ex: Parallélisme de Pipeline, Parallélisme de Tenseur) divise le modèle entre les disposités pour faire tenir des architectures plus larges, mais introduit des coûts de communication importants pour les activations et souffre souvent de bulles de pipeline et de déséquilibres de charge.

Les solutions actuelles telles que le Parallélisme de Données avec Partage Complet (FSDP) et ZeRO réduisent la mémoire par dispositif en partageant les paramètres, les gradients et les états de l'optimiseur, mais elles engendrent encore un surcoût de communication substantiel, particulièrement lors de la synchronisation des gradients. Il reste un défi ouvert pour réduire la bande passante des activations et l'utilisation de la mémoire par nœud sans sacrifier la qualité de la convergence ou nécessiter des interconnexions à haute bande passante.

Méthodologie : Parallélisme de Sous-réseau de Données (SDP)

Les auteurs proposent le Parallélisme de Sous-réseau de Données (SDP), un cadre d'entraînement distribué qui partitionne un modèle en sous-réseaux structurés entraînés à travers des travailleurs sans échanger d'activations. Contrairement aux approches de pipeline qui divisent le calcul de manière séquentielle, le SDP assigne à chaque travailleur une portion structurellement complète du modèle (un sous-réseau) qui préserve un chemin complet de l'entrée à la perte, permettant un calcul de gradient indépendant.

Mécanismes Fondamentaux

1. Construction du Sous-réseau :

   • Au niveau des Blocs (B-SDP) : Supprime des blocs architecturaux entiers (ex: blocs Transformer ou blocs de base ResNet).
   • Au niveau des Neurones/Canaux (N-SDP) : Supprime des neurones spécifiques dans les couches entièrement connectées ou des canaux dans les couches convolutionnelles, assurant la cohérence entre les couches adjacentes.
   • Les travailleurs synchronisent les paramètres chevauchants par un moyennage par étapes.

2. Régimes de Masquage :
   L'article étudie deux stratégies de masquage complémentaires pour définir comment les sous-réseaux sont utilisés pendant l'entraînement :

   • Masquage vers l'avant (Forward-Masking) : Le modèle est évalué en utilisant des paramètres masqués ($m_{fwd} = m$). Cela supprime les paramètres, les activations et les gradients pour les composants masqués, offrant des économies substantielles de mémoire et de calcul. Cependant, les gradients reflètent la passe avant masquée.
   • Masquage vers l'arrière (Backward-Masking) : La passe avant utilise le modèle complet ($m_{fwd} = m_{uni}$), tandis que la parcimonie est appliquée uniquement lors de la rétropropagation et de l'agrégation ($m_{bwd} = m$). Cela conserve des estimations de gradient non biaisées (fondées théoriquement) et offre une base de référence rigoureuse, bien que cela économise la mémoire principalement sur les gradients et les états de l'optimiseur plutôt que sur les activations.

3. Base Théorique :
   Les auteurs fournissent une analyse de convergence pour le régime de masquage vers l'arrière dans des contextages $L$-smooth. Ils démontrent que le taux de convergence est régi par l'« erreur de masquage », définie comme la distance entre le masquage uniforme et la distribution de masque spécifique. L'analyse lie la qualité de la convergence au gap spectral du graphe de masque, montrant que les assignations de paramètres chevauchantes avec un moyennage périodique maintiennent une synchronisation partielle.

Contributions Clés

1. Nouveau Paradigme : Introduction du SDP comme un paradigme d'entraînement distribué efficace en mémoire qui distribue l'entraînement de sous-réseaux à travers les nœuds, réduisant la mémoire par dispositif de 28 % à 60 % tout en maintenant ou en améliorant les performances.
2. Garanties Théoriques : Une base théorique liant la convergence du SDP à masquage vers l'arrière à une condition de gap spectral sur le graphe de masque, prouvant que des gradients non biaisés sont maintenus sous des conditions de masquage spécifiques.
3. Composabilité : Démonstration que le SDP est orthogonal aux techniques existantes d'économie de mémoire. Il supprime physiquement les composants du modèle, réduisant simultanément les paramètres, les gradients, les états de l'optimiseur et les activations. Cela permet au SDP de se composer proprement avec le FSDP (ZeRO-3) et le Checkpointing d'Activation (AC).
4. Généralisabilité : Validation du SDP à travers diverses architectures, incluant les grands modèles de langage (LLaMA) et les modèles de classification d'images (ResNet, Swin Transformer), montrant qu'il s'agit d'une alternative "prête à l'emploi" au parallélisme de données standard.

Résultats Expérimentaux

Grands Modèles de Langage (LLaMA)

Les expériences ont été menées sur FineWeb avec des modèles LLaMA de 134M, 500M et 1B de paramètres.

• Efficacité Mémoire : À l'échelle 1B, le B-SDP avec 69 % de composants actifs ($C=0.69$) a réduit la mémoire GPU de pointe de 28 % par rapport au DDP tout en atteignant une perte de validation égale ou inférieure. À l'échelle 500M, le B-SDP avec 50 % de composants actifs a réduit la mémoire de 40 %.
• Performance Aval : Sur cinq benchmarks en aval (ARC-E, BoolQ, HellaSwag, OBQA, SciQ), toutes les variantes SDP ont égalé les performances du DDP à environ 1 point près sur les moyennes macro, certaines variantes (B-SDP) montrant même de légères améliorations sur des tâches spécifiques.
• Composabilité : La combinaison du B-SDP avec le FSDP et le Checkpointing d'Activation a atteint une utilisation de mémoire de pointe de 9,7 Go par dispositif pour le modèle 1B, une réduction de 85 % par rapport au DDP (70,9 Go) et strictement inférieure au FSDP+AC seul (11,0 Go).
• Communication : Le SDP a réduit le volume de communication par étape d'environ 28 % (déplaçant 2 253 Mo contre 3 143 Mo pour le DDP). Dans les configurations inter-nœuds, la combinaison du B-SDP et du FSDP a réduit le temps d'entraînement (wall-clock) de 19,9 % par rapport au DDP.

Classification d'Images

Les expériences ont été réalisées sur CIFAR-10 et CIFAR-100 en utilisant ResNet-18, WideResNet-18 et Swin-Tiny.

• ResNet-18 : Les configurations SDP ont atteint une précision compétitive ou supérieure au DDP en utilisant seulement 40 % du budget mémoire du DDP. Le N-SDP (niveau neurone) a montré un fort effet de régularisation, surpassant le DDP à 64 % d'utilisation de la mémoire.
• Swin Transformer : Le B-SDP a égalé la précision du DDP sur CIFAR-10 avec 68 % de la mémoire et a amélioré la précision de 2 % sur CIFAR-100 avec une réduction de mémoire de 32 %.
• Limites de Parcimonie : Bien que le masquage vers l'avant (B-SDP/N-SDP) ait performé à une parcimonie modérée, le Masquage vers l'Arrière (Bb-SDP) s'est avéré supérieur à des niveaux de parcimonie extrême (ex: $C=3/8$), empêchant l'effondrement du modèle là où les autres variantes échouaient.

Signification et Revendications

L'article affirme que le SDP offre un troisième axe pratique de parallélisme, complémentaire au parallélisme de données et de modèle. Sa principale importance réside dans sa capacité à :

1. Découpler la Mémoire de la Bande Passante : En éliminant la nécessité d'échanger des activations entre les travailleurs, le SDP atténue les goulots d'étranglement de la bande passante qui limitent souvent le parallélisme de pipeline et de tenseur.
2. Permettre une Mise à l'Échelle Adaptée aux FLOPs : Le SDP permet d'entraîner des modèles plus grands ou des séquences plus longues sous le même budget matériel. La légère réduction du taux de convergence par itération est compensée par la capacité d'augmenter le nombre d'itérations (FLOP-matched) grâce à la réduction du calcul et de la communication par étape.
3. Fournir une Solution Prête à l'Emploi : Le SDP fonctionne comme un remplacement direct du parallélisme de données standard, ne nécessitant aucun changement dans l'infrastructure distribuée sous-jacente, hormis la logique de masquage.

Les auteurs soulignent que le SDP est particulièrement précieux dans les scénarios où la mémoire et la bande passante sont les contraintes principales, offrant une voie évolutive pour entraîner de grands modèles sur des accélérateurs limités en mémoire sans la complexité de l'ordonnancement de pipeline complexe ou le surcoût de communication d'un partage complet.