Gỡ lỗi rò rỉ bộ nhớ vLLM: Bí ẩn của UCX và mmap vượt ra ngoài heap

Tóm tắt:

• Tình huống sự cố: Trong môi trường phục vụ tách rời Prefill/Decode của vLLM, đã xảy ra rò rỉ bộ nhớ hệ thống (RSS) 400MB mỗi phút, nhưng các profiler Python thông thường không phát hiện được.
• Phân tích nguyên nhân: Heaptrack và pmap xác nhận rò rỉ không nằm ở heap mà ở các ánh xạ bộ nhớ ẩn danh (mmap), sau đó nguyên nhân được lần ra bằng BPFtrace và các script GDB tự động hóa.
• Thủ phạm: UCX, thư viện giao tiếp hiệu năng cao, đã chặn các lời gọi mmap/munmap để tối ưu hóa, và nguyên nhân là nó không trả lại ngay bộ nhớ đã giải phóng mà giữ chúng trong hàng đợi vô thời hạn.
• Cách khắc phục: Vấn đề được giải quyết bằng cách đặt biến môi trường UCX_MEM_MMAP_HOOK_MODE=none để vô hiệu hóa tính năng hook bộ nhớ của UCX.

Tóm tắt chi tiết:

1. Rò rỉ bộ nhớ đầy bí ẩn

Nhóm Mistral AI đã phát hiện hiện tượng bộ nhớ hệ thống tăng tuyến tính 400MB mỗi phút trong môi trường phục vụ tách rời Prefill/Decode dùng vLLM (dựa trên NIXL).

• Triệu chứng: Bộ nhớ heap của Python vẫn ổn định, nhưng RSS (Resident Set Size) ở cấp hệ điều hành tiếp tục tăng và cuối cùng dẫn tới OOM (Out of Memory).
• Những thử nghiệm ban đầu thất bại: Các công cụ Python như Memray, Guppy 3 đều cho kết quả bình thường, GDB tiêu chuẩn làm tiến trình bị crash, còn Valgrind thì quá chậm để sử dụng.

2. Phân tích sâu ở cấp kernel

Nhận thấy nguyên nhân không nằm ở tầng ứng dụng (Python/C++) mà ở lớp thấp hơn, nhóm đã dùng các công cụ hệ thống.

• Heaptrack: Xác nhận trực quan rằng cấp phát heap (malloc/free) vẫn ổn định trong khi RSS tăng lên. Điều này cho thấy rò rỉ xảy ra ở ánh xạ bộ nhớ ẩn danh (anonymous memory mappings), nằm ngoài cơ chế quản lý heap của glibc.
• pmap: Giám sát /proc/<pid>/maps và xác nhận một vùng ánh xạ ẩn danh cụ thể liên tục tăng kích thước và thay đổi địa chỉ. Điều đó ngụ ý một chu kỳ mremap hoặc munmap rồi mmap đang lặp lại.
• BPFtrace: Để theo dõi các system call không bị bắt bởi LD_PRELOAD (vì bỏ qua glibc), nhóm đã dùng BPFtrace. Kết quả cho thấy các lời gọi mmap đang diễn ra thông qua syscall trực tiếp.

3. Bắt thủ phạm: Script GDB tự động hóa

Sau khi xác định được địa chỉ system call gây vấn đề bằng BPFtrace, nhóm đã viết script GDB để chỉ dừng tại địa chỉ đó (SYS_mmap).

Ví dụ script GDB đã dùng:

# Thiết lập breakpoint có điều kiện cho system call mmap (số 9)  
break syscall if $rdi == 9  
commands  
  silent  
  # Thiết lập breakpoint tạm tại điểm trả về của system call  
  tbreak *0x00007ffff7d9525d  
  commands  
    silent  
    # In stack trace và địa chỉ được trả về  
    bt  
    printf "Syscall returned: rax = 0x%012lx\n", $rax  
    continue  
  end  
  continue  
end  
  

Thông qua stack trace này, nhóm đã tìm được bằng chứng quyết định rằng thư viện UCX (Unified Communication X) đang chặn (intercept) các lời gọi mmap/munmap của Python ở giữa đường.

4. Nguyên nhân: Sự tối ưu hóa quá mức của UCX

UCX hook việc cấp phát/giải phóng bộ nhớ để tăng hiệu năng truyền InfiniBand.

• Cơ chế: Khi munmap được gọi, UCX không trả lại bộ nhớ ngay cho hệ điều hành mà đưa nó vào một 'hàng đợi vô hiệu hóa (invalidation queue)' để tái sử dụng hoặc dọn dẹp sau.
• Bug: Với cấu hình mặc định (UCX_RCACHE_MAX_UNRELEASED=inf), hàng đợi này có thể tăng vô hạn; trong một số mẫu sử dụng cụ thể của vLLM, logic dọn dẹp (ucp_worker_progress) không hoạt động đúng, khiến bộ nhớ chỉ tích tụ mà không được giải phóng.

5. Cách giải quyết

Trong trường hợp của vLLM, chỉ cần đăng ký một vùng bộ nhớ KVCache lớn, nên không thực sự cần tới tính năng hook bộ nhớ phức tạp của UCX.

• Khắc phục ngay lập tức: Ngăn rò rỉ bằng cách đặt biến môi trường UCX_MEM_MMAP_HOOK_MODE=none để tắt hoàn toàn cơ chế hook bộ nhớ của UCX.
• Phương án khác: Cũng có thể giới hạn kích thước hàng đợi, ví dụ UCX_RCACHE_MAX_UNRELEASED=1024, để buộc quá trình dọn dẹp xảy ra.
• Hành động: Bản sửa lỗi đã được merge để phục vụ cộng đồng vLLM, và hành vi mặc định trong các bản phát hành NIXL sắp tới cũng sẽ được cải thiện.

[jongyeans]

Phải sống cuộc đời kiểu gì mới có thể… đạt tới
cảnh giới này

[ng0301]

Thật khó mà hình dung được nội lực của những người như thế này ở mức nào.