modeling_trm_hrm.py

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