| 149 | |
| 150 | |
| 151 | class TransformerBlockSkipHook(ModelHook): |
| 152 | def __init__(self, dropout: float): |
| 153 | super().__init__() |
| 154 | self.dropout = dropout |
| 155 | |
| 156 | def initialize_hook(self, module): |
| 157 | self._metadata = TransformerBlockRegistry.get(unwrap_module(module).__class__) |
| 158 | return module |
| 159 | |
| 160 | def new_forward(self, module: torch.nn.Module, *args, **kwargs): |
| 161 | if math.isclose(self.dropout, 1.0): |
| 162 | original_hidden_states = self._metadata._get_parameter_from_args_kwargs("hidden_states", args, kwargs) |
| 163 | if self._metadata.return_encoder_hidden_states_index is None: |
| 164 | output = original_hidden_states |
| 165 | else: |
| 166 | original_encoder_hidden_states = self._metadata._get_parameter_from_args_kwargs( |
| 167 | "encoder_hidden_states", args, kwargs |
| 168 | ) |
| 169 | output = (original_hidden_states, original_encoder_hidden_states) |
| 170 | else: |
| 171 | output = self.fn_ref.original_forward(*args, **kwargs) |
| 172 | output = torch.nn.functional.dropout(output, p=self.dropout) |
| 173 | return output |
| 174 | |
| 175 | |
| 176 | def apply_layer_skip(module: torch.nn.Module, config: LayerSkipConfig) -> None: |
no outgoing calls
no test coverage detected
searching dependent graphs…