ดีบักหน่วยความจำรั่วใน vLLM: ปริศนาของ UCX และ mmap ที่ไกลกว่า Heap

สรุป:

• สถานการณ์ปัญหา: ในสภาพแวดล้อมการเสิร์ฟแบบแยก Prefill/Decode ของ vLLM (disaggregated) เกิดการรั่วของหน่วยความจำระบบ (RSS) 400MB ต่อนาที แต่ไม่สามารถตรวจพบได้ด้วย Python profiler ทั่วไป
• การวิเคราะห์สาเหตุ: ใช้ Heaptrack และ pmap ยืนยันว่าการรั่วเกิดจาก anonymous memory mappings (mmap) ที่ไม่ใช่ heap และติดตามต้นตอด้วย BPFtrace กับสคริปต์ GDB แบบอัตโนมัติ
• การระบุตัวการ: ไลบรารีสื่อสารประสิทธิภาพสูง UCX ดักจับการเรียก mmap/munmap เพื่อเพิ่มประสิทธิภาพ แต่ไม่คืนหน่วยความจำที่ปล่อยแล้วกลับทันที และปล่อยให้ค้างสะสมอยู่ในคิวได้ไม่จำกัด
• วิธีแก้: แก้ปัญหาได้ด้วยการตั้งค่าตัวแปรแวดล้อม UCX_MEM_MMAP_HOOK_MODE=none เพื่อปิดฟังก์ชัน memory hooking ของ UCX

สรุปรายละเอียด:

1. การรั่วของหน่วยความจำที่เป็นปริศนา

ทีม Mistral AI พบว่าระบบเสิร์ฟแบบแยก Prefill/Decode (อิง NIXL) ที่ใช้ vLLM มีหน่วยความจำระบบเพิ่มขึ้นแบบเส้นตรง 400MB ต่อนาที

• อาการ: หน่วยความจำ Python heap คงที่ แต่ RSS (Resident Set Size) ในระดับระบบปฏิบัติการเพิ่มขึ้นต่อเนื่อง และสุดท้ายจบลงที่ OOM (Out of Memory)
• ความพยายามแรกที่ล้มเหลว: เครื่องมือฝั่ง Python อย่าง Memray, Guppy 3 แสดงผลว่าปกติ, GDB มาตรฐานทำให้โปรเซสล่ม และ Valgrind ช้าเกินไปจนใช้งานไม่ได้

2. การวิเคราะห์เชิงลึกในระดับเคอร์เนล

ทีมคาดว่าสาเหตุของปัญหาไม่ได้อยู่ที่ระดับแอปพลิเคชัน (Python/C++) แต่ลึกลงไปกว่านั้น จึงหันไปใช้เครื่องมือระดับระบบ

• Heaptrack: ยืนยันในเชิงภาพได้ว่า heap allocation (malloc/free) คงที่ แต่ RSS เพิ่มขึ้น ซึ่งบ่งชี้ว่าการรั่วเกิดใน anonymous memory mappings ที่อยู่นอกการจัดการ heap ของ glibc
• pmap: เฝ้าดู /proc/<pid>/maps แล้วพบว่า mapping แบบ anonymous บางช่วงเติบโตต่อเนื่องและเปลี่ยนที่อยู่ไปเรื่อย ๆ ซึ่งหมายความว่ามีวงจร mremap หรือ munmap แล้วตามด้วย mmap เกิดซ้ำ ๆ
• BPFtrace: ใช้ BPFtrace เพื่อติดตาม system call ที่ LD_PRELOAD จับไม่ได้ (เพราะข้าม glibc) และพบว่าการเรียก mmap เกิดผ่าน syscall โดยตรง

3. จับตัวการได้ด้วย GDB scripting แบบอัตโนมัติ

หลังจากยืนยันที่อยู่ของ system call ที่เป็นปัญหาด้วย BPFtrace แล้ว ทีมจึงเขียนสคริปต์ GDB ให้หยุดเฉพาะที่ address นั้น (SYS_mmap)

ตัวอย่างสคริปต์ GDB ที่ใช้:

# ตั้งค่า conditional breakpoint สำหรับ mmap system call (หมายเลข 9)  
break syscall if $rdi == 9  
commands  
  silent  
  # ตั้งค่า temporary breakpoint ที่จุด return ของ system call  
  tbreak *0x00007ffff7d9525d  
  commands  
    silent  
    # แสดง stack trace และ address ที่คืนกลับมา  
    bt  
    printf "Syscall returned: rax = 0x%012lx\n", $rax  
    continue  
  end  
  continue  
end  
  

จาก stack trace นี้ ทีมพบหลักฐานชี้ขาดว่าไลบรารี UCX (Unified Communication X) กำลังดักจับการเรียก mmap/munmap ของ Python อยู่ตรงกลาง

4. สาเหตุ: การเพิ่มประสิทธิภาพที่มากเกินไปของ UCX

UCX ทำการ hook การจอง/คืนหน่วยความจำเพื่อเพิ่มประสิทธิภาพการส่งข้อมูลผ่าน InfiniBand

• กลไก: เมื่อมีการเรียก munmap UCX จะไม่คืนหน่วยความจำให้ระบบปฏิบัติการทันที แต่เก็บไว้ใน 'invalidation queue' เพื่อรอนำกลับมาใช้ซ้ำหรือจัดการภายหลัง
• บั๊ก: ค่าตั้งต้น (UCX_RCACHE_MAX_UNRELEASED=inf) ทำให้คิวนี้สามารถโตได้ไม่จำกัด และในรูปแบบการใช้งานบางอย่างของ vLLM ลอจิกการล้างคิว (ucp_worker_progress) กลับไม่ทำงานตามที่ควร ทำให้หน่วยความจำสะสมอย่างต่อเนื่อง

5. วิธีแก้ปัญหา

ในกรณีของ vLLM จำเป็นต้องลงทะเบียนเพียง memory region ขนาดใหญ่ของ KVCache ชุดเดียว ดังนั้นฟีเจอร์ memory hooking ที่ซับซ้อนของ UCX จึงไม่จำเป็นนัก

• วิธีแก้ทันที: ตั้งค่าตัวแปรแวดล้อม UCX_MEM_MMAP_HOOK_MODE=none เพื่อปิด memory hooking ของ UCX ทั้งหมด และหยุดการรั่วได้
• ทางเลือก: สามารถจำกัดขนาดคิวด้วยค่าอย่าง UCX_RCACHE_MAX_UNRELEASED=1024 เพื่อบังคับให้มีการล้างคิว
• การดำเนินการ: การแก้ไขนี้ถูกรวมเข้ากับคอมมูนิตี้ vLLM แล้ว และค่าเริ่มต้นในรีลีส NIXL รุ่นถัดไปก็มีกำหนดจะได้รับการปรับปรุง

[jongyeans]

ต้องใช้ชีวิตแบบไหนกันนะ…ถึงจะไปถึง ระดับนี้ได้

[ng0301]

ไม่อาจเดาได้เลยว่าคนแบบนี้มีฝีมือลึกซึ้งขนาดไหน