Les heaps mentent bien : déboguer une fuite mémoire dans vLLM.

Il y a quelques mois, notre équipe a enquêté sur une fuite mémoire suspectée dans vLLM. Au départ, nous pensions que le problème serait facile à repérer, limité aux couches hautes de la codebase. Mais plus nous creusions, plus il devenait complexe. Cet article ouvre notre nouvelle série d’analyses techniques approfondies, dans laquelle nous expliquerons comment nous menons nos investigations techniques et construisons des solutions chez Mistral AI.

Le problème est apparu pour la première fois lors de tests de préproduction du serving désagrégé avec l’un de nos modèles de pointe. L’utilisation mémoire augmentait régulièrement, mais seulement dans des conditions précises : avec vLLM, avec notre modèle Mistral Medium 3.1, et avec la compilation de graphe activée. Aucun crash, aucune erreur : seulement une lente augmentation linéaire de la mémoire système, de 400 Mo par minute, sur un trafic proche de la production. Au bout de quelques heures, cela menait à un état « out of memory ».

Nous avons ensuite mené une recherche méthodique, en partant d’outils Python de haut niveau avant de descendre jusqu’au tracing au niveau du kernel, jusqu’à identifier la véritable source du problème. Voici comment nous l’avons trouvée, et ce qu’elle révèle sur les risques cachés des couches de dépendances dans les logiciels actuels.

Ce n’est pas le type de graphique en hausse que nous aimons voir dans Grafana 

Une fuite qui jouait à cache-cache.

Au départ, notre approche suivait un chemin de dépannage classique : isoler la source de la fuite en reproduisant le problème sur un modèle plus petit, avec moins d’optimisations de production activées. Mais après avoir testé différents paramètres et modèles, nous n’avons pas réussi à le reproduire dans une autre configuration. L’erreur n’apparaissait que dans une configuration désagrégée Prefill/Decode avec NIXL. 

Comme la désagrégation Prefill/Decode (P/D) joue un rôle central dans cette histoire, commençons par expliquer les mécanismes généraux de cette configuration d’inférence. Une configuration P/D désagrégée sépare le traitement d’une requête en deux phases, prises en charge par des instances différentes :

1. D’abord, le router envoie une « requête de prefill » (en définissant max_tokens=1 et un ensemble vide de métadonnées KV Transfer) à une instance vLLM de prefill pour calculer le KVCache de la requête.

2. Une fois l’opération terminée, le router transfère les métadonnées du KVCache avec une « requête de decode » vers une instance vLLM de decode.

3. Le transfert du KVCache est initié via NIXL, et la génération de tokens s’effectue sur l’instance vLLM de decode en utilisant et en étendant le KVCache transféré. 

La fuite n’était observée que côté decode dans cette configuration désagrégée, ce qui suggérait fortement que le transfert du KV Cache via NIXL était la cause racine. Dans notre configuration, NIXL s’appuie sur UCX (Unified Communication X), une bibliothèque de communication hautes performances conçue pour l’échange de données dans les systèmes distribués. UCX permet des transferts de données optimisés sur un large ensemble de technologies, dont Infiniband, une technologie d’interconnexion à faible latence et haut débit, couramment utilisée en HPC et dans les data centers.

Vue d’ensemble d’un déploiement de serving P/D désagrégé.

Pour la suite de notre investigation, nous avons travaillé dans cette configuration et commencé avec des outils de profilage mémoire Python pour localiser la source de la fuite.

Nous avons essayé Memray et Guppy 3, mais aucun des deux n’a montré de fuite, et tout ce qu’ils nous permettaient d’observer semblait normal. Tenter d’utiliser GDB faisait crasher l’ensemble du processus. Notre configuration vLLM était aussi trop lourde pour des outils comme Valgrind, ce qui les rendait extrêmement lents, voire inutilisables.

Il était clair qu’un outil plus puissant était nécessaire pour suivre la fuite. Mais avant d’y consacrer plus de temps, nous avons voulu vérifier que cette fuite était reproductible par d’autres. Nous avons contacté l’équipe vLLM en ouvrant une issue sur leur dépôt GitHub, ce qui a confirmé que nous n’étions pas les seuls à observer ce problème et qu’une investigation plus poussée était justifiée.

Compter les malloc et les free avec Heaptrack.

Pour mieux suivre ce qui se passait, nous nous sommes tournés vers Heaptrack : un profileur mémoire qui remplace les opérations mémoire comme malloc ou free et enregistre ces événements avec les stack traces associées.

Millian Wolf, le créateur de Heaptrack, a écrit un excellent article de blog d’introduction pour prendre en main l’outil. Le processus se déroule en deux étapes : exécuter d’abord le programme avec tracing, puis interpréter les dumps de données.

Pour suivre les allocations limitées au processus worker de vLLM, nous avons défini LD_PRELOAD sur libheaptrack_preload.so via vLLM, afin que cette bibliothèque soit chargée avant toutes les autres et remplace les comportements des fonctions d’allocation mémoire, ce qui nous fournit le dump de données.

Nous avons ensuite pu visualiser ces données avec heaptrack_interpret :

$ git clone https://github.com/KDE/heaptrack.git
$ cd heaptrack && mkdir build && cd build && cmake .. && sudo make install

# Setting LD_PRELOAD=/path/to/libheaptrack_preload.so creates a temporary file named heaptrack.<pid>, here the pid is 2028233

$ /usr/local/lib/heaptrack/libexec/heaptrack_interpret < heaptrack.2028233 | gzip > heaptrack.vllm.2028233.gz

Heaptrack fournit un graphe interactif détaillé de toutes les allocations du tas, jusqu’au niveau des fonctions. Nous pouvons suivre chaque malloc et chaque free, avec une répartition claire de l’utilisation mémoire.

L’utilisation mémoire de notre worker vLLM affichée avec Heaptrack

À ce stade, une question se pose : où est la fuite mémoire ? En réalité, la seule augmentation mémoire visible provenait d’une initialisation paresseuse de NIXL.

Pour vérifier que la fuite se produisait bien dans cette configuration, nous avons exécuté un benchmark vLLM et créé deux snapshots Heaptrack avec heaptrack_interpret : l’un au début, l’autre vers la fin. Même si la mémoire du tas restait stable, le pic de mémoire résidente (RSS), que nous aborderons dans la section suivante, différait entre les deux snapshots. Cet écart était visible dans l’onglet de synthèse de Heaptrack.

Écart de pic RSS dans Heaptrack : avant (1) et après (2) le benchmark

Cela signifiait que la fuite se produisait en dehors du tas, et donc en dehors de la mémoire analysée par Heaptrack. Nous devions changer d’outil pour suivre les allocations hors du tas.

Au-delà du tas : comprendre la mémoire résidente et les allocations système.

Pour comprendre pourquoi Heaptrack ne détectait pas la fuite, il faut d’abord préciser ce que le Resident Set Size (RSS) inclut réellement. Le RSS représente la partie de la mémoire d’un processus conservée en RAM, et il ne se limite pas au tas. Il inclut notamment :

• Le tas est traditionnellement géré avec les anciens appels système sbrk et brk, qui ajustent ou définissent l’adresse de rupture du programme, c’est-à-dire le pointeur qui marque la fin du segment de tas.

• La pile, qui stocke les variables locales et les frames d’appel de fonction.

• Les mappings mémoire anonymes, c’est-à-dire des régions de mémoire allouées directement via l’appel système mmap, sans fichier sous-jacent. Ils sont souvent utilisés par des allocateurs personnalisés ou par malloc pour des blocs mémoire plus volumineux. Les adresses des mappings anonymes se trouvent généralement entre l’espace d’adressage du tas et l’espace d’adressage de la pile, dans une région appelée segment de mapping mémoire.

Même si malloc peut utiliser sbrk pour les petites allocations, les implémentations modernes préfèrent généralement mmap avec des mappings anonymes, car cette approche est plus flexible et permet d’allouer des huge pages (pages mémoire de 2 Mo ou 1 Go selon votre configuration).

Heaptrack ne s’insère que dans les fonctions malloc et free de glibc. Cela signifie qu’il peut suivre directement toutes les allocations classiques du tas et les mappings anonymes alloués par malloc, mais qu’il ne voit pas la mémoire allouée via des appels directs à mmap ou d’autres mécanismes système hors du contrôle de glibc.

Heureusement, tout n’est pas perdu pour suivre ce qui se passait. Le système de fichiers /proc est un dossier spécial de Linux qui sert d’API du noyau, expose une interface virtuelle pour interagir avec les processus en cours d’exécution et fournit un accès en temps réel aux détails des processus, par exemple :

• /proc/<pid>/fd, qui liste tous les descripteurs de fichiers ouverts pour un processus.

• /proc/<pid>/maps, qui affiche une carte détaillée des régions mémoire du processus, dont le tas, la pile, les bibliothèques partagées et les mappings anonymes.

• Et beaucoup d’autres éléments (un simple ls dans /proc/<pid>/ liste ce qui est disponible).

Pour poursuivre notre investigation, nous avons utilisé la commande pmap, qui lit /proc/<pid>/maps et présente l’utilisation mémoire dans un format lisible. Notre objectif était de suivre l’évolution des régions mémoire dans le temps. Nous avons donc exécuté :

$ watch -n 1 "pmap -X $pid | (head -n 2 && tail -n +3 | sort -k7 -nr)"

Cette commande exécute pmap toutes les secondes pour afficher des informations mémoire étendues pour un PID donné, ignore l’en-tête et trie la sortie par taille mémoire, ce qui nous permet de nous concentrer sur les régions mémoire les plus volumineuses.

Avec cette commande, nous avons observé un motif intéressant : seuls certains mappings mémoire anonymes grossissaient au fil du temps, et leurs adresses de début changeaient. La taille de ces allocations devenait énorme avec le temps, tandis que la plupart des autres ne bougeaient tout simplement pas.

Les pages mémoire listées par notre pmap, triées par taille RSS, rendaient les allocations suspectes faciles à repérer. Le point orange met en évidence un exemple.

Ce comportement est caractéristique de mremap, un appel système utilisé pour redimensionner ou déplacer des régions mémoire existantes sans les libérer. Contrairement à realloc, qui opère dans le tas et repose sur la gestion mémoire de glibc, mremap fonctionne à un niveau plus bas et est souvent utilisé par des allocateurs personnalisés, des bibliothèques, ou même du code de gestion mémoire manuel pour ajuster dynamiquement l’organisation de la mémoire.

Ce motif pouvait aussi provenir de cycles répétés de munmap suivis de mmap, où la mémoire est libérée puis réallouée, mais où l’utilisation totale continue d’augmenter, en raison d’une fragmentation, de fuites dans des allocateurs personnalisés ou d’une logique de redimensionnement incorrecte. Dans notre cas, les adresses changeantes et la taille croissante suggéraient fortement que la mémoire était réallouée sans jamais être correctement libérée.

C’était notre premier indice concret indiquant que la fuite ne se trouvait pas dans le tas, mais dans des régions mémoire anonymes redimensionnées sans libération correcte.

Tracer la fuite avec BPFtrace.

Notre investigation avait réduit l’origine de la fuite aux appels bruts mmap ou mremap, mais nous devions confirmer lequel était responsable. Notre première tentative a consisté à utiliser LD_PRELOAD avec une petite bibliothèque C personnalisée qui enregistrait chaque appel mmap et mremap, puisque Heaptrack ne le faisait pas, dans l’espoir d’intercepter ceux qui nous intéressaient. Cependant, cette approche avait des limites : les appels mmap/mremap ne passent pas tous par glibc. Nos hooks personnalisés voyaient certaines allocations, mais elles ne correspondaient pas aux adresses qui fuyaient dans notre sortie pmap. Les régions concernées continuaient de croître sans être suivies par notre hook LD_PRELOAD. Cela suggérait que ces allocations étaient effectuées via un syscall manuel, ou qu’un autre mécanisme de hook était en place.

Pour obtenir une vue complète, nous nous sommes tournés vers BPFtrace, un outil de traçage en temps réel des appels système et des événements du noyau. Il s’appuie sur la machine virtuelle eBPF du noyau Linux, qui exécute un bytecode léger et prévalidé attaché à des tracepoints ou à des probes, ce qui permet une analyse sûre avec peu de surcoût. BPFtrace est aussi utilisé par certains outils Kubernetes pour détecter des comportements anormaux ou dangereux dans les clusters, comme des accès non autorisés ou des abus de ressources, sans mettre en risque la stabilité du noyau.

Nous avons aussi envisagé d’utiliser strace, mais sa dépendance à PTRACE le rendait trop lent pour analyser efficacement le problème à ce stade. À la place, nous avons écrit un script BPFtrace pour consigner chaque appel mmap et mremap avec ses arguments et ses traces de pile, y compris les appels qui ne passent pas par glibc. Voici le script que nous avons écrit avec l’aide de Le Chat :

tracepoint:syscalls:sys_enter_mmap /pid == (uint64)$1/ {
    printf("Stack trace:\n%s\n", ustack(perf));
    printf("PID/TID: %d %d | ", pid, tid);
    printf("ENTER mmap(addr=%p, len=%d, prot=%d, flags=%d, fd=%d, off=%d)\n",
           args->addr, args->len, args->prot, args->flags, args->fd, args->off);
}
tracepoint:syscalls:sys_exit_mmap /pid == (uint64)$1/ {
    printf("PID/TID: %d %d | ", pid, tid);
    printf("EXIT mmap: ret=%p\n", args->ret);

}
tracepoint:syscalls:sys_enter_munmap /pid == (uint64)$1/ {
    printf("PID/TID: %d %d | ", pid, tid);
    printf("ENTER munmap(addr=%p, len=%d)\n", args->addr, args->len);
}
tracepoint:syscalls:sys_exit_munmap /pid == (uint64)$1/ {
    printf("PID/TID: %d %d | ", pid, tid);
    printf("EXIT munmap: ret=%d\n", args->ret);
}
tracepoint:syscalls:sys_enter_mremap /pid == (uint64)$1/ {
    printf("PID/TID: %d %d | mremap", pid, tid);
    printf("old_addr=%p, old_len=%d, new_len=%d, flags=%d, new_addr=%p)\n",
           args->addr, args->old_len, args->new_len, args->flags, args->new_addr);
}
tracepoint:syscalls:sys_exit_mremap /pid == (uint64)$1/ {
    printf("PID/TID: %d %d | ", pid, tid);
    printf("EXIT mremap: ret=%p\n", args->ret);
}

Nous avons exécuté le script en tant que root avec cette commande, en remplaçant $pid par le PID du processus worker vLLM qui fuyait selon pmap :

bpftrace /host/script_bpftrace.txt $pid > out_$pid.txt

En résumé, ce script :

• Trace les appels système mmap, munmap et mremap lorsqu’ils entrent dans le noyau.

• Affiche l’ID de thread, le type d’appel, l’adresse demandée et la longueur.

• Affiche la trace de pile en espace utilisateur (ustack dans BPFtrace) pour identifier l’origine de l’appel.

Voici un exemple de sortie de ce script :

Stack trace:
  7ffff7d6b88d syscall+29 (/usr/lib/x86_64-linux-gnu/libc.so.6)
PID/TID: 441359 441359 | ENTER mmap(addr=(nil), len=151552, prot=3, flags=34, fd=-1, off=0)
PID/TID: 441359 441359 | EXIT mmap: ret=0x7fd8a78ee000

Extrait de la sortie du script BPFtrace. Le syscall+29 est important.

À ce stade, synthétisons les informations collectées :

• pmap nous a montré les allocations à la croissance suspecte et leurs adresses de base.

• BPFtrace nous a permis de comprendre que ces adresses étaient obtenues avec des appels mmap, et non mremap. C’était surprenant, car mremap semblait être le suspect idéal pour une allocation mémoire qui continuait de croître.

• Plus intriguant encore, les appels provenaient de syscall+29, alias le wrapper de syscall brut de glibc, qui permet aux utilisateurs d’effectuer un syscall brut via une API comme : syscall(SYS_mmap, ...)

C’était une avancée importante, mais une investigation plus poussée restait nécessaire. BPFtrace nous montrait la trace de pile utilisateur de l’appel système lui-même (ustack dans sa documentation), mais ne nous donnait que le premier élément de la pile d’appels utilisateur, pas le contexte complet en espace utilisateur qui menait à l’allocation. Nous pouvions voir où mmap était appelé, mais pas les appelants précédents. Cette information nous était essentielle. Nous avons donc essayé de comprendre pourquoi nous ne pouvions pas l’obtenir. Au départ, nous pensions que c’était lié aux frame pointers désactivés dans notre configuration, et nous avons décidé d’explorer cette piste.

Pour le contexte, un frame pointer est une fonctionnalité qui stocke l’adresse de retour des appels de fonction dans un registre ou un emplacement mémoire, ce qui permet aux outils de reconstruire la pile d’appels complète. Les frame pointers ont historiquement été désactivés comme optimisation dans la plupart des bibliothèques, car ils ajoutaient un léger surcoût. Les distributions modernes ont toutefois commencé à les réactiver, les gains de performance étant désormais négligeables au regard des bénéfices pour le débogage.

Malheureusement, passer d’Ubuntu 22.04 LTS à Ubuntu 24.04 LTS, où les frame pointers sont activés pour les bibliothèques natives, n’a pas suffi. Cela suggérait qu’une dépendance Python déjà optimisée, qui désactivait les frame pointers, était en cause.

Nous avons réfléchi à la situation pendant un moment : quel package Python pouvait effectuer des syscalls aussi directs, en contournant les appels habituels de la bibliothèque standard ? À ce stade, nous avions deux suspects possibles :

1. UCX (Unified Communication X), une bibliothèque de communication haute performance utilisée pour les réseaux accélérés et RDMA (Remote Direct Memory Access). UCX est une dépendance de NIXL, que vLLM utilise pour le serving désagrégé. UCX est connue pour ses optimisations mémoire bas niveau, y compris des allocateurs mémoire personnalisés.

2. PyTorch, qui effectue sa propre gestion mémoire et ses propres optimisations, en contournant souvent les allocateurs standard pour des raisons de performance. Les allocations personnalisées et la compilation JIT de PyTorch pouvaient aussi être responsables de la fuite.

Avec ces deux suspects en tête, nous devions trouver un moyen d’aller plus loin. À ce stade, nous nous sommes tournés vers l’automatisation de GDB.

L’automatisation de GDB à la rescousse.

Utiliser GDB plus tôt dans l’investigation n’était pas envisageable pour une raison simple : GDB s’attache à un processus entier. Avec vLLM, attacher GDB au processus principal aurait arrêté tous les workers, rendant impossible l’observation de la fuite en temps réel. Comme la fuite ne provoquait pas de crash, et que les dumps mémoire étaient impraticables en raison de la taille du processus, nous étions bloqués.

Cependant, à partir de nos logs BPFtrace, nous avons remarqué que chaque appel mmap qui fuyait provenait de la même adresse. Comme indiqué dans la section précédente, cette adresse ne pointait pas directement vers mmap, mais se trouvait dans le thin-wrapper syscall de glibc. Cette fonction contourne le wrapper mmap habituel de glibc, ce qui explique pourquoi nos hooks LD_PRELOAD l’avaient manquée plus tôt. Nous avons essayé d’utiliser LD_PRELOAD pour intercepter le syscall de glibc, mais curieusement, cela ne fonctionnait pas non plus… Il ne nous restait presque aucun moyen de tracer ces appels dynamiquement.

Comme les appels qui fuyaient provenaient toujours de la même instruction syscall, nous pouvions automatiser GDB pour interrompre l’exécution uniquement lorsque cette adresse précise était atteinte. Voici comment nous avons procédé :

1. Nous avons défini un point d’arrêt conditionnel sur l’adresse syscall, déclenché uniquement si le numéro d’appel système correspondait à SYS_mmap.

2. Nous avons interrompu temporairement l’exécution à la sortie de l’appel système mmap pour inspecter la valeur de retour, c’est-à-dire l’adresse allouée, et afficher la trace de pile complète.

3. Nous avons exporté tout le contexte disponible en une seule fois : l’adresse de retour, la pile d’appels et tout autre registre ou emplacement mémoire pertinent.

4. Nous avons exécuté ce script conditionnel pendant quelques secondes durant le benchmark, puis comparé les adresses de retour capturées avec notre surveillance pmap pour confirmer si elles correspondaient aux régions connues comme fuyantes.

Voici le script, qui peut être exécuté via gdb -x gdb_script.txt :

# Attach to the process
attach 2199304

# Open a log file for output
set logging file syscall_output.txt
set logging on

# Set a conditional breakpoint on syscall for rdi == 9 (mmap)
break syscall if $rdi == 9

# Commands for the syscall breakpoint
commands
  silent
  # Set a temporary breakpoint at the return point
  tbreak *0x00007ffff7d9525d
  # Commands for the temporary breakpoint
  commands $bpnum + 1
    silent
    set $ret_val = $rax
    bt
    printf "Syscall returned: rax = 0x%012lx\n", $ret_val
    continue
  end
  continue
end

# Run the process
continue

Cette approche nous a apporté deux bénéfices majeurs :

• Nous avons capturé la trace de pile complète en espace utilisateur au moment de l’allocation, ce que BPFtrace ne pouvait pas fournir de manière fiable.

• Nous pouvions aussi recouper les adresses renvoyées avec notre sortie pmap pour confirmer si elles correspondaient aux régions fuyantes.

En résumé, nous avons transformé GDB en observateur ciblé et non intrusif : il ne s’arrêtait que lorsque la fuite se produisait, affichait tout ce dont nous avions besoin, puis laissait le processus continuer. Cela nous a enfin permis de relier les appels mmap aux régions anonymes en croissance que nous avions vues dans pmap.

La première trace de pile montrait Python (ligne #5) invoquant mmap via UCX (ligne #4), ce qui était inattendu, car Python devrait normalement appeler directement le mmap de glibc.

#0  syscall () at ../sysdeps/unix/sysv/linux/x86_64/syscall.S:29
#1  0x00007ffc61759ac2 in ucm_orig_mmap_syscall () from /.venv/lib/python3.12/site-packages/.nixl.mesonpy.libs/plugins/../../nixl.libs/libucm-e091ff91.so.0.0.0
#2  0x00007ffc61753bd1 in ?? () from /.venv/lib/python3.12/site-packages/.nixl.mesonpy.libs/plugins/../../nixl.libs/libucm-e091ff91.so.0.0.0
#3  0x00007ffc61753e3b in ucm_event_dispatch () from /.venv/lib/python3.12/site-packages/.nixl.mesonpy.libs/plugins/../../nixl.libs/libucm-e091ff91.so.0.0.0
#4  0x00007ffc61754009 in ucm_mmap () from /.venv/lib/python3.12/site-packages/.nixl.mesonpy.libs/plugins/../../nixl.libs/libucm-e091ff91.so.0.0.0
#5  0x0000000000674ac0 in _PyMem_ArenaAlloc (_unused_ctx=<optimized out>, size=<optimized out>) at ../Objects/obmalloc.c:138
...

La deuxième trace de pile était encore plus déroutante : Python (ligne #8) appelait munmap via UCX (ligne #7), tout en déclenchant d’une manière ou d’une autre une allocation mmap (ligne #1) pendant l’opération :

#0  syscall () at ../sysdeps/unix/sysv/linux/x86_64/syscall.S:38
#1  0x00007ffc60f58ac2 in ucm_orig_mmap_syscall () from /.venv/lib/python3.12/site-packages/.nixl.mesonpy.libs/plugins/../../nixl.libs/libucm-e091ff91.so.0.0.0
#2  0x00007ffc60fae47c in ?? () from /.venv/lib/python3.12/site-packages/.nixl.mesonpy.libs/plugins/../../nixl.libs/libucs-311e600f.so.0.0.0
#3  0x00007ffc60f9b9c4 in ucs_mpool_grow () from /.venv/lib/python3.12/site-packages/.nixl.mesonpy.libs/plugins/../../nixl.libs/libucs-311e600f.so.0.0.0
#4  0x00007ffc60f9bbf5 in ucs_mpool_get_grow () from /.venv/lib/python3.12/site-packages/.nixl.mesonpy.libs/plugins/../../nixl.libs/libucs-311e600f.so.0.0.0
#5  0x00007ffc60fafe2c in ?? () from /.venv/lib/python3.12/site-packages/.nixl.mesonpy.libs/plugins/../../nixl.libs/libucs-311e600f.so.0.0.0
#6  0x00007ffc60f52e3b in ucm_event_dispatch () from /.venv/lib/python3.12/site-packages/.nixl.mesonpy.libs/plugins/../../nixl.libs/libucm-e091ff91.so.0.0.0
#7  0x00007ffc60f5313b in ucm_munmap () from /.venv/lib/python3.12/site-packages/.nixl.mesonpy.libs/plugins/../../nixl.libs/libucm-e091ff91.so.0.0.0
#8  0x0000000000607d15 in _PyThreadState_PopFrame (tstate=0xba6ac8 <_PyRuntime+459656>, frame=<optimized out>) at ../Python/pystate.c:2992
...

C’était inattendu, car munmap sert à libérer de la mémoire, pas à en allouer. Observer un appel mmap pendant une opération munmap dans UCM, le module de gestion mémoire d’UCX, suggérait qu’un problème se produisait dans la gestion du pool mémoire d’UCX.

Identifier le rôle des hooks mémoire d’UCX.

Nous avons partagé ces résultats avec l’équipe vLLM, qui nous a aidés à valider et affiner notre compréhension du problème. Ensemble, nous avons découvert qu’UCX utilise un mécanisme de hook mmap pour optimiser les opérations mémoire avec InfiniBand, notamment en prémettant en cache les données pour les transferts, via une fonctionnalité appelée Registration Cache, ou RCache. La gestion mémoire avec InfiniBand est souvent coûteuse, car elle nécessite un enregistrement mémoire au niveau matériel.

Cependant, ce mécanisme intercepte tous les appels mmap par défaut, pas seulement ceux liés aux opérations UCX ou InfiniBand. Cette interception large explique pourquoi nos tentatives précédentes de tracer les allocations avec notre hook basé sur LD_PRELOAD pour mmap ont échoué : UCX patch dynamiquement les entrées de la Global Offset Table (GOT) utilisées par les applications pour appeler des fonctions comme mmap et munmap. C’est pourquoi il contournait entièrement nos hooks précédents. Nous avons aussi découvert que ce mécanisme de hook est automatiquement désactivé lorsque Valgrind est détecté, ce qui nous aurait empêchés d’utiliser Valgrind pour une analyse plus approfondie.

La GOT est une structure de données utilisée par l’éditeur de liens dynamique pour résoudre les appels de fonctions dans les bibliothèques liées dynamiquement. Lorsqu’un programme démarre, l’éditeur de liens dynamique remplit la GOT avec les adresses réelles des fonctions, comme mmap, provenant de bibliothèques partagées.

Modifier la GOT à l’exécution est généralement considéré comme une mauvaise pratique, car cela peut introduire de l’instabilité, compliquer le débogage et casser des hypothèses sur lesquelles s’appuient d’autres parties du programme ou des bibliothèques. Mais UCX le fait pour une bonne raison. UCX utilise ce mécanisme pour gérer son cache d’enregistrement, qui suit la mémoire « enregistrée » ou « verrouillée ». Il s’agit de mémoire verrouillée en place, ce qui garantit que sa correspondance entre adresses virtuelles et physiques reste fixe. Les adaptateurs réseau peuvent ainsi transférer les données directement entre la fabric réseau et la RAM, sans intervention du CPU. C’est essentiel pour les performances, mais cela exige une gestion prudente, car la mémoire enregistrée est une ressource limitée.

Résoudre le problème.

Une fois que nous avons compris que le mécanisme de hook de mmap était en jeu, nous avons réalisé que c’était une bonne nouvelle : nous pouvions le désactiver entièrement en définissant cette variable d’environnement UCX_MEM_MMAP_HOOK_MODE=none. Cela a bien désactivé ce comportement, ce qui a résolu la fuite mémoire sans impact sur les performances. Les hooks mmap sont utiles lorsqu’il faut envoyer plusieurs fragments de mémoire via RDMA, mais dans le cas d’usage de vLLM, nous devons seulement gérer une grande région mémoire contiguë : toute la mémoire du gestionnaire KVCache de vLLM. NIXL devait enregistrer la mémoire une seule fois pour ses transferts. Dans le cas d’usage de vLLM, désactiver le mécanisme de hook était donc sûr et n’avait aucun impact négatif sur les performances du serving désagrégé.

Après en avoir discuté avec l’équipe UCX, nous avons appris qu’UCX ne libère pas immédiatement la mémoire lorsque munmap est appelé. UCX déplace plutôt la région vers une file d’invalidation pour nettoyage ultérieur. Cette file est gérée par le pool mémoire d’UCX, qui s’étend dynamiquement pour accueillir davantage d’entrées si nécessaire. Résultat : les régions mémoire s’accumulaient sans être libérées, et la file en croissance nécessitait des allocations supplémentaires, ce qui expliquait pourquoi mmap était appelé pendant les opérations munmap. Une autre solution à la fuite mémoire consiste à définir la variable d’environnement UCX_RCACHE_MAX_UNRELEASED=1024 (la valeur par défaut est inf) afin de limiter le nombre de régions mémoire non libérées dans la file, ce qui force UCX à lancer le nettoyage une fois le seuil atteint.

Le problème, c’est que cela n’aurait pas dû se produire. NIXL et vLLM appelaient bien la fonction ucp_worker_progress(), qui aurait dû déclencher le nettoyage du pool mémoire. Nous ne savons toujours pas pourquoi ce nettoyage ne s’est pas déclenché dans ce cas limite précis. Mais cela a montré que définir la valeur par défaut de UCX_RCACHE_MAX_UNRELEASED à l’infini n’était pas correct. Les équipes UCX et NIXL ont décidé de changer ce comportement dans une future version de NIXL. En attendant, nous avons effectué le merge d’un correctif dans le dépôt vLLM pour aider la communauté à éviter la même fuite. 

Résumé de l’enquête.

L’enquête en bref :

1. Nous avons constaté une fuite mémoire qui augmentait rapidement dans notre environnement de production lors du déploiement de l’un de nos modèles de dernière génération avec serving désagrégé.

2. Nous avons essayé de réduire l’environnement et d’obtenir un reproducteur minimal. Malheureusement, le bug ne se reproduisait que sur une configuration complexe, avec un grand modèle et la désagrégation activée.

3. Nous nous sommes tournés vers Memray, Guppy 3 et Heaptrack pour analyser la fuite. Ces outils n’ont rien révélé d’évident. Nous avons toutefois remarqué un point particulier dans les métriques de Heaptrack. La mémoire résidente était anormalement élevée, nous avons donc décidé d’explorer cette piste.

4. Avec pmap, nous avons pu voir des allocations RSS qui continuaient à augmenter, avec leurs pointeurs de base associés.

5. Nous voulions obtenir plus d’indices sur l’origine de ces appels. Avec un peu de scripting et BPFtrace, nous avons découvert que les fuites provenaient d’appels mmap. Malgré tous nos efforts, nous n’avons pas réussi à collecter des stack traces complètes, qui nous auraient permis d’identifier précisément le site d’appel responsable, mais cela nous a conduits à penser que les appels système étaient effectués par un package fortement optimisé.

6. À partir des informations recueillies, nous avons pu configurer des breakpoints GDB très spécifiques qui ne se déclenchaient que sur les appels responsables.

7. Nous avons découvert qu’UCX effectuait ces appels. Leur objectif était légitime : améliorer les performances des transferts InfiniBand. Mais ils créaient des allocations RSS en croissance continue et nous ont empêchés de déployer le serving désagrégé pendant plusieurs jours.

8. Une fois la source de la fuite identifiée, la solution de contournement était simple à trouver : le fait de définir UCX_MEM_MMAP_HOOK_MODE=none a corrigé notre problème. Nous avons discuté de cette enquête dans le dépôt vLLM et effectué le merge d’un patch pour la communauté.

Ce que nous avons appris.

Les stacks logicielles modernes reposent sur des couches de dépendances, chacune ajoutant de la complexité et des points de défaillance potentiels. Ces abstractions améliorent fortement la productivité des développeurs, mais elles ne les isolent pas totalement des problèmes sous-jacents de la stack. Il est donc essentiel d’être prêt à aller en profondeur lors du débogage. Dans ces environnements, c’est toutefois rarement simple, en particulier lorsque les optimisations de performance introduisent des cas limites subtils. UCX en est un bon exemple. Sa conception privilégie les performances, mais sa façon d’intercepter les appels mmap peut créer des risques difficiles à tracer. Cette expérience a montré une nouvelle fois à quel point le diagnostic peut être difficile dans des systèmes fortement interconnectés.

Cette enquête montre aussi l’importance de la transparence et de la collaboration lorsqu’il faut travailler avec des dépendances critiques pour les performances. Nous remercions les équipes vLLM, NIXL et UCX pour leur collaboration dans la confirmation et la correction de ce comportement. Leur expertise a été déterminante pour parvenir à une résolution, et nous sommes heureux de poursuivre ce travail ensemble.

Nous remercions les personnes suivantes pour leur aide et leur collaboration dans la résolution de ce problème :

• Robert Shaw (Red Hat, mainteneur de vLLM et llm-d)

• Will Eaton (Red Hat, mainteneur de vLLM et llm-d)

• Nicolò Lucchesi (Red Hat, mainteneur de vLLM et llm-d)

• Mikhail Brinskii (NVIDIA, mainteneur de NIXL)

• Leonid Genkin (NVIDIA, mainteneur d’UCX)

• Nathan Bellalou (NVIDIA, mainteneur d’UCX et NIXL)

Nous recrutons

Ces défis vous intéressent ? Rejoignez Mistral AI et contribuez à construire l’avenir de l’infrastructure IA. Nous recherchons toujours des ingénieurs et chercheurs talentueux pour collaborer sur des projets de dernière génération.