|
|
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) |
|
|
|
