Upload folder using huggingface_hub

Cosmos-T2A-low.best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:432493235c7821108adfc57c7642fd81f8aa035b6ffa3a8af3030855c8de0689
+size 42266815

Cosmos-T2A-low.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:23cdbd4510ef1b8cc56910fc76a6886f61deaaf7144db54b1396ad05f5f1d69c
+size 42300893

README.md

@@ -0,0 +1,189 @@
+---
+license: apache-2.0
+language:
+- en
+library_name: transformers
+pipeline_tag: text-generation
+tags:
+- chain-of-thought
+- reasoning
+- instruct
+- pretrained-from-scratch
+- decoder-only
+- transformer
+- qwen-tokenizer
+- rope
+- rmsnorm
+- swiglu
+- gqa
+- engram
+datasets:
+- wop/minitron-dataset
+model-index:
+- name: Cosmos T2A-low
+  results:
+  - task:
+      type: text-generation
+      name: Causal Language Modeling
+    dataset:
+      name: wop/minitron-dataset
+      type: wop/minitron-dataset
+      split: train
+    metrics:
+    - type: loss
+      name: Final training loss (cross-entropy)
+      value: 5.6389
+    - type: perplexity
+      name: Final training perplexity
+      value: 281.15
+    - type: loss
+      name: Final validation loss (cross-entropy)
+      value: 5.4913
+    - type: perplexity
+      name: Final validation perplexity
+      value: 242.56
+---
+
+<img src="https://calm-heart-d697.mmmmmm505090.workers.dev?text=Cosmos T2A-low" width="900" alt="Cosmos T2A-low" />
+
+# Cosmos T2A-low
+
+Universal Kaggle-ready training notebook for the Cosmos T2A-low series.
+
+> Notebook-generated card. Final metrics are filled after the Kaggle training run.
+> This notebook is designed to stay Kaggle-friendly on 2x T4 GPUs. The goal is a reusable training recipe, not a production assistant.
+
+## Model Details
+
+| | |
+|---|---|
+| **Model class** | `CosmosT2_Accelerate_LLM` |
+| **Architecture** | Decoder-only Transformer with RoPE, RMSNorm, SwiGLU, GQA, and a configurable Engram memory path |
+| **Parameters** | `~9.96 M` |
+| **Layers** | `4` |
+| **Attention heads** | `4` |
+| **KV heads** | `1` |
+| **d_model** | `64` |
+| **FFN hidden** | `256` |
+| **Positional encoding** | RoPE (`rope_base=10000`) |
+| **Normalization** | RMSNorm |
+| **MLP** | SwiGLU |
+| **Memory** | Engram (`use_engram=True`, every `2` blocks) |
+| **Context length** | `1028` |
+| **Training block size** | `1028` |
+| **Tokenizer** | [`Qwen/Qwen2.5-0.5B`](https://huggingface.co/Qwen/Qwen2.5-0.5B) |
+| **Dataset** | [`wop/minitron-dataset`](https://huggingface.co/datasets/wop/minitron-dataset) |
+| **License** | Apache-2.0 |
+
+### Why these choices
+
+- **RoPE** keeps positional handling compact and avoids learned absolute embeddings.
+- **RMSNorm** is cheaper and more stable than LayerNorm for this small decoder-only model.
+- **SwiGLU** usually gives a better quality/compute tradeoff than a plain GELU MLP.
+- **GQA** reduces KV cost while keeping multi-head query capacity.
+- **Engram** gives the stack a lightweight explicit memory path for repeated reasoning patterns.
+- **Dynamic isolated batching** keeps conversations separate while padding and masking each batch on CPU.
+- **KV-cache generation** avoids recomputing the full prompt for every generated token in the app.
+
+## Training Summary
+
+| Metric | Value |
+|---|---|
+| Rows used | `288,350` |
+| Loss tokens seen | `46,651,194` |
+| Epochs | `1` |
+| Batch size | `2` |
+| Peak LR | `3.00e-04` |
+| Weight decay | `0.1` |
+| Gradient clipping | `1.0` |
+| Wall-clock time | `1h 37m 23s` |
+| Final training loss | `5.6389` |
+| Final training perplexity | `281.15` |
+| Final validation loss | `5.4913` |
+| Final validation perplexity | `242.56` |
+| Best validation loss | `4.9614` |
+| Best epoch | `1` |
+
+### Loss and perplexity
+
+The notebook shows live loss and perplexity plots every `5000` epochs and does not save the graph to disk.
+
+## How to Use
+
+### Quick start
+
+~~~python
+import torch
+from transformers import AutoTokenizer
+
+from app import CosmosT2_Accelerate_LLM
+
+tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B")
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+ckpt = torch.load("$CHECKPOINT_NAME", map_location="cpu")
+model = CosmosT2_Accelerate_LLM(**ckpt["config"])
+model.load_state_dict(ckpt["model_state"])
+model.eval()
+
+prompt = tokenizer.apply_chat_template(
+    [
+        {"role": "system", "content": "Enable thinking features: INTUITION"},
+        {"role": "user", "content": "What is 12 * 7?"},
+    ],
+    tokenize=False,
+    add_generation_prompt=True,
+)
+ids = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).input_ids
+out = model.generate(ids, max_new_tokens=120, temperature=0.8, top_k=50)
+print(tokenizer.decode(out[0], skip_special_tokens=False))
+~~~
+
+### Prompt format
+
+Use the Qwen2.5 chat template. The default system prompt is:
+
+~~~text
+Enable thinking features: INTUITION
+~~~
+
+The model will then emit a `<think>` block followed by an answer when it has enough signal.
+
+The model is trained to end its turn with the `<|im_end|>` token (ChatML), so generation stops there. During data prep, any example longer than the `1028`-token context has its `<think>` reasoning replaced by a short placeholder (or is dropped) so every training sequence ends cleanly - the model is never trained on a mid-thought truncation.
+
+## Limitations
+
+- The model is intentionally small and is still a research/demo artifact.
+- Training on chain-of-thought data can overfit quickly if the corpus is tiny.
+- Long-context behavior is limited by the configured block size.
+- The model is not safety-aligned and should not be exposed as a public assistant without additional work.
+
+## Intended Use
+
+- Research into small-scale pretraining and reasoning-style formatting
+- Educational demos for decoder-only Transformer training
+- Hugging Face Spaces or local inference demos
+- Not for production use
+
+## Cosmos T2A-low Series
+
+This notebook is designed to train future Cosmos T2A-low variants by changing only the config block at the top.
+
+## Citation
+
+~~~bibtex
+@misc{cosmos-t2,
+  author    = {wop},
+  title     = {Cosmos-T2: A small from-scratch chain-of-thought Transformer},
+  year      = {2026},
+  publisher = {Hugging Face},
+  url       = {https://huggingface.co/wop/Cosmos-T2A-low}
+}
+~~~
+
+## Acknowledgements
+
+- Tokenizer from Qwen2.5 by Alibaba Cloud
+- Training data from wop/minitron-dataset
+- Trained on Kaggle T4 GPUs

demo.py

@@ -0,0 +1,382 @@
+"""Cosmos T2A-low — Standalone CLI Demo
+
+Interactive command-line demo for Cosmos T2A-low.
+No Gradio, no web server — just pure Python inference.
+
+Usage:
+    python demo.py                          # interactive chat
+    python demo.py --prompt "What is 12*7?" # single-prompt mode
+
+Architecture: RoPE + RMSNorm + SwiGLU + GQA + Engram memory
+"""
+
+import argparse
+import sys
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import AutoTokenizer
+
+# ---------------------------------------------------------------------------
+# Constants
+# ---------------------------------------------------------------------------
+CHECKPOINT_DIR = Path(__file__).parent
+FINAL_CKPT = CHECKPOINT_DIR / "Cosmos-T2A-low.pt"
+BEST_CKPT = CHECKPOINT_DIR / "Cosmos-T2A-low.best.pt"
+
+TOKENIZER_NAME = "Qwen/Qwen2.5-0.5B"
+DEFAULT_SYSTEM_PROMPT = "Enable thinking features: INTUITION"
+STOP_IDS = {151645, 151643}
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
+
+
+# ---------------------------------------------------------------------------
+# Architecture
+# ---------------------------------------------------------------------------
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+    def forward(self, x):
+        rms = x.pow(2).mean(dim=-1, keepdim=True)
+        return x * torch.rsqrt(rms + self.eps) * self.weight
+
+
+def rotate_half(x):
+    x1, x2 = x[..., ::2], x[..., 1::2]
+    return torch.stack((-x2, x1), dim=-1).flatten(-2)
+
+
+def apply_rope(q, k, cos, sin):
+    return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)
+
+
+class GQAAttention(nn.Module):
+    def __init__(self, d_model, n_heads, n_kv_heads, rope_base=10000, dropout=0.0):
+        super().__init__()
+        assert d_model % n_heads == 0 and n_heads % n_kv_heads == 0
+        self.n_heads = n_heads
+        self.n_kv_heads = n_kv_heads
+        self.head_dim = d_model // n_heads
+        self.dropout = dropout
+        self.q_proj = nn.Linear(d_model, n_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(d_model, d_model, bias=False)
+
+    def forward(self, x, rope_cos, rope_sin, past_kv=None, use_cache=False):
+        batch, seq_len, _ = x.shape
+        q = self.q_proj(x).view(batch, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(batch, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(batch, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        q, k = apply_rope(q, k, rope_cos, rope_sin)
+        if past_kv is not None:
+            pk, pv = past_kv
+            k = torch.cat([pk, k], dim=2)
+            v = torch.cat([pv, v], dim=2)
+        present_kv = (k, v) if use_cache else None
+        if self.n_kv_heads != self.n_heads:
+            k = k.repeat_interleave(self.n_heads // self.n_kv_heads, dim=1)
+            v = v.repeat_interleave(self.n_heads // self.n_kv_heads, dim=1)
+        out = F.scaled_dot_product_attention(
+            q, k, v, is_causal=(past_kv is None),
+            dropout_p=self.dropout if self.training else 0.0
+        )
+        out = out.transpose(1, 2).contiguous().view(batch, seq_len, -1)
+        out = self.o_proj(out)
+        return (out, present_kv) if use_cache else out
+
+
+class SwiGLUMLP(nn.Module):
+    def __init__(self, d_model, hidden_dim, dropout=0.0):
+        super().__init__()
+        self.gate = nn.Linear(d_model, hidden_dim, bias=False)
+        self.up = nn.Linear(d_model, hidden_dim, bias=False)
+        self.down = nn.Linear(hidden_dim, d_model, bias=False)
+        self.dropout = nn.Dropout(dropout)
+    def forward(self, x):
+        return self.down(self.dropout(F.silu(self.gate(x)) * self.up(x)))
+
+
+class EngramMemory(nn.Module):
+    def __init__(self, d_model, bucket_count, memory_dim, order, pad_id=0, dropout=0.0):
+        super().__init__()
+        self.bucket_count = bucket_count
+        self.order = order
+        self.pad_id = pad_id
+        self.bucket = nn.Embedding(bucket_count, memory_dim)
+        self.query = nn.Linear(d_model, memory_dim, bias=False)
+        self.project = nn.Linear(memory_dim, d_model, bias=False)
+        self.gate = nn.Linear(d_model, d_model, bias=True)
+        self.dropout = nn.Dropout(dropout)
+        primes = [1, 1315423911, 2654435761, 97531, 433494437]
+        self.register_buffer("primes", torch.tensor(primes[:order], dtype=torch.long), persistent=False)
+
+    def hash_tokens(self, idx):
+        B, S = idx.shape
+        pad = torch.full((B, self.order - 1), self.pad_id, device=idx.device, dtype=idx.dtype)
+        history = torch.cat([pad, idx], dim=1)
+        hashed = torch.zeros((B, S), device=idx.device, dtype=torch.long)
+        for off in range(self.order):
+            sl = history[:, off: off + S].long()
+            hashed = (hashed * 1315423911 + sl * self.primes[off]) % self.bucket_count
+        return hashed
+
+    def forward(self, x, idx):
+        hashed = self.hash_tokens(idx)
+        if hashed.size(1) != x.size(1):
+            hashed = hashed[:, -x.size(1):]
+        query = torch.tanh(self.query(x))
+        mem = self.bucket(hashed) * query
+        mem = self.project(mem)
+        gate = torch.sigmoid(self.gate(x))
+        return self.dropout(gate * mem)
+
+
+class Block(nn.Module):
+    def __init__(self, d_model, n_heads, n_kv_heads, d_ff, rope_base, dropout=0.0,
+                 use_engram=False, engram_bucket_count=128, engram_dim=16, engram_order=3, pad_id=0):
+        super().__init__()
+        self.norm1 = RMSNorm(d_model)
+        self.attn = GQAAttention(d_model, n_heads, n_kv_heads, rope_base=rope_base, dropout=dropout)
+        self.norm2 = RMSNorm(d_model)
+        self.engram = EngramMemory(d_model, engram_bucket_count, engram_dim, engram_order,
+                                   pad_id=pad_id, dropout=dropout) if use_engram else None
+        self.norm3 = RMSNorm(d_model)
+        self.mlp = SwiGLUMLP(d_model, d_ff, dropout=dropout)
+
+    def forward(self, x, idx, rope_cos, rope_sin):
+        x = x + self.attn(self.norm1(x), rope_cos, rope_sin)
+        if self.engram is not None:
+            x = x + self.engram(self.norm2(x), idx)
+        return x + self.mlp(self.norm3(x))
+
+
+class CosmosT2_Accelerate_LLM(nn.Module):
+    def __init__(self, vocab_size, d_model=64, n_layers=4, n_heads=4, n_kv_heads=1,
+                 d_ff=256, max_len=1028, rope_base=10000, dropout=0.05, use_engram=True,
+                 engram_every=2, engram_bucket_count=128, engram_dim=16, engram_order=3, pad_id=0):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.n_kv_heads = n_kv_heads
+        self.head_dim = d_model // n_heads
+        self.max_len = max_len
+        self.rope_base = rope_base
+        self.pad_id = pad_id
+        self.tok_emb = nn.Embedding(vocab_size, d_model)
+        self.blocks = nn.ModuleList()
+        for layer_index in range(n_layers):
+            block_uses_engram = use_engram and ((layer_index + 1) % engram_every == 0)
+            self.blocks.append(Block(
+                d_model=d_model, n_heads=n_heads, n_kv_heads=n_kv_heads, d_ff=d_ff,
+                rope_base=rope_base, dropout=dropout, use_engram=block_uses_engram,
+                engram_bucket_count=engram_bucket_count, engram_dim=engram_dim,
+                engram_order=engram_order, pad_id=pad_id,
+            ))
+        self.norm_f = RMSNorm(d_model)
+
+    def build_rope(self, seq_len, device, dtype, start_pos=0):
+        inv_freq = 1.0 / (self.rope_base ** (
+            torch.arange(0, self.head_dim, 2, device=device, dtype=torch.float32) / self.head_dim
+        ))
+        positions = torch.arange(start_pos, start_pos + seq_len, device=device, dtype=torch.float32)
+        freqs = torch.outer(positions, inv_freq)
+        cos = freqs.cos().repeat_interleave(2, dim=-1).to(dtype)[None, None, :, :]
+        sin = freqs.sin().repeat_interleave(2, dim=-1).to(dtype)[None, None, :, :]
+        return cos, sin
+
+    def forward(self, idx, targets=None):
+        if idx.size(1) > self.max_len:
+            idx = idx[:, -self.max_len:]
+        seq_len = idx.size(1)
+        rope_cos, rope_sin = self.build_rope(seq_len, idx.device, self.tok_emb.weight.dtype)
+        x = self.tok_emb(idx)
+        for block in self.blocks:
+            x = block(x, idx, rope_cos, rope_sin)
+        x = self.norm_f(x)
+        logits = F.linear(x, self.tok_emb.weight)
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1))
+        return logits, loss
+
+    @torch.no_grad()
+    def generate(self, input_ids, max_new_tokens=128, temperature=0.1, top_k=1):
+        self.eval()
+        cur_ids = input_ids.to(DEVICE)
+        for _ in range(max_new_tokens):
+            logits, _ = self.forward(cur_ids[:, -self.max_len:])
+            next_logit = logits[:, -1, :] / max(temperature, 1e-8)
+            if top_k > 1:
+                top_vals, _ = torch.topk(next_logit, top_k, dim=-1)
+                next_logit[next_logit < top_vals[:, -1:]] = float('-inf')
+            probs = F.softmax(next_logit, dim=-1)
+            nxt = torch.multinomial(probs, num_samples=1)
+            if nxt.item() in STOP_IDS:
+                break
+            cur_ids = torch.cat([cur_ids, nxt], dim=1)
+        return cur_ids
+
+
+# ---------------------------------------------------------------------------
+# Loading
+# ---------------------------------------------------------------------------
+
+def load_model(checkpoint_path=None):
+    if checkpoint_path is None:
+        ckpt_path = FINAL_CKPT if FINAL_CKPT.exists() else BEST_CKPT
+    else:
+        ckpt_path = Path(checkpoint_path)
+    if not ckpt_path.exists():
+        raise FileNotFoundError(f"Checkpoint not found: {ckpt_path}")
+
+    print(f"Loading checkpoint: {ckpt_path}")
+    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+    config = ckpt.get("config", {})
+
+    model = CosmosT2_Accelerate_LLM(
+        vocab_size=config.get("vocab_size", 151936),
+        d_model=config.get("d_model", 64),
+        n_layers=config.get("n_layers", 4),
+        n_heads=config.get("n_heads", 4),
+        n_kv_heads=config.get("n_kv_heads", 1),
+        d_ff=config.get("d_ff", 256),
+        max_len=config.get("max_len", 1028),
+        rope_base=config.get("rope_base", 10000),
+        dropout=config.get("dropout", 0.05),
+        use_engram=config.get("use_engram", True),
+        engram_every=config.get("engram_every", 2),
+        engram_bucket_count=config.get("engram_bucket_count", 128),
+        engram_dim=config.get("engram_dim", 16),
+        engram_order=config.get("engram_order", 3),
+        pad_id=config.get("pad_id", 0),
+    )
+    model.load_state_dict(ckpt["model_state"], strict=False)
+    model = model.to(DEVICE, dtype=DTYPE)
+    model.eval()
+    params = sum(p.numel() for p in model.parameters())
+    print(f"Model loaded on {DEVICE}: {params/1e6:.2f}M parameters")
+    return model
+
+
+# ---------------------------------------------------------------------------
+# Prompt helpers
+# ---------------------------------------------------------------------------
+
+def build_prompt(user_message, system_prompt=None):
+    if system_prompt is None:
+        system_prompt = DEFAULT_SYSTEM_PROMPT
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": user_message},
+    ]
+    return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+
+
+def generate_response(model, prompt, max_new=256, temperature=0.1, top_k=1):
+    ids = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).input_ids.to(DEVICE)
+    out = model.generate(ids, max_new_tokens=max_new, temperature=temperature, top_k=top_k)
+    response = tokenizer.decode(out[0], skip_special_tokens=False)
+    if "<|im_start|>assistant" in response:
+        parts = response.split("<|im_start|>assistant")
+        if len(parts) > 1:
+            response = parts[-1].split("<|im_end|>")[0].strip()
+    return response
+
+
+# ---------------------------------------------------------------------------
+# Interactive chat
+# ---------------------------------------------------------------------------
+
+def interactive_chat(model):
+    print("\n" + "=" * 60)
+    print("  Cosmos T2A-low — Interactive Demo")
+    print("  Type /quit to exit, /clear to reset context")
+    print("=" * 60 + "\n")
+
+    history = []
+    while True:
+        try:
+            user_msg = input("You: ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print("\nGoodbye!")
+            break
+        if not user_msg:
+            continue
+        if user_msg.lower() in ("/quit", "/exit", "/q"):
+            print("Goodbye!")
+            break
+        if user_msg.lower() == "/clear":
+            history = []
+            print("[Context cleared]\n")
+            continue
+
+        if history:
+            prompt = tokenizer.apply_chat_template(
+                history + [{"role": "user", "content": user_msg}],
+                tokenize=False, add_generation_prompt=True
+            )
+        else:
+            prompt = build_prompt(user_msg)
+
+        print("Cosmos: ", end="", flush=True)
+        ids = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).input_ids.to(DEVICE)
+        out = model.generate(ids, max_new_tokens=256, temperature=0.1, top_k=1)
+        full_response = tokenizer.decode(out[0], skip_special_tokens=False)
+
+        if "<|im_start|>assistant" in full_response:
+            parts = full_response.split("<|im_start|>assistant")
+            response = parts[-1].split("<|im_end|>")[0].strip() if len(parts) > 1 else full_response.strip()
+        else:
+            response = full_response.strip()
+
+        print(response)
+        print()
+        history.append({"role": "user", "content": user_msg})
+        history.append({"role": "assistant", "content": response})
+        if len(history) > 20:
+            history = history[-20:]
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Cosmos T2A-low CLI Demo")
+    parser.add_argument("--checkpoint", type=str, default=None,
+                        help="Path to checkpoint .pt file")
+    parser.add_argument("--prompt", type=str, default=None,
+                        help="Single prompt (non-interactive mode)")
+    parser.add_argument("--max-new", type=int, default=256,
+                        help="Max new tokens to generate")
+    parser.add_argument("--temperature", type=float, default=0.1,
+                        help="Sampling temperature")
+    parser.add_argument("--top-k", type=int, default=1,
+                        help="Top-K sampling")
+    args = parser.parse_args()
+
+    print(f"Device: {DEVICE} | Dtype: {DTYPE}")
+    print(f"Loading tokenizer: {TOKENIZER_NAME}")
+    tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    model = load_model(args.checkpoint)
+
+    if args.prompt:
+        print(f"\nPrompt: {args.prompt}\n")
+        prompt = build_prompt(args.prompt)
+        response = generate_response(model, prompt, max_new=args.max_new,
+                                      temperature=args.temperature, top_k=args.top_k)
+        print(f"Response: {response}")
+    else:
+        interactive_chat(model)

model_config.json

@@ -0,0 +1,49 @@
+{
+  "model_family": "Cosmos T2A-low",
+  "model_name": "Cosmos T2A-low",
+  "model_class_name": "CosmosT2_Accelerate_LLM",
+  "hf_repo_id": "wop/Cosmos-T2A-low",
+  "tokenizer_name": "Qwen/Qwen2.5-0.5B",
+  "dataset_name": "wop/minitron-dataset",
+  "dataset_split": "train",
+  "dataset_row_limit": 288350,
+  "rows_used": 288350,
+  "stream_dataset": true,
+  "shuffle_buffer_size": 2048,
+  "train_val_fraction": 0.1,
+  "seed": 42,
+  "block_size": 1028,
+  "max_len": 1028,
+  "d_model": 64,
+  "n_layers": 4,
+  "n_heads": 4,
+  "n_kv_heads": 1,
+  "d_ff": 256,
+  "rope_base": 10000,
+  "dropout": 0.05,
+  "use_engram": true,
+  "engram_every": 2,
+  "engram_buckets": 128,
+  "engram_dim": 16,
+  "engram_order": 3,
+  "epochs": 1,
+  "batch_size": 2,
+  "lr": 0.0003,
+  "weight_decay": 0.1,
+  "warmup_steps": 50,
+  "grad_clip": 1.0,
+  "log_every_steps": 5000,
+  "eval_every_steps": 5000,
+  "plot_every_epochs": 5000,
+  "val_max_batches": 10,
+  "early_stopping_patience": 2,
+  "tokenization_batch_size": 256,
+  "num_workers": 1,
+  "pin_memory": true,
+  "dataloader_prefetch_factor": 4,
+  "persistent_workers": true,
+  "use_kv_cache": true,
+  "train_steps_per_epoch": 129757,
+  "loss_tokens_seen": 46651194,
+  "samples_seen": 222956
+}