Contexte technique

Le AsyncGRPOTrainer de TRL v1.14 a intégré la prise en charge des adaptateurs LoRA. Un adaptateur de rang 1 pour un modèle de 1,5 M paramètres ne dépasse que quelques mégaoctets, contre environ 3 Go pour le modèle complet. Cette différence de taille rend possible le transfert de l’adaptateur entre machines sans recourir à NCCL, qui est habituellement requis pour synchroniser plusieurs gigaoctets dans un cluster dense.

Hugging Face Jobs exécute chaque conteneur sur une VM distincte. Les jobs ne partagent ni disque local ni réseau direct, ce qui empêche la transmission de gros poids. En revanche, les Storage Buckets peuvent être montés comme système de fichiers POSIX via hf-mount, offrant un point de partage accessible à tous les jobs.

Architecture mise en œuvre

Le projet se compose de trois éléments : un job trainer exécutant AsyncGRPOTrainer avec LoRA et FSDP, deux jobs vLLM servant le modèle de base et les adaptateurs, et un proxy qui dirige chaque rollout vers le replica contenant le cache KV approprié et diffuse chaque mise à jour d’adaptateur à l’ensemble des replicas.

Le flux de synchronisation s’articule ainsi : toutes les n étapes d’optimisation, le trainer écrit l’adaptateur dans /lora//.vllm_lora/trl-policy-v{N}, effectue un renommage atomique, puis invoque l’endpoint /v1/load_lora_adapter de vLLM avec le chemin du répertoire. Les serveurs vLLM lisent directement le fichier depuis le bucket monté, sans échange de tenseurs.

# montage du bucket identique dans chaque job
hf jobs run ... -v hf://buckets/aminediroHF/asyncgrpo-lora-buckets:/lora ...

# lancement d'un replica vLLM avec support LoRA
hf jobs run --detach --flavor h200 -v "hf://buckets/${BUCKET}:/lora:ro" \
    -e VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 -e VLLM_SERVER_DEV_MODE=1 \
    -- vllm/vllm-openai:v0.27.1 vllm serve Qwen/Qwen2.5-Math-1.5B \
    --enable-lora --max-lora-rank 1 --max-loras 6

Analyse des performances et de la scalabilité

Le passage du transfert complet du modèle (≈3 GB) à l’échange d’un adaptateur (≈5 MB) a réduit le temps d’une exécution de 500 étapes de 3 h 27 min à 53 min, soit une accélération de plus de six fois. Le goulot d’étranglement identifié par les métriques AsyncGRPO se situe désormais au niveau du calcul du gradient plutôt qu’au transfert réseau.

Le paramètre max_staleness=4 impose que vLLM conserve les quatre versions précédentes de l’adaptateur, plus la version courante, d’où le besoin de --max-loras 6. Cette stratégie garantit que les rollouts initiés avec une version antérieure peuvent se terminer sans interruption, tout en limitant la consommation de mémoire GPU.

Limites et perspectives

Le modèle repose sur la persistance du bucket ; une perte de connectivité ou une latence élevée du système de fichiers FUSE pourrait ralentir le chargement d’adaptateurs. De plus, le proxy constitue un point unique de défaillance : s’il ne parvient pas à router correctement les rollouts, les caches KV peuvent être sous‑exploités, augmentant le temps de génération.

Enfin, la solution reste liée à la capacité du bucket à gérer des écritures atomiques rapides. Aucun mécanisme de réplication intra‑bucket n’est décrit, ce qui pourrait devenir un facteur limitant à plus grande échelle.