Gli heap mentono: debug di una perdita di memoria in vLLM.

Qualche mese fa, il nostro team ha indagato su una sospetta perdita di memoria in vLLM. Inizialmente pensavamo che il problema sarebbe stato facile da individuare, qualcosa di confinato agli strati superiori del codebase. Ma più approfondivamo, più diventava complesso. Questo articolo inaugura la nostra nuova serie Engineering Deep Dive, in cui condivideremo come affrontiamo le indagini tecniche e costruiamo soluzioni in Mistral AI.

Il problema è emerso per la prima volta durante i test di pre-produzione del serving disaggregato con uno dei nostri modelli all'avanguardia. L'utilizzo della memoria aumentava costantemente, ma solo in condizioni specifiche: con vLLM, con il nostro modello Mistral Medium 3.1 e con la compilazione del grafo abilitata. Non si verificavano crash o errori, solo un lento incremento lineare della memoria di sistema di 400 MB al minuto su traffico simile alla produzione. Dopo alcune ore, questo avrebbe portato a uno stato di "out of memory".

Ne è seguita una ricerca metodica, partita da strumenti Python di alto livello e arrivata fino al tracing a livello di kernel, finché non abbiamo finalmente scoperto la vera origine. Ecco come l’abbiamo individuata e cosa ha rivelato sui rischi nascosti dei livelli di dipendenze nel software di oggi.

Non il tipo di grafico in salita che ci piace vedere su Grafana

Una perdita che giocava a nascondino.

Inizialmente, il nostro approccio ha seguito un percorso standard di troubleshooting: puntavamo a isolare l'origine della perdita replicando il problema su un modello più piccolo, con meno ottimizzazioni di produzione attivate. Ma dopo aver provato impostazioni e modelli diversi non siamo riusciti a riprodurlo su un'altra configurazione. L'errore era presente solo in una configurazione Prefill/Decode disaggregata con NIXL.

Dato il ruolo centrale che la disaggregazione Prefill/Decode (P/D) svolge nella nostra storia, illustriamo i meccanismi di alto livello con cui funziona questa configurazione di inferenza. P/D Disaggregated suddivide l'elaborazione di una query in due fasi, gestite da istanze diverse:

1. Per prima cosa il router invia una “prefill request” (impostando max_tokens=1 e impostando un insieme vuoto di metadati KV Transfer) a un'istanza vLLM di prefill per calcolare la KVCache della richiesta.

2. Al completamento, il router trasferisce i metadati KVCache insieme a una “decode request” a un'istanza vLLM di decode.

3. Il trasferimento della KVCache viene avviato tramite NIXL e la generazione dei token avviene sull'istanza vLLM di decode utilizzando ed estendendo la KVCache trasferita.

La perdita è stata osservata solo sul lato decode di questa configurazione disaggregata, suggerendo fortemente che il trasferimento della KV Cache tramite NIXL fosse la causa principale. Nella nostra configurazione NIXL si basa su UCX (Unified Communication X), una libreria di comunicazione ad alte prestazioni progettata per lo scambio di dati nei sistemi distribuiti. UCX consente un trasferimento dati ottimizzato su un ampio insieme di tecnologie, tra cui Infiniband, una tecnologia di interconnessione a bassa latenza e alta larghezza di banda comunemente utilizzata in HPC e nei data center.

Panoramica di un deployment di serving P/D Disaggregated.

Per il resto della nostra indagine abbiamo lavorato in questa configurazione e abbiamo iniziato con strumenti di profiling della memoria Python per individuare l'origine della perdita.

Abbiamo provato Memray e Guppy 3, ma nessuno dei due ha mostrato una perdita e tutto ciò che ci hanno permesso di osservare era normale. Il tentativo di usare GDB ha fatto crashare l'intero processo. Anche la nostra configurazione vLLM era troppo pesante per strumenti come Valgrind, rendendoli impraticabilmente lenti o persino impossibili da usare.

Era chiaro che serviva uno strumento più potente per tracciare la perdita. Ma prima di investire altro tempo, abbiamo deciso di assicurarci che questa perdita fosse riproducibile da altri. Abbiamo contattato il team vLLM aprendo una issue nel loro repository GitHub, il che ha aiutato a confermare che non eravamo gli unici a riscontrare questo problema e che era giustificata un'indagine più approfondita.

Conteggio di malloc e free con Heaptrack.

Per tracciare meglio ciò che stava accadendo, ci siamo rivolti a Heaptrack: un profiler di memoria che sovrascrive operazioni di memoria come malloc o free e registra questi eventi insieme agli stack trace.

Millian Wolf, il creatore di Heaptrack, ha scritto un eccellente blog post introduttivo per iniziare a usare lo strumento. È un processo in due passaggi: prima si esegue il programma con il tracing, poi si interpretano i dump dei dati.

Per tracciare le allocazioni confinate nel processo worker di vLLM, abbiamo impostato LD_PRELOAD su libheaptrack_preload.so tramite vLLM per garantire che questa libreria venisse caricata prima di qualsiasi altra e sovrascrivesse i comportamenti delle funzioni di allocazione della memoria, fornendoci il dump dei dati.

Abbiamo quindi potuto visualizzare questi dati tramite 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 fornisce un grafico dettagliato e interattivo di tutte le allocazioni heap, fino al livello di funzione. Possiamo tracciare ogni malloc e free, con una scomposizione chiara dell'utilizzo della memoria.

L'utilizzo della memoria mostrato con Heaptrack per il nostro worker vLLM

A questo punto, ci si potrebbe chiedere: dov'è la perdita di memoria? In effetti, l'unico aumento di memoria visibile era dovuto a un'inizializzazione lazy di NIXL.

Per verificare che la perdita si stesse effettivamente verificando in questa configurazione, abbiamo eseguito un benchmark vLLM e creato due snapshot Heaptrack usando heaptrack_interpret: uno all'inizio e uno verso la fine. Sebbene la memoria heap stessa sia rimasta stabile, il picco di memoria residente (RSS), che tratteremo nella sezione successiva, differiva tra i due snapshot. Questa discrepanza era visibile nella scheda di riepilogo di Heaptrack.

Discrepanza del picco RSS in Heaptrack: prima (1) e dopo (2) il benchmark

Questo significava che la perdita si verificava al di fuori dell'heap e quindi non faceva parte della memoria analizzata da Heaptrack. Dovevamo cambiare strumenti per tracciare le allocazioni esterne all'heap.

Oltre l'heap: comprendere la memoria residente e le allocazioni di sistema.

Per capire perché Heaptrack non riuscisse a rilevare la perdita, dobbiamo prima chiarire cosa include effettivamente il Resident Set Size (RSS). L'RSS rappresenta la porzione di memoria di un processo mantenuta in RAM e contiene più del solo heap. In particolare, include:

• L'heap è tradizionalmente gestito usando le chiamate di sistema legacy sbrk e brk, che regolano o impostano il program break address, il puntatore che segna la fine del segmento heap.

• Lo stack, che memorizza variabili locali e frame di chiamata delle funzioni.

• I mapping di memoria anonimi, ovvero regioni di memoria allocate direttamente tramite la chiamata di sistema mmap senza un file di backing. Sono spesso usati da allocatori personalizzati o da malloc per blocchi di memoria più grandi. Gli indirizzi dei mapping anonimi risiedono di solito tra lo spazio di indirizzi dell'heap e lo spazio di indirizzi dello stack, in una regione nota come segmento di mapping della memoria.

Sebbene malloc possa usare sbrk per piccole allocazioni, le implementazioni moderne in genere preferiscono usare mmap con mapping anonimi, perché è più flessibile e consente di allocare huge pages (pagine di memoria da 2 MB o 1 GB a seconda della configurazione).

Heaptrack esegue hook solo sulle funzioni malloc e free di glibc. Questo significa che può tracciare tutte le allocazioni tradizionali su heap e i mapping anonimi allocati direttamente da malloc, ma non rileva la memoria allocata tramite chiamate dirette a mmap o altri meccanismi a livello di sistema fuori dal controllo di glibc.

Fortunatamente, non tutto è perduto quando si tratta di tracciare ciò che stava accadendo. Il filesystem /proc è una cartella speciale in Linux che funge da API del kernel, esponendo un'interfaccia virtuale per interagire con i processi in esecuzione e fornendo accesso in tempo reale ai dettagli dei processi, come:

• /proc/<pid>/fd, che elenca tutti i file descriptor aperti per un processo.

• /proc/<pid>/maps, che mostra una mappa dettagliata delle regioni di memoria del processo, inclusi heap, stack, librerie condivise e mapping anonimi.

• E molto altro (un semplice ls in /proc/<pid>/ elenca ciò che è disponibile).

Per continuare la nostra indagine, abbiamo usato il comando pmap, che legge /proc/<pid>/maps e presenta l'utilizzo della memoria in un formato leggibile. Il nostro obiettivo era tracciare i cambiamenti nelle regioni di memoria nel tempo, quindi abbiamo eseguito:

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

Questo comando esegue pmap ogni secondo per mostrare informazioni di memoria estese per un PID specificato, salta l'header e ordina l'output per dimensione della memoria, permettendoci di concentrarci sulle regioni di memoria più grandi.

Con questo comando abbiamo osservato uno schema interessante: solo alcuni mapping di memoria anonimi crescevano nel tempo e i loro indirizzi iniziali cambiavano. La dimensione di queste allocazioni è diventata enorme nel tempo, mentre la maggior parte delle altre semplicemente non si muoveva.

Le pagine di memoria elencate dal nostro pmap comando, ordinate per dimensione RSS, hanno reso facile individuare allocazioni sospette. Il punto arancione evidenzia un esempio.

Questo comportamento è caratteristico di mremap, una chiamata di sistema usata per ridimensionare o rilocare regioni di memoria esistenti senza liberarle. A differenza di realloc, che opera all'interno dell'heap e si basa sulla gestione della memoria di glibc, mremap funziona a un livello più basso ed è spesso usata da allocatori personalizzati, librerie o persino codice di gestione manuale della memoria per regolare dinamicamente i layout di memoria.

Questo schema potrebbe anche derivare da cicli ripetuti di munmap seguiti da mmap, in cui la memoria viene liberata e riallocata ma l'utilizzo totale continua a crescere, a causa di frammentazione, perdite in allocatori personalizzati o logica di ridimensionamento impropria. Nel nostro caso, gli indirizzi variabili e la dimensione crescente suggerivano fortemente che la memoria venisse riallocata ma mai rilasciata correttamente.

Questa è stata la nostra prima indicazione concreta che la perdita non era nell'heap, ma in regioni di memoria anonime ridimensionate senza un rilascio adeguato.

Tracing della perdita con BPFtrace.

La nostra indagine aveva ristretto l’origine del leak alle chiamate raw mmap o mremap, ma dovevamo confermare quale delle due fosse responsabile. Il primo tentativo è stato usare LD_PRELOAD con una piccola libreria C custom che registrava ogni chiamata mmap e mremap, poiché Heaptrack non lo faceva, sperando di intercettare quelle che ci interessavano. Tuttavia, questo approccio aveva dei limiti: non tutte le mmap/mremap passano attraverso glibc. I nostri hook custom vedevano alcune allocazioni, ma non corrispondevano agli indirizzi che risultavano in leak nell’output di pmap. Le regioni interessate dal leak continuavano a crescere, senza essere tracciate dal nostro hook LD_PRELOAD. Questo suggeriva che tali allocazioni fossero effettuate eseguendo manualmente una syscall o che fosse in atto un altro meccanismo di hooking.

Per ottenere una visione completa, ci siamo rivolti a BPFtrace, uno strumento per il tracing in tempo reale delle chiamate di sistema e degli eventi del kernel. Si basa sulla macchina virtuale eBPF del kernel Linux, che esegue bytecode leggero e pre-validato collegato a tracepoint o probe, consentendo analisi sicure e a basso overhead. BPFtrace è usato anche da alcuni strumenti Kubernetes per rilevare comportamenti anomali o pericolosi nei cluster, come accessi non autorizzati o abuso di risorse, senza rischiare la stabilità del kernel.

Abbiamo considerato anche l’uso di strace, ma la sua dipendenza da PTRACE lo rendeva troppo lento per analizzare efficacemente il problema in questa fase. Abbiamo quindi scritto uno script BPFtrace per fare log di ogni chiamata mmap e mremap con i relativi argomenti e stack trace, incluse le chiamate che non passano attraverso glibc. Ecco lo script che abbiamo scritto con l’aiuto di 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);
}

Abbiamo eseguito lo script come root con questo comando, sostituendo $pid con il PID del processo worker vLLM che risultava in leak secondo pmap:

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

In sostanza, questo script:

• Traccia le chiamate di sistema mmap, munmap e mremap mentre entrano nel kernel.

• Stampa l’ID del thread, il tipo di chiamata, l’indirizzo richiesto e la lunghezza.

• Stampa lo stack trace user-space (ustack in BPFtrace) per identificare da dove ha avuto origine la chiamata.

Ecco un esempio di output di questo 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

Estratto dell’output dello script BPFtrace. Il syscall+29 è importante.

A questo punto, sintetizziamo le informazioni che abbiamo raccolto:

• pmap ci ha mostrato le allocazioni che crescevano in modo sospetto e i relativi indirizzi di base.

• BPFtrace ci ha permesso di capire che questi indirizzi erano ottenuti con chiamate mmap, non mremap. Per noi è stata una sorpresa, perché mremap sembrava il sospetto ideale per un’allocazione di memoria che continuava a crescere.

• Ancora più interessante, le chiamate avevano origine da syscall+29, ovvero il wrapper raw syscall di glibc, che consente agli utenti di eseguire una syscall raw tramite un’API come: syscall(SYS_mmap, ...)

Questo rappresentava un progresso significativo, ma erano necessarie ulteriori indagini. Sebbene BPFtrace ci mostrasse lo stack trace utente della chiamata di sistema stessa (ustack nella relativa documentazione), forniva solo il primo elemento dello stack di chiamate utente, non il contesto user-space completo che portava all’allocazione. Potevamo vedere dove veniva chiamata mmap, ma non i chiamanti precedenti. Quell’informazione per noi sarebbe stata essenziale, quindi abbiamo cercato di capire perché non riuscissimo a ottenerla e inizialmente abbiamo pensato che fosse legata ai frame pointer disabilitati nel nostro setup, decidendo di indagare in quella direzione.

Per contestualizzare, il frame pointer è una funzionalità che memorizza l’indirizzo di ritorno delle chiamate di funzione in un registro o in una posizione di memoria, consentendo agli strumenti di ricostruire l’intero call stack. Storicamente, i frame pointer erano disabilitati come ottimizzazione per la maggior parte delle librerie, perché introducevano un overhead minimo, ma le distribuzioni moderne hanno iniziato a riabilitarli, dato che i guadagni prestazionali sono ormai trascurabili rispetto ai benefici per il debugging.

Purtroppo, passare da Ubuntu 22.04 LTS a Ubuntu 24.04 LTS, che ha i frame pointer abilitati per le librerie native, non è stato sufficiente. Questo suggeriva che il colpevole fosse una dipendenza Python già ottimizzata, con frame pointer disabilitati.

Abbiamo riflettuto per un po’ sulla situazione: quale pacchetto Python poteva eseguire syscall così dirette, aggirando le consuete chiamate della libreria standard? A questo punto, avevamo due potenziali colpevoli:

1. UCX (Unified Communication X), una libreria di comunicazione ad alte prestazioni usata per networking accelerato e RDMA (Remote Direct Memory Access). UCX è una dipendenza di NIXL, che vLLM usa per il serving disaggregato. UCX è nota per le sue ottimizzazioni di memoria a basso livello, inclusi allocator di memoria custom.

2. PyTorch, che esegue la propria gestione e le proprie ottimizzazioni della memoria, spesso bypassando gli allocator standard per motivi di performance. Anche le allocazioni custom e la compilazione JIT di PyTorch potevano essere responsabili del leak.

Con entrambi i sospetti sul tavolo, ci serviva un modo per andare più a fondo. A questo punto ci siamo invece rivolti all’automazione di GDB.

Automazione GDB in soccorso.

Usare GDB nelle fasi precedenti dell’indagine non era praticabile per un motivo semplice: GDB si collega a un intero processo. Con vLLM, collegare GDB al processo principale avrebbe bloccato tutti i worker, rendendo impossibile osservare il leak in tempo reale. Poiché il leak non causava un crash e i dump di memoria erano impraticabili per via delle dimensioni enormi del processo, eravamo bloccati.

Tuttavia, dai nostri log BPFtrace abbiamo notato che ogni chiamata mmap interessata dal leak aveva origine dallo stesso indirizzo. Come detto nella sezione precedente, questo indirizzo non puntava direttamente a mmap, ma si trovava all’interno del thin wrapper syscall di glibc. Questa funzione bypassa il consueto wrapper mmap di glibc, il che spiega perché in precedenza i nostri hook LD_PRELOAD non l’avessero intercettata. Abbiamo provato a usare LD_PRELOAD per intercettare la syscall di glibc, ma stranamente nemmeno questo funzionava, lasciandoci quasi senza alcun modo di tracciare queste chiamate dinamicamente.

Poiché le chiamate interessate dal leak provenivano sempre dalla stessa istruzione syscall, potevamo automatizzare GDB affinché si interrompesse solo quando veniva raggiunto quello specifico indirizzo. Ecco come abbiamo proceduto:

1. Abbiamo impostato un breakpoint condizionale sull’indirizzo syscall, attivato solo se il numero della chiamata di sistema corrispondeva a SYS_mmap.

2. Abbiamo eseguito temporaneamente un break all’uscita della chiamata di sistema mmap per ispezionare il valore di ritorno, cioè l’indirizzo allocato, e stampato lo stack trace completo.

3. Abbiamo eseguito il dump di tutto il contesto disponibile in una sola volta: l’indirizzo di ritorno, il call stack e qualsiasi altro registro o memoria rilevante.

4. Abbiamo eseguito questo script condizionale per alcuni secondi durante il benchmark e confrontato gli indirizzi di ritorno acquisiti con il nostro monitoraggio pmap per confermare se corrispondessero alle regioni note interessate dal leak.

Ecco lo script, che può essere eseguito tramite 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

Questo approccio ci ha dato due vantaggi principali:

• Abbiamo acquisito lo stack trace user-space completo al momento dell’allocazione, qualcosa che BPFtrace non riusciva a fornire in modo affidabile.

• Potevamo inoltre confrontare gli indirizzi restituiti con il nostro output pmap per confermare se corrispondessero alle regioni interessate dal leak.

In breve, abbiamo trasformato GDB in un osservatore mirato e non intrusivo, che si interrompeva solo quando si verificava il leak, stampava tutto ciò di cui avevamo bisogno e lasciava proseguire il processo. Questo ci ha finalmente permesso di collegare i punti tra le chiamate mmap e le regioni anonime in crescita che avevamo visto in pmap.

Il primo stack trace mostrava Python (riga #5) invocare mmap tramite UCX (riga #4), cosa inattesa poiché Python, in circostanze normali, dovrebbe chiamare direttamente la mmap di 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
...

Ancora più confuso era il secondo stack trace, in cui Python (riga #8) chiamava munmap tramite UCX (riga #7), ma in qualche modo attivava un’allocazione mmap (riga #1) durante il processo:

#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
...

Questo era inatteso, poiché munmap serve a liberare memoria, non ad allocarla. Vedere una chiamata mmap durante un’operazione munmap all’interno di UCM, il modulo di gestione della memoria di UCX, suggeriva che qualcosa non stesse funzionando correttamente nella gestione del memory pool di UCX.

Identificazione del ruolo dei memory hook di UCX.

Abbiamo condiviso questi risultati con il team vLLM, che ci ha aiutato a validare e affinare la nostra comprensione del problema. Insieme, abbiamo scoperto che UCX usa un meccanismo di hooking di mmap per ottimizzare le operazioni di memoria per InfiniBand, inclusa la pre-cache dei dati per i trasferimenti, una funzionalità chiamata Registration Cache, o RCache. La gestione della memoria per InfiniBand è spesso costosa a causa della necessità di registrare la memoria a livello hardware.

Tuttavia, questo meccanismo intercetta tutte le chiamate mmap per impostazione predefinita, non solo quelle relative a operazioni UCX o InfiniBand. Questa ampia intercettazione spiega perché i nostri tentativi precedenti di tracciare le allocazioni usando il nostro hook LD_PRELOAD basato su mmap fossero falliti: UCX applica patch dinamicamente alle voci della Global Offset Table (GOT) usate dalle applicazioni per chiamare funzioni come mmap e munmap. Per questo bypassava completamente i nostri hook precedenti. Abbiamo anche scoperto che questo meccanismo di hooking viene disabilitato automaticamente quando viene rilevato Valgrind, cosa che ci avrebbe impedito di usare Valgrind per un’analisi più approfondita.

La GOT è una struttura dati usata dal linker dinamico per risolvere le chiamate di funzione nelle librerie collegate dinamicamente. Quando un programma si avvia, il linker dinamico popola la GOT con gli indirizzi effettivi delle funzioni, come mmap, provenienti dalle librerie condivise.

Modificare la GOT a runtime è generalmente considerato una cattiva pratica, perché può introdurre instabilità, rendere il debugging più difficile e violare assunzioni su cui si basano altre parti del programma o le librerie. Ma UCX lo fa per una buona ragione. Infatti, UCX lo fa per gestire la propria Registration Cache, che tiene traccia della memoria "registrata" (o "pinned"). Si tratta di memoria bloccata in posizione, in modo che la sua mappatura da indirizzo virtuale a fisico rimanga fissa. Questo consente agli adattatori di rete di trasferire dati direttamente tra il network fabric e la RAM senza coinvolgere la CPU, un aspetto critico per le prestazioni ma che richiede una gestione attenta, poiché la memoria registrata è una risorsa limitata.

Risoluzione del problema.

Una volta compreso che era in gioco il meccanismo di hooking di mmap, ci siamo resi conto che in realtà era una buona notizia, perché potevamo disabilitarlo completamente impostando questa variabile d’ambiente UCX_MEM_MMAP_HOOK_MODE=none. Il comportamento è stato disabilitato correttamente, risolvendo la perdita di memoria senza impattare sulle prestazioni. Gli hook di mmap sono utili quando si inviano vari blocchi di memoria tramite RDMA, ma nel caso d’uso di vLLM dobbiamo gestire solo un’unica grande regione di memoria contigua: l’intera memoria del vLLM KVCache Manager. NIXL doveva registrare la memoria una sola volta per i propri trasferimenti. Pertanto, nel caso d’uso di vLLM, disabilitare il meccanismo di hooking era sicuro e non aveva alcun impatto negativo sulle prestazioni del serving disaggregato.

Dopo averne discusso con il team UCX, abbiamo appreso che UCX non libera immediatamente la memoria quando viene chiamato munmap. Sposta invece la regione in una coda di invalidazione per una pulizia successiva. Questa coda è gestita dal memory pool di UCX, che si espande dinamicamente per accogliere più voci quando necessario. Di conseguenza, le regioni di memoria si accumulavano senza essere rilasciate e la coda in crescita richiedeva allocazioni aggiuntive, spiegando perché mmap veniva chiamato durante le operazioni munmap. Come soluzione alternativa alla perdita di memoria, l’impostazione della variabile d’ambiente UCX_RCACHE_MAX_UNRELEASED=1024 (il valore predefinito è inf) limita il numero di regioni di memoria non rilasciate nella coda, forzando UCX ad avviare la pulizia una volta raggiunta la soglia.

Il punto è che questo non avrebbe dovuto verificarsi fin dall’inizio. NIXL e vLLM stavano effettivamente chiamando la funzione ucp_worker_progress(), che avrebbe dovuto attivare la pulizia del memory pool. Non è ancora chiaro perché non sia stata attivata in questo specifico edge case. Ma ha mostrato che impostare a infinito il valore predefinito di UCX_RCACHE_MAX_UNRELEASED non era corretto. I team UCX e NIXL hanno deciso di modificare questo comportamento per una futura release di NIXL. Nel frattempo, abbiamo integrato una correzione nel repository vLLM per aiutare la community a evitare di incorrere nella stessa perdita.

Riepilogo dell’indagine.

L’indagine in breve:

1. Abbiamo notato una perdita di memoria in rapida crescita nel nostro ambiente di produzione durante il deployment di uno dei nostri modelli all’avanguardia con serving disaggregato.

2. Abbiamo provato a ridurre l’ambiente e a ottenere un riproduttore minimo. Purtroppo, il bug si riproduceva solo su un setup complesso, con un modello di grandi dimensioni e la disaggregazione abilitata.

3. Ci siamo rivolti a Memray, Guppy 3 e Heaptrack per analizzare la perdita. Da questi strumenti non è emerso nulla di evidente. Tuttavia, abbiamo notato qualcosa di peculiare nelle metriche di Heaptrack. La memoria residente era insolitamente elevata, quindi abbiamo deciso di approfondire in quella direzione.

4. Utilizzando pmap, siamo riusciti a vedere allocazioni RSS che continuavano a crescere, con i relativi puntatori di base associati.

5. Volevamo ottenere più indizi su chi stesse effettuando quelle chiamate. Con un po’ di scripting e l’uso di BPFtrace, abbiamo scoperto che le perdite originavano da chiamate a mmap. Nonostante tutti i nostri sforzi, non siamo riusciti a raccogliere stack trace completi (che ci avrebbero permesso di individuare con precisione il call site responsabile), ma questo ci ha portato a ritenere che le syscall fossero effettuate da un pacchetto fortemente ottimizzato.

6. Grazie alle informazioni raccolte, siamo riusciti a configurare breakpoint GDB molto specifici, che sarebbero stati attivati solo dalle chiamate incriminate.

7. Abbiamo scoperto che era UCX a effettuare quelle chiamate. Sebbene lo scopo delle chiamate sia legittimo, ovvero migliorare le prestazioni dei trasferimenti InfiniBand, ciò creava allocazioni RSS in continua crescita e ci ha impedito per giorni di effettuare il deployment del serving disaggregato.

8. Una volta nota l’origine della perdita, la soluzione alternativa è stata facile da trovare: impostare UCX_MEM_MMAP_HOOK_MODE=none ha risolto il problema. Abbiamo discusso questa indagine nel repository vLLM e integrato una patch per la community.

Cosa abbiamo imparato.

Gli stack software moderni sono costruiti su livelli di dipendenze, ognuno dei quali aggiunge complessità e potenziali punti di errore. Sebbene queste astrazioni aumentino notevolmente la produttività dei programmatori, non li isolano completamente dai problemi sottostanti nello stack. Per questo è essenziale essere preparati ad andare in profondità durante il debugging. Tuttavia, farlo in questi ambienti è raramente semplice, soprattutto quando le ottimizzazioni delle prestazioni introducono edge case sottili. UCX ne è un ottimo esempio. Il suo design privilegia le prestazioni, ma il modo in cui intercetta le chiamate mmap può creare rischi difficili da tracciare. Questa esperienza ha dimostrato ancora una volta quanto possa essere complesso diagnosticare problemi in sistemi profondamente interconnessi.

Questa indagine mostra anche l’importanza della trasparenza e della collaborazione quando si lavora con dipendenze critiche per le prestazioni. Siamo grati per la collaborazione con i team vLLM, NIXL e UCX nel confermare e affrontare questo comportamento. La loro competenza è stata fondamentale per arrivare a una risoluzione e non vediamo l’ora di continuare a lavorare insieme.

Esprimiamo la nostra gratitudine alle seguenti persone per l’assistenza e la collaborazione nella risoluzione di questo problema:

• Robert Shaw (Red Hat, Maintainer di vLLM e llm-d)

• Will Eaton (Red Hat, Maintainer di vLLM e llm-d)

• Nicolò Lucchesi (Red Hat, Maintainer di vLLM e llm-d)

• Mikhail Brinskii (NVIDIA, Maintainer di NIXL)

• Leonid Genkin (NVIDIA, Maintainer di UCX)

• Nathan Bellalou (NVIDIA, Maintainer di UCX e NIXL)

Cerchiamo nuove persone

Le interessa affrontare sfide come queste? Entri in Mistral AI e ci aiuti a plasmare il futuro dell’infrastruttura AI. Siamo sempre alla ricerca di ingegneri e ricercatori di talento con cui collaborare su progetti all’avanguardia.