Batching continu : optimiser l'inférence des LLM efficacement

Le problème du batching statique

Dans le domaine de l'inférence des modèles de langage de grande taille (LLM), le batching statique est souvent utilisé pour traiter plusieurs requêtes simultanément. Cette méthode consiste à regrouper les requêtes en lots de taille fixe, traitant chaque lot ensemble. Cependant, cette approche présente des limitations significatives en termes d'efficacité.

Prenons l'exemple de trois requêtes distinctes :

• A : "La capitale de la France est" qui génère 6 tokens supplémentaires.
• B : "La météo d'aujourd'hui est si" qui génère 50 tokens supplémentaires.
• C : "En apprentissage automatique, un transformateur est" qui génère 300 tokens supplémentaires.

Dans ce scénario, les requêtes A et B se terminent bien avant la requête C. Pourtant, leurs emplacements restent occupés et inactifs pendant que le GPU continue de traiter les tokens restants pour C. Cela entraîne une consommation inutile de cycles de calcul.

Lors du processus de décodage, la situation est la suivante :

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 …300

A <BOS> <La> <capitale> <de> <la> <France> <est> <la> <capitale> <de> <la> <République> <EOS> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> …<PAD>

B <BOS> <La> < météo> <d’> <aujourd’hui> <est> <si> <froide> <que> <c’> <est> <difficile> <de> <voir> <le> <soleil> <.> <Mais> <ce> < n’> <est> <pas> <comme> <si> <nous> <allions> <…>

C <BOS> <En> <apprentissage> <automatique> <,> <un> <transformateur> <est> <un> <type> <de> <machine> <learning> <algorithme> <qui> <peut> <être> <utilisé> <pour> <…>

Les tokens spéciaux tels que <BOS> (début de la phrase), <EOS> (fin de la phrase) et <PAD> (remplissage) illustrent comment la requête A est remplie jusqu'à la fin des 300 tokens dans ce lot. Le GPU consomme des cycles de calcul qui n'apportent rien à aucun résultat avec les tokens de remplissage. De plus, la réponse à la requête A n'est probablement pas livrée avant que la requête C ne soit terminée.

Exemple de code de batching statique

Pour illustrer le fonctionnement du batching statique, considérons un modèle GPT-2 exécutant six requêtes de longueurs différentes avec une taille de lot de 3. Chaque requête est associée à un prompt et un nombre maximum de tokens à générer.

MODEL_ID = "[openai](/dossier/openai)-community/gpt2"
BATCH_SIZE = 3
requests = [
    ("La capitale de la France est", 6),
    ("La météo d'aujourd'hui est si", 50),
    ("En apprentissage automatique, un transformateur est", 300),
    ("Il était une fois dans un pays lointain,", 30),
    ("L'informatique quantique diffère de l'informatique classique parce que", 180),
    ("L'histoire de l'Empire romain a commencé", 45),
]

La fonction de batching statique traite les requêtes par vagues de taille définie, chaque vague s'exécutant jusqu'à ce que la requête la plus longue soit terminée. Les requêtes plus courtes restent inactives, et aucun nouvel emplacement ne peut être utilisé tant que toute la vague n'est pas terminée.

def static_batching(requests: list[tuple[str, int]], tokenizer, model) -> list[str]:
    if not requests:
        return []
    tokenizer.padding_side = "left"
    results: dict[int, str] = {}
    indexed = list(enumerate(requests))
    for wave_start in range(0, len(indexed), BATCH_SIZE):
        wave = indexed[wave_start: wave_start + BATCH_SIZE]
        wave_max = max(cap for _req_id, (_prompt, cap) in wave)
        prompts = [p for _, (p, _) in wave]
        inputs = tokenizer(
            prompts, return_tensors="pt", padding=True, truncation=True
        ).to(model.device)
        with torch.no_grad():
            output_ids = model.generate(
                **inputs,
                max_new_tokens=wave_max,
                pad_token_id=tokenizer.eos_token_id,
                do_sample=False,
            )
        width = inputs.input_ids.shape[1]
        for slot, ((req_id, (prompt, cap)), row) in enumerate(zip(wave, output_ids)):
            text = prompt + tokenizer.decode(row[width:width + cap], skip_special_tokens=True)
            results[req_id] = text
    return [results[k] for k in sorted(results)]

L'exécution de ce code montre que les requêtes plus courtes attendent inutilement la fin de la requête la plus longue, illustrant ainsi les inefficacités du batching statique.

Vers une solution plus efficace

Le batching continu, avec sa planification dynamique et son batching irrégulier, promet d'améliorer ces inefficacités. En adaptant la taille des lots et en optimisant la gestion des requêtes, cette approche pourrait offrir des gains significatifs en termes de performance et de rapidité de traitement.

Conclusion

Le batching statique présente des limitations en termes d'efficacité, notamment en raison de l'attente des requêtes plus courtes. Dans les sections suivantes, nous explorerons comment le batching continu et la planification dynamique peuvent améliorer cette situation.