ModernBERT packed ONNX for ruri-v3-30m

cl-nagoya/ruri-v3-30m を packed batch 対応の ONNX として使うための最小構成である。

目的は、通常の [batch, seq] padded ONNX ではなく、batch 内の pad token を除去して [1, total_tokens] に連結した ONNX 推論経路を使うことだ。cl-nagoya/ruri-v3-30m は SentenceTransformer 形式の日本語 embedding model で、mean pooling を使う構成として公開されている。

Quick Start

依存関係、Hugging Face cache、ONNX 出力は repo 内に閉じる。

UV_CACHE_DIR=.uv-cache HF_HOME=.hf-home uv run python -m modernbert_batch_onnx.fast export

これで onnx/packed.onnx ができる。文章を embedding にする。

UV_CACHE_DIR=.uv-cache HF_HOME=.hf-home uv run python -m modernbert_batch_onnx.fast encode \
  --text "検索したい日本語の文章" \
  --text "比較対象になる別の文章" \
  --output results/embeddings.npy

出力は float32 の NumPy 配列で、shape は [num_texts, 256]。デフォルトで L2 normalize 済みである。

テキストファイルを使う場合は、1行1文章で渡す。

UV_CACHE_DIR=.uv-cache HF_HOME=.hf-home uv run python -m modernbert_batch_onnx.fast encode \
  --input texts.txt \
  --output results/embeddings.npy

Python API

組み込み側はこのクラスだけ使えばよい。

from modernbert_batch_onnx import PackedModernBertEncoder

encoder = PackedModernBertEncoder("onnx/packed.onnx", threads=4)
embeddings = encoder.encode(
    [
        "検索したい日本語の文章",
        "比較対象になる別の文章",
    ],
    batch_size=16,
)

print(embeddings.shape)  # (2, 256)

normalize しない raw embedding が必要な場合だけ normalize=False を指定する。

encoder = PackedModernBertEncoder("onnx/packed.onnx", normalize=False)

ONNX Interface

onnx/packed.onnx の入力は固定 batch 次元を持つ packed 形式である。

input_ids              [1, tokens]
position_ids           [1, tokens]
full_attention_mask    [1, 1, tokens, tokens]
sliding_attention_mask [1, 1, tokens, tokens]

fast encode と PackedModernBertEncoder は次を行う。

texts
  -> tokenize
  -> pad token なしで横に連結
  -> sequence ごとに position_ids を 0 から振り直す
  -> sequence をまたぐ attention を 4D mask で遮断
  -> ONNX Runtime
  -> mean pooling
  -> L2 normalize

ModernBERT 本体を NumPy などで再実装しているわけではない。Hugging Face の AutoModel を torch.onnx.export() し、packed 入力と 4D additive mask を ONNX Runtime に渡している。

Benchmark

PoC の比較 benchmark は別入口にしてある。

UV_CACHE_DIR=.uv-cache HF_HOME=.hf-home HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 \
uv run python -m modernbert_batch_onnx.poc --threads 4 --report results/report.md

現在の macOS CPU 実測は results/report.md に残している。多様な日本語文 16 件、4096 padded tokens 中 638 real tokens の条件で、packed ONNX は median latency が 454.67 ms から 72.01 ms になった。

Linux GPU

Linux GPU では CUDA 対応の ONNX Runtime を使い、provider を差し替える。

UV_CACHE_DIR=.uv-cache HF_HOME=.hf-home uv run python -m modernbert_batch_onnx.fast encode \
  --provider CUDAExecutionProvider \
  --input texts.txt \
  --output results/embeddings.npy

この repo の macOS 検証は CPUExecutionProvider で行っている。GPU での速度は環境依存なので、同じ poc コマンドで再測定すること。

Notes

• torch.onnx.export() で PyTorch / Transformers の FlashAttention2 varlen kernel がそのまま ONNX に保存されるわけではない。
• この repo の高速化は、ONNX Runtime に渡す入力を packed にし、pad token の embedding/MLP 計算と attention の無駄を減らすことで成立している。
• 4D mask は tokens x tokens なので、1 batch に詰め込みすぎると mask 生成とメモリが重くなる。まずは --batch-size 16 前後から測るのが現実的である。
• onnx/packed.onnx は生成物なので git 管理しない。