Debugging kebocoran memori vLLM: misteri UCX dan mmap di luar heap

Ringkasan:

• Situasi masalah: Dalam lingkungan serving vLLM dengan pemisahan Prefill/Decode (disaggregated), terjadi kebocoran memori sistem (RSS) sebesar 400MB per menit, tetapi tidak terdeteksi oleh profiler Python biasa.
• Analisis penyebab: Dengan Heaptrack dan pmap, dipastikan bahwa kebocoran terjadi bukan di heap melainkan pada pemetaan memori anonim (mmap), lalu penyebabnya dilacak menggunakan BPFtrace dan skrip GDB otomatis.
• Pelaku teridentifikasi: Library komunikasi berkinerja tinggi UCX mencegat panggilan mmap/munmap untuk optimasi, dan penyebab masalahnya adalah memori yang sudah dilepas tidak segera dikembalikan, melainkan terus menumpuk tanpa batas di antrean.
• Solusi: Masalah diselesaikan dengan menonaktifkan fitur memory hooking UCX melalui pengaturan variabel lingkungan UCX_MEM_MMAP_HOOK_MODE=none.

Ringkasan detail:

1. Kebocoran memori yang misterius

Tim Mistral AI menemukan fenomena memori sistem yang meningkat secara linear sebesar 400MB per menit dalam lingkungan serving dengan pemisahan Prefill/Decode berbasis NIXL menggunakan vLLM.

• Gejala: Memori heap Python tetap stabil, tetapi RSS (Resident Set Size) pada level sistem operasi terus meningkat dan akhirnya menyebabkan OOM (Out of Memory).
• Kegagalan upaya awal: Alat Python seperti Memray dan Guppy 3 menunjukkan kondisi normal, GDB standar membuat proses crash, dan Valgrind terlalu lambat untuk digunakan.

2. Analisis mendalam hingga level kernel

Karena terindikasi bahwa akar masalah bukan berada di level aplikasi (Python/C++) melainkan di level yang lebih rendah, digunakanlah alat sistem.

• Heaptrack: Memastikan secara visual bahwa alokasi heap (malloc/free) stabil sementara RSS terus naik. Ini menunjukkan kebocoran terjadi pada pemetaan memori anonim (anonymous memory mappings) di luar pengelolaan heap glibc.
• pmap: Dengan memantau /proc/<pid>/maps, terlihat area pemetaan anonim tertentu terus membesar dan alamatnya berubah. Ini berarti siklus mremap atau munmap lalu mmap terus berulang.
• BPFtrace: Untuk melacak system call yang tidak tertangkap oleh LD_PRELOAD (karena melewati glibc), digunakan BPFtrace. Hasilnya menunjukkan bahwa panggilan mmap terjadi melalui syscall langsung.

3. Pelaku ditemukan: scripting GDB otomatis

Setelah alamat system call bermasalah dikonfirmasi lewat BPFtrace, dibuat skrip GDB agar berhenti hanya di alamat tersebut (SYS_mmap).

Contoh skrip GDB yang digunakan:

# Tetapkan breakpoint kondisional pada system call mmap (nomor 9)  
break syscall if $rdi == 9  
commands  
  silent  
  # Tetapkan breakpoint sementara pada titik return system call  
  tbreak *0x00007ffff7d9525d  
  commands  
    silent  
    # Cetak stack trace dan alamat yang dikembalikan  
    bt  
    printf "Syscall returned: rax = 0x%012lx\n", $rax  
    continue  
  end  
  continue  
end  

Melalui stack trace ini, diperoleh petunjuk krusial bahwa library UCX (Unified Communication X) sedang mencegat panggilan mmap/munmap Python di tengah jalan.

4. Penyebab: optimasi UCX yang berlebihan

UCX melakukan hooking terhadap alokasi/pelepasan memori untuk meningkatkan performa transfer InfiniBand.

• Mekanisme: Saat munmap dipanggil, memori tidak langsung dikembalikan ke sistem operasi, tetapi dimasukkan ke dalam 'invalidation queue' agar bisa digunakan kembali atau dibersihkan nanti.
• Bug: Karena konfigurasi default (UCX_RCACHE_MAX_UNRELEASED=inf), antrean ini bisa tumbuh tanpa batas. Pada pola penggunaan tertentu di vLLM, logika pembersihan (ucp_worker_progress) tidak berjalan semestinya, sehingga memori hanya terus menumpuk.

5. Cara mengatasinya

Dalam kasus vLLM, cukup mendaftarkan satu area memori KVCache yang besar, sehingga fitur memory hooking UCX yang kompleks sebenarnya tidak diperlukan.

• Solusi langsung: Kebocoran dihentikan dengan sepenuhnya mematikan memory hooking UCX melalui variabel lingkungan UCX_MEM_MMAP_HOOK_MODE=none.
• Alternatif: Ukuran antrean juga dapat dibatasi, misalnya dengan UCX_RCACHE_MAX_UNRELEASED=1024, agar pembersihan dipaksa terjadi.
• Tindakan: Perbaikan terkait telah di-merge untuk komunitas vLLM, dan perilaku default juga direncanakan akan diperbaiki pada rilis NIXL mendatang.

[jongyeans]

Harus menjalani hidup seperti apa... sampai bisa mencapai
tingkat setinggi ini

[ng0301]

Sulit membayangkan seberapa tinggi tingkat keahlian orang-orang seperti ini.