mixture_of_recursion.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PretrainedConfig, PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast
import math
class RecursiveLanguageModelConfig(PretrainedConfig):
    model_type="recursive_language_model"
    def __init__(
        self,
        vocab_size=50260,
        embedding_dim=768,
        num_layers=16,
        num_attention_heads=12,
        max_recursion_steps=5,
        max_position_embeddings=512,
        hidden_dropout_prob=0.1,
        attention_dropout_prob=0.1,
        intermediate_size=3072,
        layer_norm_eps=1e-5,
        pad_token_id=50257,
        bos_token_id=50256,
        eos_token_id=50256,
        simple_recursion_steps=1,
        medium_recursion_steps=3,
        complex_recursion_steps=5,
        initializer_range=0.02,
        **kwargs
    ):
        super().__init__(
            pad_token_id=pad_token_id,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            **kwargs
        )
        self.vocab_size=vocab_size
        self.embedding_dim=embedding_dim
        self.num_layers=num_layers
        self.num_attention_heads=num_attention_heads
        self.max_recursion_steps=max_recursion_steps
        self.max_position_embeddings=max_position_embeddings
        self.hidden_dropout_prob=hidden_dropout_prob
        self.attention_dropout_prob=attention_dropout_prob
        self.intermediate_size=intermediate_size
        self.layer_norm_eps=layer_norm_eps
        self.simple_recursion_steps=simple_recursion_steps
        self.medium_recursion_steps=medium_recursion_steps
        self.complex_recursion_steps=complex_recursion_steps
        self.initializer_range=initializer_range
class RotaryPositionalEmbedding(nn.Module):
    def __init__(self,dim,max_seq_len=2048,base=10000):
        super().__init__()
        inv_freq=1.0/(base**(torch.arange(0,dim,2).float()/dim))
        self.register_buffer('inv_freq',inv_freq)
    def forward(self,seq_len,device):
        t=torch.arange(seq_len,device=device).float()
        freqs=torch.outer(t,self.inv_freq)
        emb=torch.cat([freqs,freqs], dim=-1)
        return emb.cos(),emb.sin()
def rotate_half(x):
    x1, x2=x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:]
    return torch.cat([-x2, x1],dim=-1)
def apply_rotary_pos_emb(q, k, cos, sin):
    cos=cos.unsqueeze(0).unsqueeze(0)
    sin=sin.unsqueeze(0).unsqueeze(0)
    return(q*cos)+(rotate_half(q)*sin),(k*cos)+(rotate_half(k)*sin)
class MultiHeadAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.num_heads=config.num_attention_heads
        self.head_dim=config.embedding_dim // config.num_attention_heads
        self.embed_dim=config.embedding_dim
        assert self.embed_dim % self.num_heads==0
        self.q_proj=nn.Linear(self.embed_dim, self.embed_dim, bias=False)
        self.k_proj=nn.Linear(self.embed_dim, self.embed_dim, bias=False)
        self.v_proj=nn.Linear(self.embed_dim, self.embed_dim, bias=False)
        self.out_proj=nn.Linear(self.embed_dim, self.embed_dim, bias=False)
        self.attn_drop=nn.Dropout(config.attention_dropout_prob)
        self.rope=RotaryPositionalEmbedding(self.head_dim, config.max_position_embeddings)
    def forward(self,x,causal_mask=None):
        B,T,C=x.shape
        q=self.q_proj(x).view(B,T,self.num_heads, self.head_dim).transpose(1, 2)
        k=self.k_proj(x).view(B,T,self.num_heads, self.head_dim).transpose(1, 2)
        v=self.v_proj(x).view(B,T,self.num_heads, self.head_dim).transpose(1, 2)
        cos,sin=self.rope(T,x.device)
        q,k=apply_rotary_pos_emb(q, k, cos, sin)
        scale=math.sqrt(self.head_dim)
        scores=torch.matmul(q, k.transpose(-2, -1))/scale 
        if causal_mask is not None:
            scores=scores+causal_mask
        scores=scores.clamp(min=-1e4, max=1e4)
        attn=F.softmax(scores, dim=-1)
        attn=torch.nan_to_num(attn, nan=0.0, posinf=0.0, neginf=0.0)
        attn=self.attn_drop(attn)
        out=torch.matmul(attn, v)                             
        out=out.transpose(1, 2).contiguous().view(B, T, C)
        return self.out_proj(out)
class FeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.fc1=nn.Linear(config.embedding_dim, config.intermediate_size, bias=False)
        self.fc2=nn.Linear(config.intermediate_size, config.embedding_dim, bias=False)
        self.drop=nn.Dropout(config.hidden_dropout_prob)
    def forward(self, x):
        return self.drop(self.fc2(self.drop(F.gelu(self.fc1(x)))))
class TransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.attn=MultiHeadAttention(config)
        self.ff=FeedForward(config)
        self.ln1=nn.LayerNorm(config.embedding_dim,eps=config.layer_norm_eps)
        self.ln2=nn.LayerNorm(config.embedding_dim,eps=config.layer_norm_eps)
    def forward(self, x, mask=None):
        x=x+self.attn(self.ln1(x), mask)
        x=x+self.ff(self.ln2(x))
        return x
class SequenceLevelRouter(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.pooler=nn.Linear(config.embedding_dim, config.embedding_dim)
        self.act=nn.Tanh()
        self.head=nn.Sequential(
            nn.Linear(config.embedding_dim, config.embedding_dim // 2),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(config.embedding_dim // 2, 3)
        )
        self.register_buffer('steps_map', torch.tensor([
            config.simple_recursion_steps,
            config.medium_recursion_steps,
            config.complex_recursion_steps,
        ], dtype=torch.long))
    def forward(self, x, valid_mask=None):
        if valid_mask is not None:
            m=valid_mask.unsqueeze(-1).float()
            pooled=(x * m).sum(1)/m.sum(1).clamp(min=1e-9)
        else:
            pooled=x.mean(1)
        pooled=self.act(self.pooler(pooled))
        logits=self.head(pooled)
        cls=logits.argmax(dim=-1)
        return logits, cls, self.steps_map[cls]
class RecursionLayer(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.block=TransformerBlock(config)
    def forward(self,x,mask=None):
        return self.block(x, mask)
class RecursiveLanguageModel(PreTrainedModel):
    config_class = RecursiveLanguageModelConfig
    def __init__(self, config):
        super().__init__(config)
        self.config=config
        self.embed=nn.Embedding(config.vocab_size, config.embedding_dim,
                                     padding_idx=config.pad_token_id)
        self.layers=nn.ModuleList([TransformerBlock(config) for _ in range(config.num_layers)])
        self.router=SequenceLevelRouter(config)
        self.rec_layer=RecursionLayer(config)
        self.ln_f=nn.LayerNorm(config.embedding_dim, eps=config.layer_norm_eps)
        self.lm_head=nn.Linear(config.embedding_dim, config.vocab_size, bias=False)
        self.post_init()
    def get_input_embeddings(self): return self.embed
    def set_input_embeddings(self, v): self.embed = v
    def get_output_embeddings(self): return self.lm_head
    def set_output_embeddings(self, v): self.lm_head = v
    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight,mean=0.0,std=self.config.initializer_range)
            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=self.config.initializer_range)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.LayerNorm):
            nn.init.zeros_(module.bias)
            nn.init.ones_(module.weight)
    def _make_causal_mask(self, input_ids):
        B,T=input_ids.shape
        device=input_ids.device
        mask=torch.zeros(T, T, device=device)
        mask=mask.masked_fill(
            torch.triu(torch.ones(T, T, device=device, dtype=torch.bool), diagonal=1),
            -1e4
        )
        mask=mask.unsqueeze(0).unsqueeze(0)
        pad_mask=(input_ids==self.config.pad_token_id)
        valid_mask=~pad_mask
        if pad_mask.any():
            pad_key_mask=pad_mask.unsqueeze(1).unsqueeze(2).float()*-1e4
            mask=mask+pad_key_mask
        return mask, valid_mask
    def forward(self,input_ids,labels=None,attention_mask=None,**kwargs):
        B,T=input_ids.shape
        x=self.embed(input_ids)
        causal_mask,valid_mask=self._make_causal_mask(input_ids)
        for layer in self.layers:
            x=layer(x,causal_mask)
        router_logits, cls, steps=self.router(x,valid_mask)
        max_steps=int(steps.max().item())
        for s in range(max_steps):
            gate=(steps > s).float().view(B, 1, 1)
            x=gate*self.rec_layer(x, causal_mask)+(1-gate)*x
        x=self.ln_f(x)
        logits=self.lm_head(x)
        loss=None
        if labels is not None:
            shift_logits=logits[:, :-1, :].contiguous()
            shift_labels=labels[:, 1:].contiguous()
            lm_loss=F.cross_entropy(
                shift_logits.view(-1, self.config.vocab_size),
                shift_labels.view(-1),
                ignore_index=-100
            )
            with torch.no_grad():
                per_tok= F.cross_entropy(
                    shift_logits.view(-1, self.config.vocab_size),
                    shift_labels.view(-1),
                    ignore_index=-100, reduction='none'
                ).view(B, -1)
                valid_tok=(shift_labels!=-100).sum(1).clamp(min=1).float()
                ppl=torch.exp((per_tok.sum(1)/valid_tok).clamp(max=20))
                pseudo=torch.zeros(B,dtype=torch.long,device=input_ids.device)
                pseudo[(ppl>=20)&(ppl<50)]=1
                pseudo[ppl>= 50] = 2
            router_loss=F.cross_entropy(router_logits, pseudo)
            loss=lm_loss+0.1*router_loss
        return CausalLMOutputWithPast(loss=loss,logits=logits)
    @torch.no_grad()
    def generate(self,input_ids,max_new_tokens=100,temperature=0.8,
                 top_p=0.9, do_sample=True, **kwargs):
        self.eval()
        gen=input_ids
        for _ in range(max_new_tokens):
            ctx=gen[:,-self.config.max_position_embeddings:]
            logits=self.forward(ctx).logits[:, -1, :]
            if temperature!=1.0:
                logits=logits/temperature
            if do_sample:
                probs=F.softmax(logits, dim=-1)
                sorted_probs,sorted_idx=torch.sort(probs,descending=True)
                cum_probs=torch.cumsum(sorted_probs, dim=-1)
                remove=cum_probs-sorted_probs>top_p
                sorted_probs[remove]=0.0
                sorted_probs=sorted_probs/sorted_probs.sum(dim=-1,keepdim=True)
                next_tok=torch.gather(sorted_idx,-1,torch.multinomial(sorted_probs,1))
            else:
                next_tok=logits.argmax(dim=-1,keepdim=True)
            gen=torch.cat([gen,next_tok],dim=-1)
            if (next_tok==self.config.eos_token_id).all():
                break
        return gen