Transformers documentation
モデルの中間出力のトレース
モデルの中間出力のトレース
すべてのモデルのforward()メソッドは、output_attentionsのようなNoneフラグを設定のデフォルト値から手動で解決し、レイヤーごとのアテンション重みと隠れ状態をタプルに蓄積し、return_dict=Falseの場合にModelOutputデータクラスをプレーンなタプルに変換していました。これらのボイラープレートはすべて、2つのデコレーターに置き換えられます。
@capture_outputsは出力フラグを解決し、中間値を収集して、return_dict変換を処理します。@merge_with_config_defaultsは config からuse_cacheを解決します。キャッシュを使用しないモデル（CLIPModel など）では、これを省略してください。
これらのデコレーターは、新しいモデルを統合する際に主に遭遇します。ステップバイステップのガイドについては、🤗 Transformers にモデルを追加するを参照してください。
キャプチャするサブモジュールを宣言
ベースモデルのforward()メソッドに@capture_outputsを適用します。これにより、以下の処理を行うforwardフックがアタッチされます。
- forward pass 中に指定されたサブモジュールクラスからの出力をインターセプトします。この際、監視されていることをサブモジュール側が認識する必要はありません。
- 各レイヤーのattention weightsとhidden statesをタプルに収集します。
- 収集した値を、返される
ModelOutputデータクラスに注入します。 return_dict=Falseの場合、dataclass を通常の tuple に変換します。Noneの場合、output_attentionsとoutput_hidden_statesを kwargs またはself.configから解決します。
出力フィールドをサブモジュールにマッピングする
@capture_outputs は、どのサブモジュールがどの出力を生成するかを把握する必要があります。PreTrainedModel のサブクラスに、_can_record_outputs というクラスレベルの辞書を宣言してください。各キーは出力フィールド名（"hidden_states"、"attentions"、"cross_attentions"）であり、各値はモジュールクラス、クラス名の文字列、OutputRecorder インスタンス、またはそれらのリスト（1 つのキーで複数のモジュールタイプを記録する場合）です。
OutputRecorderによる詳細な制御
OutputRecorder は、収集対象の出力を持つ nn.Module のサブクラスである target_class と、モジュールの出力タプルから取得する要素を選択するための任意の index を受け取ります。特定の属性名を持つモジュールにのみフックを付与するには、layer_name を渡します。自己注意と相互注意のように、2 つのレイヤーが同じクラスを共有している場合は、layer_name を使用してください。
以下の例では、異なる出力制御レベルで両方のデコレーターを実際に使用する方法を示しています。実際の使用例については、LlamaModelを参照してください。
from ...processing_utils import Unpack
from ...utils import TransformersKwargs
from ...utils.generic import merge_with_config_defaults
from ...utils.output_capturing import capture_outputs, OutputRecorder
class MyPreTrainedModel(PreTrainedModel):
_can_record_outputs = {
# Capture hidden_states: hook fires after each MyDecoderBlock forward,
# grabbing its first output (index 0 by default).
"hidden_states": MyDecoderBlock,
# Capture self-attention weights: hook fires after each MyAttention
# forward, grabbing its second output (index=1 by default).
"attentions": MyAttention,
# Capture cross-attention weights: same class, different submodule.
# layer_name targets the attribute `self.crossattention` inside the block.
# Captures second output as requested (index=1)
"cross_attentions": OutputRecorder(
MyAttention, layer_name="crossattention", index=1
),
}
# Now in base model's forward we need the decorators and `Unpack` `kwargs`
class MyModel(MyPreTrainedModel):
@merge_with_config_defaults # ← resolves use_cache
@capture_outputs # ← handles output collection + return_dict
def forward(
self,
input_ids: torch.LongTensor | None = None,
past_key_values: Cache = None,
**kwargs: Unpack[TransformersKwargs],
) -> BaseModelOutputWithPast:
# No manual collection needed. Just run your layers normally.
hidden_states = self.embed_tokens(input_ids)
for layer in self.layers:
hidden_states = layer(hidden_states, **kwargs)
# Return the primary outputs. The decorator will fill in
# hidden_states/attentions/cross_attentions automatically.
return BaseModelOutputWithPast(
last_hidden_state=hidden_states,
past_key_values=past_key_values
)Patchレイヤーのクラス
Output tracing は、_can_record_outputs がモデルのレイヤーがインスタンス化する正確なクラスを指していることに依存します。レイヤーの実装をカスタムなアテンションカーネル、量子化されたエキスパートレイヤー、またはアーキテクチャのバリアントに置き換えた場合、これらのポインタは同期を保たなければなりません。パッチング API は、_can_record_outputs の整合性を保つためのクリーンなグローバルレジストリを提供します。
register_patch_mapping は、元のクラス名を置き換え用の nn.Module サブクラスにマッピングします。キーには、完全一致するクラス名または正規表現パターンを指定できます。完全一致が優先されます。パターンは re.search() でテストされるため、アンカーなしのパターンはクラス名のどこにでも一致します。同じキーを2回登録すると、overwrite=True を渡さない限り ValueError が発生します。
unregister_patch_mapping を使用してエントリを削除します。
from transformers.monkey_patching import register_patch_mapping, unregister_patch_mapping
# Exact name – replaces only Qwen2MoeExperts
register_patch_mapping({"Qwen2MoeExperts": SequentialExperts})
# Regex – replaces every class whose name ends in "Attention"
register_patch_mapping({".*Attention$": FusedAttention})
# Anchored version – only matches Llama2Attention, Llama3Attention, …
register_patch_mapping({"^Llama\\d+Attention$": CustomLlamaAttention})
# Same way, custom keys can be removed from the registry by passing the name that was registered
unregister_patch_mapping(["Qwen2MoeExperts", ".*Attention$"])マッピングが登録されると、patch_output_recordersはすべてのサブモジュールを走査し、各OutputRecorder.target_classを登録された置換先に更新します。
from_pretrained() メソッドは
patch_output_recordersを自動的に呼び出します。モデルを直接構築する場合にのみ、自分で呼び出す必要があります。
from transformers.monkey_patching import patch_output_recorders
# Built manually, outside from_pretrained
model = Qwen2MoeModel(config)
# Without this, _can_record_outputs still points at the original Qwen2MoeExperts class
# and hooks will never fire on CustomExperts instances.
patch_output_recorders(model)