Rastreamento das saídas intermediárias do modelo

O método forward() de cada modelo costumava resolver manualmente as bandeiras None, como output_attentions, a partir dos padrões de configuração, acumular os pesos de atenção e os estados ocultos por camada em tuplas e converter os dataclasses ModelOutput em tuplas simples quando return_dict=False. Dois decoradores substituem toda essa boilerplate.

• @capture_outputs resolve as flags de saída, coleta valores intermediários e trata a conversão de return_dict.
• @merge_with_config_defaults resolve use_cache a partir da config. Omita-o para modelos que não fazem cache, como CLIPModel.

Você encontrará esses decoradores principalmente ao integrar um novo modelo. Consulte adicionando um modelo ao 🤗 Transformers para um guia passo a passo.

Declare quais submódulos capturar

Aplique @capture_outputs ao método forward() do modelo base. Ele anexa hooks de forward que:

• Intercepta as saídas de classes de submódulos especificados durante a passagem de frente, sem que esses submódulos precisem saber que estão sendo observados.
• Colete os pesos de atenção e os estados ocultos por camada em tuplas.
• Injete os valores coletados na dataclass ModelOutput retornada.
• Converta o dataclass em uma tupla simples quando return_dict=False.
• Resolve output_attentions e output_hidden_states a partir de kwargs ou self.config quando None.

Mapear campos de saída para submódulos

@capture_outputs precisa saber qual submódulo produz qual saída. Declare um dicionário de nível de classe _can_record_outputs em sua subclasse de PreTrainedModel. Cada chave é um nome de campo de saída ("hidden_states", "attentions", "cross_attentions") e cada valor é uma classe de módulo, uma string com o nome da classe, uma instância de OutputRecorder ou uma lista desses para registrar vários tipos de módulo sob uma única chave.

Controle fino com OutputRecorder

OutputRecorder aceita um target_class (uma subclasse de nn.Module cujas saídas devem ser coletadas) e um index opcional para selecionar qual elemento da tupla de saída do módulo deve ser obtido. Passe layer_name para anexar hooks apenas a módulos com um nome de atributo específico. Use layer_name quando duas camadas compartilham a mesma classe, por exemplo, autoatenção versus atenção cruzada.

O exemplo abaixo mostra os dois decoradores em prática, com diferentes níveis de controle de saída. Consulte LlamaModel para uma referência do mundo real.

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 camadas de patch

O rastreamento de saída depende de _can_record_outputs apontar para as classes exatas que as camadas de um modelo instanciam. Se você substituir uma implementação de camada por um kernel de atenção personalizado, uma camada de especialista quantizada ou uma variante de arquitetura, esses ponteiros precisam permanecer sincronizados. A API de patching fornece um registro global limpo para manter _can_record_outputs consistente.

register_patch_mapping mapeia nomes de classes originais para subclasses de nn.Module de substituição. As chaves podem ser nomes de classes exatos ou padrões de expressão regular. Os correspondentes exatos têm prioridade. Os padrões são testados com re.search(), portanto, padrões sem âncora correspondem em qualquer parte do nome da classe. Registrar a mesma chave duas vezes levanta ValueError, a menos que você passe overwrite=True.

Remova as entradas com 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$"])

Quando os mapeamentos são registrados, patch_output_recorders percorre todos os submódulos e atualiza cada OutputRecorder.target_class para o substituto registrado.

O método from_pretrained() chama patch_output_recorders automaticamente. Você só precisa chamá-lo manualmente ao construir um modelo diretamente.

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)