| |
| """ |
| HF-compatible single-language Hawk / RG-LRU model, for lm-eval. |
| |
| Self-contained `trust_remote_code` modeling file. The building blocks are the |
| SAME code used at training time (De et al., 2024, Griffin/Hawk; arXiv:2402.19427), |
| so the per-language export state_dict maps 1:1 onto this module's parameters |
| (top-level attribute names wte / layers / norm_f / lm_head match the export keys |
| exactly -- no renaming, no transpose). Exposes the standard |
| forward(input_ids, labels=None) -> CausalLMOutputWithPast that lm-eval expects. |
| |
| Register via config.json: |
| "model_type": "hawk_rglru", |
| "architectures": ["HawkForCausalLM"], |
| "auto_map": { |
| "AutoConfig": "modeling_hawk.HawkConfig", |
| "AutoModelForCausalLM": "modeling_hawk.HawkForCausalLM" |
| } |
| """ |
|
|
| from typing import Optional |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from transformers import PreTrainedModel, PretrainedConfig |
| from transformers.generation import GenerationMixin |
| from transformers.modeling_outputs import ( |
| CausalLMOutputWithPast, SequenceClassifierOutput, |
| ) |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, dim: int, eps: float = 1e-6): |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(dim)) |
| self.eps = eps |
|
|
| def forward(self, x): |
| norm = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) |
| return self.weight * (x * norm) |
|
|
|
|
| class RGLRU(nn.Module): |
| """Real-Gated Linear Recurrent Unit (De et al., 2024).""" |
|
|
| def __init__(self, width: int, c: float = 8.0): |
| super().__init__() |
| self.width = width |
| self.c = c |
| self.input_gate = nn.Linear(width, width) |
| self.recur_gate = nn.Linear(width, width) |
| lam = torch.empty(width).uniform_(2.197, 6.907) |
| self.log_lambda = nn.Parameter(lam) |
|
|
| def forward(self, x): |
| B, T, W = x.shape |
| r = torch.sigmoid(self.recur_gate(x)) |
| i = torch.sigmoid(self.input_gate(x)) |
| log_a = -F.softplus(-self.log_lambda) |
| log_a_t = self.c * r * log_a |
| a_t = torch.exp(log_a_t) |
| mult = torch.sqrt(torch.clamp(-torch.expm1(2.0 * log_a_t), min=1e-8)) |
| gated_x = mult * (i * x) |
| h = torch.zeros(B, W, device=x.device, dtype=x.dtype) |
| outs = [] |
| for t in range(T): |
| h = a_t[:, t] * h + gated_x[:, t] |
| outs.append(h) |
| return torch.stack(outs, dim=1) |
|
|
|
|
| class RecurrentBlock(nn.Module): |
| def __init__(self, d_model: int, d_rnn: int, conv_kernel: int = 4, rglru_c: float = 8.0): |
| super().__init__() |
| self.conv_kernel = conv_kernel |
| self.in_gate = nn.Linear(d_model, d_rnn) |
| self.in_recur = nn.Linear(d_model, d_rnn) |
| self.conv = nn.Conv1d(d_rnn, d_rnn, conv_kernel, groups=d_rnn, |
| padding=conv_kernel - 1) |
| self.rglru = RGLRU(d_rnn, rglru_c) |
| self.out = nn.Linear(d_rnn, d_model) |
|
|
| def forward(self, x): |
| gate = F.gelu(self.in_gate(x)) |
| rec = self.in_recur(x).transpose(1, 2) |
| rec = self.conv(rec)[..., : x.size(1)] |
| rec = self.rglru(rec.transpose(1, 2)) |
| return self.out(gate * rec) |
|
|
|
|
| class MLPBlock(nn.Module): |
| def __init__(self, d_model: int, expansion: int = 3): |
| super().__init__() |
| hidden = expansion * d_model |
| self.gate = nn.Linear(d_model, hidden) |
| self.up = nn.Linear(d_model, hidden) |
| self.down = nn.Linear(hidden, d_model) |
|
|
| def forward(self, x): |
| return self.down(F.gelu(self.gate(x)) * self.up(x)) |
|
|
|
|
| class HawkLayer(nn.Module): |
| def __init__(self, d_model, d_rnn, conv_kernel, mlp_expansion, eps, rglru_c=8.0): |
| super().__init__() |
| self.norm1 = RMSNorm(d_model, eps) |
| self.recur = RecurrentBlock(d_model, d_rnn, conv_kernel, rglru_c) |
| self.norm2 = RMSNorm(d_model, eps) |
| self.mlp = MLPBlock(d_model, mlp_expansion) |
|
|
| def forward(self, x): |
| x = x + self.recur(self.norm1(x)) |
| x = x + self.mlp(self.norm2(x)) |
| return x |
|
|
|
|
| class HawkConfig(PretrainedConfig): |
| model_type = "hawk_rglru" |
|
|
| def __init__(self, vocab_size: int = 16384, n_layer: int = 12, n_embd: int = 768, |
| rnn_width: Optional[int] = None, conv_kernel: int = 4, |
| mlp_expansion: int = 3, rmsnorm_eps: float = 1e-6, rglru_c: float = 8.0, |
| max_position_embeddings: int = 1024, tie_word_embeddings: bool = True, |
| bos_token_id: int = 2, eos_token_id: int = 3, pad_token_id: int = 1, |
| **kwargs): |
| self.vocab_size = vocab_size |
| self.n_layer = n_layer |
| self.n_embd = n_embd |
| self.rnn_width = rnn_width |
| self.conv_kernel = conv_kernel |
| self.mlp_expansion = mlp_expansion |
| self.rmsnorm_eps = rmsnorm_eps |
| self.rglru_c = rglru_c |
| self.max_position_embeddings = max_position_embeddings |
| self.auto_map = { |
| "AutoConfig": "modeling_hawk.HawkConfig", |
| "AutoModelForCausalLM": "modeling_hawk.HawkForCausalLM", |
| } |
| super().__init__(tie_word_embeddings=tie_word_embeddings, |
| bos_token_id=bos_token_id, eos_token_id=eos_token_id, |
| pad_token_id=pad_token_id, **kwargs) |
|
|
| @property |
| def d_rnn(self): |
| return self.rnn_width if self.rnn_width is not None else self.n_embd |
|
|
|
|
| class HawkForCausalLM(PreTrainedModel, GenerationMixin): |
| config_class = HawkConfig |
| _tied_weights_keys = {"lm_head.weight": "wte.weight"} |
|
|
| def __init__(self, config: HawkConfig): |
| super().__init__(config) |
| d_rnn = config.d_rnn |
| self.wte = nn.Embedding(config.vocab_size, config.n_embd) |
| self.layers = nn.ModuleList([ |
| HawkLayer(config.n_embd, d_rnn, config.conv_kernel, |
| config.mlp_expansion, config.rmsnorm_eps, config.rglru_c) |
| for _ in range(config.n_layer)]) |
| self.norm_f = RMSNorm(config.n_embd, config.rmsnorm_eps) |
| self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) |
| self.post_init() |
|
|
| def get_input_embeddings(self): |
| return self.wte |
|
|
| def set_input_embeddings(self, new): |
| self.wte = new |
|
|
| def get_output_embeddings(self): |
| return self.lm_head |
|
|
| def forward(self, input_ids: torch.LongTensor, |
| attention_mask: Optional[torch.Tensor] = None, |
| labels: Optional[torch.LongTensor] = None, |
| **kwargs) -> CausalLMOutputWithPast: |
| x = self.wte(input_ids) |
| for layer in self.layers: |
| x = layer(x) |
| x = self.norm_f(x) |
| logits = self.lm_head(x) |
| loss = None |
| if labels is not None: |
| shift_logits = logits[:, :-1, :].contiguous() |
| shift_labels = labels[:, 1:].contiguous() |
| loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), |
| shift_labels.view(-1), ignore_index=-100) |
| return CausalLMOutputWithPast(loss=loss, logits=logits) |
|
|
|
|
| class HawkForSequenceClassification(PreTrainedModel): |
| """ |
| Sequence-classification head on top of the SAME Hawk backbone used for |
| causal LM. The backbone attribute names (wte / layers / norm_f) are |
| IDENTICAL to HawkForCausalLM, so a CausalLM export state_dict maps 1:1 onto |
| the backbone with no renaming. Only `score` is newly initialised, which is |
| the expected behaviour when starting a fine-tuning run. |
| |
| The pooled representation is read from the hidden state at the last |
| non-padding position (right padding, as produced by the BabyLM finetune |
| tokenizer), matching the GPT-2 / Mamba sequence-classification convention. |
| """ |
|
|
| config_class = HawkConfig |
|
|
| def __init__(self, config: HawkConfig): |
| super().__init__(config) |
| self.num_labels = config.num_labels |
| d_rnn = config.d_rnn |
| self.wte = nn.Embedding(config.vocab_size, config.n_embd) |
| self.layers = nn.ModuleList([ |
| HawkLayer(config.n_embd, d_rnn, config.conv_kernel, |
| config.mlp_expansion, config.rmsnorm_eps, config.rglru_c) |
| for _ in range(config.n_layer)]) |
| self.norm_f = RMSNorm(config.n_embd, config.rmsnorm_eps) |
| self.score = nn.Linear(config.n_embd, self.num_labels, bias=False) |
| self.post_init() |
|
|
| def get_input_embeddings(self): |
| return self.wte |
|
|
| def set_input_embeddings(self, new): |
| self.wte = new |
|
|
| def forward(self, input_ids: torch.LongTensor, |
| attention_mask: Optional[torch.Tensor] = None, |
| labels: Optional[torch.LongTensor] = None, |
| **kwargs) -> SequenceClassifierOutput: |
| x = self.wte(input_ids) |
| for layer in self.layers: |
| x = layer(x) |
| x = self.norm_f(x) |
| logits = self.score(x) |
|
|
| B, T = input_ids.shape[:2] |
| |
| if attention_mask is not None: |
| last_idx = attention_mask.long().sum(-1) - 1 |
| elif self.config.pad_token_id is not None: |
| last_idx = (input_ids != self.config.pad_token_id).int().sum(-1) - 1 |
| else: |
| last_idx = torch.full((B,), T - 1, device=input_ids.device) |
| last_idx = last_idx.clamp(min=0) |
| pooled_logits = logits[torch.arange(B, device=input_ids.device), last_idx] |
|
|
| loss = None |
| if labels is not None: |
| if self.config.problem_type is None: |
| if self.num_labels == 1: |
| self.config.problem_type = "regression" |
| elif self.num_labels > 1 and labels.dtype in (torch.long, torch.int): |
| self.config.problem_type = "single_label_classification" |
| else: |
| self.config.problem_type = "multi_label_classification" |
|
|
| if self.config.problem_type == "regression": |
| loss_fct = nn.MSELoss() |
| loss = (loss_fct(pooled_logits.squeeze(), labels.squeeze()) |
| if self.num_labels == 1 |
| else loss_fct(pooled_logits, labels)) |
| elif self.config.problem_type == "single_label_classification": |
| loss_fct = nn.CrossEntropyLoss() |
| loss = loss_fct(pooled_logits.view(-1, self.num_labels), |
| labels.view(-1)) |
| else: |
| loss_fct = nn.BCEWithLogitsLoss() |
| loss = loss_fct(pooled_logits, labels.float()) |
|
|
| return SequenceClassifierOutput(loss=loss, logits=pooled_logits) |
