(self, input_ids, attention_mask: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None)
| 93 | ) |
| 94 | |
| 95 | def forward(self, input_ids, attention_mask: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None): |
| 96 | if attention_mask is not None: |
| 97 | attention_mask = get_extended_attention_mask(attention_mask, input_ids.shape, self.cfg.attention.causal_attention) |
| 98 | hidden_states = self.input_projection(self.embedding(input_ids)) |
| 99 | |
| 100 | if self.seq_first: |
| 101 | hidden_states = hidden_states.transpose(0, 1).contiguous() |
| 102 | |
| 103 | # Main transformer blocks: |
| 104 | if self.gradient_checkpointing and self.training: |
| 105 | # Hide this away from any jit-ing... |
| 106 | hidden_states = self.forward_checkpointed(hidden_states, attention_mask) |
| 107 | else: |
| 108 | if self.layer_drop_theta is None: |
| 109 | for i, layer_module in enumerate(self.layers): |
| 110 | hidden_states = layer_module(hidden_states, attention_mask, self.p) |
| 111 | else: |
| 112 | p = self.p.clone() |
| 113 | step = (1 - self.layer_drop_theta) / len(self.layers) |
| 114 | for i, layer_module in enumerate(self.layers): |
| 115 | p = p - step |
| 116 | if torch.bernoulli(p): |
| 117 | hidden_states = layer_module(hidden_states, attention_mask, res_scale=1 / p) |
| 118 | if self.seq_first: |
| 119 | hidden_states = hidden_states.transpose(0, 1).contiguous() |
| 120 | |
| 121 | return self.final_norm(hidden_states) |
| 122 | |
| 123 | @torch.jit.ignore |
| 124 | def forward_checkpointed(self, hidden_states, attention_mask: Optional[torch.Tensor] = None): |
nothing calls this directly
no test coverage detected