Résumé Technique : Batching Dynamique en Ligne avec Garanties Formelles pour l'Entraînement de LLM

1. Énoncé du Problème

Les pipelines modernes d'entraînement et de fine-tuning multimodal de Large Language Models (LLM) souffrent d'un problème fondamental d'observabilité concernant le batching au moment de l'entraînement. Dans les échantillonneurs de batchs hors ligne (offline) standards, le coût réel d'entraînement (longueur de séquence) d'un échantillon est inconnu jusqu'à ce qu'un pipeline de prétraitement complexe impliquant l'augmentation, le templating de chat, la tokenisation et l'expansion des tokens visuels multimodaux soit terminé.

Par conséquent, la construction des batchs est « aveugle » à la variable qui détermine le padding, l'utilisation de la mémoire et la saturation du GPU. Les solutions existantes font face à un compromis :

• Échantillonnage à batch fixe : Pour éviter les erreurs de type Out-of-Memory (OOM) sur les données à longue traîne, des tailles de batch de petite taille sont utilisées, ce qui conduit à une sous-utilisation des GPU. Des batchs fixes de grande taille provoquent un padding excessif ou des OOM.
• Caches de longueur hors ligne : Les méthodes qui précalculent les longueurs (ex: oracles GMT/BMT) nécessitent des caches statiques coûteux qui doivent être reconstruits chaque fois que les politiques d'augmentation, les templates ou les longueurs de coupure changent.
• Sequence packing : Bien qu'efficace pour supprimer le padding, cela nécessite souvent des interventions au niveau du modèle ou du kernel (ex: attention à longueur variable) plutôt qu'une solution purement liée au DataLoader.

De plus, déplacer la formation des batchs au point d'une observabilité précise (post-prétraitement) introduit un Problème d'Alignement de Groupe Distribué (DGAP). Dans un cadre de Distributed Data Parallel (DDP), tous les rangs doivent exécuter le même nombre d'étapes de réduction de gradient. Si chaque rang forme indépendamment des groupes de tailles variables basés sur les longueurs réalisées localement, le nombre de groupes différera d'un rang à l'autre, brisant le contrat DDP et provoquant des deadlocks ou des pertes d'échantillons.

2. Méthodologie : Batching Dynamique en Ligne (ODB)

Les auteurs proposent l'Online Dynamic Batching (ODB), un système prêt à l'emploi implémenté à la frontière du DataLoader qui observe les longueurs réalisées après le prétraitement et forme des batchs de taille variable dynamiquement sans modifier le modèle, l'optimiseur ou les kernels d'attention.

Architecture

• Placement : ODB enveloppe l'itérateur du DataLoader PyTorch. Il laisse le Dataset et le Model intacts.
• Flux de travail :
  1. Workers : S'exécutent avec une fonction de collation nulle, passant des échantillons uniques à un Processus de Collation dédié.
  2. Groupement : Le Processus de Collation vide les échantillons dans un buffer, les trie par longueur et forme des batchs de taille variable de manière gloutonne pour respecter un budget de tokens spécifié par l'utilisateur ($L_{max}$).
  3. Alignement : Un groupe Gloo dédié (isolé du groupe NCCL principal) synchronise le nombre de groupes entre tous les rangs.
  4. Sortie : Les groupes alignés sont transmis au trainer. Les slots sous-remplis sont complétés par des sentinelles IDLE_DATA, que le processus principal ignore, assurant ainsi l'alignement des étapes.

Le Problème d'Alignement de Groupe Distribué (DGAP)

ODB formalise l'exigence de synchronisation comme le DGAP. Il introduit un Protocole d'Alignement de Groupe Bidirectionnel basé sur le Maximum pour garantir que tous les rangs effectuent le même nombre d'opérations AllReduce sans deadlock ni perte d'échantillons.

• Calcul de la cible : Le nombre cible de groupes ($T_{grp}$) est calculé comme le maximum entre la capacité de sortie positive minimale et le nombre de samples mis en tampon positif minimal à travers les rangs actifs.
• Ajustement :
  • Division (Split) : Si un rang possède moins de groupes que $T_{grp}$, il divise les groupes plus larges en singletons.
  • Dépassement (Overflow) : Si un rang possède plus de groupes, il conserve les $T_{grp}$ meilleurs groupes et réinjecte les échantillons restants dans le buffer.
• Modes de terminaison :
  • Mode Join par défaut : Assure une couverture d'identité stricte en vidant les vues d'échantillon en attente avant la complétion globale.
  • Mode Non-Join optionnel : Assure la clôture du quota d'échantillons (compte cumulatif) sans garanties strictes d'identité par itération, adapté aux temps d'exécution contraints.

Mise à l'échelle de la perte (Loss Scaling)

Puisque les batchs ODB diffèrent en nombre de tokens par rang, un moyennage DDP naïf produit une estimation biaisée de la perte. ODB implémente une mise à l'échelle de la perte au niveau du token, en pondérant la perte de chaque rang par sa fraction de tokens totaux ($w_r = t_r / T_{tok}$) pour retrouver la perte de référence par token de précision bit-par-bit.

3. Contributions Clés

1. Conception Système : Introduction d'ODB, un système côté DataLoader qui observe les longueurs post-pipeline et forme des batchs de taille variable sûrs pour le DDP sans réécriture de modèle ou de cache de longueur précalculé.
2. Garanties Formelles : Formalisation du DGAP et preuve de terminaison bornée sans deadlock, de couverture d'identité stricte (en mode join) et de clôture de quota d'échantillons (en mode non-join).
3. Performance Empirique : Évaluation sur des modèles Qwen3-VL 2B/8B sur UltraChat, LLaVA et ShareGPT4o, démontrant des gains de débit significatifs avec une qualité comparable au Standard.
4. Open Source : Publication du package online-dynamic-batching avec des adaptateurs de trainer légers.

4. Résultats Expérimentaux

Les expériences ont été menées sur des nœuds 8×H20 (DeepSpeed ZeRO-2, bf16) utilisant les modèles Qwen3-VL-2B/8B.

Gains de Débit (Throughput)

ODB surpasse significativement les baselines à batch fixe (Standard, Sorted) et approche les méthodes hors ligne (GMT/BMT) sans leur surcharge de cache :

• Single-Node (Full FT/LoRA) : Amélioration de 1,58× – 2,51× par rapport au Standard.
• Two-Node (Full FT) : Amélioration de 1,71× – 3,78× par rapport au Standard.
• Haute Hétérogénéité (ShareGPT4o, CV=1.00) : ODB atteint un gain de vitesse de 2,46× (8B) et 2,47× (2B), tandis que Sorted est limité à environ 1,03× en raison des contraintes de la longue traîne imposant de petites tailles de batch.
• Cas de Production (MM-Mix) : Sur un mélange multimodal de production avec une forte densité de samples courts, ODB a obtenu un gain de vitesse de 4,43× par rapport au Standard.

Métriques de Qualité

• Perte de Validation & Benchmarks : ODB maintient la qualité dans la « bande comparable au Standard ». Par exemple, sur 8B UltraChat, le score MMLU d'ODB (74,75 %) est comparable à celui de Packing (75,18 %) et de GMT (75,14 %).
• Comparaison avec Sorted : Bien que Sorted atteigne souvent un débit supérieur au Standard, il souffre fréquemment de régression de la perte de validation et de dégradation du format de réponse (ex: sur LLaVA, le score MMMU de la réponse générée de Sorted s'est effondré à 5,52 % contre 22,30 % pour le Standard en raison de la dérive de format). ODB évite cela en groupant plutôt qu'en triant.
• Comparaison avec les Oracles : Les oracles hors ligne (GMT/BMT) atteignent souvent un débit brut légèrement supérieur mais nécessitent des caches statiques coûteux devant être reconstruits lors des changements de politique. ODB opère dans le régime « online/drop-in », offrant une qualité comparable sans précalcul de cache.

Études d'Ablation

• Budget de Tokens ($L_{max}$) : Le débit augmente avec $L_{max}$ jusqu'à ce que la pression mémoire ou le temps par étape devienne dominant. L'$L_{max}$ optimal varie selon le dataset (ex: 12288 pour UltraChat, 14334 pour ShareGPT4o).
• Profondeur d'Attente (Outstanding Depth $D$) : Une fois que le chevauchement du pipeline est saturé (typiquement >95 %), augmenter $D$ produit des rendements décroissants.
• Levier des Samples Courts : Le gain de vitesse est corrélé au coefficient de variation (CV), mais est amplifié par la fraction de samples courts ($f_s$). MM-Mix (CV plus faible mais $f_s$ élevée) a vu des gains plus importants que ShareGPT4o (CV élevé, $f_s$ faible).

5. Signification et Revendications

L'article affirme qu'ODB occupe avec succès le régime online/drop-in pour le fine-tuning de LLM à haute hétérogénéité. Sa signification réside dans :

• Combler le fossé : Il réduit l'écart de débit entre l'entraînement à batch fixe et les méthodes plus fortes basées sur le modèle ou hors ligne (comme le packing ou le batching par oracle) tout en évitant le besoin de caches de longueur scalaires ou de réécritures de kernels côté modèle.
• Sécurité Formelle : Il fournit les premières garanties formelles (absence de deadlock, couverture d'identité, clôture de quota) pour le batching de taille variable à l'exécution en DDP, un régime précédemment évité par les systèmes statiques.
• Viabilité de Déploiement : En opérant purement à la frontière du DataLoader, il est compatible avec les stacks d'entraînement existantes (DeepSpeed, LLaMA-Factory) et robuste aux changements dynamiques de politiques d'augmentation ou de templating qui invalideraient les caches hors ligne.

Les auteurs restent modestes, notant que les gains dépendent des données (meilleurs sur les datasets à haut CV/haut taux de samples courts) et qu'ODB complète plutôt qu'il ne remplace le packing côté modèle lorsqu'ils sont disponibles. Ils ne prétendent pas qu'ODB résout toutes les inefficacités de batching, mais qu'il fournit une solution robuste et formellement garantie pour le goulot d'étranglement spécifique de l'observabilité dans l'entraînement moderne des modèles multimodaux et augmentés.