import torch import torch.nn as nn import torch.nn.functional as F from transformers import PreTrainedModel, GenerationMixin from transformers.modeling_outputs import CausalLMOutputWithPast from .configuration_trm_hrm import TRMHRMConfig class RMSNorm(nn.Module): def __init__(self, dim, eps=1e-8): super().__init__() self.weight = nn.Parameter(torch.ones(dim)) self.eps = eps def forward(self, x): return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) * self.weight class RoPE(nn.Module): def __init__(self, dim, max_seq_len=4096, base=10000): super().__init__() assert dim % 2 == 0 half = dim // 2 inv_freq = 1.0 / (base ** (torch.arange(half).float() / half)) pos = torch.arange(max_seq_len).float() angles = torch.einsum("i,j->ij", pos, inv_freq) self.register_buffer("sin", angles.sin()[None, :, :], persistent=False) self.register_buffer("cos", angles.cos()[None, :, :], persistent=False) def forward(self, x): seq_len = x.size(1) sin = self.sin[:, :seq_len, :].to(device=x.device, dtype=x.dtype) cos = self.cos[:, :seq_len, :].to(device=x.device, dtype=x.dtype) half = x.size(-1) // 2 x1 = x[..., :half] x2 = x[..., half:] return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1) class CausalDepthwiseConv1d(nn.Module): def __init__(self, dim, kernel_size, dilation): super().__init__() self.left_padding = dilation * (kernel_size - 1) self.conv = nn.Conv1d( dim, dim, kernel_size=kernel_size, dilation=dilation, groups=dim, bias=True, ) def forward(self, x): h = x.transpose(1, 2) h = F.pad(h, (self.left_padding, 0)) h = self.conv(h) return h.transpose(1, 2) class CausalDilatedDepthwiseMixer(nn.Module): def __init__(self, dim, kernel_size, dilations): super().__init__() self.convs = nn.ModuleList([ CausalDepthwiseConv1d(dim, kernel_size, d) for d in dilations ]) self.proj = nn.Linear(dim * len(dilations), dim) def forward(self, x): return self.proj(torch.cat([conv(x) for conv in self.convs], dim=-1)) class TinyRecursiveBlock(nn.Module): def __init__(self, config, dilations, rope): super().__init__() self.norm = RMSNorm(config.dim) self.rope = rope self.mixer = CausalDilatedDepthwiseMixer( config.dim, config.kernel_size, dilations, ) self.up = nn.Linear(config.dim, config.hidden_dim * 2) self.down = nn.Linear(config.hidden_dim, config.dim) self.dropout = nn.Dropout(config.dropout) self.res_scale = nn.Parameter(torch.tensor(0.1)) def forward(self, x): h = self.norm(x) h = self.rope(h) h = self.mixer(h) a, gate = self.up(h).chunk(2, dim=-1) h = a * F.silu(gate) h = self.down(h) h = self.dropout(h) return x + self.res_scale * h class TRMUnit(nn.Module): def __init__(self, config, n_steps, dilations, rope): super().__init__() self.block = TinyRecursiveBlock(config, dilations, rope) self.n_steps = n_steps def forward(self, x): for _ in range(self.n_steps): x = self.block(x) return x class TRMHRMForCausalLM(PreTrainedModel, GenerationMixin): config_class = TRMHRMConfig base_model_prefix = "trm_hrm" supports_gradient_checkpointing = False _tied_weights_keys = {} all_tied_weights_keys = {} def __init__(self, config): super().__init__(config) self.token_embedding = nn.Embedding( config.vocab_size, config.dim, padding_idx=config.pad_token_id, ) self.rope = RoPE(config.dim, config.max_seq_len) self.low = TRMUnit(config, config.low_steps, config.low_dilations, self.rope) self.mid = TRMUnit(config, config.mid_steps, config.mid_dilations, self.rope) self.high = TRMUnit(config, config.high_steps, config.high_dilations, self.rope) self.low_to_mid = nn.Linear(config.dim, config.dim) self.mid_to_high = nn.Linear(config.dim, config.dim) self.high_to_mid = nn.Linear(config.dim, config.dim) self.mid_to_low = nn.Linear(config.dim, config.dim) self.final_norm = RMSNorm(config.dim) self.lm_head = nn.Linear(config.dim, config.vocab_size, bias=False) self.post_init() @classmethod def _can_set_experts_implementation(cls): return False def _init_weights(self, module): if isinstance(module, nn.Linear): nn.init.normal_(module.weight, mean=0.0, std=0.02) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): nn.init.normal_(module.weight, mean=0.0, std=0.02) def get_input_embeddings(self): return self.token_embedding def set_input_embeddings(self, value): self.token_embedding = value def get_output_embeddings(self): return self.lm_head def set_output_embeddings(self, value): self.lm_head = value def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs): if input_ids.size(1) > self.config.max_seq_len: input_ids = input_ids[:, -self.config.max_seq_len:] if attention_mask is not None: attention_mask = attention_mask[:, -self.config.max_seq_len:] return { "input_ids": input_ids, "attention_mask": attention_mask, } def forward( self, input_ids=None, attention_mask=None, labels=None, return_dict=True, **kwargs, ): low = self.token_embedding(input_ids) if attention_mask is not None: low = low * attention_mask.unsqueeze(-1).to(low.dtype) mid = torch.zeros_like(low) high = torch.zeros_like(low) for _ in range(self.config.cycles): low = self.low(low) mid = mid + 0.25 * self.low_to_mid(low) mid = self.mid(mid) high = high + 0.25 * self.mid_to_high(mid) high = self.high(high) mid = mid + 0.25 * self.high_to_mid(high) low = low + 0.25 * self.mid_to_low(mid) logits = self.lm_head(self.final_norm(low)) loss = None if labels is not None: loss = F.cross_entropy( logits[:, :-1, :].contiguous().view(-1, self.config.vocab_size), labels[:, 1:].contiguous().view(-1), ignore_index=-100, ) if not return_dict: return (loss, logits) if loss is not None else (logits,) return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=None, hidden_states=None, attentions=None, )