Tăng cường khả năng cạnh tranh của GPU AMD trong suy luận LLM (2023)

• MLC-LLM cho phép biên dịch và triển khai LLM trên GPU AMD bằng ROCm, giúp Radeon RX 7900 XTX đạt 80% tốc độ của RTX 4090 và 94% tốc độ của RTX 3090 Ti trong suy luận Llama 2 7B/13B
• Ba GPU được so sánh đều có 24GB bộ nhớ và băng thông tương tự nhau, nên trong suy luận LLM nhạy với độ trễ, nghẽn cổ chai bộ nhớ có tác động lớn hơn hiệu năng tính toán FP16
• Khoảng cách của AMD chủ yếu đến từ thiếu hỗ trợ phần mềm và tối ưu hóa hơn là bản thân phần cứng; đầu tư vào ROCm và biên dịch machine learning đóng vai trò như một trục giúp giảm chi phí hỗ trợ nhiều backend
• ROCm được áp dụng cho Radeon 7900 XTX, còn Vulkan được áp dụng cho SteamDeck AMD APU, cho thấy có thể chạy Llama-7B lượng tử hóa 4-bit ngay cả trên các thiết bị AMD phổ thông
• Kết quả dựa trên benchmark batch đơn; tối ưu hóa attention có thể thay đổi số liệu của cả AMD lẫn NVIDIA, nên cần diễn giải với biên độ sai số 10%

---

Suy luận LLM trên GPU AMD dựa trên ROCm

• MLC-LLM cho phép biên dịch và triển khai LLM trên GPU AMD bằng ROCm
• Hiệu năng Llama 2 7B/13B của Radeon RX 7900 XTX đạt mức sau
  • 80% tốc độ của NVIDIA GeForce RTX 4090
  • 94% tốc độ của NVIDIA GeForce RTX 3090 Ti
• Ngoài ROCm, công cụ này còn hỗ trợ Vulkan, mở rộng phạm vi triển khai LLM sang các thiết bị như SteamDeck dùng AMD APU

So sánh phần cứng và nghẽn cổ chai thực tế

• Về thông số, AMD RX 7900 XTX ở vị trí có thể so sánh với RTX 4090 và RTX 3090 Ti
  • Cả ba GPU đều có 24GB bộ nhớ, nên có thể tải các mô hình cùng kích thước
  • Băng thông bộ nhớ cũng tương tự nhau
  • RTX 4090 có hiệu năng FP16 cao gấp 2 lần RX 7900 XTX, còn RTX 3090 Ti cao gấp 1,3 lần
• Suy luận LLM nhạy với độ trễ phần lớn bị nghẽn cổ chai bộ nhớ, nên khác biệt hiệu năng FP16 ở đây không phải nút thắt chính
• RX 7900 XTX rẻ hơn 40% so với RTX 4090
• RTX 3090 Ti là sản phẩm thế hệ trước nên khó so sánh giá hơn, và được dùng như một điểm tham chiếu

Nguyên nhân khoảng cách của AMD và cách tiếp cận của MLC

• Lý do AMD bị tụt lại gần với việc thiếu hỗ trợ phần mềm và tối ưu hóa cho các mô hình liên quan hơn là phần cứng
• Hai thay đổi đang đóng vai trò thu hẹp khoảng cách
  • AMD đầu tư vào stack ROCm để bắt kịp
  • Biên dịch machine learning giúp giảm chi phí hỗ trợ phần mềm phổ quát cho nhiều backend
• Biên dịch machine learning biên dịch tác vụ machine learning và tự động hóa tối ưu hóa, thay vì viết trực tiếp từng kernel riêng cho ROCm hay CUDA
• MLC-LLM là giải pháp triển khai LLM dựa trên biên dịch machine learning, được xây dựng trên Apache TVM Unity
  • Cung cấp luồng phát triển lấy Python làm trung tâm
  • Bao gồm biến đổi đồ thị tính toán, tối ưu hóa layout/lịch trình kernel GPU, và API triển khai native
  • Hỗ trợ CUDA, Metal, ROCm, Vulkan, OpenCL
  • Mở rộng phạm vi từ GPU cấp máy chủ tới thiết bị di động như iPhone và Android

Cách hỗ trợ GPU và APU AMD

• Các đường hỗ trợ GPU AMD gồm ROCm, OpenCL, Vulkan và WebGPU
  • ROCm là stack mà AMD gần đây đang đẩy mạnh, có các thành phần tương tự CUDA
  • Vulkan là tiêu chuẩn đồ họa hiện đại và được hỗ trợ rộng rãi trên nhiều thiết bị GPU
  • WebGPU là tiêu chuẩn web mới cho phép chạy tính toán trong trình duyệt web
• Có ít phần mềm machine learning xây dựng giải pháp ngoài CUDA, và chi phí kỹ thuật để nhân bản stack cho từng phần cứng mới hoặc mô hình lập trình GPU mới là rất lớn
• MLC hỗ trợ nhiều đường bằng sinh mã tự động, thay vì viết lại kernel GPU cho từng backend
• Hiệu năng thực tế phụ thuộc vào chất lượng của runtime GPU cấp thấp và khả năng sử dụng trên từng nền tảng

Triển khai tối ưu hóa ROCm

• ROCm được chọn cho Radeon 7900 XTX, còn Vulkan được chọn cho SteamDeck APU
• Stack ROCm hoạt động ngay, và nhờ pipeline phát triển dựa trên Python của TVM Unity, phiên bản tối ưu hóa có thể được thêm vào trong vài giờ
• Hỗ trợ ROCm tái sử dụng các thành phần hiện có
  • Toàn bộ pipeline MLC của các target hiện có như CUDA và Metal
    • Lập kế hoạch bộ nhớ
    • Fusion toán tử
  • Không gian tối ưu hóa kernel GPU phổ quát được viết bằng TVM TensorIR
  • Luồng sinh mã ROCm của TVM, tạo kernel ROCm cấp thấp thông qua LLVM
• Mã được sinh ra được xuất dưới dạng thư viện dùng chung hoặc thư viện tĩnh, và có thể gọi từ CLI, Python, REST API

Điều kiện và cách diễn giải benchmark Llama 2

• Benchmark đo Llama 2 7B và 13B với lượng tử hóa 4-bit
• Hiệu năng giải mã được đo bằng cách đưa vào một token prompt đơn và tạo 512 token
• Tất cả kết quả đều dựa trên suy luận batch đơn
• Theo bản phát hành ROCm 5.6, hiệu năng suy luận batch đơn đạt 80% tốc độ của NVIDIA 4090
• Baseline CUDA được xem là mức hàng đầu tại thời điểm đó cho tác vụ này
• Vẫn còn dư địa cải thiện, chẳng hạn tối ưu hóa attention tốt hơn; nếu tối ưu hóa đó được đưa vào MLC, số liệu của cả AMD và NVIDIA đều có thể cải thiện
• Nếu tối ưu hóa chỉ được triển khai ở phía NVIDIA, khoảng cách có thể tăng từ 20% lên 30%, nên khuyến nghị diễn giải các số liệu với biên độ sai số 10%

Điều kiện chạy trực tiếp và ví dụ

• Có cung cấp wheel dựng sẵn và hướng dẫn chạy để tái lập benchmark
• Điều kiện chạy là GPU AMD hoạt động với ROCm 5.6 trở lên trên Linux
• Việc cài đặt gói MLC dựng sẵn đã bật ROCm làm theo tài liệu dùng thử MLC-LLM
• Ví dụ Python tải mô hình Llama-2-7b-chat-hf-q4f16_1 bằng mlc_chat.ChatModule và đo hiệu năng bằng benchmark_generate("Hi", generate_length=512)
• MLC-LLM cũng cung cấp CLI tương tác, nhưng trên ROCm cần build CLI từ source và làm theo tài liệu build CLI

Sử dụng Vulkan và bộ nhớ hợp nhất trên SteamDeck

• SteamDeck được dùng như ví dụ cho nhóm thiết bị AMD rộng hơn có trang bị AMD APU
• GPU VRAM có thể dùng trong ROCm bị BIOS giới hạn ở 4GB
• Driver Mesa Vulkan cho phép buffer sử dụng vượt giới hạn này, lên tới 16GB, thông qua bộ nhớ hợp nhất
• Dung lượng bộ nhớ này đủ để chạy Llama-7B lượng tử hóa 4-bit
• Kết quả là khả năng hỗ trợ LLM cũng được xác nhận trên nhiều thiết bị AMD phổ thông

Công việc trong tương lai và liên kết dự án

• Trong kỷ nguyên AI tạo sinh, tính sẵn có của phần cứng trở thành vấn đề quan trọng
• Biên dịch machine learning có thể giảm nhẹ vấn đề này bằng cách cho phép triển khai phổ quát hiệu năng cao trên nhiều backend phần cứng
• Nghiên cứu hiện tại tập trung vào GPU cấp người dùng phổ thông
• Theo kinh nghiệm trước đây, tối ưu hóa MLC cho GPU phổ thông đôi khi có thể tổng quát hóa sang GPU đám mây, theo hướng từ RTX 4090 sang A100, A10g
• Các lĩnh vực công việc trong tương lai gồm
  • Hỗ trợ batching và đa GPU
  • Tích hợp hệ sinh thái PyTorch
  • Hỗ trợ nhiều phương thức lượng tử hóa và kiến trúc mô hình hơn
  • Mở rộng tối ưu hóa tự động trên nhiều backend phần cứng hơn
• NVIDIA vẫn giữ vị thế dẫn đầu nhờ đổi mới liên tục, và môi trường có thể thay đổi theo phần cứng mới như H100 cùng sự tiến hóa của phần mềm
• Hướng dẫn MLC-LLM có tại trang dự án, và mã nguồn nằm trong kho GitHub