Intégrer les kernels TileGym à Transformers

L'objectif principal du projet TileGym est de fournir des kernels performants pour l'entraînement et l'inférence des LLM. Nous allons intégrer les kernels appropriés disponibles dans le projet TileGym aux modèles LLM fournis par la bibliothèque transformers de Hugging Face pour valider la correction fonctionnelle end-to-end et les améliorations de performance. Au lieu de modifier le code source de transformers, nous adopterons une approche de monkey-patch non-intrusive : nous remplacerons certains modules/classes/méthodes de la bibliothèque transformers qui implémentent le modèle Transformer que nous souhaitons intégrer, de sorte qu'à l'instanciation du modèle, les composants principaux du modèle seront remplacés par des implémentations TileGym. À l'exécution, le modèle invoquera réellement les kernels TileGym en arrière-plan.

Instructions

Le processus d'intégration suit un flux de travail « rechercher les exigences des kernels et fournir les offres -> proposer les candidats d'intégration de kernels -> implémenter les intégrations de kernels et vérifier -> agréger les intégrations valides ». Consultez le diagramme ci-dessous pour comprendre le processus global, puis vérifiez le texte numéroté ci-dessous pour plus de détails. Si vous avez du mal à interpréter le script Mermaid intégré, consultez l'image PNG rendue qui représente le diagramme de flux de travail exactement identique : <details>

</details>

flowchart TD
  %% Nodes are labeled ONLY by step number; read the numbered text below for details.
  %% Styling encodes who executes the step (orchestrator vs subagent).

  classDef orch fill:#E3F2FD,stroke:#1E88E5,stroke-width:2px,color:#0D47A1;
  classDef sub fill:#FFF3E0,stroke:#FB8C00,stroke-width:2px,color:#E65100;
  classDef decision fill:#E8F5E9,stroke:#43A047,stroke-width:2px,color:#1B5E20;
  classDef terminal fill:#F3E5F5,stroke:#8E24AA,stroke-width:2px,color:#4A148C;

  S([Start]):::terminal
  E([End]):::terminal

  step1[Step 1]:::orch
  join1(( )):::orch

  step2{Step 2}:::decision
  step2_1[Step 2.1]:::orch
  step2_2[Step 2.2]:::orch

  step3[Step 3]:::orch
  step3_1[Step 3.1]:::orch
  step3_2[Step 3.2]:::orch
  step3_3{Step 3.3}:::decision
  step3_4[Step 3.4]:::orch
  step3_5[Step 3.5]:::orch

  %% Explore subagents for Step 1 (delegated research; two distinct agents)
  subgraph subagent_explore_1
    %% Explore subagent A: Step 1.1
    step1_1[Step 1.1]:::sub
  end

  subgraph subagent_explore_2
    %% Explore subagent B: Step 1.2
    step1_2[Step 1.2]:::sub
  end

  %% Explore subagents (research phase)
  S --> step1
  step1 -->|parallel| step1_1
  step1 -->|parallel| step1_2
  step1_1 --> join1
  step1_2 --> join1
  join1 --> step2

  %% Plan phase branching
  step2 -->|already patched| E
  step2 -->|needs patching| step2_1 --> step2_2 --> step3

  %% Execute-and-verify phase (orchestrator)
  step3 --> step3_1 --> step3_2

  %% Code subagent per integration-plan item (sub-workflow for Step 3.2)
  subgraph subagent_code
    %% Code subagent loop (runs once per integration-plan item)
    step3_2_1[Step 3.2.1]:::sub
    step3_2_2[Step 3.2.2]:::sub
    step3_2_3[Step 3.2.3]:::sub
    step3_2_4{Step 3.2.4}:::decision
    step3_2_5{Step 3.2.5}:::decision
    step3_2_6[Step 3.2.6]:::sub
    step3_2_7[Step 3.2.7]:::sub

    step3_2_1 --> step3_2_2 --> step3_2_3 --> step3_2_4
    step3_2_4 -->|no candidates| step3_2_7
    step3_2_4 -->|candidate selected| step3_2_5
    step3_2_5 -->|mismatch: invalidate candidate| step3_2_4
    step3_2_5 -->|match| step3_2_6 --> step3_2_7
  end

  %% Orchestrator iterates items; accepts/rejects subagent output
  step3_2 --> step3_2_1
  step3_2_7 -->|next plan item| step3_2
  step3_2 -->|all items attempted| step3_3

  %% Aggregate + finalize, or exit if nothing viable
  step3_3 -->|none verified| E
  step3_3 -->|some verified| step3_4 --> step3_5 --> E

• Note de correspondance : Step 1.1/1.2 correspondent aux deux points des explore-subagent sous Step 1 ; Step 2.1/2.2 correspondent aux deux sous-étapes du plan sous Step 2 ; Step 3.2.1-3.2.7 correspondent aux sous-étapes du code-subagent sous Step 3.2.

Étapes détaillées

1. Phase de recherche : Étudiez le modèle Transformer cible et les implémentations disponibles de kernels et de monkey-patch dans TileGym. Lancez 2 subagents explore en parallèle :

   • Recherchez l'ID du modèle sur HuggingFace pour connaître les architectures qu'il utilise. Puis recherchez sur GitHub le code pour obtenir l'implémentation de cette architecture. Passez en revue les détails pour comprendre les calculs effectués sur chaque composant. Résumez une liste complète des exigences avec tous les détails nécessaires. Concentrez-vous sur les détails. Certains modèles pourraient utiliser des variantes d'Attention/MoE/normalisation standard et/ou utiliser des types de données distincts dans différentes parties des calculs ;
   • Parcourez @src/tilegym/ pour inventorier les implémentations de kernels disponibles, les interfaces OP et les monkey-patches de modèles Transformer. Faites attention aux mappages @dispatch("<OP name>") et @register("<OP name>"), et aux schémas de patch apply_tilegym_kernel_to_<transformer_module>. Résumez un manifeste qui liste toutes les fonctions de monkey-patch disponibles, les interfaces OP, les implémentations de kernels avec suffisamment de détails pour distinguer les variantes d'opérations. Consultez mais ne vous appuyez pas sur les docstrings/commentaires ; concentrez-vous sur les détails qui distinguent les kernels similaires. En cas de doute sur la sémantique du kernel cuda.tile, consultez https://docs.nvidia.com/cuda/cutile-python/operations.html.

2. Phase de planification : Vérifiez si l'architecture du modèle cible est déjà patchée. Si c'est le cas, informez l'utilisateur et quittez ; Sinon, proposez un plan d'intégration en suivant ces sous-étapes :

   1. Consultez la liste des exigences et le manifeste pour déterminer quel ensemble de calculs pourrait être patchée par les implémentations TileGym. Soyez optimiste puisque les étapes/subagents suivants élimineront les propositions inadéquates ;
   2. Pour chaque calcul sélectionné à la sous-étape précédente, proposez des interfaces OP TileGym correspondantes et/ou des implémentations de kernels concrets. Vous pouvez proposer plusieurs candidats en cas d'incertitude, mais gardez le pool de candidats petit selon votre meilleur jugement.

3. Phase d'exécution et de vérification : Préparez l'environnement de développement, lancez les subagents pour implémenter le monkey-patch pour chacun des éléments du plan d'intégration une fois à la fois, vérifiez-le dans l'environnement de développement et acceptez/rejetez ce monkey-patch. Sous-étapes spécifiques :

   1. L'agent orchestrateur (c'est-à-dire vous) prépare un environnement de développement GPU en suivant environment-setup.md. Cet environnement sera utilisé par les subagents suivants.

   2. Pour chaque élément du plan d'intégration non vérifié (c'est-à-dire un mappage d'un calcul du modèle Transformer <-> un ou plusieurs candidats d'implémentation TileGym), lancez un subagent de code en parallèle avec une très forte capacité de suivi de plan et de débogage. Dites à ce subagent le nom du nœud alloué à la sous-étape 3.1. Le flux de travail de ce subagent de code est :

      1. Étudiez src/tilegym/transformers/monkey_patch.py et modeling/transformers/infer.py pour comprendre comment patcher un calcul transformer avec une implémentation TileGym. Étudiez docker-gpu-guide.md. Si une étape ultérieure nécessite l'exécution d'un script, utilisez toujours SSH vers le nom du nœud fourni par l'agent orchestrateur, construisez ou reconstruisez l'image par modeling/transformers/Dockerfile si nécessaire, et exécutez ce script dans l'environnement Docker NVIDIA ;

      2. Localisez le point d'intégration dans la bibliothèque transformers. Par exemple, il pourrait s'agir d'une sous-classe nn.Module qui correspond à une couche du modèle transformer, ou d'une fonction utilitaire qui applique certaines modifications aux variables/tenseurs intermédiaires du modèle transformer. Utilisez la recherche de code GitHub MCP pour les imports non vus ;

      3. Collectez les entrées et sorties autour du point d'intégration pour servir de références pour les vérifications ultérieures. Vous pouvez créer un simple script Python de débogage qui appelle l'API .generate() de la bibliothèque transformers pour inviter le modèle Transformer à générer « The capital of France is », et ajouter du code avant et après le point d'intégration pour enregistrer les tenseurs PyTorch intermédiaires et autres variables nécessaires sur le disque comme références futures. Critique : .generate() non optimisé est lent, collectez aussi peu de données que possible ;

      4. Sélectionnez le prochain candidat d'implémentation TileGym non vérifié. Si aucun candidat non vérifié n'est disponible, quittez le subagent actuel et informez l'agent orchestrateur que le calcul Transformer actuel ne convient pas pour TileGym à patcher ; Sinon, implémentez une fonction de monkey-patch en suivant la convention étudiée à la sous-étape 3.2.1. La fonction de patch du calcul actuel va à src/tilegym/transformers/<submodule_name>/monkeypatch<compute_name>.py. Si des modifications supplémentaires sont nécessaires pour le modèle transformer actuel (similaire au scénario de src/tilegym/transformers/deepseek2/modeling_deepseek.py), vérifiez l'existence (créée par d'autres subagents) ou créez un sous-module Python autonome src/tilegym/transformers/<submodulename>/modeling<submodule_name>.py et placez-y les modifications ;

      5. Vérifiez l'implémentation du monkey-patch à la sous-étape 3.2.4 en créant un script Python qui instancie un sous-module contenant le point d'intégration, applique le monkey-patch, alimente les données d'entrée collectées à la sous-étape 3.2.3 et collecte les données de sortie. Les données de sortie doivent correspondre aux données de sortie de référence collectées à la sous-étape 3.2.3 dans une tolérance d'erreur raisonnable. Faites de votre mieux pour corriger les erreurs causées par l'intégration et pour résoudre les décalages. Si vous ne pouvez pas corriger, marquez le candidat d'implémentation TileGym actuel comme invalide et retournez à la sous-étape 3.2.4 ; Sinon, continuez à la sous-étape suivante ;

      6. Consolidez le code de débogage et de test que vous avez implémenté vers src/tilegym/transformers/<submodule_name>/test_monkeypatch<compute_name>.py et organisez-le en style pytest et supprimez tous les autres fichiers/scripts/documents/fichiers de données binaires que vous avez créés pendant le débogage. Assurez-vous que seul un cas de test reste qui vérifie que l'entrée-sortie autour du point d'intégration correspond à celles de l'implémentation d'origine et assurez-vous que le cas de test passe. À ce stade, le répertoire src/tilegym/transformers/<submodule_name>/ devrait ressembler à :

         src/tilegym/transformers/<submodule_name>/
         |- monkey_patch_<compute_name>.py  # Patch function for compute assigned to current subagent.
         |- test_monkey_patch_<compute_name>.py  # Test logic specific to <compute_name> patching.
         |- # Optional [monkey_patch_<other_compute_name>.py, test_monkey_patch_<other_compute_name>.py] pairs created by other subagents assigned with <other_compute_name>s.
         |- modeling_<submodule_name>.py  # Optional if need to modify submodule or function, could be initially created by other subagents.

      7. Quittez le subagent actuel et informez l'agent orchestrateur que le calcul Transformer assigné peut être patchée par l'implémentation TileGym vérifiée aux sous-étapes 3.2.5 et 3.2.6 et que la fonction de patch est disponible à src/tilegym/transformers/<submodule_name>/monkeypatch<compute_name>.py.

   3. Agrégez tous les calculs vérifiés et les patches correspondants. Si aucun des calculs ne peut être fidèlement intégré, quittez le flux de travail et informez les utilisateurs ; Sinon, agrégez toute la logique de patching à une fonction principale de monkey-patch def apply_tilegym_kernel_to_<submodule_name>(...) et placez-la à src/tilegym/transformers/monkey_patch.py. Chaque calcul a un indicateur booléen correspondant comme argument de fonction ;

   4. Mettez à jour modeling/transformers/infer.py pour inclure la fonction principale de monkey-patch dans le flux d'inférence et de benchmark. Créez un script Bash modeling/transformers/bench_<submodule_name>.sh similaire aux autres scripts bench de ce répertoire. Assurez-vous d'utiliser --use_cutile à l'appel de 2e infer.py, car nous nous concentrons principalement sur le backend cuTile ;

   5. Étudiez docker-gpu-guide.md, utilisez SSH vers le nom du nœud alloué à la sous-étape 3.1, construisez l'image par modeling/transformers/Dockerfile et vérifiez le script d'inférence end-to-end créé à la sous-étape 3.4. Une fois tous les tests réussis, libérez le nœud GPU.