README.md

---
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