Utiliser des LLM locaux pour découvrir des algorithmes hautes performances

Depuis que je suis enfant, je suis fasciné par le dessin. Ce qui m’a frappé n’était pas seulement l’acte de dessiner lui-même, mais aussi l’idée que chaque dessin pouvait être amélioré de plus en plus. Je me souviens avoir atteint des niveaux très élevés avec mon style de dessin. Cependant, une fois atteint le sommet de la perfection, j’essayais de voir comment je pourrais améliorer encore le dessin – hélas, avec des résultats désastreux.

À partir de là, je garde toujours en tête le même mantra : « affinez et itérez et vous atteindrez la perfection ». À l’université, mon approche consistait à lire des livres plusieurs fois, élargissant mes connaissances en recherchant d’autres sources, en trouvant des couches cachées de sens dans chaque concept. Aujourd’hui, j’applique cette même philosophie à l’IA/ML et au codage.

Nous savons que la multiplication matricielle (matmul pour plus de simplicité ici) est la partie essentielle de tout processus d’IA. Dans le passé, j’ai développé LLM.rust, un miroir Rust de Karpathy. LLM.c. Le point le plus difficile de l’implémentation de Rust a été la multiplication matricielle. Puisque nous devons effectuer des milliers d’itérations pour affiner un modèle basé sur GPT, nous avons besoin d’une opération matmul efficace. Pour cela, j’ai dû utiliser la bibliothèque BLAS, implémentant un unsafe stratégie pour surmonter les limites et les obstacles. L’utilisation de unsafe in Rust est contraire à la philosophie de Rust, c’est pourquoi je recherche toujours des méthodes plus sûres pour améliorer matmul dans ce contexte.

Ainsi, en m’inspirant de la déclaration de Sam Altman – « demandez à GPT comment créer de la valeur » – j’ai décidé de demander aux LLM locaux de générer, de comparer et d’itérer sur leurs propres algorithmes pour créer une meilleure implémentation native de Rust matmul.

Le défi comporte certaines contraintes :

• Nous devons utiliser notre environnement local. Dans mon cas, un MacBook Pro, M3, 36 Go de RAM ;
• Dépassez les limites des jetons ;
• Chronométrez et comparez le code dans la boucle de génération elle-même

Je sais qu’il est presque impossible d’atteindre des performances de niveau BLAS avec cette méthode, mais je tiens à souligner comment nous pouvons exploiter l’IA pour répondre à des besoins personnalisés, même avec nos « petits » ordinateurs portables, afin de pouvoir débloquer des idées et repousser les limites dans n’importe quel domaine. Cet article veut être une source d’inspiration pour les praticiens et les personnes qui souhaitent se familiariser avec Microsoft Autogen et le déploiement LLM local.

Toute l’implémentation du code peut être trouvée dans ceci Dépôt Github. Il s’agit d’une expérience en cours et de nombreux changements/améliorations seront apportés.

Idée générale

L’idée générale est d’avoir une table ronde d’agents. Le point de départ est le MrAderMacher Mixtral 8x7B modèle Q4 K_M modèle local. A partir du modèle nous créons 5 entités :

• le Proposer propose un nouvel algorithme de type Strassen, pour trouver un moyen meilleur et plus efficace d’effectuer le matmul ;
• le Verifier examine la formulation matmul à l’aide de mathématiques symboliques ;
• le Coder crée le code Rust sous-jacent ;
• le Tester l’exécute et enregistre toutes les informations dans la base de données vectorielles ;
• le Manager agit silencieusement, contrôlant le flux de travail global.

Agent	Fonction de rôle
Proposant	Analyse les temps de référence et propose de nouveaux paramètres de réglage et formulations matmul.
Vérificateur	(Actuellement désactivé dans le code). Il vérifie la formulation mathématique du proposant par une vérification symbolique.
Codeur	Il prend les paramètres et élabore le code du modèle Rust.
Testeur	Il exécute le code Rust, enregistre le code et calcule le timing de référence.
Directeur	Contrôle global du flux de travail.

Le flux de travail global peut être orchestré via Microsoft Autogen, comme illustré sur la fig.1.

Préparer les données d’entrée et la base de données vectorielles

Les données d’entrée sont collectées à partir de tous les articles académiques, axés sur l’optimisation de la multiplication matricielle. Beaucoup de ces articles sont référencés et liés à, Article Strassen de DeepMind. Je veux commencer simplement, j’ai donc rassemblé 50 articles, publiés de 2020 à 2025, qui traitent spécifiquement de la multiplication matricielle.

Ensuite, j’ai utilisé chroma pour créer la base de données vectorielles. L’aspect critique dans la génération d’une nouvelle base de données vectorielles est la façon dont les PDF sont fragmentés. Dans ce contexte, j’ai utilisé un sémantique morceaur. Contrairement aux méthodes de texte fractionné, le chunker sémantique utilise la signification réelle du texte pour déterminer où couper. L’objectif est de conserver les phrases associées ensemble en un seul morceau, ce qui rend la base de données vectorielles finale plus cohérente et précise. Cela se fait en utilisant le modèle local BAAI/bge-base-en-v1.5. L’essentiel de Github ci-dessous montre la mise en œuvre complète.

Le code de base : autogen-core et modèles GGML

J’ai utilisé Microsoft Autogen, en particulier le autogen-core variante (version 0.7.5). Contrairement au chat de niveau supérieur, dans autogen-core nous pouvons avoir accès à des éléments de base pilotés par des événements de bas niveau, qui sont nécessaires pour créer un flux de travail piloté par une machine à états selon nos besoins. En fait, le défi consiste à maintenir un flux de travail strict. Tous les agents agissant doivent agir dans un ordre précis : Proposer -> Verifier -> Coder -> Tester.

La partie centrale est la BaseMatMulAgentqui hérite de celui d’AutoGen RoutedAgent. Cette classe de base nous permet de standardiser la manière dont les agents LLM participeront au chat et se comporteront.

À partir du code ci-dessus, nous pouvons voir que la classe est conçue pour participer à une discussion de groupe asynchrone, en gérant l’historique des conversations, les appels à des outils externes et en générant des réponses via le LLM local.

Le composant principal est @message_handlerun décorateur qui enregistre une méthode comme listener ou subscriber en fonction du type de message. Le décorateur détecte automatiquement l’indice de type de l’argument de la première méthode – dans notre cas, c’est message: GroupChatMessage. Il abonne ensuite l’agent pour recevoir tous les événements de ce type envoyés au sujet de l’agent. Le handle_message La méthode async se charge alors de mettre à jour la mémoire interne de l’agent, sans générer de réponse.

Une fois le mécanisme écouteur-abonné en place, nous pouvons nous concentrer sur la classe Manager. Le MatMulManager hérite RoutedAgent et orchestre le flux global des agents.

Le code ci-dessus gère tous les agents. Nous sautons le Verifier partie, pour le moment. Le Coder publier le code final, et le Tester se charge de sauvegarder à la fois le code et l’ensemble du contexte dans la base de données vectorielles. De cette façon, nous pouvons éviter de consommer tous les tokens de notre modèle local. A chaque nouvelle exécution, le modèle rattrapera les derniers algorithmes générés à partir de la base de données vectorielles et proposera une nouvelle solution.

Une mise en garde très importante, pour être sûr autogen-core peut travailler avec llama modèles sur MacOS, utilisez l’extrait suivant :

#!/bin/bash 

CMAKE_ARGS="-DGGML_METAL=on" FORCE_CMAKE=1 pip install --upgrade --verbose --force-reinstall llama-cpp-python --no-cache-dir

La figure 2 résume l’intégralité du code. Nous pouvons grossièrement subdiviser le code en 3 blocs principaux :

• Le BaseAgentqui gère les messages via les agents de LLM, en évaluant la formulation mathématique et en générant du code ;
• Le MatMulManager orchestre l’ensemble du flux des agents ;
• autogen_core.SingleThreadedAgentRuntime nous permet de faire de l’ensemble du flux de travail une réalité.

Résultats et benchmark

Tout le code Rust a été révisé et réexécuté manuellement. Bien que le flux de travail soit robuste, travailler avec des LLM nécessite un œil critique. Plusieurs fois, le modèle a fabulé^*générant un code qui semblait optimisé mais qui n’a pas réussi à effectuer le travail matmul réel.

La toute première itération génère une sorte d’algorithme de type Strassen (code « Run 0 » dans la fig.3) :

Le modèle pense à de meilleures implémentations, plus de type Rust-NEON, de sorte qu’après 4 itérations il donne le code suivant (« Run 3 » sur la fig.3) :

Nous pouvons voir l’utilisation de fonctions comme vaddq_f32instruction CPU spécifique pour les processeurs ARM, provenant de std::arch::aarch64. Le modèle parvient à utiliser rayon pour diviser le flux de travail sur plusieurs cœurs de processeur et, à l’intérieur des threads parallèles, il utilise les éléments intrinsèques de NEON. Le code lui-même n’est pas totalement correct, de plus, j’ai remarqué que nous rencontrons une erreur de mémoire insuffisante lorsque nous traitons des matrices 1024×1024. J’ai dû retravailler manuellement le code pour le faire fonctionner.

Cela nous ramène à notre mantra « itérer à la perfection », et nous pouvons nous demander : « un agent local peut-il affiner de manière autonome le code Rust au point de maîtriser les intrinsèques complexes de NEON ? Les résultats montrent que oui, même sur du matériel grand public, ce niveau d’optimisation est réalisable.

La figure 3 montre les résultats finaux que j’ai obtenus après chaque itération.

Les 0ème et 2ème benchmark comportent quelques erreurs, car il est physiquement impossible d’obtenir de tels résultats sur un matmul 1024×1024 sur un CPU :

• le premier code souffre d’une erreur diagonale, donc le code calcule uniquement les blocs diagonaux de la matrice et ignore le reste ;
• le deuxième code a un tampon cassé, car il écrase à plusieurs reprises un petit tampon chaud de cache de 1028 flotteurs, plutôt que de parcourir l’intégralité d’un million d’éléments.

Cependant, le code a produit deux vrais codes, le run 1 et le run 3. La première itération atteint 760 ms, et elle constitue une véritable baseline. Il souffre de manques de cache et de manque de vectorisation SIMD. Le run 3 enregistre 359 ms, l’amélioration est la mise en œuvre du parallélisme NEON SIMD et Rayon.

* : J’ai écrit « le modèle confabule » à des fins. D’un point de vue médical, tous les LLM ne sont pas hallucinants, mais fabulateurs. Les hallucinations sont une situation totalement différente de ce que font les LLM lorsqu’ils babillent et génèrent de « fausses » réponses.

Conclusions

Cette expérience a commencé avec une question qui semblait un défi impossible : « pouvons-nous utiliser des LLM locaux grand public pour découvrir des algorithmes Rust hautes performances capables de rivaliser avec l’implémentation de BLAS ?

Nous pouvons dire oui, ou du moins nous avons une expérience valide et solide, où nous pouvons créer un meilleur code pour obtenir un code complet de type BLAS dans Rust.

L’article montrait comment interagir avec Microsoft Autogen, autogen-coreet comment créer une table ronde d’agents.

Le modèle de base utilisé provient de GGUF et peut fonctionner sur un MacBook Pro M3 de 36 Go.

Bien sûr, nous n’avons pas (encore) trouvé quelque chose de mieux que BLAS dans un seul code simple. Cependant, nous avons prouvé que le flux de travail agent local, sur un MacBook Pro, peut réaliser ce que l’on pensait auparavant nécessiter un cluster et des modèles massifs. Finalement, le modèle a réussi à trouver une implémentation Rust-NEON raisonnable, « Run 3 ci-dessus », qui a une accélération de plus de 50 % par rapport à l’implémentation Rayon standard. Nous devons souligner que la mise en œuvre du backbone a été générée par l’IA.

La frontière est ouverte. J’espère que cet article de blog pourra vous inspirer pour essayer de voir quelles limites nous pouvons surmonter avec le déploiement local de LLM.

---

J’écris ceci à titre personnel ; ces opinions sont les miennes.