Theo dõi các đầu ra trung gian của mô hình

Phương thức forward() của mỗi model trước đây phải tự động phân giải các cờ None như output_attentions từ các giá trị mặc định trong config, tích lũy trọng số attention và hidden states của từng layer vào các tuple, và chuyển đổi các dataclass ModelOutput thành các tuple thông thường khi return_dict=False. Hai decorator thay thế toàn bộ phần boilerplate đó.

• @capture_outputs giải quyết các cờ đầu ra, thu thập các giá trị trung gian và xử lý chuyển đổi return_dict.
• @merge_with_config_defaults giải quyết use_cache từ config. Bỏ qua nó đối với các model không cache, chẳng hạn như CLIPModel.

Bạn sẽ chủ yếu gặp các decorator này khi tích hợp một mô hình mới. Xem cách thêm một mô hình vào 🤗 Transformers để có hướng dẫn từng bước.

Khai báo các submodule cần bắt

Áp dụng @capture_outputs vào phương thức forward() của mô hình cơ sở. Nó gắn các forward hook để:

• Chặn các đầu ra từ các lớp con module được chỉ định trong quá trình truyền thuận, mà không cần các con module đó phải biết chúng đang được quan sát.
• Thu thập trọng số chú ý và trạng thái ẩn cho từng lớp vào các bộ (tuples).
• Tiêm các giá trị đã thu thập vào dataclass ModelOutput được trả về.
• Chuyển đổi dataclass thành tuple thông thường khi return_dict=False.
• Giải quyết output_attentions và output_hidden_states từ kwargs hoặc self.config khi None.

Ánh xạ các trường đầu ra sang các module con

@capture_outputs cần biết submodule nào tạo ra output nào. Khai báo một từ điển cấp lớp _can_record_outputs trên lớp con PreTrainedModel của bạn. Mỗi khóa là tên trường output ("hidden_states", "attentions", "cross_attentions"), và mỗi giá trị là một lớp module, một chuỗi tên lớp, một instance OutputRecorder, hoặc một danh sách các đối tượng đó để ghi lại nhiều loại module dưới một khóa.

Kiểm soát chi tiết với OutputRecorder

OutputRecorder nhận một target_class (một lớp con của nn.Module mà bạn muốn thu thập đầu ra) và một index tùy chọn để chọn phần tử nào trong bộ đầu ra của module. Truyền layer_name để gắn hook chỉ vào các module có tên thuộc tính cụ thể. Sử dụng layer_name khi hai lớp chia sẻ cùng một class, ví dụ như self-attention so với cross-attention.

Ví dụ dưới đây minh họa cả hai decorator trong thực tế với các mức độ kiểm soát đầu ra khác nhau. Xem LlamaModel để tham khảo trong thực tế.

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
        )

Các lớp lớp vá

Việc theo dõi đầu ra phụ thuộc vào _can_record_outputs trỏ đến các lớp chính xác mà các lớp của mô hình khởi tạo. Nếu bạn thay thế một triển khai lớp bằng một nhân chú ý tùy chỉnh, một lớp chuyên gia lượng tử hóa, hoặc một biến thể kiến trúc, các con trỏ đó phải được đồng bộ. API vá cung cấp một bộ đăng ký toàn cục sạch sẽ để giữ cho _can_record_outputs nhất quán.

register_patch_mapping ánh xạ tên lớp gốc sang các lớp con thay thế của nn.Module. Các khóa có thể là tên lớp chính xác hoặc các mẫu regex. Các khớp chính xác có độ ưu tiên cao hơn. Các mẫu được kiểm tra bằng re.search(), do đó các mẫu không neo sẽ khớp ở bất kỳ vị trí nào trong tên lớp. Việc đăng ký cùng một khóa hai lần sẽ gây ra ValueError trừ khi bạn truyền overwrite=True.

Xóa các mục bằng 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$"])

Sau khi các ánh xạ được đăng ký, patch_output_recorders sẽ duyệt qua mọi submodule và cập nhật mỗi OutputRecorder.target_class thành phần thay thế đã được đăng ký.

Phương thức from_pretrained() sẽ tự động gọi patch_output_recorders. Bạn chỉ cần tự gọi nó khi xây dựng mô hình trực tiếp.

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)