Truy vấn 3 tỷ vector

• Ghi chép về một thử nghiệm kỹ thuật xoay quanh quá trình tự triển khai giải pháp map-reduce tối ưu, xuất phát từ bài toán truy vấn 3 tỷ vector của Jeff Dean
• Bắt đầu từ cách triển khai naive để tính dot product giữa 3 tỷ vector float32 768 chiều và 1.000 vector truy vấn, rồi đạt được cải thiện hiệu năng theo từng bước nhờ vector hóa với numpy và chuyển sang float32
• Với 3.000 vector, cách naive mất khoảng 2 giây; sau khi vector hóa còn 0,01 giây, và khi dùng float32 thì giảm xuống 0,0045 giây
• Khi mở rộng lên quy mô 3 tỷ vector, ma trận kết quả cần khoảng 8,6 TB bộ nhớ, gây ra vấn đề OOM, nên cần thêm các tối ưu như xử lý theo lô, memory mapping, viết lại bằng Rust/C, hoặc dùng SimSIMD
• Nhấn mạnh rằng thay vì giải pháp kỹ thuật, bài toán cốt lõi khó hơn là định nghĩa yêu cầu (dạng trả về, cấu hình máy, có cho phép nén hay không, v.v.)

---

Điểm khởi đầu của bài toán

• Sau khi thấy cuộc trao đổi về Jeff Dean và bài toán truy vấn 3 tỷ vector, tác giả bắt đầu thử tự triển khai một giải pháp map-reduce tối ưu
• Vector là mảng các số thực dấu phẩy động có n chiều, và tìm kiếm vector được dùng để tìm các từ hoặc mục có độ tương đồng ngữ nghĩa
• Đây là một mẫu phổ biến khi dùng embedding trong các ứng dụng tìm kiếm sinh sinh như tìm kiếm, gợi ý, hoặc Cursor

Triển khai naive

• Giả định ban đầu: có 3 tỷ vector tài liệu cần tìm kiếm và khoảng 1.000 vector truy vấn, tất cả đều được lưu trên đĩa ở định dạng .npy
• Số chiều của vector là 768, một kích thước phổ biến trong nhiều mô hình truy vấn embedding dựa trên độ tương đồng
• Dùng cách lặp hai tầng để so sánh từng vector truy vấn với toàn bộ vector tài liệu, tính dot product và trả về kết quả
• Khi thử trước với 3.000 vector, kết quả là mất khoảng 2 giây trên M2 MacBook (quy mô nhỏ hơn 1 triệu lần so với 3 tỷ)

Tối ưu bằng vector hóa

• Áp dụng cách vector hóa bằng phép nhân ma trận của numpy (vectors_file @ query_vectors.T) để thay thế vòng lặp đôi
• Với 3.000 vector, thời gian giảm mạnh còn 0,0107 giây
• Khi mở rộng lên 3 triệu vector, thời gian là 12,85 giây

Chuyển sang float32

• Thực hiện tối ưu thêm bằng cách chuyển dữ liệu sang np.float32
• Với 3.000 vector, thời gian tiếp tục giảm xuống 0,0045 giây
• Vì 3 triệu vector mất khoảng 13 giây, nên ước tính với 3 tỷ vector sẽ gấp 1.000 lần, tức khoảng 3.216 phút

Tóm tắt so sánh hiệu năng

• Cách naive (3.000 vector): 1,9877 giây
• Cách vector hóa (3.000 vector): 0,0107 giây
• Cách vector hóa (3 triệu vector): 12,8491 giây
• Cách float32 (3.000 vector): 0,0045 giây

Vấn đề OOM và hướng tối ưu bổ sung

• Khi xử lý 3 tỷ đối tượng ở dạng float32 (4 byte), bộ nhớ cần thiết là khoảng 8,6 TB, và khi chạy sẽ gặp OOM
• Cần thêm các tối ưu theo hướng Jeff Dean gợi ý:
  • Chuyển mã sang generator và phép so sánh theo lô
  • Ghi xuống đĩa sau mỗi lượng tính toán nhất định hoặc dùng memory mapping
  • Viết lại mã tối ưu ở mức hệ thống bằng Rust hoặc C
  • Tận dụng các thư viện chuyên cho so sánh độ tương đồng vector quy mô lớn như SimSIMD

Tầm quan trọng của việc định nghĩa yêu cầu

• Trước khi tối ưu thêm, tác giả nhận ra nhiều điểm mơ hồ ngay trong chính bài toán:
  • Có phải chạy một vector truy vấn qua toàn bộ 3 tỷ vector rồi trả về toàn bộ kết quả, hay là tìm kiếm top-k ANN
  • Có cần trả về cả vector gốc hay không, và có cần re-ranking theo độ tương đồng hay không
  • Có phải dùng 1.000 vector truy vấn để quét toàn bộ tập dữ liệu một lần hay không
  • Vector đã nằm sẵn trong bộ nhớ chưa, đang được đọc từng cái từ đĩa, hay theo kiểu streaming
  • Chạy cục bộ hay không, cấu hình máy là gì, có GPU hay không
  • Kích thước embedding ảnh hưởng thế nào đến cách biểu diễn kết quả và kích thước vector đầu vào/đầu ra
  • Việc nén từ float64 xuống float32 có chấp nhận được về mặt độ chính xác hay không
  • Đây là dự án làm một lần hay có bao nhiêu thời gian cho việc xây dựng
• Điều khó nhất không phải bản thân giải pháp kỹ thuật, mà là định nghĩa chính xác yêu cầu của bài toán