Encodeur multimodal

Le troisième problème critique du pré-entraînement à l'échelle 3 T : le calcul très variable de l'encodeur visuel est exposé sur le chemin critique.

Pourquoi le ViT est un problème

L'encodeur visuel ne pèse que 401 M de paramètres — 0,014 % du modèle. Il devrait être négligeable. Il ne l'est pas, pour deux raisons.

1. Sa charge est extrêmement variable

Une vignette de 224 × 224 coûte 256 carreaux. Une image de 3 584 × 3 584 en coûte 65 536 — un facteur 256. Une vidéo longue davantage encore.

Contrairement au texte, dont chaque jeton coûte la même chose, le coût visuel d'un échantillon varie de plusieurs ordres de grandeur.

2. Il est sur le chemin critique

Le ViT doit terminer avant que le backbone puisse traiter les jetons visuels. Un GPU qui traite une grosse image fait attendre tous les autres.

Le rapport résume : large images and long videos substantially increase the computation time of the vision encoder and cause significant load imbalance across devices.

Solution 1 : le parallélisme de contexte dynamique

Deux niveaux de partitionnement.

Partitionner une grande image

Une image unique est partitionnée le long de la dimension des carreaux entre plusieurs appareils, et l'attention est calculée en rassemblant les paires clé–valeur (gather-KV) entre rangs CP.

image 3584×3584 = 65 536 carreaux
        │
        ├──► GPU 0 : carreaux 0 … 16 383
        ├──► GPU 1 : carreaux 16 384 … 32 767
        ├──► GPU 2 : carreaux 32 768 … 49 151
        └──► GPU 3 : carreaux 49 152 … 65 535
                     │
                     └──► gather-KV entre rangs pour l'attention

Équilibrer plusieurs grandes images

Chaque groupe CP est divisé en plusieurs sous-groupes CP, et plusieurs grandes images sont distribuées entre eux de façon équilibrée.

Pourquoi cette seconde étape est nécessaire

Le rapport en donne la raison précise : cela empêche la fraction de communication de croître avec l'échelle.

Si l'on répartissait une seule image sur tous les rangs disponibles, la communication gather-KV croîtrait avec le nombre de rangs, jusqu'à dominer le calcul. En sous-groupant, chaque image reste sur un petit nombre de rangs, et l'on gagne en parallélisme par le nombre d'images traitées simultanément.

Effet : réduit à la fois la latence de l'encodeur sur les grands échantillons visuels et le déséquilibre de charge entre appareils, ce qui permet de cacher le calcul restant dans les bulles de pipeline.

Solution 2 : le calcul dans les bulles de pipeline

C'est l'optimisation la plus élégante de cette section.

Le point de départ : DEP

Kimi K2.5 avait introduit le Decoupled Encoder Process (DEP), qui sépare l'entraînement du ViT et du texte en étages distincts, et équilibre les passes avant et arrière de la vision entre étages PP.

L'observation de K3

Ce que l'équipe a remarqué

Sous l'ordonnancement de pipeline 1F1B entrelacé :

• les passes avant texte des premiers micro-lots PP sont toutes ordonnancées tout au début ;
• les passes arrière texte des derniers micro-lots PP ne se terminent qu'à la toute fin.

Il y a donc des trous structurels dans l'occupation du pipeline — les bulles.

La décomposition

L'équipe décompose davantage le calcul du ViT :

Partie	Ordonnancement
Passes avant des premiers micro-lots PP	Exécutées synchroniquement en amont
Passes avant restantes	Ordonnancées dans les bulles de pipeline
Passes arrière	Traitées de façon analogue

Le résultat

As a result, most of the ViT computation is hidden within pipeline bubbles, largely eliminating the effective overhead of the vision encoder.

Le coût de l'encodeur visuel devient quasi nul en temps de mur : il est payé avec du temps GPU qui, autrement, serait perdu.

Sans optimisation :
  GPU0 ████ txt ░░░░ bulle ░░░░ ████ txt   +  ▓▓ ViT (chemin critique)
  GPU1 ░░░░ ████ txt ░░░░ ████ txt

Avec optimisation :
  GPU0 ████ txt ▓▓▓▓ ViT ▓▓▓▓ ████ txt      ← les bulles sont remplies
  GPU1 ▓▓▓▓ ████ txt ▓▓▓▓ ████ txt

L'idée générale à retenir

Un pipeline a nécessairement des bulles. Plutôt que de tenter de les supprimer (ce qui est difficile), on peut y loger un travail indépendant. L'encodeur visuel est le candidat idéal : son calcul ne dépend pas du reste du pipeline à cet instant.

Le lien avec l'architecture

Cohérence de conception

Cette optimisation n'est possible que parce que Kimi K3 est nativement multimodal, avec l'encodeur entraîné conjointement dans la même boucle.

Dans une approche « greffe et alignement », le ViT serait pré-entraîné séparément, puis gelé ou affiné dans une phase distincte — il n'y aurait pas de pipeline texte dans lequel cacher son calcul.

L'architecture et l'infrastructure se justifient mutuellement.

Vérification de compréhension

Pourquoi ne pas simplement redimensionner toutes les images à une taille fixe ?

On perdrait la précision sur les documents, captures d'écran et graphiques — exactement les cas d'usage visés par Kimi K3 (OmniDocBench : 91,1 %, meilleur score du panel). La résolution native est un choix de qualité, dont le déséquilibre de charge est le prix ; ces optimisations le paient.

Le gather-KV entre rangs CP n'est-il pas très coûteux ?

Il l'est, ce qui explique la seconde étape : en sous-groupant les rangs CP, on borne le nombre de participants au gather pour une image donnée. C'est exactement le sens de « empêcher la fraction de communication de croître avec l'échelle ».

Pourquoi exécuter synchroniquement les passes avant des premiers micro-lots ?

Parce qu'au tout début du pipeline, il n'y a pas encore de bulles : tous les étages sont en phase d'échauffement et les premiers micro-lots avancent sans attendre. Il n'y a donc pas de trou à remplir, et il faut bien faire ce travail. Les bulles n'apparaissent qu'ensuite.

---

Chapitre précédent : Mémoire et parallélismes · Chapitre suivant : Infra pour le RL à 1 M