Transformers 4.40 のパフォーマンス向上(OpenAI GPT‑OSS 向け):ゼロビルドカーネル、MXFP4 量子化、並列化、ロード高速化

TL;DR

OpenAI の GPT‑OSS モデルが Hugging Face の transformers ライブラリでフルサポートされ、ゼロビルドでダウンロード可能なカーネル、MXFP4 4 ビット量子化、テンソルおよびエキスパート並列、動的スライディングウィンドウキャッシュ、連続バッチング、モデルロード高速化といった一連のパフォーマンス向上が提供されました。これにより、開発者は単一 GPU または複数 GPU 上でこれらのモデルをより効率的にロード、実行、ファインチューニングできるようになります。

ゼロビルドカーネル(Hub から)

結論: 事前コンパイル済みのカスタムカーネルを自動でダウンロードでき、ビルド時の依存関係が不要になり、一般的な LLM 操作で最大 10 倍の速度向上が期待できます。

transformerskernels パッケージを統合し、初回使用時に Hugging Face Hub からバイナリカーネル(例: Liger RMSNorm、MegaBlocks MoE、FlashAttention 3)を取得します。モデル読み込み時に use_kernels=True を渡すだけで有効化できます。

from transformers import AutoTokenizer, AutoModelForCausalLM
import logging
logging.basicConfig(level=logging.INFO)

model_id = "openai/gpt-oss-20b"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    dtype="auto",
    device_map="auto",
    use_kernels=True,
)

ログ出力でロードされたカーネル(例: LigerRMSNormMegaBlocksMoeMLP)が確認できます。元記事のベンチマークでは、これらのカーネルは大きなバッチサイズで特に効果的ですが、実際のワークロードでベンチマークを取ることが推奨されます。

FlashAttention 3 と Attention Sinks

結論: Attention Sinks をサポートする FlashAttention 3 カーネルを有効にすると、Hopper クラス GPU でレイテンシが低減します。

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    dtype="auto",
    device_map="auto",
    attn_implementation="kernels-community/vllm-flash-attn3",
)

このカーネルは NVIDIA Hopper GPU と互換性があり、GPT‑OSS で導入された attention‑sink 機能を活用します。

MXFP4 4‑ビット量子化

結論: MXFP4 により、120 B パラメータの GPT‑OSS モデルの VRAM 使用量が最大 80 GB 削減され、消費電力の低い単一 GPU でも実行可能になり、専用の Triton カーネルで推論速度を維持します。

MXFP4 とは

MXFP4 は E2M1 4 ビット浮動小数点形式(符号 1 ビット、指数 2 ビット、仮数 1 ビット)を採用し、ブロック単位スケーリング(32 要素ブロックが同一スケールを共有)と組み合わせています。この設計により、粗い仮数でも動的範囲を保持できます。

transformers でのネイティブサポート

ライブラリには量子化器 (quantizer_mxfp4.py) と統合フック (integrations/mxfp4.py) が同梱されています。モデル設定に "quant_method": "mxfp4" が含まれると、MXFP4 パスが自動的に選択されます。

from transformers import GptOssConfig
cfg = GptOssConfig.from_pretrained("openai/gpt-oss-120b")
print(cfg.quantization_config)

必要な環境(acceleratekernelstriton>=3.4、Compute Capability ≥ 7.5 の GPU)が整っていれば MXFP4 モードで実行され、そうでなければ bfloat16 にフォールバックし、約 4 倍のメモリを消費します。

メモリ削減効果

元記事の図 3 では、量子化された 20 B モデルが約 16 GB、非量子化時は約 64 GB を使用することが示されています。同様のスケールが 120 B モデルでも当てはまり、約 80 GB 対 >300 GB です。

テンソル並列(TP)

結論: TP はテンソルを GPU 間で分割し、単一 GPU のメモリ上限を超えるモデルでもマルチ GPU ノードでスループットを向上させます。

transformersfrom_pretrainedtp_plan="auto" 引数を受け付け、組み込みのシャーディングレシピを自動選択します。使用例:

from transformers import PreTrainedTokenizerFast, GptOssForCausalLM
model_id = "openai/gpt-oss-120b"
tokenizer = PreTrainedTokenizerFast.from_pretrained(model_id)
model = GptOssForCausalLM.from_pretrained(
    model_id,
    tp_plan="auto",
    dtype="auto",
).eval()

TP は高速なノード内リンクを持つ単一ノードで最も効果的で、device_map="auto"(メモリ配置のみ)とは別概念です。

エキスパート並列(EP)

結論: EP は MoE のエキスパートを GPU 間で分散し、TP を補完して MoE モデルの GPU 当たり計算負荷をさらに低減します。

from transformers import DistributedConfig
model = GptOssForCausalLM.from_pretrained(
    model_id,
    distributed_config=DistributedConfig(enable_expert_parallel=True),
    dtype="auto",
).eval()

EP が有効になると TP が自動的に有効化され、両者の効果が合算されます。

動的スライディングウィンドウ層 & キャッシュ

結論: 新しい DynamicSlidingWindowLayerDynamicCache により、アテンションウィンドウ到達後に KV キャッシュが増大しなくなり、ハイブリッドアテンションモデル(例: GPT‑OSS)のキャッシュメモリが半減します。

デフォルトで有効化されており、必要に応じて明示的にキャッシュを作成できます:

from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache
model = AutoModelForCausalLM.from_pretrained(
    "openai/gpt-oss-20b",
    dtype="auto",
    device_map="auto",
).eval()
cache = DynamicCache(config=model.config)

ベンチマーク(図 6)では、長い生成時のメモリ削減とレイテンシ改善が大きく示されています。

連続バッチング & ページドアテンション

結論: generate_batch は動的(連続)バッチングを実装し、完了したスロットを新しいリクエストで即座に埋めることで GPU の稼働率を高め、静的バッチングに比べてトークン/秒を向上させます。

この API は実験的で、主に研究・評価用途向けです(本番サービングには vLLM や SGLang が適しています)。元記事のスクリプトとベンチマークでは、静的バッチングに対して最大約 2 倍の速度向上が報告されています。

モデルロード高速化

結論: transformers は GPU ごとに大きなメモリブロックを事前確保し、重みコピー時に数千回に及ぶ小規模割り当て呼び出しを削減することで、数十億パラメータモデルのロード時間を大幅に短縮します。

この挙動は device_map="auto" または任意の明示的デバイスマップ使用時に自動的に適用され、TP 有効時にも効果があります。

全体的なインパクト

結論: コミュニティ主導のカーネル、MXFP4 量子化、先進的な並列化戦略を transformers に直接統合したことで、最先端 LLM のハードウェアハードルが大幅に下がり、推論・ファインチューニングの速度が向上し、他ツールキット(MLX、llama.cpp、vLLM)向けの統一リファレンス実装が提供されました。

開発者は現在以下が可能です:

  • MXFP4 で無料枠の Colab GPU 上で GPT‑OSS 20 B をロード
  • 単一 torchrun コマンドで 4 GPU に跨って GPT‑OSS 120 B をスケール
  • 手動コンパイル不要で自動カーネルダウンロードを活用
  • 長文コンテキスト向けに KV キャッシュメモリを削減

これらの機能はすべて transformers リポジトリでオープンソースとして公開され、PR 参照やサンプルスクリプトが本文中にリンクされています。


主なリソース


開始手順

  1. 最新の transformers(≥ 4.40)をオプションの extras と共にインストール:
    pip install "transformers[torch,accelerate,triton,kernels]"
    
  2. モデルを選択(例: openai/gpt-oss-20b)。
  3. 必要な機能を有効化(use_kernels=Truetp_plan="auto"DistributedConfig(enable_expert_parallel=True))。
  4. 提供されたベンチマークスクリプトを実行し、ハードウェア上で速度とメモリ削減を確認。

これらの手順に従うことで、2025 年 9 月の Hugging Face ブログ記事で発表されたパフォーマンス向上をすぐに体感できます。


今後の方向性 ブログでは、今回の統合はスナップショットに過ぎず、コミュニティ貢献により新しい量子化形式、追加カーネルバックエンド、vLLM などのサービングスタックとの連携が今後も拡充されると述べられています。

Sources

関連

  • Dispatch
  • Dispatch
  • Dispatch
  • Dispatch
  • Dispatch