Fixed model.py

model.py

@@ -0,0 +1,541 @@
+from typing import Tuple
+
+import torch
+from torch import Tensor
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import PreTrainedModel, GenerationMixin
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, MaskedLMOutput
+from transformers.cache_utils import Cache, DynamicCache
+
+from rotary_embedding_torch import RotaryEmbedding
+from .config import FSTConfig
+
+# === Util ===
+
+class Residual(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: Tensor, delta: Tensor):
+        return x + delta
+
+# === MLP ===
+
+class MLP(nn.Module):
+    def __init__(
+        self, 
+        hidden_size: int,
+        intermediate_size: int
+    ):
+        super().__init__()
+        
+        self.fc_up = nn.Linear(hidden_size, intermediate_size)
+        self.activation = nn.GELU()
+        self.fc_down = nn.Linear(intermediate_size, hidden_size)
+
+    def forward(self, x: Tensor):
+        return self.fc_down(self.activation(self.fc_up(x)))
+
+# === Attention ===
+
+class MHAttention(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_attention_heads: int,
+        use_causal_attention: bool = True,
+        layer_idx: int | None = None
+    ):
+        super().__init__()
+        
+        self.hidden_size = hidden_size
+        self.num_attention_heads = num_attention_heads
+        self.head_dim = hidden_size // num_attention_heads
+
+        assert self.head_dim * self.num_attention_heads == self.hidden_size
+
+        self.use_causal_attention = use_causal_attention
+        self.layer_idx = layer_idx
+        
+        self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=True)
+        self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=True)
+
+        self.rotary_emb = RotaryEmbedding(dim=self.head_dim)
+        self.scale = self.head_dim ** -0.5
+
+    def forward(
+        self,
+        q: Tensor,
+        k: Tensor | None = None,
+        v: Tensor | None = None,
+        attention_mask: Tensor | None = None,
+        past_key_values: Cache | None = None
+    ):
+        B, T, _ = q.size()
+
+        if k is None:
+            k = q
+        if v is None:
+            v = q
+
+        q = self.q_proj(q)
+        k = self.k_proj(k)
+        v = self.v_proj(v)
+
+        q = q.view(B, T, self.num_attention_heads, self.head_dim).transpose(1, 2)
+        k = k.view(B, T, self.num_attention_heads, self.head_dim).transpose(1, 2)
+        v = v.view(B, T, self.num_attention_heads, self.head_dim).transpose(1, 2)
+
+        if past_key_values is None:
+            
+            q = self.rotary_emb.rotate_queries_or_keys(q)
+            k = self.rotary_emb.rotate_queries_or_keys(k)
+
+        else:
+            
+            cache_position = past_key_values.get_seq_length(self.layer_idx)
+            
+            q = self.rotary_emb.rotate_queries_or_keys(q, offset=cache_position)
+            k = self.rotary_emb.rotate_queries_or_keys(k, offset=cache_position)
+            
+            k, v = past_key_values.update(k, v, self.layer_idx)
+
+        is_causal = self.use_causal_attention and attention_mask is None
+        attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask, scale=self.scale, is_causal=is_causal)
+
+        attn_output = attn_output.transpose(1, 2).contiguous().view(B, T, self.hidden_size)
+        out = self.o_proj(attn_output)
+            
+        return out
+
+# === Blocks ===
+
+class FeatureBlock(nn.Module):
+
+    def __init__(
+        self, 
+        config: FSTConfig, 
+        layer_idx: int = None
+    ):
+        super().__init__()
+        
+        self.attn = MHAttention(
+            hidden_size=config.hidden_size,
+            num_attention_heads=config.num_attention_heads,
+            use_causal_attention=config.use_causal_attention,
+            layer_idx=layer_idx,
+        )
+        
+        self.mlp = MLP(
+            config.hidden_size, 
+            config.intermediate_size
+        )
+        
+        self.norm_attn = nn.LayerNorm(config.hidden_size)
+        self.norm_mlp = nn.LayerNorm(config.hidden_size)
+
+        self.resid_attn = Residual()
+        self.resid_mlp = Residual()
+
+    def forward(
+        self, 
+        x: Tensor, 
+        attention_mask: Tensor | None = None, 
+        past_key_values: Cache | None = None
+    ):
+
+        attn_out = self.attn(self.norm_attn(x), attention_mask=attention_mask, past_key_values=past_key_values)
+        x = self.resid_attn(x, attn_out)
+
+        mlp_out = self.mlp(self.norm_mlp(x))
+        x = self.resid_mlp(x, mlp_out)
+
+        return x
+
+class PredictiveBlock(nn.Module):
+
+    def __init__(
+        self, 
+        config: FSTConfig, 
+        layer_idx: int = None
+    ):
+        super().__init__()
+        
+        self.attn = MHAttention(
+            hidden_size=config.hidden_size,
+            num_attention_heads=config.num_attention_heads,
+            use_causal_attention=config.use_causal_attention,
+            layer_idx=layer_idx,
+        )
+        
+        self.mlp = MLP(
+            config.hidden_size, 
+            config.intermediate_size
+        )
+        
+        self.norm_attn_qk = nn.LayerNorm(config.hidden_size)
+        self.norm_attn_v = nn.LayerNorm(config.hidden_size)
+        self.norm_mlp = nn.LayerNorm(config.hidden_size)
+
+        self.resid_attn = Residual()
+        self.resid_mlp = Residual()
+
+    def forward(
+        self, 
+        phi: Tensor,
+        f: Tensor, 
+        e: Tensor,
+        attention_mask: Tensor | None = None, 
+        past_key_values: Cache | None = None
+    ):
+        
+        qk = self.norm_attn_qk(phi)
+        v = self.norm_attn_v(e)
+
+        attn_out = self.attn(qk, qk, v, attention_mask=attention_mask, past_key_values=past_key_values)
+        f = self.resid_attn(f, attn_out)
+
+        mlp_out = self.mlp(self.norm_mlp(f))
+        f = self.resid_mlp(f, mlp_out)
+
+        return f
+
+# === Base Model ===
+
+class FSTPreTrainedModel(PreTrainedModel):
+    
+    config_class = FSTConfig
+    base_model_prefix = "model"
+    _no_split_modules = ["FSTBlock"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn_2 = True
+    _supports_cache_class = True
+
+    # Initialization taken from Deepseek and Falcon
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+class FSTModel(FSTPreTrainedModel):
+    
+    def __init__(
+        self, 
+        config: FSTConfig
+    ):
+        super().__init__(config)
+
+        self.config = config
+        self.embedding = nn.Embedding(config.vocab_size, config.hidden_size)
+
+        self.feature_blocks = nn.ModuleList([FeatureBlock(config, layer_idx) for layer_idx in range(0, config.num_hidden_layers, 2)])
+        self.predictive_blocks = nn.ModuleList([PredictiveBlock(config, layer_idx) for layer_idx in range(1, config.num_hidden_layers, 2)])
+        self.norm_out = nn.LayerNorm(config.hidden_size)
+
+        self.post_init()
+
+    def _prepare_attention_mask(
+        self, 
+        x: Tensor,
+        attention_mask: Tensor | None = None,
+        past_key_values: Cache | None = None,
+        use_causal_attention: bool = True
+    ):
+        
+        device = x.device
+        B = x.shape[0]
+        T = x.shape[1]
+
+        T_past = past_key_values.get_seq_length() if past_key_values is not None else 0
+        T_total = T + T_past
+
+        if use_causal_attention:
+            causal_mask = ~torch.triu(
+                torch.ones((T, T_total), dtype=torch.bool, device=device), 
+                diagonal=(1 + T_past)
+            ).unsqueeze(0).unsqueeze(0)
+        
+        if attention_mask is not None:
+            attn_len = attention_mask.shape[-1]
+
+            if attn_len < T_total:
+                pad = torch.ones(B, T_past, device=device, dtype=attention_mask.dtype)  # Fixed: ones instead of zeros
+                attention_mask = torch.cat([pad, attention_mask], dim=-1)
+            elif attn_len > T_total:
+                attention_mask = attention_mask[:, -T_total:]
+            
+            expanded_mask = (attention_mask == 1).view(B, 1, 1, T_total)
+        
+        if use_causal_attention and attention_mask is not None:
+            return causal_mask & expanded_mask
+        elif use_causal_attention:
+            return causal_mask
+        elif attention_mask is not None:  # Added: handle non-causal with custom mask
+            return expanded_mask
+        else:
+            return torch.ones((1, 1, T, T_total), dtype=torch.bool, device=device)
+        
+    def forward(
+        self,
+        input_ids: Tensor | None = None,
+        attention_mask: Tensor | None = None,
+        inputs_embeds: Tensor | None = None,
+        past_key_values = None,
+        use_cache: bool | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs,
+    ):
+        
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        assert not (input_ids is not None and inputs_embeds is not None), "You cannot specify both input_ids and inputs_embeds"
+        assert not (input_ids is None and inputs_embeds is None), "You must specify either input_ids or inputs_embeds"
+
+        e = self.embedding(input_ids) if input_ids is not None else inputs_embeds
+
+        B, T, _ = e.shape
+        device = e.device
+        dtype = e.dtype
+
+        if not use_cache:
+            past_key_values=None
+        elif past_key_values is None:
+            past_key_values = DynamicCache()
+
+        # Note that we must use an attention mask when caching- otherwise, SDPA uses is_casual and breaks
+        if attention_mask is not None or past_key_values is not None:
+            attention_mask = self._prepare_attention_mask(e, attention_mask=attention_mask, use_causal_attention=self.config.use_causal_attention, past_key_values=past_key_values)
+
+        hidden_states = [] if output_hidden_states else None
+
+        phi = e
+        f = torch.zeros(B, T, self.config.hidden_size, dtype=dtype, device=device) # Initialize f as zero for purity, but f=e also works fine
+
+        for feature_block, predictive_block in zip(self.feature_blocks, self.predictive_blocks):
+            
+            phi = feature_block(phi, attention_mask=attention_mask, past_key_values=past_key_values)
+            f = predictive_block(phi, f, e, attention_mask=attention_mask, past_key_values=past_key_values)
+
+            if output_hidden_states:
+                hidden_states.append(phi)
+                hidden_states.append(f)
+
+        if hidden_states is not None:
+            hidden_states = tuple(hidden_states)
+
+        f = self.norm_out(f)
+
+        if return_dict:
+            return BaseModelOutputWithPast(
+                last_hidden_state=f,
+                past_key_values=past_key_values,
+                hidden_states=hidden_states
+            )
+
+        return f, past_key_values, hidden_states
+
+# === Applied Models ===
+
+class FSTForCausalLM(GenerationMixin, FSTPreTrainedModel):
+    
+    accepts_loss_kwargs = False
+
+    def __init__(
+            self, 
+            config: FSTConfig
+    ):
+        super().__init__(config)
+    
+        self.model = FSTModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        if config.tie_word_embeddings:
+            self.tie_weights()
+            self._dynamic_tied_weights_keys = {"lm_head.weight": "model.embedding.weight"} # Avoids safetensor naming issues
+
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embedding
+
+    def set_input_embeddings(self, new_embeddings):
+        self.model.embedding = new_embeddings
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def tie_weights(self):
+        self._tie_or_clone_weights(self.lm_head, self.get_input_embeddings())
+
+    def forward(
+        self,
+        input_ids: Tensor | None = None,
+        attention_mask: Tensor | None = None,
+        past_key_values = None,
+        inputs_embeds: Tensor | None = None,
+        labels: Tensor | None = None,
+        use_cache: bool | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs,
+    ):
+
+        if labels is not None:
+            return_dict = True
+        else:
+            return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+    
+        model_output = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_hidden_states=output_hidden_states
+        )
+
+        logits = self.lm_head(model_output[0])
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100
+            )
+
+        if not return_dict:
+            output = (logits,) + model_output[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=model_output.past_key_values,
+            hidden_states=model_output.hidden_states
+        )
+
+    def _prepare_inputs_for_generation(
+        self,
+        input_ids: Tensor,
+        past_key_values: Cache | None = None,
+        attention_mask: Tensor | None = None,
+        **kwargs
+    ):
+        if past_key_values is not None:
+            input_ids = input_ids[:, -1:]
+
+        model_inputs = {"input_ids": input_ids, "past_key_values": past_key_values, "use_cache": True}
+
+        if attention_mask is not None:
+            model_inputs["attention_mask"] = attention_mask
+
+        for key, value in kwargs.items():
+            model_inputs[key] = value
+
+        return model_inputs
+    
+    def _reorder_cache(self, past_key_values: Cache, beam_idx: Tensor):
+        return past_key_values.reorder_cache(beam_idx)
+    
+class FSTForMaskedLM(FSTPreTrainedModel):
+    
+    accepts_loss_kwargs = False
+
+    def __init__(
+            self, 
+            config: FSTConfig
+    ):
+        super().__init__(config)
+        
+        assert not config.use_causal_attention, "FSTForMaskedLM requires use_causal_attention=False"
+        assert not config.use_cache, "FSTForMaskedLM requires use_cache=False (caching not supported for bidirectional models)"
+    
+        self.model = FSTModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        if config.tie_word_embeddings:
+            self.tie_weights()
+            self._dynamic_tied_weights_keys = {"lm_head.weight": "model.embedding.weight"} # Avoids safetensor naming issues
+
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embedding
+
+    def set_input_embeddings(self, new_embeddings):
+        self.model.embedding = new_embeddings
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def tie_weights(self):
+        self._tie_or_clone_weights(self.lm_head, self.get_input_embeddings())
+
+    def forward(
+        self,
+        input_ids: Tensor | None = None,
+        attention_mask: Tensor | None = None,
+        inputs_embeds: Tensor | None = None,
+        labels: Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs,
+    ):
+
+        if labels is not None:
+            return_dict = True
+        else:
+            return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+    
+        model_output = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            past_key_values=None,
+            use_cache=False,
+            output_hidden_states=output_hidden_states
+        )
+
+        logits = self.lm_head(model_output[0])
+
+        loss = None
+        if labels is not None:
+            
+            loss = F.cross_entropy(
+                logits.view(-1, logits.size(-1)),
+                labels.view(-1),
+                ignore_index=self.config.pad_token_id if self.config.pad_token_id is not None else -100
+            )
+
+        if not return_dict:
+            output = (logits,) + model_output[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return MaskedLMOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=model_output.hidden_states
+        )