yufuin/plamo-2-translate-AWQ-INT4-ONNX-WEBGPU

PLaMo 2 Translate AWQ-INT4 ONNX for WebGPU

pfnet/plamo-2-translate にAWQ (INT4) 量子化を適用し、さらにWebGPU推論向けの調整を施したONNXモデルです。

2026年3月時点のWebGPUバックエンドでの利用を前提としてopグラフ構造を調整しています。Pythonランタイムなど非WebGPU環境で利用するとむしろ計算効率が悪い構造となっているので注意してください。また、2026年3月現在はonnxruntime-web wasm (CPU) バックエンドではメモリサイズの制約のため使用できません。

動作確認はChromeとFireFox (ver 146) の2環境で行ってます。FireFox ver 147, 148では途中の層から演算結果が崩壊し計算できない挙動を確認しており使用できません。 またブラウザのWebGPUのbufferSizeが1GB以上の環境でしか動作しません。

想定する利用経路は、transformers.jsによりモデルをロードし、同梱のplamo2_onnx_for_webgpu_module.mjsをヘルパとしてmodel.generate関数にパッチをあてて翻訳を実行する構成です。

デフォルトのeosは<|plamo:op|>に設定しています。

使い方

基本の利用手順は次のとおりです。

1. 本レポジトリに含まれるplamo2_onnx_for_webgpu_module.mjsをあなたのコード内でインポートする
2. Transformers.js v4によりtokenizerをロードする
3. 提供ヘルパ関数getPlamo2OnnxForWebGPUModelLoadOptions()を使って必要オプションを用意しTransformers.js v4によりモデルをロードする
4. 提供ヘルパ関数attachPlamo2GenerateFastPath(...)によりモデルインスタンスにパッチをあてて専用generateへ差し替える
5. 提供ヘルパ関数buildPlamo2TranslationPrompt(...)で翻訳プロンプトを組み立て、tokenizerでinput_idsへエンコードする
6. model.generate(...)で翻訳する

import * as tx from '@huggingface/transformers';
import { Tensor } from '@huggingface/transformers/src/utils/tensor.js';
import {
  attachPlamo2GenerateFastPath,
  buildPlamo2TranslationPrompt,
  getPlamo2OnnxForWebGPUModelLoadOptions,
} from './plamo2_onnx_for_webgpu_module.mjs';

const tokenizer = await tx.AutoTokenizer.from_pretrained('yufuin/plamo-2-translate-AWQ-INT4-ONNX-WEBGPU');

const model = await tx.AutoModel.from_pretrained('yufuin/plamo-2-translate-AWQ-INT4-ONNX-WEBGPU', getPlamo2OnnxForWebGPUModelLoadOptions());

attachPlamo2GenerateFastPath(model, { Tensor });

const prompt = buildPlamo2TranslationPrompt({
  sourceLang: 'English',
  targetLang: 'Japanese',
  text: 'Please translate this sentence.',
});

const encoded = tokenizer(prompt, {
  add_special_tokens: false,
  return_attention_mask: false,
});

const output = await model.generate({
  inputs: encoded.input_ids,
  max_new_tokens: 256,
  do_sample: false,
  temperature: 1.0,
  top_k: 50,
  top_p: 1.0,
  num_beams: 1,
  num_return_sequences: 1,
});

プロンプトについて

plamo-2-translateは想定する翻訳プロンプトフォーマットが決まっています。 フォーマットに合うプロンプトを構成するヘルパとしてbuildPlamo2TranslationPromptを提供しています。

低レベル情報

ONNX 入力

名前	DType	Shape	説明
input_ids	int64	[batch, sequence]	現在の prompt または decode step の token ids
past_attn_key	float32	[16, batch, 4, cache_sequence, 128]	attention KV cache
past_attn_value	float32	[16, batch, 4, cache_sequence, 128]	attention KV cache
past_mamba_conv_state	float32	[16, batch, 8192, 3]	Mamba convolution state cache
past_mamba_ssm_state	float32	[16, batch, 64, 128, 64]	Mamba SSM state cache

ONNX 出力

名前	DType	Shape	説明
logits	float32	[batch, 1, 100032]	最後の位置に対応する logits
next_past_attn_key	float32	[16, batch, 4, cache_sequence_out, 128]	更新後 attention KV cache
next_past_attn_value	float32	[16, batch, 4, cache_sequence_out, 128]	更新後 attention KV cache
next_past_mamba_conv_state	float32	[16, batch, 8192, 3]	更新後 Mamba convolution state
next_past_mamba_ssm_state	float32	[16, batch, 64, 128, 64]	更新後 Mamba SSM state

Shape 補足

• input_idsのsequenceサイズは[1,8]の範囲に限られます。
  • このモデルは高速化のためにScanを8ステップ固定のopグラフへ展開しており、それによる制約です。
• 想定する使用方法通りattachPlamo2GenerateFastPath適用後のmodel.generate(...)を実行した場合はinput_idsのスライスは自動で処理されるため、通常はONNX入出力形状を直接意識する必要はありません。
  • batch = 1前提です。

ライセンス

Built with PLaMo。 このリポジトリは pfnet/plamo-2-translate の派生配布物であり、ライセンスはupstreamと同じPLaMo Community Licenseです。

• 日本語版ライセンス: LICENSE/ja
• 英語版ライセンス: LICENSE/en
• upstream model page: https://huggingface.co/pfnet/plamo-2-translate

このモデルを使用する際も依然 PLaMo Community License は適用されます。 再配布や派生物の公開時はPLaMo Community Licenseライセンス条件に従ってください。特に、関連する UI・ドキュメント・配布ページ等で Built with PLaMo を明示する必要があります。