Suivi des sorties intermédiaires du modèle

La méthode forward() de chaque modèle résolvait manuellement les drapeaux None tels que output_attentions à partir des valeurs par défaut de la configuration, accumulait les poids d'attention et les états cachés de chaque couche dans des tuples, et convertissait les dataclasses ModelOutput en tuples simples lorsque return_dict=False. Deux décorateurs remplacent l'ensemble de ce code répétitif.

@capture_outputs résout les drapeaux de sortie, collecte les valeurs intermédiaires et gère la conversion return_dict.
@merge_with_config_defaults résout use_cache à partir de la configuration. Omettez-le pour les modèles qui ne mettent pas en cache, comme CLIPModel.

Vous rencontrerez principalement ces décorateurs lors de l’intégration d’un nouveau modèle. Consultez ajouter un modèle à 🤗 Transformers pour un guide étape par étape.

Déclarer les sous-modules à capturer

Appliquez @capture_outputs à la méthode forward() du modèle de base. Elle attache des hooks avant qui :

Intercepter les sorties des classes de sous-modules spécifiés pendant la passe avant, sans que ces sous-modules aient besoin de savoir qu’ils sont observés.
Collecter les poids d'attention et les états cachés par couche dans des tuples.
Injectez les valeurs collectées dans la dataclass ModelOutput retournée.
Convertir le dataclass en tuple simple lorsque return_dict=False.
Résoudre output_attentions et output_hidden_states à partir des kwargs ou de self.config lorsque la valeur est None.

Associer les champs de sortie aux sous-modules

@capture_outputs doit savoir quel sous-module produit quelle sortie. Déclarez un dictionnaire de classe _can_record_outputs sur votre sous-classe de PreTrainedModel. Chaque clé est le nom d'un champ de sortie ("hidden_states", "attentions", "cross_attentions"), et chaque valeur est une classe de module, une chaîne de caractères contenant le nom de la classe, une instance de OutputRecorder, ou une liste de ceux-ci pour enregistrer plusieurs types de modules sous une même clé.

Contrôle fin avec OutputRecorder

OutputRecorder accepte un target_class (une sous-classe de nn.Module dont les sorties doivent être collectées) et un index facultatif pour sélectionner l’élément du tuple de sortie du module à récupérer. Passez layer_name pour n’attacher des hooks qu’aux modules portant un nom d’attribut spécifique. Utilisez layer_name lorsque deux couches partagent la même classe, par exemple l’auto-attention et l’attention croisée.

L’exemple ci-dessous montre les deux décorateurs en pratique, avec différents niveaux de contrôle de la sortie. Consultez LlamaModel pour une référence issue du monde réel.

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
        )

Classes de couches de patch

Le traçage des sorties dépend de _can_record_outputs qui doit pointer vers les classes exactes instanciées par les couches d’un modèle. Si vous remplacez une implémentation de couche par un noyau d’attention personnalisé, une couche d’expert quantifiée ou une variante d’architecture, ces pointeurs doivent rester synchronisés. L’API de patching fournit un registre global propre pour maintenir la cohérence de _can_record_outputs.

register_patch_mapping associe les noms de classes d’origine aux sous-classes de remplacement de nn.Module. Les clés peuvent être des noms de classes exacts ou des motifs regex. Les correspondances exactes ont la priorité. Les motifs sont testés avec re.search(), de sorte que les motifs non ancrés correspondent n’importe où dans le nom de la classe. Enregistrer deux fois la même clé lève une ValueError sauf si vous passez overwrite=True.

Supprimez les entrées avec 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$"])

Une fois les correspondances enregistrées, patch_output_recorders parcourt chaque sous-module et met à jour chaque OutputRecorder.target_class avec le remplacement enregistré.

La méthode from_pretrained() appelle patch_output_recorders automatiquement. Vous n'avez besoin de l'appeler vous-même que lors de la construction directe d'un modèle.

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)

Update on GitHub