ZorBA: Zeroth-order Federated Fine-tuning of LLMs with Heterogeneous Block Activation

Ce papier présente ZorBA, un cadre de fine-tuning fédéré pour les grands modèles de langage qui utilise l'optimisation d'ordre zéro et une activation hétérogène de blocs pour réduire significativement l'usage de la VRAM et les coûts de communication tout en accélérant la convergence.

L'essentiel

🌍 Le Contexte : Une Grande Cuisine Collaborative

Imaginez que vous voulez apprendre à cuisiner un plat complexe (un Grand Modèle de Langage ou LLM, comme ceux qui écrivent des poèmes ou répondent à des questions).

Pour apprendre, vous avez besoin de beaucoup de pratique. Mais au lieu d'avoir un seul chef dans une immense cuisine, vous avez 50 cuisiniers (les clients) dispersés dans le monde, chacun avec sa propre petite cuisine et ses propres ingrédients (données). C'est ce qu'on appelle l'Apprentissage Fédéré : on apprend ensemble sans jamais partager les ingrédients bruts (les données privées).

🚧 Le Problème : La Cuisine est Trop Petite

Le problème, c'est que ce "plat" (le modèle) est gigantesque.

1. L'espace de travail (VRAM) : Pour cuisiner ce plat, il faut une énorme table de travail. La plupart des cuisiniers ont de petites tables. Si tout le monde essaie de cuisiner le plat entier en même temps, les tables débordent, et le travail s'arrête.
2. Le bruit dans les couloirs (Communication) : Pour s'entraider, les cuisiniers doivent constamment s'envoyer des notes sur ce qu'ils ont fait. Comme le plat est énorme, ces messages sont si gros qu'ils bloquent les couloirs de l'immeuble.

Les méthodes actuelles demandent aux cuisiniers de faire des calculs très complexes (rétropropagation) qui nécessitent encore plus d'espace sur la table.

💡 La Solution : ZorBA (Le Chef Intelligemment Organisé)

Les auteurs proposent ZorBA, une nouvelle méthode pour organiser cette cuisine collaborative. Voici comment ça marche, avec trois astuces principales :

1. L'Art de "Sentir" sans Calculer (Optimisation d'Ordre Zéro)

Normalement, pour savoir si une recette est bonne, un chef goûte, ajuste, et recalcule tout le processus chimique (c'est la rétropropagation). C'est lent et demande beaucoup de papier brouillon (mémoire).

ZorBA change la règle :

• Au lieu de calculer tout, le chef fait une petite expérience : il ajoute une pincée de sel (une perturbation), goûte, puis enlève la pincée et goûte à nouveau.
• En comparant les deux goûts, il devine dans quelle direction aller.
• L'avantage : Il n'a plus besoin de garder tout le papier brouillon des calculs précédents. Il économise énormément d'espace sur sa table (réduction de la VRAM).

2. La Répartition Inégale des Tâches (Activation de Blocs Hétérogène)

Le modèle est comme un livre de cuisine divisé en chapitres (les blocs de transformation).

• L'ancienne méthode : Tout le monde lit et travaille sur tous les chapitres. C'est trop lourd pour les petites tables.
• La méthode ZorBA : Le chef central (le serveur) dit :
  • "Toi, tu travailles sur les chapitres 1, 2 et 5."
  • "Toi, avec ta petite table, tu ne travailles que sur les chapitres 3 et 4."
  • "Toi, tu prends les chapitres 6 à 10."
• L'astuce : Chaque cuisinier ne travaille que sur une partie du livre qui rentre sur sa table. Mais ensemble, ils couvrent tout le livre. Cela permet d'utiliser des tables plus petites tout en allant vite.

3. Le Code Secret (Graines Aléatoires Partagées)

Pour que les cuisiniers puissent se comprendre sans s'envoyer des livres entiers, ils utilisent un code secret (une graine aléatoire partagée).

• Au lieu de s'envoyer les résultats complets, ils s'envoient juste le résultat de la "pincée de sel" (la différence de goût).
• Comme ils ont tous la même graine secrète, ils peuvent recréer exactement les mêmes expériences de leur côté.
• Résultat : Les messages envoyés sont minuscules. Les couloirs ne sont plus bloqués.

📊 Les Résultats : Pourquoi c'est génial ?

Les chercheurs ont testé ZorBA avec des modèles réels (comme OPT-125M et OPT-1.3B) sur des tâches de classification de texte.

• Économie d'espace : ZorBA utilise jusqu'à 62 % moins d'espace sur les tables (VRAM) que les méthodes classiques. C'est comme si vous pouviez cuisiner un banquet géant sur une table de cuisine de studio.
• Vitesse : Même en travaillant sur des morceaux différents, ZorBA apprend aussi vite, voire plus vite, que les autres méthodes parce que l'organisation est optimisée mathématiquement.
• Communication : Les échanges entre les cuisiniers sont réduits au minimum, évitant les embouteillages.

🏁 En Résumé

ZorBA est une méthode intelligente pour entraîner des intelligences artificielles géantes sur des appareils modestes (comme des téléphones ou des petits serveurs).

Au lieu de demander à tout le monde de soulever un éléphant (le modèle entier) en même temps, ZorBA dit : "Chacun soulève une patte, mais on le fait de manière coordonnée, sans avoir besoin de se souvenir de tout ce qu'on a soulevé avant, et en se parlant très peu."

C'est une victoire pour la vie privée, l'efficacité énergétique et l'accessibilité de l'IA.

Résumé technique

1. Problématique

Le fine-tuning (ajustement fin) collaboratif de grands modèles de langage (LLM) via l'apprentissage fédéré (FL) se heurte à deux obstacles majeurs lorsqu'il est appliqué à des modèles de grande taille :

1. Consommation excessive de VRAM : Les approches classiques de FL utilisent la rétropropagation (backpropagation) pour calculer les gradients, ce qui nécessite de stocker les activations de l'avant (forward passes) et les gradients. Pour des LLMs contenant des centaines de millions ou des milliards de paramètres, cela dépasse souvent la capacité de mémoire vidéo (VRAM) des clients aux ressources limitées.
2. Surcharge de communication : L'échange fréquent de gradients de haute dimension entre les clients et le serveur central génère une latence et une bande passante considérables.
3. Limites des méthodes d'ordre zéro existantes : Bien que l'optimisation d'ordre zéro (Zeroth-Order Optimization - ZOO) permette d'estimer les gradients sans rétropropagation (en utilisant uniquement des passes avant), les travaux antérieurs appliquent cette méthode à tous les blocs du modèle. Cela entraîne toujours une consommation de VRAM élevée (due aux activations) et une convergence lente dans les espaces de paramètres de haute dimension en raison du bruit introduit par les vecteurs de perturbation.

L'objectif est donc de concevoir un cadre de fine-tuning fédéré qui réduise la VRAM et la communication tout en maintenant une vitesse de convergence acceptable.

2. Méthodologie : Le cadre ZorBA

Les auteurs proposent ZorBA, un cadre d'optimisation d'ordre zéro pour le fine-tuning fédéré, reposant sur trois piliers principaux :

A. Optimisation d'ordre zéro (Zeroth-Order)

Au lieu de calculer les gradients par rétropropagation, ZorBA estime les gradients en utilisant des différences finies de la fonction de perte. Pour chaque itération, le client effectue des passes avant avec des vecteurs de perturbation aléatoires pour estimer la direction du gradient. Cela élimine le besoin de stocker les gradients, réduisant ainsi la charge mémoire.

B. Activation hétérogène des blocs (Heterogeneous Block Activation)

C'est l'innovation centrale. Au lieu d'activer tous les blocs de transformateurs pour tous les clients, le serveur central alloue dynamiquement des sous-ensembles différents de blocs à chaque client, en fonction de leurs contraintes de VRAM.

• Un client n'active que les blocs dont il a la capacité de gérer les activations.
• Cela permet de réduire drastiquement la VRAM utilisée par client ($\psi_{act}$) tout en assurant que l'ensemble du modèle est mis à jour collectivement.

C. Graines aléatoires partagées et réduction de la communication

Pour éviter d'échanger les vecteurs de perturbation (qui sont de haute dimension), le serveur et les clients partagent un ensemble de graines aléatoires (random seeds).

• Les clients et le serveur génèrent indépendamment les mêmes vecteurs de perturbation à partir de ces graines.
• Au lieu d'envoyer les gradients estimés, les clients n'envoient au serveur que les différences finies scalaires des pertes.
• Le serveur reconstruit les mises à jour globales en utilisant ces différences et les vecteurs de perturbation générés localement. Cela réduit considérablement la charge de communication.

3. Contributions Clés

1. Proposition de ZorBA : Un cadre intégrant l'optimisation d'ordre zéro et une stratégie d'activation sélective des blocs pour les LLMs en environnement fédéré.
2. Analyse théorique de la convergence :
   • Les auteurs dérivent une borne de convergence pour ZorBA dans un cadre non convexe.
   • Ils introduisent une métrique clé, $\Lambda(A)$, qui dépend de la matrice d'activation des blocs $A$. Ils démontrent que minimiser $\Lambda(A)$ (lié à la "popularité" minimale des blocs activés) est essentiel pour accélérer la convergence.
   • Ils montrent qu'il existe un compromis fondamental : activer plus de blocs accélère la convergence mais augmente la VRAM.
3. Algorithme d'optimisation ($\epsilon$-constraint Lexicographic) :
   • Pour résoudre le problème d'optimisation multi-objectif (minimiser $\Lambda(A)$ et la VRAM simultanément), ils proposent un algorithme en deux étapes :
     1. Maximiser la "popularité minimale" des blocs (le nombre minimum de clients activant un bloc donné) pour garantir la convergence.
     2. Ajuster l'activation des blocs supplémentaires via un algorithme gourmand (greedy) pour minimiser le nombre de clients restant à ce niveau minimal, tout en respectant les contraintes de VRAM.
   • Cela permet de trouver une solution de Pareto optimale entre performance et coût matériel.
4. Validation expérimentale : Des expériences sur des modèles OPT-125M et OPT-1.3B avec des jeux de données de classification de texte (AG-News, SST-2, SNLI).

4. Résultats Expérimentaux

Les résultats montrent que ZorBA surpasse les méthodes de base (FedIT, FedZO, DeComFL) :

• Réduction de la VRAM : ZorBA réduit l'utilisation totale de VRAM d'environ 62,41 % par rapport aux méthodes qui activent tous les blocs (FedIT, FedZO, DeComFL).
• Surcharge de communication : La communication est réduite de plusieurs ordres de grandeur (passant de $10^{11}$-$10^{13}$ paramètres transmis à environ $10^5$) grâce à l'échange de différences finies au lieu des gradients complets.
• Convergence :
  • ZorBA converge plus rapidement que FedZO et DeComFL (réduction du nombre de tours d'entraînement nécessaire pour atteindre la précision cible).
  • L'analyse de la courbe de Pareto (Fig. 4) confirme que l'algorithme d'optimisation des blocs permet de trouver un équilibre où la convergence reste rapide malgré une réduction significative de la VRAM.
  • Une étude de cas montre que l'équilibre de la répartition des blocs (la métrique $\Lambda$) est plus important que le nombre total de blocs activés pour la vitesse de convergence.

5. Signification et Impact

L'article ZorBA est significatif car il résout le goulot d'étranglement majeur du fine-tuning fédéré des LLMs : la limitation matérielle (VRAM).

• Accessibilité : Il permet à des clients avec des GPU peu puissants de participer au fine-tuning de modèles massifs, ce qui était auparavant impossible avec les méthodes de rétropropagation complètes.
• Efficacité : En combinant l'optimisation d'ordre zéro avec une allocation intelligente des ressources (activation hétérogène), il démontre qu'il n'est pas nécessaire de sacrifier la convergence pour économiser de la mémoire.
• Généralité : L'approche de partage de graines aléatoires et de différences finies offre un modèle de communication efficace applicable à d'autres scénarios de FL avec des modèles noirs ou à haute dimension.

En résumé, ZorBA fournit une solution théoriquement fondée et pratiquement efficace pour déployer le fine-tuning collaboratif de LLMs dans des environnements distribués aux ressources contraintes.