| """Attention layers used by the correctness-first scaffold.""" |
|
|
| from __future__ import annotations |
|
|
| import math |
|
|
| import torch |
| from torch import nn |
|
|
| from .padding import masked_hidden |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, hidden_size: int, eps: float = 1e-5): |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(hidden_size)) |
| self.eps = eps |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| variance = x.pow(2).mean(dim=-1, keepdim=True) |
| return x * torch.rsqrt(variance + self.eps) * self.weight |
|
|
|
|
| class FeedForward(nn.Module): |
| def __init__(self, hidden_size: int, intermediate_size: int, dropout: float, activation: str = "gelu"): |
| super().__init__() |
| self.activation_name = activation |
| if activation == "geglu": |
| self.in_proj = nn.Linear(hidden_size, intermediate_size * 2) |
| self.out_proj = nn.Linear(intermediate_size, hidden_size) |
| else: |
| self.in_proj = nn.Linear(hidden_size, intermediate_size) |
| self.out_proj = nn.Linear(intermediate_size, hidden_size) |
| self.dropout = nn.Dropout(dropout) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| x = self.in_proj(x) |
| if self.activation_name == "geglu": |
| value, gate = x.chunk(2, dim=-1) |
| x = value * torch.nn.functional.gelu(gate) |
| else: |
| x = torch.nn.functional.gelu(x) |
| return self.out_proj(self.dropout(x)) |
|
|
|
|
| class StrataBertAttentionLayer(nn.Module): |
| def __init__(self, config, layer_type: str): |
| super().__init__() |
| if config.hidden_size % config.num_attention_heads != 0: |
| raise ValueError("hidden_size must be divisible by num_attention_heads") |
| self.layer_type = layer_type |
| self.num_heads = config.num_attention_heads |
| self.head_dim = config.hidden_size // config.num_attention_heads |
| self.window = config.local_attention_window |
| self.norm = RMSNorm(config.hidden_size, config.norm_eps) |
| self.qkv = nn.Linear(config.hidden_size, config.hidden_size * 3) |
| self.out_proj = nn.Linear(config.hidden_size, config.hidden_size) |
| self.dropout = nn.Dropout(config.attention_dropout) |
| self.ffn_norm = RMSNorm(config.hidden_size, config.norm_eps) |
| self.ffn = FeedForward( |
| config.hidden_size, |
| config.intermediate_size, |
| config.hidden_dropout, |
| config.hidden_activation, |
| ) |
|
|
| def forward( |
| self, |
| hidden_states: torch.Tensor, |
| attention_mask: torch.Tensor, |
| segment_ids: torch.Tensor | None = None, |
| ) -> torch.Tensor: |
| residual = hidden_states |
| x = self.norm(hidden_states) |
| batch, length, hidden = x.shape |
| qkv = self.qkv(x).view(batch, length, 3, self.num_heads, self.head_dim) |
| q, k, v = qkv.unbind(dim=2) |
| q = q.transpose(1, 2) |
| k = k.transpose(1, 2) |
| v = v.transpose(1, 2) |
|
|
| scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(self.head_dim) |
| key_mask = attention_mask[:, None, None, :] |
| query_mask = attention_mask[:, None, :, None] |
| scores = scores.masked_fill(~key_mask, -1.0e4) |
| if segment_ids is not None: |
| same_segment = segment_ids[:, None, :, None] == segment_ids[:, None, None, :] |
| scores = scores.masked_fill(~same_segment, -1.0e4) |
| if self.layer_type == "local_attention": |
| idx = torch.arange(length, device=x.device) |
| local = (idx[None, :] - idx[:, None]).abs() <= self.window |
| scores = scores.masked_fill(~local[None, None, :, :], -1.0e4) |
| probs = torch.softmax(scores, dim=-1) |
| probs = self.dropout(probs) * query_mask.to(probs.dtype) |
| attended = torch.matmul(probs, v).transpose(1, 2).contiguous().view(batch, length, hidden) |
| hidden_states = residual + self.out_proj(attended) |
| hidden_states = masked_hidden(hidden_states, attention_mask) |
|
|
| hidden_states = hidden_states + self.ffn(self.ffn_norm(hidden_states)) |
| return masked_hidden(hidden_states, attention_mask) |
|
|
