推論用の互換性のあるモデルバックエンドの構築

Transformers のモデルは、vLLM や SGLang などの推論エンジンと互換性があります。同じ Transformers モデルをどこでも使用でき、各推論エンジンごとにモデルをゼロから再実装する必要がありません。また、エンジンがネイティブに再実装していないモデルについても、以下のセクションで説明するアテンションバックエンドインターフェースをサポートしている限り、このアプローチが有効です。

このガイドでは、Transformersでモデルを実装し、あらゆる推論エンジンのバックエンドとして動作させる方法を示します。

モデルの実装

モデルの貢献ガイドラインまたはカスタムモデルの貢献ガイドラインに従ってください。モデルのディレクトリに有効なconfig.jsonが存在し、かつ設定内のモデルクラスを指す有効なauto_mapフィールドが設定されている必要があります。

カスタムおよび最適化されたアテンション関数には、AttentionInterface クラスを使用します。このインターフェースにより、各推論エンジンのパフォーマンス機能が解放されます。

アテンション層を定義する際はALL_ATTENTION_FUNCTIONSを使用し、ベースのMyModelクラスからアテンション層へ**kwargsを伝播させます。PreTrainedModelで_supports_attention_backendをTrueに設定してください。

以下のコードを例として拡張します。

modeling_my_model.py

from transformers import PreTrainedModel
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
from torch import nn

class MyAttention(nn.Module):

    def forward(self, hidden_states, **kwargs):
        ...
        attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
        attn_output, attn_weights = attention_interface(
            self,
            query_states,
            key_states,
            value_states,
            **kwargs,
        )
        ...

class MyModel(PreTrainedModel):
    _supports_attention_backend = True

PreTrainedConfig に以下のキーを追加することで、任意のテンソル並列またはパイプライン並列を有効にできます。
- base_model_tp_plan は、完全修飾レイヤー名パターンをテンソル並列スタイルにマッピングすることで、テンソル並列を有効にします。サポートされる分割戦略は "colwise" と "rowwise" のみです。
- base_model_pp_plan は、直接の子レイヤー名を文字列のリストのタプルにマッピングすることにより、パイプライン並列性を有効にします。タプルの最初の要素には、入力引数の名前が含まれます。最後の要素には、モデリングコード内のレイヤー出力の変数名が含まれます。
以下のコードを例として拡張します。
configuration_my_model.py
```
from transformers import PreTrainedConfig

class MyConfig(PreTrainedConfig):
    base_model_tp_plan = {
        "layers.*.self_attn.k_proj": "colwise",
        "layers.*.self_attn.v_proj": "colwise",
        "layers.*.self_attn.o_proj": "rowwise",
        "layers.*.mlp.gate_proj": "colwise",
        "layers.*.mlp.up_proj": "colwise",
        "layers.*.mlp.down_proj": "rowwise",
    }
    base_model_pp_plan = {
        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
        "norm": (["hidden_states"], ["hidden_states"]),
    }
```

マルチモーダルモデル

マルチモーダルモデルには、ビジョン言語モデルの貢献ガイドラインに加えて、追加の変更が必要です。これらの変更により、マルチモーダル入力が適切に処理されることを保証します。

ProcessorMixin クラスには、self.image_token と self.image_token_ids 属性を含める必要があります。これらのプレースホルダートークンは、入力における画像の位置を示します。同じトークンは、画像の入力プロンプト内と、画像特徴を配置するためのモデルコードの両方に出現します。
ProcessorMixin クラスには、self._get_num_multimodal_tokens メソッドを含める必要があります。このメソッドは、指定されたサイズのマルチモーダル入力に必要なプレースホルダートークンの数を計算します。戻り値は MultiModalData オブジェクトです。<image> トークンの間に存在するプレースホルダー（行トークンや列トークンなど）は、画像プレースホルダーとしてカウントしません。モデリングコードの後段で画像特徴量に置き換えられるトークンのみをカウントしてください。
ProcessorMixin クラスは、return_mm_token_type_ids の値を確認し、mm_token_type_ids を返す必要があります。これは、各ポジションがテキストトークン（0）、画像プレースホルダートークン（1）、または動画プレースホルダートークン（2）のいずれかを示します。マルチモーダルトークンタイプIDのシーケンスは、連続するトークンの間に中断がないように、連続して配置する必要があります。開始、終了、行、および列のトークン用の特殊トークンは、プレースホルダーとして扱います。

以下のコードを例として拡張します。

modeling_my_multimodal_model.py

class MyMultimodalProcessor(ProcessorMixin):

    def __call__(self, images=None, text=None, **kwargs):
        if return_mm_token_type_ids:
            mm_token_type_ids = np.zeros_like(input_ids)
            mm_token_type_ids[input_ids == self.image_token_id] = 1
            text_inputs["mm_token_type_ids"] = mm_token_type_ids.tolist()
        return BatchFeature(data={**text_inputs, **image_inputs}, tensor_type=return_tensors)

    def _get_num_multimodal_tokens(self, image_sizes=None, **kwargs):
        """
        Computes the number of placeholder tokens needed for multimodal inputs with the given sizes.
        Args:
            image_sizes (`list[list[int]]`, *optional*):
                The input sizes formatted as (height, width) per each image.
        Returns:
            `MultiModalData`: A `MultiModalData` object holding number of tokens per each of the provided
            input modalities, along with other useful data.
        """
        vision_data = {}
        if image_sizes is not None:
            num_image_tokens = [256] * len(image_sizes) # 256 placeholder tokens for each image always
            num_image_patches = [1] * len(image_sizes) # no patching, thus each image is processed as a single base image
            vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches})
        return MultiModalData(**vision_data)

リソース

詳細については、Transformers バックエンドの vLLM 統合に関するブログ記事をご覧ください。
詳細については、SGLangにおけるTransformersバックエンド統合のブログ記事をご覧ください。

Update on GitHub