| | import math |
| | import torch |
| | import torch.nn as nn |
| | from einops import rearrange, repeat |
| | from einops.layers.torch import EinMix |
| | from transformers import PreTrainedModel, PretrainedConfig |
| |
|
| | |
| | |
| | |
| | class TRMConfig(PretrainedConfig): |
| | model_type = "trm" |
| |
|
| | def __init__(self, |
| | vocab_size=32000, |
| | hidden_size=256, |
| | seq_len=128, |
| | depth_L=2, |
| | depth_H=2, |
| | act_threshold=0.9, |
| | act_epsilon=1e-2, |
| | **kwargs): |
| | super().__init__(**kwargs) |
| | self.vocab_size = vocab_size |
| | self.hidden_size = hidden_size |
| | self.seq_len = seq_len |
| | self.depth_L = depth_L |
| | self.depth_H = depth_H |
| | self.act_threshold = act_threshold |
| | self.act_epsilon = act_epsilon |
| |
|
| |
|
| | |
| | |
| | |
| | class HaltingBlock(nn.Module): |
| | def __init__(self, hidden_size, act_threshold, act_epsilon): |
| | super().__init__() |
| | self.proj = nn.Linear(hidden_size, hidden_size) |
| | self.act_proj = nn.Linear(hidden_size, 1) |
| | self.act_threshold = act_threshold |
| | self.act_epsilon = act_epsilon |
| |
|
| | def forward(self, x): |
| | halting_probs = torch.sigmoid(self.act_proj(x)) |
| | remainders = torch.zeros_like(halting_probs) |
| | n_updates = torch.zeros_like(halting_probs) |
| | still_running = torch.ones_like(halting_probs, dtype=torch.bool) |
| | accumulated_output = torch.zeros_like(x) |
| | accumulated_prob = torch.zeros_like(halting_probs) |
| |
|
| | while still_running.any(): |
| | p = torch.where(still_running, halting_probs, torch.zeros_like(halting_probs)) |
| | new_accum = accumulated_prob + p |
| |
|
| | still_running = new_accum < self.act_threshold |
| | remainder = torch.where(still_running, torch.zeros_like(halting_probs), 1 - accumulated_prob) |
| |
|
| | update_weights = torch.where(still_running, p, remainder) |
| | accumulated_output += update_weights * torch.tanh(self.proj(x)) |
| | accumulated_prob += update_weights |
| | n_updates += still_running.float() |
| |
|
| | if (1 - accumulated_prob).mean() < self.act_epsilon: |
| | break |
| |
|
| | return accumulated_output, accumulated_prob.mean() |
| |
|
| |
|
| | class TRMLayer(nn.Module): |
| | def __init__(self, hidden_size, depth_H, act_threshold, act_epsilon): |
| | super().__init__() |
| | self.blocks = nn.ModuleList([ |
| | HaltingBlock(hidden_size, act_threshold, act_epsilon) for _ in range(depth_H) |
| | ]) |
| | self.norm = nn.LayerNorm(hidden_size) |
| |
|
| | def forward(self, x): |
| | for block in self.blocks: |
| | x, _ = block(x) |
| | return self.norm(x) |
| |
|
| |
|
| | class TRM(PreTrainedModel): |
| | config_class = TRMConfig |
| |
|
| | def __init__(self, config): |
| | super().__init__(config) |
| | self.emb = nn.Embedding(config.vocab_size, config.hidden_size) |
| | self.pos_emb = nn.Parameter(torch.zeros(1, config.seq_len, config.hidden_size)) |
| | self.layers = nn.ModuleList([ |
| | TRMLayer(config.hidden_size, config.depth_H, config.act_threshold, config.act_epsilon) |
| | for _ in range(config.depth_L) |
| | ]) |
| | self.norm = nn.LayerNorm(config.hidden_size) |
| | self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) |
| |
|
| | self.post_init() |
| |
|
| | def forward(self, input_ids, labels=None): |
| | x = self.emb(input_ids) + self.pos_emb[:, :input_ids.size(1), :] |
| | for layer in self.layers: |
| | x = layer(x) |
| | x = self.norm(x) |
| | logits = self.lm_head(x) |
| |
|
| | loss = None |
| | if labels is not None: |
| | shift_logits = logits[..., :-1, :].contiguous() |
| | shift_labels = labels[..., 1:].contiguous() |
| | loss_fct = nn.CrossEntropyLoss() |
| | loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) |
| |
|
| | return {"loss": loss, "logits": logits} |
| |
|
| |
|
| | |
| | |
| | |
| | from transformers import AutoConfig, AutoModel |
| |
|
| | AutoConfig.register("trm", TRMConfig) |
| | AutoModel.register(TRMConfig, TRM) |
| |
|
