Отслеживание промежуточных выходов модели

Метод forward() каждой модели ранее вручную разрешал флаги None, такие как output_attentions, из значений по умолчанию в конфигурации, накапливал веса внимания и скрытые состояния для каждого слоя в кортежи, а также преобразовывал dataclass ModelOutput в обычные кортежи, когда return_dict=False. Два декоратора заменяют весь этот шаблонный код.

• @capture_outputs разрешает флаги вывода, собирает промежуточные значения и обрабатывает преобразование return_dict.
• @merge_with_config_defaults определяет use_cache из конфигурации. Для моделей, которые не используют кэш, например CLIPModel, этот параметр можно опустить.

Вы чаще всего столкнётесь с этими декораторами при интеграции новой модели. Пошаговое руководство см. в разделе adding a model to 🤗 Transformers.

Объявите, какие подмодули нужно захватывать

Примените @capture_outputs к методу forward() базовой модели. Он добавляет forward-хуки, которые:

• Перехватывайте выходные данные указанных классов подмодулей во время прямого прохода, не требуя, чтобы эти подмодули знали о том, что за ними наблюдают.
• Собирайте веса внимания и скрытые состояния для каждого слоя в кортежи.
• Вставьте собранные значения в возвращаемый dataclass ModelOutput.
• Преобразуйте dataclass в обычный кортеж, если return_dict=False.
• Определите output_attentions и output_hidden_states из kwargs или self.config, если они равны None.

Соответствие полей выходных данных подмодулям

@capture_outputs необходимо знать, какой подмодуль генерирует какой вывод. Объявите словарь уровня класса _can_record_outputs в вашем подклассе PreTrainedModel. Каждый ключ — это имя поля вывода ("hidden_states", "attentions", "cross_attentions"), а каждое значение — класс модуля, строка с именем класса, экземпляр OutputRecorder или список таких элементов для записи нескольких типов модулей под одним ключом.

Тонкий контроль с помощью OutputRecorder

OutputRecorder принимает target_class (подкласс nn.Module, для которого нужно собрать выходные данные) и необязательный index для выбора конкретного элемента кортежа выходов модуля. Передайте layer_name, чтобы прикреплять хуки только к модулям с определённым именем атрибута. Используйте layer_name, если два слоя имеют один и тот же класс, например self-attention и cross-attention.

Приведённый ниже пример демонстрирует оба декоратора на практике с различными уровнями контроля вывода. См. 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
        )

Классы патч-слоёв

Слежение за выходом зависит от того, на какие точные классы указывает _can_record_outputs — те, которые инстанцируют слои модели. Если вы заменяете реализацию слоя на собственное ядро внимания, квантованный слой эксперта или вариант архитектуры, эти указатели должны оставаться синхронизированными. API патчинга предоставляет чистый глобальный реестр для поддержания согласованности _can_record_outputs.

register_patch_mapping сопоставляет исходные имена классов с заменяющими подклассами nn.Module. В качестве ключей можно использовать точные имена классов или регулярные выражения. Точные совпадения имеют приоритет. Паттерны проверяются с помощью re.search(), поэтому не закреплённые паттерны совпадают в любом месте имени класса. Повторная регистрация того же ключа вызывает ValueError, если не передать overwrite=True.

Удалите записи с помощью 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)