Pré‑chargement dynamique des poids des LLM pour réduire la latence d’inférence avec Workflow v6

Contexte et contraintes de latence

Dans les architectures où les LLM servent des requêtes interactives (chatbots, assistants de code, recherche sémantique), la latence perçue par l’utilisateur est un critère de succès. La majorité du temps de réponse provient du chargement des poids du modèle : même avec un GPU performant, le transfert depuis la mémoire système ou le disque vers la VRAM représente plusieurs dizaines de millisecondes, voire plus selon la taille du modèle.

Les pipelines traditionnels de Workflow v6 chargent les poids au démarrage du worker et les gardent en mémoire pendant toute la durée du pod. Cette approche garantit la simplicité, mais elle ne scale pas lorsqu’on veut héberger plusieurs modèles de tailles différentes sur le même cluster ou lorsqu’on veut exploiter les ressources de manière éphémère (serverless, auto‑scaling). Le pré‑chargement dynamique consiste à déplacer la charge de transfert hors du chemin critique en anticipant les besoins et en les satisfaisant de façon asynchrone.

Principes du pré‑chargement dynamique

1. Segmentation des poids – Les poids d’un LLM sont découpés en blocs (par ex. par couche ou par groupe de paramètres). Cette segmentation est déjà fournie par les formats de checkpoint modernes (HF safetensors, torch.save avec metadata).
2. Cache de poids partagé – Un service dédié (ex. un WeightCache déployé en tant que StatefulSet) expose une API d’accès en lecture‑seule aux blocs déjà résidents en VRAM ou en mémoire NVRAM.
3. Chargement anticipé – Le scheduler de Workflow v6 déclenche une tâche de pré‑chargement dès qu’un nouveau worker est provisionné, en fonction d’un profil de probabilité d’utilisation (ex. % de requêtes ciblant le modèle B vs C).
4. Fallback transparent – Si le cache ne contient pas le bloc demandé, le worker charge le bloc depuis le stockage persistant (S3, GCS) tout en continuant à servir les requêtes déjà prêtes.

Ces principes permettent de réduire la latence de la première requête d’un modèle « froid » de plusieurs dizaines de millisecondes, tout en limitant le gaspillage de VRAM.

Implémentation avec Workflow v6

1. Définir le stockage segmenté

# workflow/v6/weights/bigmodel.yaml
resources:
  storage:
    type: s3
    path: s3://ml-models/bigmodel/
    layout: "layers/{layer_id}.safetensors"

Le layout indique à Workflow v6 comment récupérer chaque bloc.

2. Créer le service de cache de poids

# weight_cache.py
from fastapi import FastAPI, HTTPException
import torch
import os

app = FastAPI()
CACHE_DIR = "/mnt/weight_cache"  # monté sur un SSD NVMe partagé

@app.get("/weights/{layer_id}")
async def get_layer(layer_id: str):
    path = os.path.join(CACHE_DIR, f"{layer_id}.pt")
    if not os.path.exists(path):
        raise HTTPException(status_code=404, detail="Layer not cached")
    tensor = torch.load(path, map_location="cpu")
    return {"tensor": tensor.tolist()}

Le service expose les poids sous forme de tableau JSON. En production on privilégiera le streaming binaire (gRPC) mais le principe reste le même.

3. Pipeline de pré‑chargement

# workflow/v6/pipelines/preload.yaml
name: preload_weights
steps:
  - name: list_layers
    action: python
    script: |
      import json, boto3
      s3 = boto3.client('s3')
      objects = s3.list_objects_v2(Bucket='ml-models', Prefix='bigmodel/layers/')
      layers = [obj['Key'].split('/')[-1].replace('.safetensors','') for obj in objects.get('Contents',[])]
      print(json.dumps(layers))
  - name: push_to_cache
    action: parallel
    foreach: "{{ steps.list_layers.output }}"
    steps:
      - name: download
        action: s3
        bucket: ml-models
        key: "bigmodel/layers/{{ item }}.safetensors"
        destination: "/tmp/{{ item }}.pt"
      - name: upload_cache
        action: http
        method: POST
        url: "http://weight-cache:8000/weights/{{ item }}"
        body: "@/tmp/{{ item }}.pt"

Cette pipeline s’exécute à chaque montée de nouveaux workers et pré‑charge les blocs les plus probables.

4. Modification du worker

# worker.py
import torch, requests

def load_layer(layer_id):
    resp = requests.get(f"http://weight-cache:8000/weights/{layer_id}")
    if resp.status_code == 200:
        data = resp.json()['tensor']
        return torch.tensor(data)
    # Fallback: load from S3 (omitted for brevity)

Le worker appelle d’abord le cache, ce qui rend le code identique quelle que soit la provenance des poids.

Mesures et bonnes pratiques

• Profilage préalable : utilisez le profiler intégré de Workflow v6 (workflow profile) pour identifier les couches qui dominent le temps de chargement.
• Politique d’éviction : le cache doit implémenter LRU ou LFU afin de libérer la VRAM lorsqu’une couche n’est plus sollicitée.
• Dimensionnement du cache : commencez avec un SSD NVMe d’une capacité équivalente à 20 % de la taille totale du modèle, puis ajustez en fonction des métriques de hit‑rate.
• Sécurité : le service de cache doit être isolé par un réseau interne et authentifié via les secrets de Workflow v6 pour éviter toute fuite de poids propriétaires.
• Observabilité : exposez des métriques Prometheus (taux de hit, temps moyen de chargement) et configurez des alertes si le taux de hit descend sous un seuil critique (ex. 85 %).

En appliquant ces étapes, les équipes produit peuvent réduire de façon mesurable la latence d’inférence des LLM tout en conservant la flexibilité d’un pipeline CI/CD entièrement orchestré par Workflow v6.