Upload folder using huggingface_hub

model/__init__.py

@@ -0,0 +1,4 @@
+from .config import IvmeConfig
+from .transformer import IvmeConversateV2
+
+__all__ = ["IvmeConfig", "IvmeConversateV2"]

model/attention.py

@@ -0,0 +1,48 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .rope import apply_rope
+
+
+class CausalSelfAttention(nn.Module):
+    """Full multi-head causal self-attention (Section 4.4).
+
+    Deliberately NOT using Grouped Query Attention (GQA) — the doc is explicit
+    that at this scale, GQA's memory savings are negligible and it can quietly
+    cost quality. Every head gets its own independent K/V projections.
+    """
+
+    def __init__(self, hidden_dim: int, n_heads: int, dropout: float = 0.0):
+        super().__init__()
+        assert hidden_dim % n_heads == 0
+        self.n_heads = n_heads
+        self.head_dim = hidden_dim // n_heads
+        self.dropout = dropout
+
+        # separate q, k, v projections -- no sharing across heads (full attention)
+        self.q_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.k_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.v_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.out_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
+
+    def forward(self, x: torch.Tensor, rope_freqs: torch.Tensor) -> torch.Tensor:
+        B, T, C = x.shape
+
+        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+
+        q = apply_rope(q, rope_freqs[:T])
+        k = apply_rope(k, rope_freqs[:T])
+
+        # scaled dot-product attention with causal masking (built-in flash-attention
+        # kernel when running on a CUDA GPU; falls back to a math kernel on CPU)
+        out = F.scaled_dot_product_attention(
+            q, k, v,
+            is_causal=True,
+            dropout_p=self.dropout if self.training else 0.0,
+        )
+
+        out = out.transpose(1, 2).contiguous().view(B, T, C)
+        return self.out_proj(out)

model/config.py

@@ -0,0 +1,29 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class IvmeConfig:
+    """Ivme-Conversate-v2 (Dense) architecture config.
+
+    Every field here corresponds to a decision in Section 4 of the design doc.
+    Values are chosen to match v1 wherever the doc calls for it, so any quality
+    difference between v1 and v2 is attributable to data/training, not size.
+    """
+
+    vocab_size: int = 16_000       # Section 4.9: 16k tokens, English-only
+    hidden_dim: int = 384          # Section 4.2: matches v1
+    n_layers: int = 10             # Section 4.3: matches v1
+    n_heads: int = 6               # Section 4.4: full attention, no GQA
+    context_len: int = 1024        # Section 4.10: matches v1
+    ffn_mult: float = 4.0          # SwiGLU hidden expansion (adjusted below for param parity)
+    rope_theta: float = 10_000.0   # standard RoPE base frequency
+    norm_eps: float = 1e-5         # RMSNorm epsilon
+    tie_embeddings: bool = True    # Section 4.8
+    dropout: float = 0.0           # no dropout at this data:param ratio (heavily overtrained regime)
+
+    def __post_init__(self):
+        assert self.hidden_dim % self.n_heads == 0, "hidden_dim must be divisible by n_heads"
+
+    @property
+    def head_dim(self) -> int:
+        return self.hidden_dim // self.n_heads

model/feedforward.py

@@ -0,0 +1,27 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class SwiGLU(nn.Module):
+    """SwiGLU feed-forward block (Section 4.6), as used in Llama/PaLM.
+
+    Standard formulation: down_proj(silu(gate_proj(x)) * up_proj(x))
+    The inner dim is scaled down from the naive 4x so that SwiGLU's extra
+    gate_proj matrix doesn't blow the parameter budget relative to a plain MLP
+    of the same nominal "4x" size -- this matches how Llama-style models size it.
+    """
+
+    def __init__(self, hidden_dim: int, mult: float = 4.0):
+        super().__init__()
+        # standard correction: 4 * hidden * (2/3) keeps param count comparable
+        # to a plain (non-gated) 4x MLP, rounded to a clean multiple of 8.
+        inner_dim = int(hidden_dim * mult * 2 / 3)
+        inner_dim = ((inner_dim + 7) // 8) * 8
+
+        self.gate_proj = nn.Linear(hidden_dim, inner_dim, bias=False)
+        self.up_proj = nn.Linear(hidden_dim, inner_dim, bias=False)
+        self.down_proj = nn.Linear(inner_dim, hidden_dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))

model/rmsnorm.py

@@ -0,0 +1,23 @@
+import torch
+import torch.nn as nn
+
+
+class RMSNorm(nn.Module):
+    """RMSNorm (Section 4.7): cheaper alternative to LayerNorm.
+
+    Rescales by root-mean-square of the activations instead of full
+    mean/variance normalization. No bias, single learnable scale per dim.
+    """
+
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # compute in float32 for stability regardless of input dtype (bf16 etc.)
+        dtype = x.dtype
+        x = x.float()
+        rms = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+        out = x * rms
+        return (out.to(dtype)) * self.weight

model/rope.py

@@ -0,0 +1,28 @@
+import torch
+
+
+def precompute_rope_freqs(head_dim: int, max_seq_len: int, theta: float = 10_000.0):
+    """Precompute the rotation angles used by RoPE (Section 4.5).
+
+    Returns a complex tensor of shape (max_seq_len, head_dim // 2) where each
+    entry encodes the rotation to apply at that position/frequency pair.
+    """
+    assert head_dim % 2 == 0, "RoPE requires an even head_dim"
+    freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
+    positions = torch.arange(max_seq_len).float()
+    angles = torch.outer(positions, freqs)  # (seq_len, head_dim/2)
+    return torch.polar(torch.ones_like(angles), angles)  # complex64
+
+
+def apply_rope(x: torch.Tensor, rope_freqs: torch.Tensor) -> torch.Tensor:
+    """Apply rotary position embedding to a tensor of shape (B, n_heads, T, head_dim).
+
+    rope_freqs should be pre-sliced to the current sequence length T before
+    being passed in, i.e. rope_freqs[:T].
+    """
+    B, H, T, D = x.shape
+    x_complex = torch.view_as_complex(x.float().reshape(B, H, T, D // 2, 2))
+    freqs = rope_freqs.view(1, 1, T, D // 2)
+    x_rotated = x_complex * freqs
+    out = torch.view_as_real(x_rotated).reshape(B, H, T, D)
+    return out.type_as(x)

model/transformer.py

@@ -0,0 +1,90 @@
+import torch
+import torch.nn as nn
+
+from .config import IvmeConfig
+from .rmsnorm import RMSNorm
+from .rope import precompute_rope_freqs
+from .attention import CausalSelfAttention
+from .feedforward import SwiGLU
+
+
+class TransformerBlock(nn.Module):
+    """One dense transformer layer (Section 3): pre-norm attention + pre-norm SwiGLU,
+    with residual connections around each. Identical shape repeated n_layers times --
+    no loops, no weight sharing (Section 3.1, distinguishing this from the shelved
+    Ivmetron design).
+    """
+
+    def __init__(self, cfg: IvmeConfig):
+        super().__init__()
+        self.attn_norm = RMSNorm(cfg.hidden_dim, eps=cfg.norm_eps)
+        self.attn = CausalSelfAttention(cfg.hidden_dim, cfg.n_heads, cfg.dropout)
+        self.ffn_norm = RMSNorm(cfg.hidden_dim, eps=cfg.norm_eps)
+        self.ffn = SwiGLU(cfg.hidden_dim, cfg.ffn_mult)
+
+    def forward(self, x: torch.Tensor, rope_freqs: torch.Tensor) -> torch.Tensor:
+        x = x + self.attn(self.attn_norm(x), rope_freqs)
+        x = x + self.ffn(self.ffn_norm(x))
+        return x
+
+
+class IvmeConversateV2(nn.Module):
+    """Ivme-Conversate-v2 (Dense) -- the full model described in Section 4.
+
+    ~20M parameters, 10 layers, hidden_dim 384, 6 heads, RoPE, SwiGLU, RMSNorm,
+    tied embeddings, 16k vocab, 1024 context. See config.py for the exact spec.
+    """
+
+    def __init__(self, cfg: IvmeConfig):
+        super().__init__()
+        self.cfg = cfg
+
+        self.tok_embed = nn.Embedding(cfg.vocab_size, cfg.hidden_dim)
+        self.blocks = nn.ModuleList([TransformerBlock(cfg) for _ in range(cfg.n_layers)])
+        self.final_norm = RMSNorm(cfg.hidden_dim, eps=cfg.norm_eps)
+
+        # Section 4.8: tied embeddings -- output head reuses the input embedding
+        # table instead of learning a separate one.
+        self.lm_head = nn.Linear(cfg.hidden_dim, cfg.vocab_size, bias=False)
+        if cfg.tie_embeddings:
+            self.lm_head.weight = self.tok_embed.weight
+
+        rope_freqs = precompute_rope_freqs(cfg.head_dim, cfg.context_len, cfg.rope_theta)
+        self.register_buffer("rope_freqs", rope_freqs, persistent=False)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module: nn.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 forward(self, idx: torch.Tensor, targets: torch.Tensor | None = None):
+        B, T = idx.shape
+        assert T <= self.cfg.context_len, (
+            f"sequence length {T} exceeds context_len {self.cfg.context_len}"
+        )
+
+        x = self.tok_embed(idx)
+        for block in self.blocks:
+            x = block(x, self.rope_freqs)
+        x = self.final_norm(x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if targets is not None:
+            loss = nn.functional.cross_entropy(
+                logits.view(-1, logits.size(-1)),
+                targets.view(-1),
+                ignore_index=-1,
+            )
+        return logits, loss
+
+    def num_params(self, non_embedding: bool = False) -> int:
+        n = sum(p.numel() for p in self.parameters())
+        if non_embedding:
+            n -= self.tok_embed.weight.numel()
+        return n