Skip to content

Module Wrapper

Introduction

Module Wrapper is the submodule that is responsible for managing the module wrappers. The module wrappers are essential to add custom hook where in the original transfomer codebase the hook is not available. For example, the transformer module does not have a hook to get the attention matrix of a head. The module wrapper is used to add this hook. The module_wrapper submodel is composed of the following files: - manager.py: The manager file is responsible for managing the module wrappers. It is the standard interface to add the wrap around models. - base.py: The base file is the base class for the module wrapper. Implement a base form of a Wrapper class. - model_name_attention.py: The model name attention file is the module wrapper for the attention matrix of a single model. When add a new model, add a new file with the name model_name_attention.py and implement the ModelNameAttention class. It is basically a copy of the forward pass of the attention module with the addition of the hook to get the attention matrix.

Manager Wrappers and Abstract Base Class

AttentionWrapperFactory

Maps a given model name to the correct attention wrapper class.

Source code in easyroutine/interpretability/module_wrappers/manager.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
class AttentionWrapperFactory:
    """
    Maps a given model name to the correct attention wrapper class.
    """

    AVAILABLE_MODULE_WRAPPERS: dict = {
        ChameleonAttentionWrapper.original_name(): ChameleonAttentionWrapper,
        LlamaAttentionWrapper.original_name(): LlamaAttentionWrapper,
        T5AttentionWrapper.original_name(): T5AttentionWrapper,
        MistralAttentionWrapper.original_name(): MistralAttentionWrapper,
        Gemma3AttentionWrapper.original_name(): Gemma3AttentionWrapper,
    }

    # MODEL_NAME_TO_WRAPPER = {
    #     "facebook/chameleon-7b": ChameleonAttentionWrapper,
    #     "facebook/chameleon-30b": ChameleonAttentionWrapper,
    #     "mistral-community/pixtral-12b": LlamaAttentionWrapper,
    #     "llava-hf/llava-v1.6-mistral-7b-hf": LlamaAttentionWrapper,
    #     "hf-internal-testing/tiny-random-LlamaForCausalLM": LlamaAttentionWrapper,
    #     "ChoereForAI/aya-101": T5AttentionWrapper,
    # }

    @staticmethod
    def get_wrapper_classes(
        model: nn.Module,
    ) -> Dict[
        str,
        Union[
            Type[ChameleonAttentionWrapper],
            Type[LlamaAttentionWrapper],
            Type[T5AttentionWrapper],
            Type[MistralAttentionWrapper],
            Type[Gemma3AttentionWrapper],
        ],
    ]:
        """
        Returns a dictionary mapping module names to their corresponding wrapper classes
        for all supported modules found in the model.

        Args:
            model (nn.Module): The model to analyze

        Returns:
            Dict[str, Type]: Dictionary mapping original module names to wrapper classes
        """
        all_modules = find_all_modules(model, return_only_names=True)
        found_wrappers = {}

        for (
            candidate_name,
            candidate_wrapper,
        ) in AttentionWrapperFactory.AVAILABLE_MODULE_WRAPPERS.items():
            if candidate_name in all_modules:
                logger.info(f"Found a wrapper for {candidate_name}")
                found_wrappers[candidate_name] = candidate_wrapper

        if not found_wrappers:
            logger.warning(f"No wrappers found for any module in {model}")

        return found_wrappers

get_wrapper_classes(model) staticmethod

Returns a dictionary mapping module names to their corresponding wrapper classes for all supported modules found in the model.

Parameters:

Name Type Description Default
model Module

The model to analyze

 required

Returns:

Type Description
Dict[str, Union[Type[ChameleonAttentionWrapper], Type[LlamaAttentionWrapper], Type[T5AttentionWrapper], Type[MistralAttentionWrapper], Type[Gemma3AttentionWrapper]]]

Dict[str, Type]: Dictionary mapping original module names to wrapper classes
Source code in easyroutine/interpretability/module_wrappers/manager.py 
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
@staticmethod
def get_wrapper_classes(
    model: nn.Module,
) -> Dict[
    str,
    Union[
        Type[ChameleonAttentionWrapper],
        Type[LlamaAttentionWrapper],
        Type[T5AttentionWrapper],
        Type[MistralAttentionWrapper],
        Type[Gemma3AttentionWrapper],
    ],
]:
    """
    Returns a dictionary mapping module names to their corresponding wrapper classes
    for all supported modules found in the model.

    Args:
        model (nn.Module): The model to analyze

    Returns:
        Dict[str, Type]: Dictionary mapping original module names to wrapper classes
    """
    all_modules = find_all_modules(model, return_only_names=True)
    found_wrappers = {}

    for (
        candidate_name,
        candidate_wrapper,
    ) in AttentionWrapperFactory.AVAILABLE_MODULE_WRAPPERS.items():
        if candidate_name in all_modules:
            logger.info(f"Found a wrapper for {candidate_name}")
            found_wrappers[candidate_name] = candidate_wrapper

    if not found_wrappers:
        logger.warning(f"No wrappers found for any module in {model}")

    return found_wrappers

ModuleWrapperManager

Handles the logic of replacing original modules within a given model with custom wrappers. Supports multiple module types per model. Also allows restoring the original modules if needed.

Source code in easyroutine/interpretability/module_wrappers/manager.py 
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
class ModuleWrapperManager:
    """
    Handles the logic of replacing original modules within a given model
    with custom wrappers. Supports multiple module types per model.
    Also allows restoring the original modules if needed.
    """

    def __init__(self, model: nn.Module, log_level: str = "INFO"):
        """
        Initializes the manager with a given model.

        Args:
            model (nn.Module): The model whose modules will be wrapped
            log_level (str): Logging level
        """
        # Get all available wrappers for modules in this model
        self.module_wrappers = AttentionWrapperFactory.get_wrapper_classes(model)

        # Dictionary to store {module_path: {module_type: original_module}}
        self.original_modules = {}

        # Store target module names for quick lookup
        self.target_module_names = list(self.module_wrappers.keys())

    def __contains__(self, module_name: str) -> bool:
        """
        Check if a module name is supported by this manager.

        Args:
            module_name (str): The name of the module to check

        Returns:
            bool: True if the module is supported, False otherwise
        """
        return module_name in self.target_module_names

    def substitute_attention_module(self, model: nn.Module) -> None:
        """
        Public method that performs the substitution of all supported modules in the model.
        Logs each replacement.

        Args:
            model (nn.Module): The model whose modules will be wrapped
        """
        self._traverse_and_modify(model, parent_path="", mode="substitute")

    def restore_original_attention_module(self, model: nn.Module) -> None:
        """
        Public method that restores the original modules in the model.
        Logs each restoration.

        Args:
            model (nn.Module): The model whose modules will be restored
        """
        self._traverse_and_modify(model, parent_path="", mode="restore")

    def _traverse_and_modify(
        self, module: nn.Module, parent_path: str, mode: str
    ) -> None:
        """
        Recursively traverses `module` and either substitutes or restores each matching
        submodule, depending on `mode`.

        - mode="substitute": Replaces original modules with wrappers
        - mode="restore": Replaces wrapper modules with their originals

        Args:
            module (nn.Module): The current module to inspect
            parent_path (str): A string that tracks the 'path' of this submodule in the overall model hierarchy
            mode (str): Either "substitute" or "restore"
        """
        for name, child in list(module.named_children()):
            # Identify the submodule path (e.g. "encoder.layer.0.attention")
            submodule_path = f"{parent_path}.{name}" if parent_path else name

            if mode == "substitute":
                # Look for any matching original module class names
                child_class_name = child.__class__.__name__
                if child_class_name in self.target_module_names:
                    # Get the appropriate wrapper for this module type
                    wrapper_class = self.module_wrappers[child_class_name]

                    # Store the original module
                    if submodule_path not in self.original_modules:
                        self.original_modules[submodule_path] = {}
                    self.original_modules[submodule_path][child_class_name] = child

                    # Wrap it
                    wrapped_module = wrapper_class(child)
                    setattr(module, name, wrapped_module)

                    logger.debug(
                        f"Substituted '{submodule_path}' with wrapper for {child_class_name}."
                    )
                else:
                    # Recurse
                    self._traverse_and_modify(child, submodule_path, mode="substitute")

            elif mode == "restore":
                # Check if this is any kind of wrapper we know about
                child_class_name = child.__class__.__name__
                for orig_name, wrapper_class in self.module_wrappers.items():
                    if child_class_name == wrapper_class.__name__:
                        # Found a wrapper, check if we have the original
                        if (
                            submodule_path in self.original_modules
                            and orig_name in self.original_modules[submodule_path]
                        ):
                            # Restore the original module
                            original_module = self.original_modules[submodule_path][
                                orig_name
                            ]
                            setattr(module, name, original_module)
                            logger.info(
                                f"Restored '{submodule_path}' to original {orig_name}."
                            )
                            break
                else:
                    # If not a wrapper or no original found, recurse
                    self._traverse_and_modify(child, submodule_path, mode="restore")

__contains__(module_name)

Check if a module name is supported by this manager.

Parameters:

Name Type Description Default
module_name str

The name of the module to check

 required

Returns:

Name Type Description
bool bool

True if the module is supported, False otherwise
Source code in easyroutine/interpretability/module_wrappers/manager.py 
109
110
111
112
113
114
115
116
117
118
119
def __contains__(self, module_name: str) -> bool:
    """
    Check if a module name is supported by this manager.

    Args:
        module_name (str): The name of the module to check

    Returns:
        bool: True if the module is supported, False otherwise
    """
    return module_name in self.target_module_names

__init__(model, log_level='INFO')

Initializes the manager with a given model.

Parameters:

Name Type Description Default
model Module

The model whose modules will be wrapped

 required
log_level str

Logging level

 'INFO'
Source code in easyroutine/interpretability/module_wrappers/manager.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
def __init__(self, model: nn.Module, log_level: str = "INFO"):
    """
    Initializes the manager with a given model.

    Args:
        model (nn.Module): The model whose modules will be wrapped
        log_level (str): Logging level
    """
    # Get all available wrappers for modules in this model
    self.module_wrappers = AttentionWrapperFactory.get_wrapper_classes(model)

    # Dictionary to store {module_path: {module_type: original_module}}
    self.original_modules = {}

    # Store target module names for quick lookup
    self.target_module_names = list(self.module_wrappers.keys())

restore_original_attention_module(model)

Public method that restores the original modules in the model. Logs each restoration.

Parameters:

Name Type Description Default
model Module

The model whose modules will be restored

 required
Source code in easyroutine/interpretability/module_wrappers/manager.py 
131
132
133
134
135
136
137
138
139
def restore_original_attention_module(self, model: nn.Module) -> None:
    """
    Public method that restores the original modules in the model.
    Logs each restoration.

    Args:
        model (nn.Module): The model whose modules will be restored
    """
    self._traverse_and_modify(model, parent_path="", mode="restore")

substitute_attention_module(model)

Public method that performs the substitution of all supported modules in the model. Logs each replacement.

Parameters:

Name Type Description Default
model Module

The model whose modules will be wrapped

 required
Source code in easyroutine/interpretability/module_wrappers/manager.py 
121
122
123
124
125
126
127
128
129
def substitute_attention_module(self, model: nn.Module) -> None:
    """
    Public method that performs the substitution of all supported modules in the model.
    Logs each replacement.

    Args:
        model (nn.Module): The model whose modules will be wrapped
    """
    self._traverse_and_modify(model, parent_path="", mode="substitute")

AttentionMatrixHookModule

Bases: Module

Computation of the attention matrix. Note: it has been added just for adding custom hooks.

Source code in easyroutine/interpretability/module_wrappers/base.py 
 4
 5
 6
 7
 8
 9
10
11
class AttentionMatrixHookModule(nn.Module):
    """Computation of the attention matrix. *Note*: it has been added just for adding custom hooks."""

    def forward(
            self,
            attention_matrix: torch.Tensor,
    ):
        return attention_matrix

BaseAttentionWrapper

Bases: Module

A base class for wrapping an original attention module.

Provides

_orig_module to store the real (unwrapped) attention. A robust __getattr__ that checks: 1) self.dict 2) self._modules 3) the base class 4) fallback to _orig_module

Source code in easyroutine/interpretability/module_wrappers/base.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
class BaseAttentionWrapper(nn.Module):
    """
    A base class for wrapping an original attention module.

    Provides:
        `_orig_module` to store the real (unwrapped) attention.
        A robust `__getattr__` that checks:
            1) self.__dict__
            2) self._modules
            3) the base class
            4) fallback to `_orig_module`
    """

    def __init__(self, original_module: nn.Module):
        super().__init__()
        # store the original module in a private attribute
        object.__setattr__(self, "_orig_module", original_module)

    def __getattr__(self, name: str):
        """
        If name is not in this wrapper, fall back to the original module.
        Also checks `self._modules` for submodules, because PyTorch
        automatically places them there.
        """
        # 1) get this wrapper's __dict__
        wrapper_dict = object.__getattribute__(self, "__dict__")

        # 2) if name is in our own instance dictionary, return it
        if name in wrapper_dict:
            return wrapper_dict[name]

        # 3) if name is in our submodules, return it
        modules_dict = wrapper_dict["_modules"]
        if name in modules_dict:
            return modules_dict[name]

        # 4) check if name is in our class (methods, etc.)
        cls = object.__getattribute__(self, "__class__")
        if hasattr(cls, name):
            return getattr(cls, name)

        # 5) fallback to _orig_module
        orig = wrapper_dict["_orig_module"]
        return getattr(orig, name)

    @staticmethod
    def original_name() -> str:
        """
        By default, you might override this in each derived class if you want
        your manager code to know which original class name this wrapper replaces.
        """
        return "BaseAttention"

__getattr__(name)

If name is not in this wrapper, fall back to the original module. Also checks self._modules for submodules, because PyTorch automatically places them there.

Source code in easyroutine/interpretability/module_wrappers/base.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
def __getattr__(self, name: str):
    """
    If name is not in this wrapper, fall back to the original module.
    Also checks `self._modules` for submodules, because PyTorch
    automatically places them there.
    """
    # 1) get this wrapper's __dict__
    wrapper_dict = object.__getattribute__(self, "__dict__")

    # 2) if name is in our own instance dictionary, return it
    if name in wrapper_dict:
        return wrapper_dict[name]

    # 3) if name is in our submodules, return it
    modules_dict = wrapper_dict["_modules"]
    if name in modules_dict:
        return modules_dict[name]

    # 4) check if name is in our class (methods, etc.)
    cls = object.__getattribute__(self, "__class__")
    if hasattr(cls, name):
        return getattr(cls, name)

    # 5) fallback to _orig_module
    orig = wrapper_dict["_orig_module"]
    return getattr(orig, name)

original_name() staticmethod

By default, you might override this in each derived class if you want your manager code to know which original class name this wrapper replaces.

Source code in easyroutine/interpretability/module_wrappers/base.py 
58
59
60
61
62
63
64
@staticmethod
def original_name() -> str:
    """
    By default, you might override this in each derived class if you want
    your manager code to know which original class name this wrapper replaces.
    """
    return "BaseAttention"

Specific Module Wrappers

LlamaAttentionWrapper

Bases: BaseAttentionWrapper

A wrapper around the original LlamaAttention. It has: - The same named attributes (q_proj, k_proj, etc.), which are references to the original module's submodules/parameters. - A private reference (_orig_attn) to the entire original attention, for falling back if something isn't found on the wrapper itself. - An additional attention_matrix_hook for intercepting attention.

Source code in easyroutine/interpretability/module_wrappers/llama_attention.py 
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
class LlamaAttentionWrapper(BaseAttentionWrapper):
    """
    A wrapper around the original LlamaAttention. It has:
    - The same named attributes (q_proj, k_proj, etc.), which are references
        to the original module's submodules/parameters.
    - A private reference (`_orig_attn`) to the entire original attention,
        for falling back if something isn't found on the wrapper itself.
    - An additional `attention_matrix_hook` for intercepting attention.
    """

    @staticmethod
    def original_name():
        return "LlamaAttention"

    def __init__(self, original_attention: nn.Module):
        """
        Store references to all relevant submodules so the wrapper
        "feels" the same. Also store a reference to the original module
        in a private attribute for fallback.
        """
        super().__init__(original_attention)

        # This is the private reference to the entire original attention.
        # We'll fallback to it for any attribute we haven't explicitly set.
        object.__setattr__(self, "_orig_attn", original_attention)

        # Now replicate the original attention's submodules as attributes of *this* wrapper.
        # These are direct references, not new modules:
        self.q_proj = original_attention.q_proj
        self.k_proj = original_attention.k_proj
        self.v_proj = original_attention.v_proj
        self.o_proj = original_attention.o_proj

        # Copy over any scalar attributes you need
        # self.num_heads = original_attention.num_heads
        # self.num_key_value_heads = original_attention.num_key_value_heads
        # self.num_key_value_groups = original_attention.num_key_value_groups
        self.head_dim = original_attention.head_dim
        # self.hidden_size = original_attention.hidden_size
        self.attention_dropout = original_attention.attention_dropout
        self.layer_idx = original_attention.layer_idx
        self.config = original_attention.config

        # Add your custom hook module
        self.attention_matrix_pre_softmax_hook = AttentionMatrixHookModule()
        self.attention_matrix_hook = AttentionMatrixHookModule()

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
        attention_mask: Optional[torch.Tensor],
        past_key_value: Optional[Cache] = None,
        cache_position: Optional[torch.LongTensor] = None,
        **kwargs: Unpack[FlashAttentionKwargs],
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
        input_shape = hidden_states.shape[:-1]
        hidden_shape = (*input_shape, -1, self.head_dim)

        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)

        cos, sin = position_embeddings
        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

        if past_key_value is not None:
            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
            key_states, value_states = past_key_value.update(
                key_states, value_states, self.layer_idx, cache_kwargs
            )

        # Inline eager_attention_forward logic
        key_states = repeat_kv(key_states, self.num_key_value_groups)
        value_states = repeat_kv(value_states, self.num_key_value_groups)
        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * self.scaling

        if attention_mask is not None:
            causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
            attn_weights = attn_weights + causal_mask

        # Pre-softmax hook
        attn_weights = self.attention_matrix_pre_softmax_hook(attn_weights)
        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
        attn_weights = self.attention_matrix_hook(attn_weights)
        attn_weights = nn.functional.dropout(
            attn_weights,
            p=0.0 if not self.training else self.attention_dropout,
            training=self.training,
        )

        attn_output = torch.matmul(attn_weights, value_states).transpose(1, 2).contiguous()
        # End inline

        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
        attn_output = self.o_proj(attn_output)
        return attn_output, attn_weights # type: ignore

__init__(original_attention)

Store references to all relevant submodules so the wrapper "feels" the same. Also store a reference to the original module in a private attribute for fallback.

Source code in easyroutine/interpretability/module_wrappers/llama_attention.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
def __init__(self, original_attention: nn.Module):
    """
    Store references to all relevant submodules so the wrapper
    "feels" the same. Also store a reference to the original module
    in a private attribute for fallback.
    """
    super().__init__(original_attention)

    # This is the private reference to the entire original attention.
    # We'll fallback to it for any attribute we haven't explicitly set.
    object.__setattr__(self, "_orig_attn", original_attention)

    # Now replicate the original attention's submodules as attributes of *this* wrapper.
    # These are direct references, not new modules:
    self.q_proj = original_attention.q_proj
    self.k_proj = original_attention.k_proj
    self.v_proj = original_attention.v_proj
    self.o_proj = original_attention.o_proj

    # Copy over any scalar attributes you need
    # self.num_heads = original_attention.num_heads
    # self.num_key_value_heads = original_attention.num_key_value_heads
    # self.num_key_value_groups = original_attention.num_key_value_groups
    self.head_dim = original_attention.head_dim
    # self.hidden_size = original_attention.hidden_size
    self.attention_dropout = original_attention.attention_dropout
    self.layer_idx = original_attention.layer_idx
    self.config = original_attention.config

    # Add your custom hook module
    self.attention_matrix_pre_softmax_hook = AttentionMatrixHookModule()
    self.attention_matrix_hook = AttentionMatrixHookModule()

repeat_kv(hidden_states, n_rep)

(batch, num_key_value_heads, seq_len, head_dim) -> (batch, num_attention_heads, seq_len, head_dim)

Source code in easyroutine/interpretability/module_wrappers/llama_attention.py 
37
38
39
40
41
42
43
44
45
46
47
48
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
    """
    (batch, num_key_value_heads, seq_len, head_dim)
        -> (batch, num_attention_heads, seq_len, head_dim)
    """
    bsz, num_kv_heads, slen, head_dim = hidden_states.shape
    if n_rep == 1:
        return hidden_states
    hidden_states = hidden_states[:, :, None, :, :].expand(
        bsz, num_kv_heads, n_rep, slen, head_dim
    )
    return hidden_states.reshape(bsz, num_kv_heads * n_rep, slen, head_dim)

ChameleonAttentionWrapper

Bases: BaseAttentionWrapper

Attention wrapper for the Chameleon model.

Source code in easyroutine/interpretability/module_wrappers/chameleon_attention.py 
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
class ChameleonAttentionWrapper(BaseAttentionWrapper):
    """
    Attention wrapper for the Chameleon model.
    """

    @staticmethod
    def original_name():
        return "ChameleonAttention"

    def __init__(self, original_attention: nn.Module):
        super().__init__(original_attention)

        self.q_proj = original_attention.q_proj
        self.k_proj = original_attention.k_proj
        self.v_proj = original_attention.v_proj
        self.q_norm = original_attention.q_norm
        self.k_norm = original_attention.k_norm
        self.o_proj = original_attention.o_proj
        # self.softmax = original_attention.softmax
        self.attention_dropout = original_attention.attention_dropout
        self.training = original_attention.training
        self.layer_idx = original_attention.layer_idx
        self.num_heads = original_attention.num_heads
        self.num_key_value_heads = original_attention.num_key_value_heads
        self.num_key_value_groups = original_attention.num_key_value_groups
        self.head_dim = original_attention.head_dim
        self.hidden_size = original_attention.hidden_size
        self.rotary_emb = original_attention.rotary_emb

        self.attention_matrix_pre_softmax_hook = AttentionMatrixHookModule()        
        self.attention_matrix_hook = AttentionMatrixHookModule()

        self.original_attention = original_attention


    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_value: Optional[Cache] = None,
        output_attentions: bool = False,
        use_cache: bool = False,
        cache_position: Optional[torch.LongTensor] = None,
        **kwargs,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
        bsz, q_len, _ = hidden_states.size()

        query_states = self.q_proj(hidden_states)
        key_states = self.k_proj(hidden_states)
        value_states = self.v_proj(hidden_states)

        query_states = query_states.reshape(-1, self.num_heads, self.head_dim)
        query_states = self.q_norm(query_states)

        key_states = key_states.reshape(-1, self.num_key_value_heads, self.head_dim)
        key_states = self.k_norm(key_states)

        query_states = query_states.reshape(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
        key_states = key_states.reshape(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)

        cos, sin = self.rotary_emb(value_states, position_ids)
        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

        if past_key_value is not None:
            # sin and cos are specific to RoPE models; position_ids needed for the static cache
            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)

        key_states = repeat_kv(key_states, self.num_key_value_groups)
        value_states = repeat_kv(value_states, self.num_key_value_groups)

        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)

        if attention_mask is not None:  # no matter the length, we just slice it
            causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
            attn_weights = attn_weights + causal_mask

        # upcast attention to fp32
        # attn_weights = nn.functional.softmax(attn_weights, dim=-1).to(query_states.dtype)
        attn_weights = self.attention_matrix_pre_softmax_hook(attn_weights)
        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
        attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
        attn_weights = self.attention_matrix_hook(attn_weights)
        attn_output = torch.matmul(attn_weights, value_states)

        if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
            raise ValueError(
                f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
                f" {attn_output.size()}"
            )

        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
        attn_output = self.o_proj(attn_output)

        if not output_attentions:
            attn_weights = None

        return attn_output, attn_weights, past_key_value # type: ignore

apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1)

Applies Rotary Position Embedding to the query and key tensors.

Parameters:

Name Type Description Default
q `torch.Tensor`

The query tensor.

 required
k `torch.Tensor`

The key tensor.

 required
cos `torch.Tensor`

The cosine part of the rotary embedding.

 required
sin `torch.Tensor`

The sine part of the rotary embedding.

 required
position_ids `torch.Tensor`, *optional*

Deprecated and unused.

 None
unsqueeze_dim `int`, *optional*, defaults to 1

The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.

 1

Returns: tuple(torch.Tensor) comprising of the query and key tensors rotated using the Rotary Position Embedding.

Source code in easyroutine/interpretability/module_wrappers/chameleon_attention.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
def apply_rotary_pos_emb( q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
    """Applies Rotary Position Embedding to the query and key tensors.

    Args:
        q (`torch.Tensor`): The query tensor.
        k (`torch.Tensor`): The key tensor.
        cos (`torch.Tensor`): The cosine part of the rotary embedding.
        sin (`torch.Tensor`): The sine part of the rotary embedding.
        position_ids (`torch.Tensor`, *optional*):
            Deprecated and unused.
        unsqueeze_dim (`int`, *optional*, defaults to 1):
            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
    Returns:
        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
    """

    cos = cos.unsqueeze(unsqueeze_dim)
    sin = sin.unsqueeze(unsqueeze_dim)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed

repeat_kv(hidden_states, n_rep)

This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)

Source code in easyroutine/interpretability/module_wrappers/chameleon_attention.py 
43
44
45
46
47
48
49
50
51
52
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
    """
    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
    """
    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
    if n_rep == 1:
        return hidden_states
    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)

rotate_half(x)

Rotates half the hidden dims of the input.

Source code in easyroutine/interpretability/module_wrappers/chameleon_attention.py 
 9
10
11
12
13
def rotate_half(x):
    """Rotates half the hidden dims of the input."""
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)

T5AttentionWrapper

Bases: BaseAttentionWrapper

Source code in easyroutine/interpretability/module_wrappers/T5_attention.py 
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
class T5AttentionWrapper(BaseAttentionWrapper):
    @staticmethod
    def original_name() -> str:
        return "T5Attention"

    def __init__(self, original_attention: nn.Module):
        super().__init__(original_attention)
        self.q = original_attention.q
        self.k = original_attention.k
        self.v = original_attention.v
        self.o = original_attention.o
        self.dropout = original_attention.dropout
        self.layer_idx = original_attention.layer_idx
        self.n_heads = original_attention.n_heads
        self.key_value_proj_dim = original_attention.key_value_proj_dim
        self.inner_dim = original_attention.inner_dim
        self.has_relative_attention_bias = (
            original_attention.has_relative_attention_bias
        )
        self.pruned_heads = original_attention.pruned_heads
        self.attention_matrix_pre_softmax_hook = AttentionMatrixHookModule()
        self.attention_matrix_hook = AttentionMatrixHookModule()
        self.original_attention = original_attention

    def forward(
        self,
        hidden_states,
        mask=None,
        key_value_states=None,
        position_bias=None,
        past_key_value=None,
        layer_head_mask=None,
        query_length=None,
        use_cache=False,
        output_attentions=False,
        cache_position=None,
    ):
        """
        Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states).
        """
        # Input is (batch_size, seq_length, dim)
        # Mask is (batch_size, 1, 1, key_length) (non-causal encoder) or (batch_size, 1, seq_length, key_length) (causal decoder)
        batch_size, seq_length = hidden_states.shape[:2]

        # if key_value_states are provided this layer is used as a cross-attention layer for the decoder
        is_cross_attention = key_value_states is not None

        query_states = self.q(hidden_states)
        query_states = query_states.view(
            batch_size, -1, self.n_heads, self.key_value_proj_dim
        ).transpose(1, 2)

        if past_key_value is not None:
            is_updated = past_key_value.is_updated.get(self.layer_idx)
            if is_cross_attention:
                # after the first generated id, we can subsequently re-use all key/value_states from cache
                curr_past_key_value = past_key_value.cross_attention_cache
            else:
                curr_past_key_value = past_key_value.self_attention_cache

        current_states = key_value_states if is_cross_attention else hidden_states
        if is_cross_attention and past_key_value is not None and is_updated:  # type: ignore
            # reuse k,v, cross_attentions
            key_states = curr_past_key_value.key_cache[self.layer_idx]  # type: ignore
            value_states = curr_past_key_value.value_cache[self.layer_idx]  # type: ignore
        else:
            key_states = self.k(current_states)
            value_states = self.v(current_states)
            key_states = key_states.view(
                batch_size, -1, self.n_heads, self.key_value_proj_dim
            ).transpose(1, 2)
            value_states = value_states.view(
                batch_size, -1, self.n_heads, self.key_value_proj_dim
            ).transpose(1, 2)

            if past_key_value is not None:
                # save all key/value_states to cache to be re-used for fast auto-regressive generation
                cache_position = cache_position if not is_cross_attention else None
                key_states, value_states = curr_past_key_value.update(  # type: ignore
                    key_states,
                    value_states,
                    self.layer_idx,
                    {"cache_position": cache_position},
                )
                # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls
                if is_cross_attention:
                    past_key_value.is_updated[self.layer_idx] = True

        # compute scores, equivalent of torch.einsum("bnqd,bnkd->bnqk", query_states, key_states), compatible with onnx op>9
        scores = torch.matmul(query_states, key_states.transpose(3, 2))

        if position_bias is None:
            key_length = key_states.shape[-2]
            # cache position is 0-indexed so we add 1 to get the real length of queries (aka with past)
            real_seq_length = (
                query_length if query_length is not None else cache_position[-1] + 1 # type: ignore
            )  # type: ignore
            if not self.has_relative_attention_bias:
                position_bias = torch.zeros(
                    (1, self.n_heads, seq_length, key_length),
                    device=scores.device,
                    dtype=scores.dtype,
                )
                if self.gradient_checkpointing and self.training:
                    position_bias.requires_grad = True
            else:
                position_bias = self.compute_bias(
                    real_seq_length,
                    key_length,
                    device=scores.device,
                    cache_position=cache_position,
                )
                position_bias = position_bias[:, :, -seq_length:, :]

            if mask is not None:
                causal_mask = mask[:, :, :, : key_states.shape[-2]]
                position_bias = position_bias + causal_mask

        if self.pruned_heads:
            mask = torch.ones(position_bias.shape[1])
            mask[list(self.pruned_heads)] = 0
            position_bias_masked = position_bias[:, mask.bool()]
        else:
            position_bias_masked = position_bias

        scores += position_bias_masked

        # (batch_size, n_heads, seq_length, key_length)
        attn_weights = self.attention_matrix_pre_softmax_hook(scores)
        attn_weights = nn.functional.softmax(scores.float(), dim=-1).type_as(scores)
        attn_weights = self.attention_matrix_hook(attn_weights)
        attn_weights = nn.functional.dropout(
            attn_weights, p=self.dropout, training=self.training
        )

        # Mask heads if we want to
        if layer_head_mask is not None:
            attn_weights = attn_weights * layer_head_mask

        attn_output = torch.matmul(attn_weights, value_states)

        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, -1, self.inner_dim)
        attn_output = self.o(attn_output)

        outputs = (attn_output, past_key_value, position_bias)

        if output_attentions:
            outputs = outputs + (attn_weights,)
        return outputs

forward(hidden_states, mask=None, key_value_states=None, position_bias=None, past_key_value=None, layer_head_mask=None, query_length=None, use_cache=False, output_attentions=False, cache_position=None)

Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states).

Source code in easyroutine/interpretability/module_wrappers/T5_attention.py 
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
def forward(
    self,
    hidden_states,
    mask=None,
    key_value_states=None,
    position_bias=None,
    past_key_value=None,
    layer_head_mask=None,
    query_length=None,
    use_cache=False,
    output_attentions=False,
    cache_position=None,
):
    """
    Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states).
    """
    # Input is (batch_size, seq_length, dim)
    # Mask is (batch_size, 1, 1, key_length) (non-causal encoder) or (batch_size, 1, seq_length, key_length) (causal decoder)
    batch_size, seq_length = hidden_states.shape[:2]

    # if key_value_states are provided this layer is used as a cross-attention layer for the decoder
    is_cross_attention = key_value_states is not None

    query_states = self.q(hidden_states)
    query_states = query_states.view(
        batch_size, -1, self.n_heads, self.key_value_proj_dim
    ).transpose(1, 2)

    if past_key_value is not None:
        is_updated = past_key_value.is_updated.get(self.layer_idx)
        if is_cross_attention:
            # after the first generated id, we can subsequently re-use all key/value_states from cache
            curr_past_key_value = past_key_value.cross_attention_cache
        else:
            curr_past_key_value = past_key_value.self_attention_cache

    current_states = key_value_states if is_cross_attention else hidden_states
    if is_cross_attention and past_key_value is not None and is_updated:  # type: ignore
        # reuse k,v, cross_attentions
        key_states = curr_past_key_value.key_cache[self.layer_idx]  # type: ignore
        value_states = curr_past_key_value.value_cache[self.layer_idx]  # type: ignore
    else:
        key_states = self.k(current_states)
        value_states = self.v(current_states)
        key_states = key_states.view(
            batch_size, -1, self.n_heads, self.key_value_proj_dim
        ).transpose(1, 2)
        value_states = value_states.view(
            batch_size, -1, self.n_heads, self.key_value_proj_dim
        ).transpose(1, 2)

        if past_key_value is not None:
            # save all key/value_states to cache to be re-used for fast auto-regressive generation
            cache_position = cache_position if not is_cross_attention else None
            key_states, value_states = curr_past_key_value.update(  # type: ignore
                key_states,
                value_states,
                self.layer_idx,
                {"cache_position": cache_position},
            )
            # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls
            if is_cross_attention:
                past_key_value.is_updated[self.layer_idx] = True

    # compute scores, equivalent of torch.einsum("bnqd,bnkd->bnqk", query_states, key_states), compatible with onnx op>9
    scores = torch.matmul(query_states, key_states.transpose(3, 2))

    if position_bias is None:
        key_length = key_states.shape[-2]
        # cache position is 0-indexed so we add 1 to get the real length of queries (aka with past)
        real_seq_length = (
            query_length if query_length is not None else cache_position[-1] + 1 # type: ignore
        )  # type: ignore
        if not self.has_relative_attention_bias:
            position_bias = torch.zeros(
                (1, self.n_heads, seq_length, key_length),
                device=scores.device,
                dtype=scores.dtype,
            )
            if self.gradient_checkpointing and self.training:
                position_bias.requires_grad = True
        else:
            position_bias = self.compute_bias(
                real_seq_length,
                key_length,
                device=scores.device,
                cache_position=cache_position,
            )
            position_bias = position_bias[:, :, -seq_length:, :]

        if mask is not None:
            causal_mask = mask[:, :, :, : key_states.shape[-2]]
            position_bias = position_bias + causal_mask

    if self.pruned_heads:
        mask = torch.ones(position_bias.shape[1])
        mask[list(self.pruned_heads)] = 0
        position_bias_masked = position_bias[:, mask.bool()]
    else:
        position_bias_masked = position_bias

    scores += position_bias_masked

    # (batch_size, n_heads, seq_length, key_length)
    attn_weights = self.attention_matrix_pre_softmax_hook(scores)
    attn_weights = nn.functional.softmax(scores.float(), dim=-1).type_as(scores)
    attn_weights = self.attention_matrix_hook(attn_weights)
    attn_weights = nn.functional.dropout(
        attn_weights, p=self.dropout, training=self.training
    )

    # Mask heads if we want to
    if layer_head_mask is not None:
        attn_weights = attn_weights * layer_head_mask

    attn_output = torch.matmul(attn_weights, value_states)

    attn_output = attn_output.transpose(1, 2).contiguous()
    attn_output = attn_output.view(batch_size, -1, self.inner_dim)
    attn_output = self.o(attn_output)

    outputs = (attn_output, past_key_value, position_bias)

    if output_attentions:
        outputs = outputs + (attn_weights,)
    return outputs