Update architecture and tokenizer

README.md

@@ -1,3 +1,64 @@
 ---
+language:
+- en
 license: mit
+library_name: transformers
+pipeline_tag: text-generation
+tags:
+- pytorch
+- research
+- llama
 ---
+
+# advanced-transformers-lib -- Llama 3 Baseline
+
+A Llama 3-style decoder-only transformer architecture for research. No pretrained
+weights -- pull the architecture from the Hub and instantiate a freshly initialised
+model from config. Override any parameter at instantiation time.
+
+> **Important:** `trust_remote_code=True` is required. It downloads the architecture
+> source files from the Hub and imports them into your Python process. Review the
+> source at [smithblack-0/llama3_baseline_dev](https://huggingface.co/smithblack-0/llama3_baseline_dev) before use.
+
+## Usage
+
+```python
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+
+# Pull architecture config -- override any parameter at instantiation time
+config = AutoConfig.from_pretrained(
+    "smithblack-0/llama3_baseline_dev",
+    trust_remote_code=True,
+    num_hidden_layers=16,  # example override
+)
+
+# Instantiate with fresh random weights -- no checkpoint required
+model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)
+
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained("smithblack-0/llama3_baseline_dev")
+
+# Save and reload after training
+model.save_pretrained("./checkpoint")
+model = AutoModelForCausalLM.from_pretrained("./checkpoint", trust_remote_code=True)
+```
+
+## Default Configuration
+
+| Parameter | Default |
+|-----------|---------|
+| `vocab_size` | 50277 |
+| `hidden_size` | 768 |
+| `intermediate_size` | 1568 |
+| `num_hidden_layers` | 24 |
+| `num_attention_heads` | 16 |
+| `num_key_value_heads` | 4 |
+| `head_dim` | 48 |
+| `max_position_embeddings` | 8192 |
+| `rope_theta` | 500000.0 |
+
+## License
+
+MIT. Clean-room synthesis: the human author has not read the Llama source code.
+Architectural decisions derive from the published paper. Tokenizer is GPT-NeoX
+(`EleutherAI/gpt-neox-20b`, Apache 2.0).

__init__.py

@@ -0,0 +1,9 @@
+from .attention import GroupedQueryAttention
+from .configuration import Llama3Config
+from .decoder_layer import DecoderLayer
+from .huggingface import Llama3ForCausalLM
+from .mlp import SwiGLUMLP
+from .model import Llama3Model
+from .rope import RotaryEmbedding
+
+__all__ = ["DecoderLayer", "GroupedQueryAttention", "Llama3Config", "Llama3ForCausalLM", "Llama3Model", "RotaryEmbedding", "SwiGLUMLP"]

attention.py

@@ -0,0 +1,149 @@
+"""Grouped Query Attention (GQA).
+
+GQA reduces KV cache memory by sharing key-value heads across groups of query heads.
+With G query heads per KV head, the KV cache is G× smaller than standard multi-head
+attention (MHA). This is the primary motivation for its use in Llama 3 at 128K context:
+8 KV heads shared across 32 query heads gives a 4× cache reduction.
+
+Setting num_key_value_heads == num_attention_heads recovers standard MHA.
+Setting num_key_value_heads == 1 gives multi-query attention (MQA).
+
+Attention is computed via torch.nn.functional.scaled_dot_product_attention (SDPA),
+which selects FlashAttention when hardware and dtype allow, falling back to standard
+attention otherwise. No custom kernel or additional dependency required.
+
+KV caching is handled via HuggingFace's Cache protocol. The cache owns K/V storage and
+accumulation; attention only calls cache.update() to store new projections and retrieve
+the full accumulated history. This cleanly separates attention computation from cache
+management: different Cache subclasses (DynamicCache, StaticCache, custom research
+variants) can be dropped in without touching the attention logic.
+
+No bias on any projection — a fixed architectural constant of this model.
+"""
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PretrainedConfig
+from transformers.cache_utils import Cache
+
+from .rope import RotaryEmbedding
+
+
+class GroupedQueryAttention(nn.Module):
+    """Grouped Query Attention with RoPE, causal masking, and KV cache support.
+
+    Implements GQA as used in Llama 3: Q heads are split into groups, each group
+    sharing a single KV head. Before attention is computed, K and V are expanded
+    by repeating each KV head across its group of query heads.
+
+    The forward pass is strictly causal. An optional pre-built boolean attention
+    mask can be threaded in from the caller; when absent, SDPA's native
+    ``is_causal`` mode applies — correct for full-sequence training.
+
+    Args:
+        config: Model config. Must expose ``num_attention_heads``,
+            ``num_key_value_heads``, ``head_dim``, ``hidden_size``,
+            and ``attention_dropout``.
+
+    Raises:
+        ValueError: If ``num_attention_heads`` is not divisible by
+            ``num_key_value_heads``.
+    """
+
+    def __init__(self, config: PretrainedConfig) -> None:
+        super().__init__()
+
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.head_dim = config.head_dim
+        self.num_groups = self.num_heads // self.num_kv_heads
+        self.attention_dropout = config.attention_dropout
+
+        if self.num_heads % self.num_kv_heads != 0:
+            raise ValueError(
+                f"num_attention_heads ({self.num_heads}) must be divisible by "
+                f"num_key_value_heads ({self.num_kv_heads})."
+            )
+
+        # No bias on any projection — architectural constant.
+        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)
+
+        self.rope = RotaryEmbedding(config)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        position_ids: torch.Tensor,
+        cache: Cache | None = None,
+        layer_idx: int = 0,
+        causal_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Apply grouped query attention to the input.
+
+        Args:
+            x: Input of shape (batch, seq_len, hidden_size).
+            position_ids: Absolute positions of shape (batch, seq_len). Used by
+                RoPE to rotate Q and K at the correct frequencies.
+            cache: HuggingFace Cache object for KV accumulation, or None when
+                caching is disabled (``use_cache=False``). When provided,
+                ``cache.update(k, v, layer_idx)`` stores the new K/V and returns
+                the full accumulated key and value tensors for this layer.
+            layer_idx: Which slot in the cache to read and write. Each decoder
+                layer has its own index so they accumulate independently.
+            causal_mask: Optional boolean attention mask of shape
+                (1, 1, seq_len, kv_len), where True indicates a position that
+                should be attended to. When None, SDPA's built-in ``is_causal``
+                mode is used, which is correct for full-sequence training
+                (square Q×K matrix). When provided, ``is_causal`` is disabled
+                and the explicit mask governs attention — required for any
+                generation pattern where Q and K lengths differ.
+
+        Returns:
+            Output tensor of shape (batch, seq_len, hidden_size).
+        """
+        batch, seq_len, _ = x.shape
+
+        # Project and reshape: (batch, seq_len, heads * head_dim)
+        #                    → (batch, heads, seq_len, head_dim)
+        q = self.q_proj(x).view(batch, seq_len, self.num_heads,    self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(batch, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(batch, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+
+        # Apply RoPE. attention_scaling is 1.0 for default/linear; YaRN returns a
+        # value != 1.0 that corrects attention magnitude after frequency manipulation.
+        q, k, attention_scaling = self.rope(q, k, position_ids)
+
+        if cache is not None:
+            k_full, v_full = cache.update(k, v, layer_idx)
+        else:
+            k_full, v_full = k, v
+
+        # Expand KV heads to align with query heads for GQA.
+        # Each KV head is repeated num_groups times so SDPA sees matching head counts.
+        if self.num_groups > 1:
+            k_full = k_full.repeat_interleave(self.num_groups, dim=1)
+            v_full = v_full.repeat_interleave(self.num_groups, dim=1)
+
+        attn_output = F.scaled_dot_product_attention(
+            q, k_full, v_full,
+            attn_mask=causal_mask,
+            dropout_p=self.attention_dropout if self.training else 0.0,
+            is_causal=causal_mask is None,
+            scale=attention_scaling / math.sqrt(self.head_dim),
+        )
+
+        # Merge heads and project back to hidden_size.
+        attn_output = (
+            attn_output
+            .transpose(1, 2)
+            .contiguous()
+            .view(batch, seq_len, self.num_heads * self.head_dim)
+        )
+
+        return self.o_proj(attn_output)

config.json

@@ -0,0 +1,22 @@
+{
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration.Llama3Config",
+    "AutoModelForCausalLM": "huggingface.Llama3ForCausalLM"
+  },
+  "head_dim": 48,
+  "hidden_size": 768,
+  "intermediate_size": 1568,
+  "max_position_embeddings": 8192,
+  "model_type": "llama3_baseline",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 24,
+  "num_key_value_heads": 4,
+  "rms_norm_eps": 1e-05,
+  "rope_parameters": null,
+  "rope_theta": 500000.0,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.3.0",
+  "use_cache": true,
+  "vocab_size": 50277
+}

configuration.py

@@ -0,0 +1,177 @@
+"""Configuration for the Llama 3 baseline transformer.
+
+All architectural parameters that vary across model scales or are meaningful research
+variables are expressed here. Architectural constants of Llama 3 (no bias in linear
+layers, SwiGLU activation with SiLU gate) are implemented in the relevant modules and
+documented at the point of use — they are not config parameters because they do not
+vary across Llama 3 scales and changing them produces a different architecture, not a
+different scale of this one.
+
+RoPE configuration is handled by HuggingFace's RotaryEmbeddingConfigMixin (mixed into
+PretrainedConfig in transformers 5.x). rope_theta and rope_scaling are passed through
+to the base class, which validates and standardises them into config.rope_parameters.
+Do not bypass or duplicate this system.
+"""
+
+from transformers import PretrainedConfig
+
+
+class Llama3Config(PretrainedConfig):
+    """Configuration class for the Llama 3 baseline decoder-only transformer.
+
+    This config is the single source of truth for every architectural dimension of the
+    model. Nothing in the architecture may use a literal number that belongs here —
+    doing so breaks the library's ability to express different model scales without
+    code changes.
+
+    RoPE scaling is handled by HuggingFace's rope system. Pass rope_scaling as a dict
+    using HF's format (key is ``rope_type``, not ``type``). Supported types:
+    ``"linear"``, ``"dynamic"``, ``"yarn"``, ``"longrope"``, ``"llama3"``. HF validates
+    the dict and standardises it into ``config.rope_parameters``.
+
+    Registered with HuggingFace AutoClass via ``auto_map``. Instantiate from the Hub::
+
+        config = AutoConfig.from_pretrained(
+            "your-namespace/advanced-transformers-lib",
+            trust_remote_code=True,
+            num_hidden_layers=16,  # override any parameter at instantiation time
+        )
+        model = AutoModelForCausalLM.from_config(config)
+
+    Args:
+        vocab_size: Vocabulary size. Controls the embedding table and output logits
+            dimension. Must match the tokenizer.
+        hidden_size: Model width. The central dimension from which all others are
+            derived or to which they project.
+        intermediate_size: FFN hidden dimension. Expressed directly rather than derived
+            from a formula because Llama 3 ratios vary by scale (~3.5x at 8B/70B,
+            ~3.25x at 405B). A formula would be wrong for at least some scales.
+        num_hidden_layers: Number of transformer blocks stacked in sequence.
+        num_attention_heads: Number of query heads. Determines how hidden_size is
+            partitioned per head.
+        num_key_value_heads: Number of KV heads for Grouped Query Attention. Must
+            evenly divide num_attention_heads. Setting equal to num_attention_heads
+            gives standard MHA; setting to 1 gives MQA; values between give GQA.
+            Llama 3 uses 8 at all scales, motivated by KV cache memory at 128K context.
+        head_dim: Dimension per attention head. Normally hidden_size //
+            num_attention_heads, but exposed as a parameter for architectures that
+            decouple head count from head size. Computed automatically if None.
+        rms_norm_eps: Epsilon passed to torch.nn.RMSNorm. Prevents division by zero
+            when layer activations are near zero.
+        rope_theta: Base rotation frequency for RoPE. Controls how fast position angles
+            rotate per dimension — higher values mean slower rotation, preventing
+            positional aliasing at long sequence distances. Llama 3 uses 500,000
+            (vs ~10,000 typical) as a prerequisite for 128K context support. This
+            value has physical meaning tied to the target context length and must
+            never be hardcoded in the architecture.
+        max_position_embeddings: The context length the model was trained at. Used by
+            HF's rope system as original_max_position_embeddings for scaling types that
+            need it (yarn, longrope, llama3). This is the training context length, not
+            an inference ceiling — the rope module handles longer sequences at runtime
+            via lazy cache extension. Llama 3 base training context: 8192.
+        rope_scaling: Optional RoPE scaling configuration for extending context beyond
+            max_position_embeddings without retraining. Pass as a dict in HF's format
+            with ``rope_type`` as the key. HF's RotaryEmbeddingConfigMixin validates
+            and stores this. None means no scaling (default RoPE behaviour).
+        attention_dropout: Dropout probability applied to attention weights. Default
+            0.0 for deterministic behaviour.
+        use_cache: Whether the model returns past_key_values for KV caching. Set True
+            for inference, may be set False during training to reduce memory pressure.
+        output_hidden_states: Whether the model returns the hidden state tensor after
+            each decoder layer. Useful for probing or intermediate representation
+            extraction. Default False.
+        tie_word_embeddings: Whether the input embedding table and the LM head share
+            weights. False for Llama 3.
+    """
+
+    model_type = "llama3_baseline"
+
+    # auto_map tells HuggingFace which classes to instantiate when loading this config
+    # with trust_remote_code=True. Paths are relative to the Hub repository root, not
+    # the local src/ layout — these are the paths used after HF downloads the files.
+    auto_map = {
+        "AutoConfig": "configuration.Llama3Config",
+        "AutoModelForCausalLM": "huggingface.Llama3ForCausalLM",
+    }
+
+    def __init__(
+        self,
+        vocab_size: int = 50277,
+        hidden_size: int = 768,
+        intermediate_size: int = 1568,
+        num_hidden_layers: int = 24,
+        num_attention_heads: int = 16,
+        num_key_value_heads: int = 4,
+        head_dim: int | None = None,
+        rms_norm_eps: float = 1e-5,
+        rope_theta: float = 500000.0,
+        max_position_embeddings: int = 8192,
+        rope_scaling: dict | None = None,
+        attention_dropout: float = 0.0,
+        use_cache: bool = True,
+        output_hidden_states: bool = False,
+        tie_word_embeddings: bool = False,
+        **kwargs,
+    ):
+        # Validate structural constraints before storing anything, so that an invalid
+        # config fails loudly at construction rather than silently producing wrong
+        # shapes at forward-pass time.
+        if hidden_size % num_attention_heads != 0:
+            raise ValueError(
+                f"hidden_size ({hidden_size}) must be divisible by "
+                f"num_attention_heads ({num_attention_heads})."
+            )
+        if num_attention_heads % num_key_value_heads != 0:
+            raise ValueError(
+                f"num_attention_heads ({num_attention_heads}) must be divisible by "
+                f"num_key_value_heads ({num_key_value_heads}). GQA requires query "
+                f"heads to divide evenly across KV head groups."
+            )
+
+        # RoPE rotates dimensions in pairs. An odd head_dim has no valid pairing and
+        # produces a cos/sin cache of size head_dim+1 (torch.arange(0, odd, 2) rounds
+        # up), causing a shape mismatch at runtime. Catch it here rather than at
+        # forward-pass time.
+        resolved_head_dim = head_dim if head_dim is not None else hidden_size // num_attention_heads
+        if resolved_head_dim % 2 != 0:
+            raise ValueError(
+                f"head_dim must be even (RoPE rotates dimensions in pairs). "
+                f"Got head_dim={resolved_head_dim} from hidden_size={hidden_size} "
+                f"and num_attention_heads={num_attention_heads}."
+            )
+
+        # head_dim is normally hidden_size // num_attention_heads but is exposed as a
+        # parameter for architectures that decouple head count from head size.
+        if head_dim is None:
+            head_dim = hidden_size // num_attention_heads
+
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.rms_norm_eps = rms_norm_eps
+        self.attention_dropout = attention_dropout
+        self.use_cache = use_cache
+
+        # rope_theta, max_position_embeddings, and rope_scaling are passed to HF's
+        # base class, which owns rope configuration via RotaryEmbeddingConfigMixin.
+        # HF validates rope_scaling and standardises everything into rope_parameters.
+        # Do not store or validate these ourselves.
+        super().__init__(
+            rope_theta=rope_theta,
+            max_position_embeddings=max_position_embeddings,
+            rope_scaling=rope_scaling,
+            tie_word_embeddings=tie_word_embeddings,
+            output_hidden_states=output_hidden_states,
+            **kwargs,
+        )
+
+        # Promote auto_map to an instance attribute so PretrainedConfig.to_dict()
+        # serialises it into config.json. Class-level attributes are not picked up
+        # by to_dict() — only self.__dict__ is serialised. model_type is the sole
+        # exception handled specially by HF; auto_map is not.
+        self.auto_map = type(self).auto_map

decoder_layer.py

@@ -0,0 +1,75 @@
+"""Decoder layer — a single transformer block.
+
+Each block applies pre-norm attention followed by pre-norm MLP, with residual
+connections around both sublayers:
+
+    normed = RMSNorm(x)
+    h      = x + Attention(normed, ...)
+    normed = RMSNorm(h)
+    out    = h + MLP(normed)
+
+Pre-norm keeps the residual stream unnormalised. Gradients flow more cleanly
+through unnormalised residuals at depth, and each sublayer receives a stable,
+normalised view of the signal.
+
+Two independent RMSNorm instances are used — one before attention, one before MLP.
+They learn different scalings because they precede layers with different dynamic
+ranges. Sharing them would be wrong.
+
+torch.nn.RMSNorm is used directly (available from PyTorch 2.4+). It omits mean
+subtraction, is faster than LayerNorm, and proved more stable at scale.
+"""
+
+import torch
+import torch.nn as nn
+from transformers import PretrainedConfig
+from transformers.cache_utils import Cache
+
+from .attention import GroupedQueryAttention
+from .mlp import SwiGLUMLP
+
+
+class DecoderLayer(nn.Module):
+    """A single pre-norm transformer decoder block.
+
+    Composes GroupedQueryAttention and SwiGLUMLP with residual connections and
+    independent RMSNorm instances on each sublayer input.
+
+    Args:
+        config: Model config passed through to attention and MLP. Must also expose
+            ``hidden_size`` and ``rms_norm_eps``.
+    """
+
+    def __init__(self, config: PretrainedConfig) -> None:
+        super().__init__()
+        self.attn_norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.mlp_norm  = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.attention = GroupedQueryAttention(config)
+        self.mlp       = SwiGLUMLP(config)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        position_ids: torch.Tensor,
+        cache: Cache | None = None,
+        layer_idx: int = 0,
+        causal_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Apply one decoder block to the input.
+
+        Args:
+            x: Input of shape (batch, seq_len, hidden_size).
+            position_ids: Absolute positions of shape (batch, seq_len).
+            cache: HuggingFace Cache object for KV accumulation, or None when
+                caching is disabled. Passed through to attention unchanged.
+            layer_idx: Cache slot index for this layer. Each layer has its own
+                index so they accumulate independently within the shared cache.
+            causal_mask: Optional boolean attention mask of shape
+                (1, 1, seq_len, kv_len). Passed through to attention unchanged.
+
+        Returns:
+            Output tensor of shape (batch, seq_len, hidden_size).
+        """
+        attn_out = self.attention(self.attn_norm(x), position_ids, cache, layer_idx, causal_mask)
+        h = x + attn_out
+        return h + self.mlp(self.mlp_norm(h))

huggingface.py

@@ -0,0 +1,248 @@
+"""HuggingFace wrapper for the Llama 3 baseline.
+
+Llama3ForCausalLM wraps Llama3Model with everything a researcher needs to
+train, evaluate, and generate from it through the HuggingFace ecosystem:
+token embedding, vocabulary projection, next-token loss, weight tying, and
+the full AutoClass and GenerationMixin contracts.
+
+The token embedding lives here, not on the backbone. Llama3Model is a pure
+transformer stack that accepts pre-embedded hidden states — it has no knowledge
+of tokens or vocabulary. This is the correct HF convention: the backbone is
+modality-agnostic; the token interface belongs on the task wrapper.
+
+The LM head projects the backbone's (batch, seq, hidden_size) output to
+(batch, seq, vocab_size) logits. When labels are provided, cross-entropy loss
+is computed with a one-position shift: token i predicts token i+1. The shift
+is applied here rather than expected from the caller — a causal LM always
+trains this way and there is no use case for an unshifted loss.
+
+Weight tying: when config.tie_word_embeddings is True, lm_head.weight is
+directly assigned to embed_tokens.weight after post_init(). Both matrices are
+shape (vocab_size, hidden_size) — same shape, no transpose needed.
+
+KV caching uses HuggingFace's Cache protocol. GenerationMixin creates and
+manages the DynamicCache for generate() calls, passing it as past_key_values
+on every forward call. The backbone updates the cache in place and returns the
+same object. _reorder_cache delegates to DynamicCache.reorder_cache() for beam
+search, keeping all beam-reordering logic inside the cache implementation.
+
+Returns a CausalLMOutputWithPast. ModelOutput subclasses support both attribute
+access (output.logits) and dict-style access (output["logits"]), satisfying
+GenerationMixin's attribute access requirements while keeping existing code unchanged.
+"""
+
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel, GenerationMixin
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .configuration import Llama3Config
+from .model import Llama3Model
+
+
+class Llama3ForCausalLM(PreTrainedModel, GenerationMixin):
+    """Llama 3 causal language model: token embedding, backbone, LM head, HF contract.
+
+    Owns the token embedding and LM head. Delegates all transformer computation
+    to Llama3Model. Adds loss computation for training, weight tying between the
+    LM head and the input embedding, and the full HuggingFace AutoClass and
+    GenerationMixin contracts.
+
+    Args:
+        config: Model configuration. Must be a ``Llama3Config`` instance.
+    """
+
+    config_class = Llama3Config
+    base_model_prefix = "model"
+    _no_split_modules = ["DecoderLayer"]
+    supports_gradient_checkpointing = True
+
+    def __init__(self, config: Llama3Config) -> None:
+        super().__init__(config)
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.model = Llama3Model(config)
+
+        # No bias — consistent with all other projections in this architecture.
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+        # Direct weight tying: both matrices are (vocab_size, hidden_size) — same shape,
+        # no transpose. Explicit here for visibility; post_init() → tie_weights() also
+        # performs this via get_input/output_embeddings(), but that is less readable.
+        if config.tie_word_embeddings:
+            self.lm_head.weight = self.embed_tokens.weight
+
+    def _init_weights(self, module: nn.Module) -> None:
+        # Suppress HF's default reinitialisation pass. HF's _init_weights overwrites
+        # all Linear and Embedding weights with normal(0, 0.02) after construction,
+        # silently replacing PyTorch's own defaults (kaiming_uniform_ for Linear,
+        # normal(0,1) for Embedding). PyTorch's reset_parameters() already ran at
+        # construction time and those initialisations should stand.
+        pass
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        """Return the token embedding matrix. Required by PreTrainedModel for weight tying and resize_token_embeddings."""
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        """Replace the token embedding matrix. Required by PreTrainedModel for weight tying and resize_token_embeddings."""
+        self.embed_tokens = value
+
+    def get_output_embeddings(self) -> nn.Linear:
+        """Return the LM head. Required by PreTrainedModel for weight tying and resize_token_embeddings."""
+        return self.lm_head
+
+    def set_output_embeddings(self, value: nn.Linear) -> None:
+        """Replace the LM head. Required by PreTrainedModel for weight tying and resize_token_embeddings."""
+        self.lm_head = value
+
+    def _reorder_cache(
+        self, past_key_values: Cache, beam_idx: torch.Tensor
+    ) -> Cache:
+        """Reorder the KV cache to match beam reordering during beam search.
+
+        GenerationMixin calls this after pruning and reordering beams at each
+        step. beam_idx[i] is the old batch position whose cache should move to
+        position i. DynamicCache.reorder_cache() handles the index-select on
+        every stored tensor's batch dimension, keeping the cache consistent with
+        the reordered beam hypotheses.
+
+        Args:
+            past_key_values: The active Cache object.
+            beam_idx: 1-D tensor of shape (batch * num_beams,) mapping new batch
+                positions to old ones.
+
+        Returns:
+            The same Cache object, reordered in place.
+        """
+        past_key_values.reorder_cache(beam_idx)
+        return past_key_values
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        position_ids: torch.Tensor | None = None,
+        past_key_values: Cache | None = None,
+        use_cache: bool | None = None,
+        output_hidden_states: bool | None = None,
+        labels: torch.Tensor | None = None,
+        cache_position: torch.Tensor | None = None,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        """Run the causal language model.
+
+        Args:
+            input_ids: Token indices of shape (batch, seq_len).
+            position_ids: Absolute positions of shape (batch, seq_len). Passed
+                through to the backbone. When use_cache=True and this is None,
+                derived from cache_position.
+            past_key_values: A HuggingFace Cache object from a prior step, or
+                None. When use_cache=True and this is None, a fresh DynamicCache
+                is created here before calling the backbone.
+            use_cache: Whether to accumulate and return a KV cache. When True
+                and no cache is provided, a DynamicCache is created. When False,
+                None is passed to the backbone regardless of what was provided.
+                Defaults to config.use_cache when None.
+            output_hidden_states: Whether to return per-layer hidden states.
+                Passed through to the backbone.
+            labels: Target token indices of shape (batch, seq_len) for computing
+                next-token prediction loss. The loss is computed over positions
+                1..seq_len predicting from positions 0..seq_len-1 — the shift
+                is applied internally. Positions with label value -100 are
+                ignored by cross-entropy, following the HuggingFace convention
+                for padding and masked positions.
+            cache_position: 1-D integer tensor of shape (seq_len,) giving the
+                absolute position of each input token in the full sequence.
+                Provided by GenerationMixin during generate(). When use_cache=True
+                and this is None, it is derived from the current cache length.
+            **kwargs: Additional keyword arguments passed by GenerationMixin
+                (e.g. return_dict). Accepted and ignored for forward compatibility.
+                We always return CausalLMOutputWithPast regardless of return_dict.
+
+        Returns:
+            CausalLMOutputWithPast with fields:
+            - ``logits``: vocabulary scores of shape (batch, seq_len, vocab_size).
+              Always present.
+            - ``loss``: scalar cross-entropy loss, or None if labels not provided.
+            - ``past_key_values``: the updated Cache object, or None.
+            - ``hidden_states``: per-layer hidden states, or None.
+        """
+        if kwargs.get("attention_mask") is not None:
+            raise ValueError(
+                "attention_mask is not supported. This model does not support padding masks. "
+                "For training on variable-length sequences, use right-padding with -100 labels."
+            )
+
+        # Resolve both flags against config defaults. Config sets the default;
+        # per-call arguments override it. Both fields in Llama3Config remain live.
+        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
+        )
+
+        # Cache lifecycle is owned here — the backbone only receives a cache or None
+        # and never decides whether to create one.
+        if use_cache:
+            if past_key_values is None:
+                past_key_values = DynamicCache()
+        else:
+            past_key_values = None
+
+        inputs_embeds = self.embed_tokens(input_ids)
+        batch, seq_len, _ = inputs_embeds.shape
+
+        # For training (use_cache=False), positions are always 0..seq_len-1.
+        # This is not inference from state — it is a trivial fact about a
+        # non-cached forward pass. The backbone requires explicit position_ids.
+        if not use_cache and position_ids is None:
+            position_ids = torch.arange(seq_len, device=inputs_embeds.device).unsqueeze(0).expand(batch, -1)
+
+        causal_mask = None
+        if use_cache:
+            # Derive absolute query positions from cache occupancy. cache_len is the
+            # number of tokens already stored; query positions are cache_len..cache_len+seq_len-1.
+            # cache_position is accepted in the signature for HuggingFace contract
+            # compatibility but is not used — get_seq_length() is the authoritative source.
+            cache_len = past_key_values.get_seq_length()
+            q_positions = torch.arange(cache_len, cache_len + seq_len, device=inputs_embeds.device)
+
+            if position_ids is None:
+                position_ids = q_positions.unsqueeze(0).expand(batch, -1)
+
+            # Build the causal attention mask. Each query at position p may attend to
+            # all keys at positions 0..p. k_len covers the full sequence after this step.
+            k_positions = torch.arange(cache_len + seq_len, device=inputs_embeds.device)
+            # mask[q, k] = True when key position k is within the causal horizon of query q.
+            # Shape: (1, 1, seq_len, k_len) — broadcast over batch and head dimensions.
+            causal_mask = (k_positions[None, :] <= q_positions[:, None]).unsqueeze(0).unsqueeze(0)
+
+        backbone_out = self.model(
+            inputs_embeds,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            output_hidden_states=output_hidden_states,
+            causal_mask=causal_mask,
+        )
+
+        logits = self.lm_head(backbone_out["last_hidden_state"])
+
+        loss = None
+        if labels is not None:
+            # Shift so that each position predicts the next token. The final
+            # logit has no target; the first label has no corresponding input.
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            loss = nn.functional.cross_entropy(
+                shift_logits.view(-1, self.config.vocab_size),
+                shift_labels.view(-1),
+            )
+
+        return CausalLMOutputWithPast(
+            logits=logits,
+            loss=loss,
+            past_key_values=backbone_out["past_key_values"],
+            hidden_states=backbone_out["hidden_states"],
+        )

mlp.py

@@ -0,0 +1,52 @@
+"""SwiGLU feed-forward sublayer.
+
+SwiGLU is a gated linear unit variant that multiplies a SiLU-gated projection
+element-wise against a separate up-projection:
+
+    output = W_down(SiLU(W_gate(x)) ⊙ W_up(x))
+
+The gating mechanism gives the network more expressive control over which features
+to propagate than a plain two-matrix FFN. It requires three weight matrices instead
+of two, which is why intermediate_size in Llama 3 is set lower than the 4× multiplier
+typical of two-matrix FFNs — the total parameter count remains comparable.
+
+SiLU is used as the gate activation because Llama 3 committed to SwiGLU specifically
+— a fixed architectural choice.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PretrainedConfig
+
+
+class SwiGLUMLP(nn.Module):
+    """SwiGLU feed-forward sublayer.
+
+    Implements the three-matrix SwiGLU FFN used in Llama 3:
+
+        output = W_down(SiLU(W_gate(x)) ⊙ W_up(x))
+
+    No bias on any projection. SiLU as the gate activation is an architectural
+    constant — it is what defines SwiGLU specifically.
+
+    Args:
+        config: Model config. Must expose ``hidden_size`` and ``intermediate_size``.
+    """
+
+    def __init__(self, config: PretrainedConfig) -> None:
+        super().__init__()
+        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.up_proj   = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply the SwiGLU feed-forward transformation.
+
+        Args:
+            x: Input tensor of shape (batch, seq_len, hidden_size).
+
+        Returns:
+            Output tensor of shape (batch, seq_len, hidden_size).
+        """
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))

model.py

@@ -0,0 +1,118 @@
+"""Transformer backbone for the Llama 3 baseline.
+
+Llama3Model is a pure PyTorch module: a sequence of DecoderLayer blocks followed
+by a final RMSNorm. It accepts pre-embedded hidden states and returns contextual
+representations. It has no knowledge of tokens, vocabulary, generation, or the
+HuggingFace contract — those concerns belong on Llama3ForCausalLM.
+
+Keeping the embedding out of the backbone is the correct HF convention and makes
+the backbone genuinely modality-agnostic. The token interface — embedding lookup,
+LM head, weight tying — belongs on the task wrapper (Llama3ForCausalLM), which is
+the only class that knows this backbone is being used for language modelling.
+
+The final RMSNorm is necessary because the decoder stack uses pre-norm throughout:
+each sublayer normalises its own input, leaving the residual stream itself
+unnormalised. After many layers of accumulated residuals, that stream arrives at the
+top with uncontrolled magnitude. The final norm brings it to a well-scaled state
+before any projection. Without it, the LM head would receive signals of arbitrary
+scale.
+
+KV caching is caller-managed. If a Cache object is provided as past_key_values, it
+is threaded through every decoder layer (each layer writes to its own slot via
+layer_idx) and returned in the output dict. If None is provided, no caching occurs.
+The decision of whether to create a cache and when belongs to the caller.
+
+Returns a plain dict with keys:
+- "last_hidden_state": normed backbone output, shape (batch, seq_len, hidden_size)
+- "past_key_values": the Cache object passed in (updated in place), or None
+- "hidden_states": tuple of per-layer activations if output_hidden_states=True, else None
+"""
+
+import torch
+import torch.nn as nn
+from transformers.cache_utils import Cache
+
+from .configuration import Llama3Config
+from .decoder_layer import DecoderLayer
+
+
+class Llama3Model(nn.Module):
+    """Pure transformer backbone: decoder stack and final normalisation.
+
+    Accepts pre-embedded hidden states of shape (batch, seq_len, hidden_size)
+    and returns contextual representations of the same shape. No token embedding,
+    vocabulary projection, or HuggingFace lifecycle concerns.
+
+    RoPE is applied inside each attention layer. Positional information is
+    encoded in the relationship between Q and K, not added to the residual
+    stream, so the backbone is agnostic to how positions are represented.
+
+    Args:
+        config: Model configuration. Must be a ``Llama3Config`` instance.
+    """
+
+    def __init__(self, config: Llama3Config) -> None:
+        super().__init__()
+        self.config = config
+        self.layers = nn.ModuleList(
+            [DecoderLayer(config) for _ in range(config.num_hidden_layers)]
+        )
+        # RMSNorm over LayerNorm: omits mean subtraction, faster, and proved more
+        # stable at scale. This is the final norm that stabilises the accumulated
+        # residual stream — distinct from the per-layer pre-norms inside each block.
+        self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        inputs_embeds: torch.Tensor,
+        position_ids: torch.Tensor,
+        past_key_values: Cache | None = None,
+        output_hidden_states: bool = False,
+        causal_mask: torch.Tensor | None = None,
+    ) -> dict:
+        """Run the transformer stack over a batch of pre-embedded sequences.
+
+        Args:
+            inputs_embeds: Pre-embedded input of shape (batch, seq_len, hidden_size).
+            position_ids: Absolute positions of shape (batch, seq_len). Required.
+                Must be provided explicitly by the caller — this module does not
+                infer positions from cache state. The caller owns the mapping from
+                tokens to sequence positions.
+            past_key_values: A Cache object carrying the accumulated K/V history from
+                prior forward passes, or None. When provided, each decoder layer writes
+                new K/V into its slot and reads back the full accumulated history. The
+                cache is updated in place and returned as-is. When None, no caching
+                occurs and None is returned for past_key_values.
+            output_hidden_states: When True, the output dict includes a tuple of
+                per-layer hidden states: (inputs_embeds, layer_0_out, ..., layer_N_out),
+                collected before the final norm.
+            causal_mask: Optional boolean attention mask of shape
+                (1, 1, seq_len, kv_len). Threaded unchanged into every decoder
+                layer. When None, each layer uses SDPA's native ``is_causal``
+                mode (correct for full-sequence training).
+
+        Returns:
+            Plain dict with keys:
+            - ``"last_hidden_state"``: normed backbone output,
+              shape (batch, seq_len, hidden_size).
+            - ``"past_key_values"``: the Cache object (updated in place), or None.
+            - ``"hidden_states"``: tuple of per-layer activations (including
+              inputs_embeds as position 0) if ``output_hidden_states`` is True,
+              else None. Collected before the final norm so each entry reflects the
+              unnormalised residual stream at that depth.
+        """
+        hidden_states = inputs_embeds
+        all_hidden_states = (hidden_states,) if output_hidden_states else None
+
+        for i, layer in enumerate(self.layers):
+            hidden_states = layer(hidden_states, position_ids, cache=past_key_values, layer_idx=i, causal_mask=causal_mask)
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+        hidden_states = self.norm(hidden_states)
+
+        return {
+            "last_hidden_state": hidden_states,
+            "past_key_values": past_key_values,
+            "hidden_states": all_hidden_states,
+        }

rope.py

@@ -0,0 +1,143 @@
+"""Rotary Position Embeddings (RoPE).
+
+RoPE encodes position in the *relationship* between query and key vectors rather than
+adding it to the inputs directly. When the attention dot product Q·Kᵀ is computed, the
+per-position rotations cancel to produce a score that depends only on the relative
+distance between positions — not on their absolute values. This is what gives RoPE
+better length generalisation than absolute learned embeddings.
+
+Each pair of head dimensions (d, d+1) is assigned a rotation frequency
+    1 / theta^(2d / head_dim)
+Higher theta → slower rotation per position → position encodings remain distinguishable
+further apart before wrapping. Llama 3 uses theta=500,000 as a prerequisite for
+128K context support.
+
+Supported rope types: "default" (standard unscaled RoPE), "linear", and "yarn".
+HuggingFace's ROPE_INIT_FUNCTIONS handles inv_freq computation for linear and yarn;
+the default case is not in that registry and is computed directly here.
+"""
+
+import torch
+import torch.nn as nn
+from transformers import PretrainedConfig
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
+
+_SUPPORTED_ROPE_TYPES = {"default", "linear", "yarn"}
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    """Apply the 90° rotation used in the RoPE update formula.
+
+    Splits the last dimension into two halves [x1, x2] and returns [-x2, x1].
+    Combined with ``x * cos + rotate_half(x) * sin``, this implements a 2D rotation
+    on each consecutive pair of dimensions.
+    """
+    d = x.shape[-1] // 2
+    x1, x2 = x[..., :d], x[..., d:]
+    return torch.cat([-x2, x1], dim=-1)
+
+
+class RotaryEmbedding(nn.Module):
+    """Rotary Position Embeddings as an nn.Module.
+
+    Computes position-dependent rotation frequencies from the model config, maintains
+    a lazily-extended cos/sin cache, and applies the rotations to query and key tensors.
+
+    The cos/sin cache grows automatically at runtime when a sequence longer than the
+    current cache is encountered. ``config.max_position_embeddings`` records the
+    training context length (required by HF's scaling computations) but does not cap
+    inference length.
+
+    Args:
+        config: Model config. Must expose ``rope_theta``, ``rope_parameters`` (set by
+            HF's RotaryEmbeddingConfigMixin), and ``head_dim``.
+        device: Optional device for initial buffer placement. Buffers move with the
+            model on ``.to()`` / ``.cuda()`` calls.
+
+    Raises:
+        NotImplementedError: If ``config.rope_parameters`` specifies an unsupported
+            rope type. Supported types: "default", "linear", "yarn".
+    """
+
+    def __init__(self, config: PretrainedConfig, device: torch.device | None = None) -> None:
+        super().__init__()
+        self.config = config
+
+        # rope_parameters is None when no rope_scaling was passed to the config.
+        rope_params = config.rope_parameters
+        self.rope_type = (
+            rope_params.get("rope_type", "default") if rope_params is not None else "default"
+        )
+
+        if self.rope_type not in _SUPPORTED_ROPE_TYPES:
+            raise NotImplementedError(
+                f"rope_type '{self.rope_type}' is not supported. "
+                f"Supported types: {sorted(_SUPPORTED_ROPE_TYPES)}"
+            )
+
+        if self.rope_type == "default":
+            # Standard RoPE: inv_freq = 1 / theta^(2i / head_dim).
+            # Not in ROPE_INIT_FUNCTIONS, so computed directly.
+            inv_freq = 1.0 / (
+                config.rope_theta
+                ** (torch.arange(0, config.head_dim, 2, dtype=torch.float32, device=device) / config.head_dim)
+            )
+            self.attention_scaling: float = 1.0
+        else:
+            inv_freq, self.attention_scaling = ROPE_INIT_FUNCTIONS[self.rope_type](config, device)
+
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        # Initialised as None; built on first forward call and extended lazily thereafter.
+        # Registered as buffers so they move with the model across devices.
+        self.register_buffer("_cos_cached", None, persistent=False)
+        self.register_buffer("_sin_cached", None, persistent=False)
+
+    def _extend_cache(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> None:
+        """Build the cos/sin table to cover positions [0, seq_len).
+
+        Registered as buffers so subsequent calls to ``.to()`` / ``.cuda()`` will
+        move them to the correct device. Rebuilds whenever the sequence grows or
+        the dtype changes (e.g. switching between fp32 and bf16).
+        """
+        positions = torch.arange(seq_len, device=device, dtype=torch.float32)
+        # outer product → (seq_len, head_dim // 2); duplicate → (seq_len, head_dim)
+        freqs = torch.outer(positions, self.inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("_cos_cached", emb.cos().to(dtype), persistent=False)
+        self.register_buffer("_sin_cached", emb.sin().to(dtype), persistent=False)
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        position_ids: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, float]:
+        """Apply rotary embeddings to query and key tensors.
+
+        The cos/sin cache is extended lazily when position_ids reference positions
+        beyond its current length.
+
+        Args:
+            q: Query tensor of shape (batch, num_heads, seq_len, head_dim).
+            k: Key tensor of shape (batch, num_kv_heads, seq_len, head_dim).
+            position_ids: Integer positions of shape (batch, seq_len).
+
+        Returns:
+            Tuple of (q_rotated, k_rotated, attention_scaling). attention_scaling is
+            1.0 for default and linear; YaRN returns a value != 1.0 that callers must
+            apply to attention logits to correct for frequency magnitude changes.
+        """
+        seq_len = int(position_ids.max().item()) + 1
+
+        if self._cos_cached is None or seq_len > self._cos_cached.shape[0] or self._cos_cached.dtype != q.dtype:
+            self._extend_cache(seq_len, device=q.device, dtype=q.dtype)
+
+        # Gather cos/sin for the given positions → (batch, seq_len, head_dim),
+        # then unsqueeze the head axis for broadcast over all heads.
+        cos = self._cos_cached[position_ids].unsqueeze(1)
+        sin = self._sin_cached[position_ids].unsqueeze(1)
+
+        q_rotated = q * cos + _rotate_half(q) * sin
+        k_rotated = k * cos + _rotate_half(k) * sin
+
+        return q_rotated, k_rotated, self.attention_scaling

tokenizer_config.json

@@ -0,0 +1,13 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "is_local": false,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<|padding|>",
+  "tokenizer_class": "GPTNeoXTokenizerFast",
+  "trim_offsets": true,
+  "unk_token": "<|endoftext|>"
+}