Rust+WASM으로 이기종 Edge에서 빠르고 포터블한 Llama2 추론 실행하기
(secondstate.io)- Rust+WASM 스택이 AI추론에서 Python의 강력한 대체제가 될 수 있음
- 파이썬과 비교해서 Rust+Wasm 앱은 크기가 1/100, 속도가 100배 빠르며, 바이너리 코드를 변경하지 않고도 전체 하드웨어 가속으로 모든 곳에서 안전하게 실행할 수 있음
- Rust는 AGI의 언어임
- LLama2 추론을 네이티브 속도로 실행하는 간단한 Rust 프로그램을 작성해보면
- Wasm 컴파일후 바이너리는 2MB에 불과하지만, 이기종 하드웨어 가속기에 완전히 이식 가능
- 또한 Wasm 런타임(WasmEdge)은 클라우드 환경을 위한 안전한 보안 실행 환경을 제공. 컨테이너 도구와 원활하게 작동하여 다양한 기기에서 휴대용 애플리케이션을 오케스트레이션하고 실행할 수 있음
왜 파이썬은 아닌가?
- llama2와 같은 LLM은 일반적으로 Python(예: PyTorch, Tensorflow, JAX)으로 학습
- 하지만 AI 컴퓨팅의 약 95%에 해당하는 추론 애플리케이션에 Python을 사용하는 것은 큰 실수
- Python 패키지는 복잡한 종속성을 가지고 있음. 설정하고 사용하기가 어려움
- Python 종속성은 엄청남. Python 또는 PyTorch용 Docker 이미지는 일반적으로 몇 GB 또는 수십 GB에 달하며, 이는 엣지 서버나 디바이스에서 AI 추론을 수행할 때 특히 문제가 됨
- Python은 매우 느린 언어. C, C++, Rust와 같은 컴파일 언어보다 최대 35,000배 느림
- Python은 느리기 때문에 대부분의 실제 워크로드는 Python 래퍼 아래의 네이티브 공유 라이브러리에 위임해야 함
- 따라서 Python 추론 앱은 데모용으로는 훌륭하지만, 비즈니스별 요구 사항에 맞게 내부적으로 수정하기는 매우 어려움
- 네이티브 라이브러리에 대한 의존도가 높고 종속성 관리가 복잡하기 때문에 기기의 고유한 하드웨어 기능을 활용하면서 Python AI 프로그램을 여러 기기에 포팅하기가 매우 어려움
- LLM 툴체인에서 일반적으로 사용되는 Python 패키지는 서로 직접적으로 충돌함
- LLVM, 텐서플로우, Swift 언어로 유명한 Chris Lattner가 이번 주 스타트업 팟캐스트에서 멋진 인터뷰를 진행했는데, 그는 파이썬이 모델 훈련에는 훌륭하지만 추론 애플리케이션에는 잘못된 선택인 이유에 대해 설명했음
Rust+Wasm의 장점
- Rust+Wasm 스택은 디바이스에서 엣지 클라우드, 온프레미스 서버, 퍼블릭 클라우드에 이르는 통합된 클라우드 컴퓨팅 인프라를 제공
- AI 추론 애플리케이션을 위한 Python 스택의 강력한 대안이며, 엘론 머스크가 Rust를 AGI의 언어라고 말한 것은 놀랍지 않음
- 초경량: 추론 애플리케이션은 모든 종속성을 포함하여 2MB에 불과. 이는 일반적인 PyTorch 컨테이너 크기의 1%도 되지 않음
- 매우 빠름: 전처리, 텐서 연산, 후처리 등 추론 애플리케이션의 모든 부분에서 네이티브 C/Rust 속도를 제공
- 포터블: 이기종 하드웨어 가속을 지원하는 모든 주요 컴퓨팅 플랫폼에서 동일한 Wasm 바이트코드 애플리케이션을 실행할 수 있음
- 설정, 개발 및 배포가 쉬움: 더 이상 복잡한 종속성이 없음. 노트북의 표준 도구를 사용하여 단일 Wasm 파일을 빌드하고 어디에나 배포 가능
- 안전하고 클라우드 지원: Wasm 런타임은 신뢰할 수 없는 사용자 코드를 격리하도록 설계됨. Wasm 런타임은 컨테이너 도구로 관리할 수 있으며 클라우드 네이티브 플랫폼에 쉽게 배포할 수 있음
Rust 추론 데모
- 40줄의 Rust 코드로 작성되었으며 Wasm으로 컴파일됨
- Wasm 바이트코드가 생기면 WasmEdge 런타임을 지원하는 어떤 기기에든 배포 가능
- llama.cpp를 기반으로 하는 WasmEdge GGML 플러그인은 기기의 하드웨어 가속을 자동으로 활용하여 llama2 모델을 실행
다음 단계
- WasmEdge GGML 툴링은 현재 사용할 수 있으며 실제로 클라우드 네이티브 고객들이 사용하고 있지만, 아직 초기 단계
- 오픈 소스 프로젝트에 기여하고 미래 LLM 추론 인프라의 방향을 형성하는 데 관심이 있으시다면, 잠재적으로 기여할 수 있는 몇 가지
- 더 많은 하드웨어와 OS 플랫폼을 위한 GGML 플러그인을 추가하기: Linux와 Windows의 TPU, ARM NPU 및 기타 특수 AI 칩등
- 더 많은 llama.cpp 구성을 지원하기: 현재 일부 구성 옵션을 Wasm에서 GGML 플러그인으로 전달하는 것을 지원하지만 GGML이 제공하는 모든 옵션을 지원하고자 함
- 다른 Wasm 호환 언어의 WASI NN API를 지원하기: 특히 Go, Zig, Kotlin, JavaScript, C 및 C++
현재 AI에서 열세인 애플이 취할 수 있는 가장 현실적인 접근이,
본문글과 같은 최적화된 S/W 추론엔진이 아이폰 A칩의 NPU로 가속되고 foundation 모델파일을 아이폰에 통짜로 넣어 버리는 상황 같습니다.
최근 학계/업계에서 AI성능이 사이즈보다는 데이터 및 fine-tunning이 중요하다는게 중론이고 llama2에서 fine-tunning 성능이 상당히 좋다라는 평이 많아 잘 학습시킨 모델을 지속적으로 edge device로 내려주는것도 충분히 가능하리라 생각됩니다. 그런면에서 아이폰에 모두 때려박는 전략이 꽤 현실적이겠습니다.