Upload model.py with huggingface_hub

model.py

@@ -0,0 +1,450 @@
+"""
+HybridFourierLM – Model Architecture
+=====================================
+Self-contained module that registers the custom config + model with
+HuggingFace `transformers` so that `AutoConfig` / `AutoModelForCausalLM`
+can load checkpoints transparently.
+"""
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    GenerationMixin,
+    PretrainedConfig,
+    PreTrainedModel,
+)
+from transformers.modeling_outputs import CausalLMOutput
+
+
+# ──────────────────────────────────────────────────────────────────────
+# Config
+# ──────────────────────────────────────────────────────────────────────
+
+class HybridFourierConfig(PretrainedConfig):
+    model_type = "hybrid_fourier_lm"
+
+    def __init__(
+        self,
+        vocab_size=50304,
+        latent_dim=768,
+        num_layers=12,
+        num_modes=64,
+        layer_types=None,
+        time_scale=128.0,
+        dropout=0.05,
+        pad_token_id=0,
+        bos_token_id=1,
+        eos_token_id=2,
+        tie_word_embeddings=True,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.latent_dim = latent_dim
+        self.num_layers = num_layers
+        self.num_modes = num_modes
+        self.time_scale = time_scale
+        self.dropout = dropout
+
+        if layer_types is None:
+            layer_types = [
+                "softmax" if (i % 4 == 3) else "linear"
+                for i in range(num_layers)
+            ]
+        assert len(layer_types) == num_layers, (
+            f"layer_types length ({len(layer_types)}) must equal num_layers ({num_layers})"
+        )
+        self.layer_types = layer_types
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+# ──────────────────────────────────────────────────────────────────────
+# Mixer layers
+# ──────────────────────────────────────────────────────────────────────
+
+class LinearFourierMixer(nn.Module):
+    def __init__(self, channels, num_modes=64, num_heads=12, time_scale=128, dropout=0.05):
+        super().__init__()
+        assert channels % num_heads == 0, (
+            f"channels ({channels}) must be perfectly divisible by num_heads ({num_heads})"
+        )
+        self.channels = channels
+        self.num_modes = num_modes
+        self.num_heads = num_heads
+        self.head_dim = channels // num_heads
+        self.time_scale = time_scale
+
+        freq_bands = torch.exp(torch.linspace(math.log(0.0001), math.log(num_modes), num_modes))
+        self.num_modes = freq_bands.shape[0]
+        self.register_buffer("frequencies", freq_bands)
+
+        self.q_proj = nn.Linear(channels, self.num_heads * self.num_modes)
+        self.k_proj = nn.Linear(channels, self.num_heads * self.num_modes)
+        self.v_proj = nn.Linear(channels, channels)
+        self.proj_v2 = nn.Linear(channels, channels)
+        self.out_proj = nn.Linear(channels, channels)
+        self.activation = nn.SiLU()
+        self.norm_in = nn.LayerNorm(channels)
+        self.norm_out = nn.LayerNorm(channels)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, attention_mask=None):
+        B, seq_len, C = x.shape
+        norm_x = self.norm_in(x)
+
+        Q = F.elu(self.q_proj(norm_x)).view(B, seq_len, self.num_heads, self.num_modes) + 1.0
+        K = F.elu(self.k_proj(norm_x)).view(B, seq_len, self.num_heads, self.num_modes) + 1.0
+
+        v1 = self.v_proj(norm_x)
+        v2 = self.activation(self.proj_v2(norm_x))
+
+        t = (torch.arange(seq_len, device=x.device, dtype=x.dtype) / self.time_scale).view(-1, 1)
+        omega_t = 2 * math.pi * t * self.frequencies.unsqueeze(0)
+        U = torch.cos(omega_t).unsqueeze(0).unsqueeze(2)
+        V = torch.sin(omega_t).unsqueeze(0).unsqueeze(2)
+
+        Q_cos = Q * U
+        Q_sin = Q * V
+        K_cos = K * U
+        K_sin = K * V
+
+        Q_rot = torch.cat([Q_cos, Q_sin], dim=-1)
+        K_rot = torch.cat([K_cos, K_sin], dim=-1)
+
+        # Linear-attention normalization in fp32 for numerical stability
+        orig_dtype = Q_rot.dtype
+        Q_rot = Q_rot.float()
+        K_rot = K_rot.float()
+
+        if seq_len > 512:
+            v1_heads = v1.view(B, seq_len, self.num_heads, self.head_dim)
+            out_chunks = []
+            chunk_size = 256 if seq_len > 1024 else 512
+            for i_start in range(0, seq_len, chunk_size):
+                i_end = min(i_start + chunk_size, seq_len)
+                Q_chunk = Q_rot[:, i_start:i_end, :, :]       # [B, C, H, 2M]
+                K_past = K_rot[:, :i_end, :, :]               # [B, j_max, H, 2M]
+                v1_past = v1_heads[:, :i_end, :, :].float()   # [B, j_max, H, D]
+
+                A_chunk = torch.einsum('b i h m, b j h m -> b h i j', Q_chunk, K_past)
+                A_chunk = A_chunk / math.sqrt(self.num_modes * 2)
+
+                i_abs = torch.arange(i_start, i_end, device=x.device).view(-1, 1)
+                j_abs = torch.arange(i_end, device=x.device).view(1, -1)
+                causal_mask = (j_abs <= i_abs).to(dtype=A_chunk.dtype)
+                A_chunk = A_chunk * causal_mask.unsqueeze(0).unsqueeze(0)
+
+                if attention_mask is not None:
+                    pad_mask = attention_mask[:, None, None, :i_end].to(dtype=A_chunk.dtype)
+                    A_chunk = A_chunk * pad_mask
+
+                row_denom = torch.abs(A_chunk.sum(dim=-1, keepdim=True)) + 1.0
+                A_chunk = A_chunk / row_denom
+
+                v1_chunk = torch.einsum('b h i j, b j h d -> b i h d', A_chunk, v1_past)
+                out_chunks.append(v1_chunk.to(orig_dtype))
+            v1_token_mixed = torch.cat(out_chunks, dim=1).reshape(B, seq_len, C)
+            v1_token_mixed = self.dropout(v1_token_mixed)
+            if attention_mask is not None:
+                v1_token_mixed = torch.nan_to_num(v1_token_mixed, nan=0.0, posinf=0.0, neginf=0.0)
+                v1_token_mixed = v1_token_mixed * attention_mask.unsqueeze(-1).to(dtype=v1_token_mixed.dtype)
+        else:
+            A = torch.einsum('b i h m, b j h m -> b h i j', Q_rot, K_rot)
+            A = A / math.sqrt(self.num_modes * 2)
+
+            causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=x.device, dtype=A.dtype))
+            A = A * causal_mask.unsqueeze(0).unsqueeze(0)
+
+            if attention_mask is not None:
+                pad_mask = attention_mask[:, None, None, :].to(dtype=A.dtype)
+                A = A * pad_mask
+
+            row_denom = torch.abs(A.sum(dim=-1, keepdim=True)) + 1.0
+            A = A / row_denom
+            A = A.to(orig_dtype)
+
+            v1_heads = v1.view(B, seq_len, self.num_heads, self.head_dim)
+            v1_token_mixed = torch.einsum('b h i j, b j h d -> b i h d', A, v1_heads)
+            v1_token_mixed = v1_token_mixed.reshape(B, seq_len, C)
+            v1_token_mixed = self.dropout(v1_token_mixed)
+            if attention_mask is not None:
+                v1_token_mixed = torch.nan_to_num(v1_token_mixed, nan=0.0, posinf=0.0, neginf=0.0)
+                v1_token_mixed = v1_token_mixed * attention_mask.unsqueeze(-1).to(dtype=v1_token_mixed.dtype)
+
+        v3 = v1_token_mixed * v2
+        return self.norm_out(self.out_proj(v3)) + x
+
+
+class SoftmaxFourierMixer(nn.Module):
+    def __init__(self, channels, num_modes=64, num_heads=12, time_scale=128.0, dropout=0.05):
+        super().__init__()
+        assert channels % num_heads == 0, (
+            f"channels ({channels}) must be perfectly divisible by num_heads ({num_heads})"
+        )
+        self.channels = channels
+        self.num_modes = num_modes
+        self.num_heads = num_heads
+        self.head_dim = channels // num_heads
+        self.time_scale = time_scale
+
+        freq_bands = torch.exp(torch.linspace(math.log(0.0001), math.log(num_modes), num_modes))
+        self.num_modes = freq_bands.shape[0]
+        self.register_buffer("frequencies", freq_bands)
+
+        self.q_proj = nn.Linear(channels, self.num_heads * self.num_modes)
+        self.k_proj = nn.Linear(channels, self.num_heads * self.num_modes)
+        self.v_proj = nn.Linear(channels, channels)
+        self.proj_v2 = nn.Linear(channels, channels)
+        self.out_proj = nn.Linear(channels, channels)
+        self.activation = nn.SiLU()
+        self.norm_in = nn.LayerNorm(channels)
+        self.norm_out = nn.LayerNorm(channels)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, attention_mask=None):
+        B, seq_len, C = x.shape
+        norm_x = self.norm_in(x)
+
+        Q = self.q_proj(norm_x).view(B, seq_len, self.num_heads, self.num_modes)
+        K = self.k_proj(norm_x).view(B, seq_len, self.num_heads, self.num_modes)
+
+        v1 = self.v_proj(norm_x)
+        v2 = self.activation(self.proj_v2(norm_x))
+
+        t = (torch.arange(seq_len, device=x.device, dtype=x.dtype) / self.time_scale).view(-1, 1)
+        omega_t = 2 * math.pi * t * self.frequencies.unsqueeze(0)
+        U = torch.cos(omega_t).unsqueeze(0).unsqueeze(2)
+        V = torch.sin(omega_t).unsqueeze(0).unsqueeze(2)
+
+        Q_cos = Q * U
+        Q_sin = Q * V
+        K_cos = K * U
+        K_sin = K * V
+
+        Q_rot = torch.cat([Q_cos, Q_sin], dim=-1)
+        K_rot = torch.cat([K_cos, K_sin], dim=-1)
+
+        v1_heads = v1.view(B, seq_len, self.num_heads, self.head_dim)
+
+        # Actual softmax attention via SDPA (B, H, L, 2M)
+        Q_b = Q_rot.transpose(1, 2)
+        K_b = K_rot.transpose(1, 2)
+        V_b = v1_heads.transpose(1, 2)
+
+        # PyTorch's MPS backend has a known bug/crash in C++ kernel
+        # (`-[__NSPlaceholderDictionary initWithObjects:forKeys:count:]`)
+        # when calling F.scaled_dot_product_attention on certain shapes or when
+        # attn_mask and is_causal are combined. We use explicit math attention
+        # on MPS (or fallback if SDPA fails) to guarantee stability across all PyTorch versions.
+        if x.device.type == "mps" or seq_len > 512:
+            scale = 1.0 / math.sqrt(Q_b.size(-1))
+            if seq_len > 256:
+                out_chunks = []
+                chunk_size = 256 if seq_len > 1024 else 512
+                for i_start in range(0, seq_len, chunk_size):
+                    i_end = min(i_start + chunk_size, seq_len)
+                    Q_chunk = Q_b[:, :, i_start:i_end, :]         # [B, H, C, 2M]
+                    K_past = K_b[:, :, :i_end, :]                 # [B, H, 2M, j_max]
+                    V_past = V_b[:, :, :i_end, :]                 # [B, H, j_max, D]
+
+                    scores_chunk = torch.matmul(Q_chunk, K_past.transpose(-2, -1)) * scale
+
+                    i_abs = torch.arange(i_start, i_end, device=x.device).view(-1, 1)
+                    j_abs = torch.arange(i_end, device=x.device).view(1, -1)
+                    causal_mask = (j_abs <= i_abs)
+                    scores_chunk = scores_chunk.masked_fill(~causal_mask.unsqueeze(0).unsqueeze(0), float("-inf"))
+
+                    if attention_mask is not None:
+                        pad_mask = attention_mask[:, None, None, :i_end].to(dtype=torch.bool)
+                        scores_chunk = scores_chunk.masked_fill(~pad_mask, float("-inf"))
+
+                    attn_weights = F.softmax(scores_chunk, dim=-1)
+                    out_chunk = torch.matmul(attn_weights, V_past)  # [B, H, C, D]
+                    out_chunks.append(out_chunk)
+                v1_token_mixed = torch.cat(out_chunks, dim=2)
+            else:
+                scores = torch.matmul(Q_b, K_b.transpose(-2, -1)) * scale
+                causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=x.device, dtype=torch.bool))
+                scores = scores.masked_fill(~causal_mask.unsqueeze(0).unsqueeze(0), float("-inf"))
+                if attention_mask is not None:
+                    pad_mask = attention_mask[:, None, None, :].to(dtype=torch.bool)
+                    scores = scores.masked_fill(~pad_mask, float("-inf"))
+                attn_weights = F.softmax(scores, dim=-1)
+                v1_token_mixed = torch.matmul(attn_weights, V_b)
+        else:
+            attn_mask = None
+            if attention_mask is not None:
+                attn_mask = attention_mask[:, None, None, :].to(dtype=Q_b.dtype)
+                attn_mask = (1.0 - attn_mask) * torch.finfo(Q_b.dtype).min
+
+            try:
+                v1_token_mixed = F.scaled_dot_product_attention(
+                    Q_b, K_b, V_b,
+                    attn_mask=attn_mask,
+                    is_causal=True,
+                )
+            except Exception:
+                scale = 1.0 / math.sqrt(Q_b.size(-1))
+                scores = torch.matmul(Q_b, K_b.transpose(-2, -1)) * scale
+                causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=x.device, dtype=torch.bool))
+                scores = scores.masked_fill(~causal_mask.unsqueeze(0).unsqueeze(0), float("-inf"))
+                if attention_mask is not None:
+                    pad_mask = attention_mask[:, None, None, :].to(dtype=torch.bool)
+                    scores = scores.masked_fill(~pad_mask, float("-inf"))
+                attn_weights = F.softmax(scores, dim=-1)
+                v1_token_mixed = torch.matmul(attn_weights, V_b)
+
+        v1_token_mixed = v1_token_mixed.transpose(1, 2).reshape(B, seq_len, C)
+        v1_token_mixed = self.dropout(v1_token_mixed)
+        if attention_mask is not None:
+            v1_token_mixed = torch.nan_to_num(v1_token_mixed, nan=0.0, posinf=0.0, neginf=0.0)
+            v1_token_mixed = v1_token_mixed * attention_mask.unsqueeze(-1).to(dtype=v1_token_mixed.dtype)
+
+        v3 = v1_token_mixed * v2
+        return self.norm_out(self.out_proj(v3)) + x
+
+
+# ──────────────────────────────────────────────────────────────────────
+# Transformer block
+# ──────────────────────────────────────────────────────────────────────
+
+class HybridSpectralBlock(nn.Module):
+    def __init__(self, latent_dim, num_modes=64, is_softmax=False,
+                 time_scale=128.0, dropout=0.05, num_heads=None):
+        super().__init__()
+        self.is_softmax = is_softmax
+        num_heads = num_heads if num_heads is not None else max(1, latent_dim // 64)
+
+        if is_softmax:
+            self.mixer = SoftmaxFourierMixer(latent_dim, num_modes, num_heads, time_scale, dropout)
+        else:
+            self.mixer = LinearFourierMixer(latent_dim, num_modes, num_heads, time_scale, dropout)
+
+        self.ffn = nn.Sequential(
+            nn.LayerNorm(latent_dim),
+            nn.Linear(latent_dim, 4 * latent_dim),
+            nn.GELU(),
+            nn.Linear(4 * latent_dim, latent_dim),
+            nn.Dropout(dropout),
+        )
+        self.gradient_checkpointing = False
+
+    def forward(self, x, attention_mask=None):
+        z = self.mixer(x, attention_mask=attention_mask)
+        return z + self.ffn(z)
+
+
+# ──────────────────────────────────────────────────────────────────────
+# Full model
+# ──────────────────────────────────────────────────────────────────────
+
+class HybridFourierPreTrainedModel(PreTrainedModel):
+    config_class = HybridFourierConfig
+    base_model_prefix = "hybrid_fourier"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["HybridSpectralBlock"]
+    _tied_weights_keys = {"lm_head.weight": "embedding.weight"}
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        elif isinstance(module, nn.LayerNorm):
+            torch.nn.init.zeros_(module.bias)
+            torch.nn.init.ones_(module.weight)
+
+
+class HybridFourierLM(HybridFourierPreTrainedModel, GenerationMixin):
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+
+        self.embedding = nn.Embedding(config.vocab_size, config.latent_dim,
+                                       padding_idx=config.pad_token_id)
+
+        blocks = []
+        for layer_type in config.layer_types:
+            blocks.append(HybridSpectralBlock(
+                config.latent_dim,
+                config.num_modes,
+                is_softmax=(layer_type == "softmax"),
+                time_scale=config.time_scale,
+                dropout=config.dropout,
+            ))
+        self.mixers = nn.ModuleList(blocks)
+
+        self.ln_f = nn.LayerNorm(config.latent_dim)
+        self.lm_head = nn.Linear(config.latent_dim, config.vocab_size, bias=False)
+
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embedding
+
+    def set_input_embeddings(self, value):
+        self.embedding = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embedding):
+        self.lm_head = new_embedding
+
+    def forward(self, input_ids=None, attention_mask=None, labels=None,
+                past_key_values=None, use_cache=None, inputs_embeds=None, **kwargs):
+        if inputs_embeds is None:
+            z = self.embedding(input_ids)
+        else:
+            z = inputs_embeds
+
+        for mixer in self.mixers:
+            if getattr(self, "gradient_checkpointing", False) and self.training:
+                z = checkpoint(mixer, z, attention_mask, use_reentrant=False)
+            else:
+                z = mixer(z, attention_mask=attention_mask)
+
+        z = self.ln_f(z)
+        logits = self.lm_head(z)
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
+            loss = loss_fct(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+            )
+
+        return CausalLMOutput(loss=loss, logits=logits)
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None,
+                                       attention_mask=None, **kwargs):
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "use_cache": False,
+        }
+
+    def _reorder_cache(self, past_key_values, beam_idx):
+        return past_key_values
+
+
+# ── Register with AutoClasses ────────────────────────────────────────
+AutoConfig.register("hybrid_fourier_lm", HybridFourierConfig)
+AutoModelForCausalLM.register(HybridFourierConfig, HybridFourierLM)