🦍

WasmでLLMが動く仕組みの解説

2023/12/16に公開

WebAssembly

LLM

tech

この記事は WebAssembly Advent Calendar 2023 16日目の記事です。

はじめに

先日、次のブログでWasmEdgeを使ってLLMを動かすことができることを知りました。

以前RustでWasm Runtimeを実装した身として、どのような仕組みで動いているのか気になって調べたので、その仕組について解説してきます。

前提知識について

解説するまえに、いくつかの前提知識について解説していきます。

Wasmでの関数呼び出し

Wasmではモジュール内に定義した関数を呼ぶことはもちろんできますが、モジュール外部の関数をimportして呼び出すこともできます。
たとえば、次のWATはcalcというモジュールからdoubleという関数をimportして使うことができます。

(module
  (import "calc" "double" ;; importするモジュールと関数
    (func $bar (param i32) (result i32)) ;; importした関数の型
  )
  ...
)

double関数の実態を知らなくても、関数を使う際のインターフェイスさえ分かれば呼び出せます。
Rustでいうと、次のようなことをやっています。

extern "C"  {
    fn double(x: i32) -> i32;
}

fn main() {
    unsafe { double(10) };
}

WASI Preview 1とは

WASI Preview 1（以降WASI）はWasmからOSのファイルや環境変数といったリソースにアクセスるためのインターフェイスです。
これらのインターフェイスはモジュールからはただの外部関数として見えます。

WASIを使う際は、さきほど解説した外部関数の呼び出し方と同じです。
たとえば、次はWATはfd_write()という関数を使って、線形メモリ上にある"Hello, World!"という文字列を標準出力に書き出しています。

(module
  (import "wasi_snapshot_preview1" "fd_write" ;; WASIのfd_write()をimport
    (func $fd_write (param i32 i32 i32 i32) (result i32)) ;; fd_write()の型を定義
  )
  (memory (export "memory") 1)
  (data (i32.const 0) "Hello, World!\n")

  (func $write_hello_world (result i32)
    (local $iovec i32)

    (i32.store (i32.const 16) (i32.const 0))
    (i32.store (i32.const 20) (i32.const 7))
    (i32.store (i32.const 24) (i32.const 7))
    (i32.store (i32.const 28) (i32.const 7))

    (local.set $iovec (i32.const 16))

    (call $fd_write
      (i32.const 1)
      (local.get $iovec)
      (i32.const 2)
      (i32.const 28)
    )
  )
  (export "_start" (func $write_hello_world))
)

fd_write()の実態はWasm Runtime側の関数です。
筆者のRuntimeではfd_write()は次のように実装しています。

このようにWASIの関数の実態はただのRuntime側の関数なので、実質なんでもできます。

`wasi-nn`について

さきほどの説明で察している方もいるかと思いますが、たとえばLLM Runtimeを使って入力を推論する関数inference()をWasm Runtime側で実装して、WasmモジュールでインポートすればWasmからLLMを実行できちゃいます。

しかし、各種Runtimeがインターフェイスを独自に定義して実装すると互換性がなくなります。
そこでwasi-nnの出番です。

wasi-nnは機械学習（以降ML）のためのWASIインターフェイスです。2019年からすでにあります。
次の関数が定義されています。

関数	概要
`load`	モデルをロードする
`load_by_name`	名前からモデルをロード
`init_execution_context`	モデルの初期化
`set_input`	入力をセット
`compute`	推論を実行
`get_output`	推論結果を取得

WasmEdgeはこのwasi-nnを実装して、WasmからLLMを動かせるようになっています。
これが冒頭のブログで紹介しているLLMが動く理由です。

あらためて、ブログで使っているllama-chat.wasmでimportしている関数を見てみると、
wasi-nnの関数を使っているのがわかりますね。

$ wasm2wat llama-chat.wasm 2>&1 | grep import
  (import "wasi_ephemeral_nn" "load_by_name" (func (;0;) (type 7)))
  (import "wasi_ephemeral_nn" "init_execution_context" (func (;1;)
  (import "wasi_ephemeral_nn" "compute" (func (;2;) (type 2)))
  (import "wasi_ephemeral_nn" "get_output" (func (;3;) (type 8)))
  (import "wasi_ephemeral_nn" "set_input" (func (;4;) (type 7)))
  ...

WasmEdgeの実装について

WasmEdgeの場合wasi-nnはプラグインという形で提供されて、WasmEdge本体には入っていません。
wasi-nnの具体的な実装はリポジトリの次の場所から見れます。

リポジトリを見てもらえれば分かると思いますが、ggmlやopenvino、torchといったRuntimeをつかったwasi-nnの実装があります。
冒頭にあったブログではggmlを使っています。

ggmlはllama.cppで使われているLLM Runtimeです。
このRuntimeはCPUでも実用レベルの速さで推論が動きます。

ggmlの詳細はこちらを参照してください。

WasmEdgeのggmlの実装はggml.cppにあるので、詳細を知りたい方はこちらを読めばどんなことをやっているのか分かると思います。

他のWasm Runtimeのwasi-nn対応

wasmtimeはopenvinoのみのようです。

wamrはtensorflowliteのみのようです。

wasmerは未サポートのようで、ディスカッションはありますが、何もコメントがついていませんでした。
WASIXやonyxlangに力を入れているかもしれないですね。

Wasm + LLMは今のところsecondstateのホームページを見ても分かるくらい、WasmEdgeが一番力を入れているなぁと感じます。
今後に期待ですね。

https://www.secondstate.io

Wasm + LLMの嬉しいところ

LLMの実行の部分はWasm Runtime側にあるので、Wasmバイナリが小さくなるというメリットがあります。
これは、アプリケーション自体がLLM Runtimeを含めてなくて済むのでWasmバイナリが小さくなるということですね。

またwasi-nnをブラウザ側で実装すれば、ブラウザだけに閉じたLLMをサービスとして提供できるようになります。
たとえばのAI文章校正機能を気軽に提供できる、サーバーのリソースを気にせずスケールできる、オフラインで実行できるしといろんなメリットがあるかなと思います。

まとめ

WasmEdgeはwasi-nnを実装している
wasi-nnはMLモデル推論するためのインターフェイス
ggmlというLLM Runtime使ったwasi-nn実装をしているため、LLMを動かせる
ggmlはllama.cppで使われていて、CPUでも高速推論できる
任意のLLM Runtimeを使ったwasi-nn実装があれば、Wasmからは任意のMLモデルを動かせる