| 108 | |
| 109 | |
| 110 | class AttentionProcessorSkipHook(ModelHook): |
| 111 | def __init__(self, skip_processor_output_fn: Callable, skip_attention_scores: bool = False, dropout: float = 1.0): |
| 112 | self.skip_processor_output_fn = skip_processor_output_fn |
| 113 | self.skip_attention_scores = skip_attention_scores |
| 114 | self.dropout = dropout |
| 115 | |
| 116 | def new_forward(self, module: torch.nn.Module, *args, **kwargs): |
| 117 | if self.skip_attention_scores: |
| 118 | if not math.isclose(self.dropout, 1.0): |
| 119 | raise ValueError( |
| 120 | "Cannot set `skip_attention_scores` to True when `dropout` is not 1.0. Please set `dropout` to 1.0." |
| 121 | ) |
| 122 | with AttentionScoreSkipFunctionMode(): |
| 123 | output = self.fn_ref.original_forward(*args, **kwargs) |
| 124 | else: |
| 125 | if math.isclose(self.dropout, 1.0): |
| 126 | output = self.skip_processor_output_fn(module, *args, **kwargs) |
| 127 | else: |
| 128 | output = self.fn_ref.original_forward(*args, **kwargs) |
| 129 | output = torch.nn.functional.dropout(output, p=self.dropout) |
| 130 | return output |
| 131 | |
| 132 | |
| 133 | class FeedForwardSkipHook(ModelHook): |
no outgoing calls
no test coverage detected
searching dependent graphs…