Suy luận Llama2 nhanh và gọn nhẹ trên môi trường edge dị chủng

 (secondstate.io)
2 điểm bởi GN⁺ 2023-11-14 | 1 bình luận | Chia sẻ qua WhatsApp
  • Để giảm gánh nặng phụ thuộc lớn và triển khai nặng nề của stack suy luận xoay quanh Python, một ứng dụng Llama2 bằng Rust+Wasm được đóng gói thành nhị phân 2MB để chạy trên nhiều thiết bị
  • Cách triển khai này là phiên bản llama.cpp được điều chỉnh cho Wasm, sử dụng tệp mô hình GGUF cùng plugin GGML của WasmEdge và API WASI NN để tận dụng tăng tốc phần cứng cục bộ
  • Ví dụ sử dụng mô hình lượng tử hóa 5-bit Llama-2-7B-Chat GGUF, đạt khoảng 25 tokens/s trên MacBook M2 giá rẻ và khoảng 50 tokens/s trên Nvidia A10G
  • Runtime WasmEdge có thể dùng cùng các công cụ container, giúp dễ triển khai cùng một nhị phân Wasm lên thiết bị edge, môi trường on-premise và cloud
  • Bộ công cụ GGML đã có thể sử dụng nhưng vẫn ở giai đoạn đầu, còn thiếu thêm plugin cho nhiều phần cứng và hệ điều hành hơn, hỗ trợ nhiều cấu hình llama.cpp hơn và hỗ trợ API WASI NN cho nhiều ngôn ngữ tương thích Wasm

Chạy suy luận Llama2 bằng Rust+Wasm

  • Stack Rust+Wasm có thể được dùng như một lựa chọn thay thế cho Python trong suy luận AI
    • Ứng dụng Rust+Wasm có thể nhỏ bằng 1/100 và nhanh hơn 100 lần so với Python, đồng thời tận dụng tăng tốc phần cứng trên nhiều môi trường khác nhau mà không cần thay đổi cùng một mã nhị phân
  • Second State đã tạo một chương trình Rust đơn giản chạy suy luận mô hình Llama2 ở tốc độ native
    • Ứng dụng nhị phân được biên dịch sang Wasm chỉ có kích thước 2MB
    • Ứng dụng Wasm này có tính di động trên các thiết bị có bộ tăng tốc phần cứng dị chủng
    • Runtime WasmEdge cung cấp môi trường thực thi an toàn trong môi trường cloud
    • WasmEdge có thể dùng cùng các công cụ container để điều phối và chạy các ứng dụng di động trên nhiều thiết bị

Dựa trên llama.cpp và GGUF

Quy trình chạy

  • Cài WasmEdge và plugin GGML trên thiết bị Linux hoặc Mac
curl -sSf https://raw.githubusercontent.com/WasmEdge/WasmEdge/… | bash -s -- --plugins wasmedge_rustls wasi_nn-ggml
  • Tải ứng dụng Wasm đã được build sẵn
curl -LO https://github.com/LlamaEdge/LlamaEdge/…
  • Ví dụ sử dụng tệp GGUF của mô hình chat-tuned Llama2 7B được lượng tử hóa với 5-bit weights
curl -LO https://huggingface.co/second-state/Llama-2-7B-Chat-GGUF/…
  • Chạy ứng dụng suy luận Wasm bằng WasmEdge và truyền mô hình GGUF để nhập câu hỏi theo chế độ tương tác
wasmedge --dir .:. --nn-preload default:GGML:AUTO:Llama-2-7b-chat-hf-Q5_K_M.gguf llama-chat.wasm

Cấu hình mô hình và ví dụ hiệu năng

  • llama-chat.wasm cấu hình cách tương tác với mô hình bằng các tùy chọn dòng lệnh
    • --ctx-size: kích thước ngữ cảnh prompt, mặc định 512
    • --n-predict: số token cần dự đoán, mặc định 1024
    • --n-gpu-layers: số layer chạy trên GPU, mặc định 100
    • --batch-size: kích thước batch xử lý prompt, mặc định 512
    • --temp: nhiệt độ lấy mẫu, mặc định 0.8
    • --repeat-penalty: hệ số phạt token lặp, mặc định 1.1
    • --prompt-template: hỗ trợ llama-2-chat, codellama-instruct, mistral-instruct-v0.1, chatml, deepseek-chat, deepseek-coder v.v.
    • --log-stat: xuất thống kê
  • Ví dụ chạy với độ dài ngữ cảnh 2048, tối đa 512 token phản hồi và bật xuất thống kê như sau
wasmedge --dir .:. --nn-preload default:GGML:AUTO:Llama-2-7b-chat-hf-Q5_K_M.gguf llama-chat.wasm -c 2048 -n 512 --log-stat
  • Trên MacBook M2 giá rẻ, phản hồi của LLM được stream ra mặc định với tốc độ khoảng 25 tokens/s
    • Ví dụ log cho thấy 25.64 tokens/s trong 82 lần chạy theo eval time
    • kv self size hiển thị 1024.00MB và compute buffer total size hiển thị 630.14MB
  • Trên thiết bị Nvidia A10G, tốc độ đạt khoảng 50 tokens/s
    • Sử dụng tăng tốc GPU CUDA và offload 35/35 layer sang GPU
    • Mức dùng VRAM hiển thị là 4474.93MB cho mô hình, 1648.02MB cho ngữ cảnh, tổng cộng 6122.95MB
    • Ví dụ log cho thấy 50.55 tokens/s trong 83 lần chạy theo eval time

Máy chủ API tương thích OpenAI và ứng dụng LLM

  • Cũng có máy chủ API tương thích OpenAI được xây dựng bằng Rust và WasmEdge
  • Dùng máy chủ này, có thể xây agent và ứng dụng LLM bằng các công cụ phát triển tương thích OpenAI như flows.network
  • Phần giải thích liên quan có tại Build a super lightweight AI agent

Giới hạn của stack suy luận Python

  • Các LLM như Llama2 thường được huấn luyện bằng các framework dựa trên Python như PyTorch, Tensorflow và JAX
  • Tuy nhiên, bài viết cho rằng Python bất lợi trong các ứng dụng suy luận, vốn chiếm khoảng 95% khối lượng tính toán AI
  • Các gói Python có thể trở nên khó thiết lập và sử dụng vì phụ thuộc phức tạp
  • Image Docker cho Python hoặc PyTorch thường có thể phình lên vài GB hoặc vài chục GB, gây gánh nặng lớn cho suy luận AI trên server edge hay thiết bị
  • Python chậm hơn nhiều so với các ngôn ngữ biên dịch như C, C++ và Rust, có thể chậm hơn tới 35.000 lần
  • Phần lớn workload thực tế được giao cho các thư viện chia sẻ native bên dưới lớp wrapper Python
    • Cấu trúc này phù hợp cho demo nhưng bị xem là khó chỉnh sửa phần bên trong theo yêu cầu kinh doanh cụ thể
  • Mức phụ thuộc cao vào thư viện native và việc quản lý phụ thuộc phức tạp làm giảm tính di động của các chương trình AI Python khi cần tận dụng khả năng phần cứng theo từng thiết bị
  • Có cả ví dụ các gói Python thường dùng trong toolchain LLM xung đột với nhau vì yêu cầu phiên bản pydantic
    • llama-cpp-python yêu cầu pydantic 2.0.1 và không hoạt động với <=2.0
    • fastapi và chromadb yêu cầu pydantic 1.9.2 và không hoạt động với >=2.0

Ưu điểm của Rust+Wasm

  • Stack Rust+Wasm có thể được dùng như hạ tầng điện toán đám mây hợp nhất, bao trùm thiết bị, edge cloud, server on-premise và public cloud
  • Trong các ứng dụng suy luận AI, đây có thể là một lựa chọn thay thế mạnh mẽ cho stack Python

  • Siêu nhẹ

    • Ứng dụng suy luận chỉ 2MB, bao gồm mọi phụ thuộc
    • Nhỏ hơn 1% kích thước của một container PyTorch điển hình

  • Thực thi nhanh

    • Có thể đạt tốc độ native C/Rust trên toàn bộ các bước tiền xử lý, tính toán tensor và hậu xử lý

  • Tính di động

    • Cùng một ứng dụng bytecode Wasm chạy được trên các nền tảng tính toán chủ chốt
    • Cũng hỗ trợ tăng tốc phần cứng dị chủng

  • Dễ thiết lập, phát triển và triển khai

    • Giảm phụ thuộc phức tạp
    • Có thể build một tệp Wasm duy nhất bằng công cụ tiêu chuẩn trên laptop rồi triển khai lên nhiều môi trường

  • An toàn và sẵn sàng cho cloud

    • Runtime Wasm được thiết kế để cô lập mã người dùng không đáng tin cậy
    • Có thể được quản lý bằng công cụ container và triển khai lên nền tảng cloud-native

Cấu trúc chương trình suy luận Rust

  • Chương trình suy luận demo được viết bằng Rust và biên dịch sang Wasm
  • Mã nguồn Rust cốt lõi chỉ khoảng 40 dòng
  • Chương trình Rust đảm nhiệm các vai trò sau
    • Quản lý đầu vào người dùng
    • Theo dõi lịch sử hội thoại
    • Chuyển văn bản sang template chat của Llama2
    • Chạy suy luận bằng API WASI NN
  • Nếu muốn tự build, hãy cài trình biên dịch Rust và target biên dịch wasm32-wasi
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
rustup target add wasm32-wasi
  • Sau đó tải dự án nguồn và build tệp Wasm bằng cargo
git clone https://github.com/second-state/llama-utils
cd llama-utils/chat/
cargo build --target wasm32-wasi --release
cp target/wasm32-wasi/release/llama-chat.wasm .

Triển khai trên cloud và edge

  • Khi đã có tệp bytecode Wasm, có thể triển khai nó lên mọi thiết bị hỗ trợ runtime WasmEdge
  • Hiện plugin GGML hỗ trợ generic Linux và Ubuntu Linux
    • CPU x86 và ARM
    • GPU Nvidia
    • Apple M1/M2/M3
  • Plugin WasmEdge GGML tự động tận dụng tăng tốc phần cứng của thiết bị dựa trên llama.cpp
    • Nếu có GPU Nvidia, trình cài đặt sẽ tự động cài plugin GGML tối ưu cho CUDA
    • Plugin GGML cho Mac OS dùng Metal API để chạy workload suy luận trên neural processing engine tích hợp của M1/M2/M3
    • Bản build Linux CPU dùng OpenBLAS để tự động phát hiện và tận dụng các tính năng tính toán CPU hiện đại như AVX và SIMD
  • Đây là cách tiếp cận nhằm đạt tính di động giữa các phần cứng AI dị chủng và các nền tảng khác nhau mà không phải đánh đổi hiệu năng

Công việc sắp tới

  • Công cụ WasmEdge GGML hiện đã có thể sử dụng và đang được khách hàng cloud-native dùng, nhưng vẫn còn ở giai đoạn đầu
  • Các mảng có thể đóng góp gồm
    • Bổ sung plugin GGML cho nhiều phần cứng và nền tảng hệ điều hành hơn
      • Nhóm quan tâm đến TPU, ARM NPU và các chip AI chuyên dụng trên Linux và Windows
    • Hỗ trợ thêm nhiều cấu hình của llama.cpp
      • Hiện mới chỉ có thể truyền một phần tùy chọn cấu hình từ Wasm sang plugin GGML
      • Mục tiêu là hỗ trợ toàn bộ tùy chọn mà GGML cung cấp
    • Hỗ trợ API WASI NN trong các ngôn ngữ tương thích Wasm khác
      • Nhóm quan tâm đến Go, Zig, Kotlin, JavaScript, C và C++

Hỗ trợ các mô hình ngoài LLM

  • WasmEdge và WASI NN cũng có thể được dùng như một lựa chọn thay thế Python nhẹ, nhanh, di động và an toàn để xây ứng dụng suy luận dựa trên các mô hình AI phổ biến ngoài LLM
  • mediapipe-rs cung cấp API Rust+Wasm cho họ mô hình Tensorflow mediapipe của Google
  • WasmEdge YOLO là dự án API Rust+Wasm để xử lý mô hình PyTorch YOLOv8
  • WasmEdge ADAS demo là ví dụ dùng mô hình Intel OpenVINO để thực hiện phân đoạn mặt đường cho xe tự lái
  • WasmEdge Document AI dự kiến sẽ cung cấp API Rust+Wasm cho một họ mô hình OCR và xử lý tài liệu
  • Có thể thảo luận và đóng góp cho WasmEdge tại WasmEdge Discord

Bài viết liên quan

1 bình luận

 
GN⁺ 2023-11-14
Các ý kiến trên Hacker News

  • Rust và WASM thì tốt, nhưng nhìn vào mã thì về cơ bản chỉ là một script dòng lệnh Rust 150 dòng
    Phần việc nặng được xử lý bởi một dòng chuyển mô hình sang backend WASI-NN, ở đây do runtime WasmEdge cung cấp
    Nhưng WasmEdge viết bằng C++ chứ không phải Rust, và trong trường hợp này lợi ích mà Rust mang lại gần như không có; backend cũng có thể được gọi từ các ngôn ngữ khác như Python

    • Lợi ích của Rust ở đây có vẻ là đóng gói và phân phối
      Việc gói Python và PyTorch thành dạng mà người dùng cuối có thể nhấp đúp để chạy hiện vẫn gần như là một mớ hỗn độn; mã hiệu năng cao thực tế ở cả hai bên đều là C++, nhưng nếu thay vì phân phối hơn 2GB cùng hàng chục nghìn tệp chỉ để gửi vài lệnh tới phần C++ đó mà có thể kết thúc bằng một tệp thực thi 2MB, thì điều đó có ý nghĩa

  • Công việc rất ấn tượng. Tuy nhiên những người muốn dùng thử vẫn phải tải xuống các tệp trọng số rất lớn
    Về bản chất, nó giống như biến llama.cpp vốn hoàn toàn portable và không phụ thuộc thành 2MB
    Nếu là nhà phát triển ứng dụng, đây có thể là cách dễ nhất để đóng gói engine suy luận vào một tệp có thể phân phối. Trọng số vốn đã portable và có thể tải xuống khi cần, nên phần thật sự muốn cố định là engine suy luận

    • Có lẽ sẽ hữu ích hơn nếu tiêu đề ghi là 2MB wasm
      Như đã nói, trọng số áp đảo kích thước đó
    • Trên máy của tôi, tệp main mà llama.cpp build ra là 1,2MB
      Bản thân kích thước 2MB không có gì đặc biệt ấn tượng; điểm chính là nhắm tới wasm để trở nên portable hơn, chứ không phải được nén đặc biệt hơn

  • Chẳng phải cuối cùng đây chỉ là bọc quanh llama.cpp sao? Thành thật mà nói tôi khá mệt với các dự án bọc quanh x.cpp
    Trong 6 tháng qua tôi đã phát triển một framework học máy Rust + WebGPU, và nhanh chóng nhận ra công việc của GG ấn tượng đến mức nào
    Vẫn còn ở giai đoạn đầu, nhưng có thể xem tại đây:
    https://www.ratchet.sh/
    https://github.com/FL33TW00D/whisper-turbo

    • Bạn có thể giải thích thêm điều gì khiến nó ấn tượng không? Tôi hoàn toàn không biết lĩnh vực này nên khó hiểu đúng giá trị của nó
    • Bạn vừa hạ thấp công việc của người khác rồi quảng bá công việc của mình trong cùng một bình luận à? Cần nghiêm túc nhìn lại về mặt đạo đức
    • GG là ai?

  • wasm-nn mà cái này phụ thuộc vào, tức https://github.com/WebAssembly/wasi-nn, là một đề xuất theo kiểu gửi các chunk tùy ý cho phần triển khai của vendor. API thực chất cũng chỉ ở mức thiết lập input, tính toán, thiết lập output
    Vì vậy nó hoàn toàn không portable
    Lý do cái này hoạt động là vì nó dựa vào phần trừu tượng hóa đã được triển khai sẵn trong llama.cpp. Có thể thấy ở đây WasmEdge đã lấy nguyên phần mã ánh xạ mô hình gguf sang nhiều target phần cứng của llama.cpp: https://github.com/WasmEdge/WasmEdge/tree/master/plugins/was...
    Vì thế cách mô tả kiểu “nhà phát triển có thể dùng binding để viết ứng dụng học máy bằng ngôn ngữ cấp cao, biên dịch sang WebAssembly rồi chạy trên runtime hỗ trợ wasi-nn như WasmEdge” là hoàn toàn sai. Thực tế không thể làm như vậy
    Cái này không portable, cũng không phải sandbox, và cũng không phải lớp trừu tượng hóa phần cứng
    Dù có binary wasm, nó chỉ chạy được khi phiên bản runtime đang dùng tình cờ triển khai đúng backend ggml cụ thể cần thiết, mà không hề có yêu cầu nào như vậy, nên phần lớn khả năng là không
    Kể cả có chạy được thì rốt cuộc cũng là gọi mã ggml của llama.cpp, nên độ an toàn phụ thuộc vào chính thư viện đó
    Bài viết nhấn mạnh quá nhiều vào “tính portable” và “Rust”, nhưng thực tế không thể hiện rõ ưu điểm của bên nào
    Giả sử có một runtime WASI trên phần cứng mới, liệu có chạy được mô hình không? Có hỗ trợ GPU không? Câu trả lời sẽ là “hãy kiểm tra xem llama.cpp có được biên dịch với hỗ trợ GPU trên nền tảng đó không, runtime bạn dùng có plugin ggml không, và trong đó có vendor đúng phiên bản ggml phù hợp không. Nếu không thì không chạy được”
    Nếu vậy thì rốt cuộc dùng WASI để làm gì?
    Đúng là hỗ trợ GPU đa nền tảng rất khó, nhưng cách này trông khá vô lý
    Hãy tưởng tượng WebGPU hoạt động theo kiểu “ném một chunk binary vào GPU, và nếu chunk đó tình cờ phù hợp với phần cứng hiện tại thì có thể nó sẽ vẽ ra thứ gì đó”; cấu trúc ở đây chính là như vậy

    • Cảm ơn phần giải thích. Tôi đã thắc mắc hỗ trợ GPU trong WASM đến từ đâu
    • Bạn có thể giải thích chi tiết hơn về tác động ở khía cạnh bảo mật không?

  • Có thể chạy cái này offline trên iPhone không? Nếu được thì khi đi cắm trại sẽ hữu ích, vì có thể dùng kiểu như tìm kiếm Internet cơ bản bất kể sóng sánh thế nào

    • Có thể chạy trên nhiều thiết bị dựa trên Linux, Mac, Windows, bao gồm cả Raspberry Pi và hầu hết laptop/server
      Tuy nhiên để nạp bản thân mô hình thì vẫn cần vài GB bộ nhớ
    • Tôi đã chạy thử dự án này trên Pixel. Có vẻ nó cũng hoạt động trên một số iPhone/iPad
      [0] https://github.com/mlc-ai/mlc-llm
    • Với mục đích đó, có lẽ tải xuống một bản Wikipedia sẽ tốt hơn. Còn có những vấn đề như entropy nữa
    • Tôi đã làm một ứng dụng iOS thương mại khá thành công cho mục đích đó
      Ban đầu làm bằng ggml, sau đó phát hiện mlc-llm và port sang nền tảng đó
      [0]: https://apps.apple.com/us/app/private-llm/id6448106860

  • Theo luồng phát triển, chắc chắn sẽ còn có thêm các cách chạy kiến trúc transformer hiệu quả và nhanh hơn ở edge, nhưng yêu cầu VRAM thì không thể giải quyết bằng Rust, nên có vẻ đang tiến gần đến giới hạn
    Đây chính là nút thắt chính khi đưa một mô hình đủ lớn lên chạy
    Có thể nói “các mô hình nhỏ đang tốt lên, hãy nhìn Mistral và llama 2”, nhưng ngay cả mô hình nhỏ cũng đang tiến sát giới hạn dung lượng. Lượng thông tin có thể nhét vào 7 tỷ tham số là có hạn
    Tôi không nghĩ cách tiếp cận AI kiểu này sẽ dẫn tới AGI. Nó có vẻ quá kém hiệu quả

    • Tôi nghĩ vẫn còn khá nhiều dư địa, chẳng hạn như hệ thống MoE ngay cả với mô hình nhỏ, hoặc cách tải động LoRA

  • Mô tả rằng “plugin GGML cho Mac OS dùng Metal API để chạy tác vụ suy luận trên bộ xử lý thần kinh tích hợp của M1/M2/M3” có vẻ không chính xác
    Việc GGML dùng Metal API nghĩa là nó chạy trên GPU của M1/2/3, chứ không có nghĩa là chạy trên Neural Engine
    Tất nhiên bản thân điều đó cũng tốt rồi, nhưng nói nghiêm ngặt thì là như vậy

  • Tôi không thích kiểu marketing câu click nói như thể chỉ vì wrapper bọc thư viện lõi dùng ngôn ngữ khác mà dự án đã giảm kích thước xuống còn 1/100 hoặc tăng tốc độ lên 100~35000 lần so với các giải pháp khác
    Nó cũng hoàn toàn bỏ qua các công cụ và chuyên môn cộng đồng đã hình thành quanh các giải pháp khác
    Trước hết, dự án này dựa trên llama.cpp[1], và công việc nặng là nạp các tệp mô hình dung lượng nhiều GB lên GPU/CPU để chạy do llama.cpp xử lý
    Tốc độ suy luận không bị giới hạn bởi việc chọn wrapper; cũng có wrapper Go, Python, Node, Rust, và cũng có thể dùng trực tiếp llama.cpp
    Kích thước binary cũng không quá quan trọng. Lý do là các tệp mô hình lượng tử hóa phổ biến nằm trong khoảng 5GB~40GB, và cần máy có GPU mạnh hoặc RAM 16~64GB
    [1] https://github.com/ggerganov/llama.cpp

  • Nếu phần lớn kích thước rốt cuộc là trọng số mô hình đã được huấn luyện, thì làm sao có thể giảm kích thước đi vài bậc độ lớn mà không mất độ chính xác?

    • Tôi nghĩ khó có thể giảm kích thước mà không mất độ chính xác. Tuy vậy GGUF đã lượng tử hóa thì rất tuyệt
      Con số 2MB ở đây có vẻ chỉ kích thước chương trình, không tính mô hình. Có vẻ là cách chạy llama.cpp bằng một máy chủ Rust chạy wasm và llama.cpp
      Tôi thích llama.cpp/examples/server nhỏ gọn nên đang nhúng nó vào FreeChat, nhưng có thêm lựa chọn công cụ thì lúc nào cũng tốt
      Kiểm tra lại thì tệp thực thi arm64/x86 mà tôi nhúng hiện là 4,2MB. FreeChat là 12,1MB, nhưng mô hình mặc định khoảng 3GB nên tôi không quá bận tâm đến chênh lệch 2MB
      [0]: https://github.com/ggerganov/llama.cpp/tree/master/examples/...
    • Nếu ý là giảm kích thước của chính mô hình, tức trọng số đã huấn luyện, đi vài bậc độ lớn mà không mất độ chính xác, thì đó là một bài toán khó riêng biệt
      Bài viết này nói về việc giảm kích thước ứng dụng suy luận xuống 100 lần

  • Tôi không hiểu lý do gì để dùng cái này thay vì dùng trực tiếp llama.cpp

    • Gợi ý: đơn vị tiền tệ của nền kinh tế Rewrite-it-in-Rust thực ra không phải là chạy được thứ gì đó
    • llama.cpp thường phải biên dịch riêng cho từng hệ điều hành và kiến trúc như Windows, macOS, Linux, nên tính di động kém
      Bài viết cũng giải thích rằng nó tận dụng tăng tốc phần cứng trên các thiết bị có bộ tăng tốc phần cứng dị chủng
      Điều này có nghĩa là chương trình được biên dịch sang wasm có thể sử dụng hiệu quả các tài nguyên phần cứng đa dạng trên nhiều thiết bị, như GPU hoặc chip AI chuyên dụng
      Nếu tự triển khai bằng C++, để đạt hiệu năng tương tự có thể sẽ cần tối ưu hóa hoặc phiên bản riêng cho từng loại phần cứng