Metadata-Version: 2.4
Name: fiveattn
Version: 0.1.0
Summary: 100M-parameter Transformer comparing 5-Component Attention (GAU) and MHA, with SmolLM tokenizer, Cut-Cross-Entropy, SwiGLU, and RMSNorm
Author: FiveAttn Team
License-Expression: Apache-2.0
Project-URL: Homepage, https://github.com/your-org/fiveattn
Project-URL: Repository, https://github.com/your-org/fiveattn
Keywords: deep-learning,transformer,attention,gated-attention,gau,cut-cross-entropy,swiglu,rmsnorm,smollm,tinystories,pytorch
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: Intended Audience :: Developers
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Operating System :: OS Independent
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.1.0
Requires-Dist: transformers>=4.40.0
Requires-Dist: cut-cross-entropy>=0.1.0
Requires-Dist: datasets>=2.18.0
Requires-Dist: pyarrow>=14.0.0
Requires-Dist: huggingface-hub>=0.20.0
Requires-Dist: tokenizers>=0.19.0
Requires-Dist: numpy
Provides-Extra: dev
Requires-Dist: pytest>=7.0.0; extra == "dev"
Requires-Dist: build>=1.0.0; extra == "dev"
Requires-Dist: twine>=4.0.0; extra == "dev"
Dynamic: license-file

# FiveAttn: 5-Component Attention & MHA 100M Transformers

A production-grade PyTorch implementation and benchmark suite comparing **5-Component Attention (Gated Attention Unit / GAU)** against standard **Multi-Head Attention (MHA)** at the **100M-parameter** scale, built for high-performance training on the **TinyStories** dataset.

---

## Highlights

- **5-Component Attention (GAU)**: Replaces standard 4-projection MHA ($W_q, W_k, W_v, W_o$) with 5 projection matrices ($W_q, W_k, W_v, W_g, W_o$), introducing a non-linear Swish gate ($G = X W_g$) that modulates the attention context prior to output projection.
- **Strict 100M Parameter Parity**: Both the standard MHA and 5-Component models are calibrated to **exactly 99,994,176 parameters** (within 0.006% of 100M), guaranteeing fair scientific ablation.
- **SmolLM-135M Tokenizer**: Uses Hugging Face's `HuggingFaceTB/SmolLM-135M` tokenizer with a 49,152 vocabulary size and tied word embeddings.
- **Apple Cut-Cross-Entropy**: Computes memory-efficient next-token cross-entropy loss directly from final hidden states without allocating the $(B \times T \times V)$ logits matrix into memory, saving gigabytes of VRAM.
- **SwiGLU & RMSNorm**: Uses Swish-Gated Linear Units (SwiGLU) for the feed-forward network and Root Mean Square Normalization (RMSNorm) with pre-normalization residuals.
- **Zero Graph Breaks in `torch.compile`**: Verified with `torch._dynamo.explain` to compile into 1 single graph with 0 graph breaks.
- **PyTorch Inductor Autotune Ready**: Supports `--compile --autotune` (`mode="max-autotune"` and coordinate descent tuning).
- **Remote-First CLI**: Download and pre-tokenize TinyStories on your remote GPU instance using memory-mapped binary files (`.bin`) for high-throughput zero-copy training.

---

## Architecture & Math

### 1. Multi-Head Attention (Standard 4 Projections)
Standard MHA projects input $X \in \mathbb{R}^{B \times T \times D}$ into queries, keys, and values:
$$Q = X W_q, \quad K = X W_k, \quad V = X W_v$$
$$A = \text{softmax}\left(\frac{Q K^T}{\sqrt{d_k}} + \text{causal\_mask}\right) V$$
$$\text{Output} = A W_o$$
Parameters per attention block: $4 \times D^2$.

### 2. 5-Component Attention (Gated Attention Unit - GAU)
5-Component Attention introduces a 5th projection matrix $W_g \in \mathbb{R}^{D \times D}$ acting as a data-dependent gating channel:
$$Q = X W_q, \quad K = X W_k, \quad V = X W_v, \quad G = X W_g$$
$$A = \text{softmax}\left(\frac{\text{RoPE}(Q) \text{RoPE}(K)^T}{\sqrt{d_k}} + \text{causal\_mask}\right) V$$
$$Z = A \odot \text{SiLU}(G)$$
$$\text{Output} = Z W_o$$
Just as SwiGLU adds a gating projection to standard MLPs, 5-Component Attention adds a gating projection to the attention mechanism.
Parameters per attention block: $5 \times D^2$.

### 3. SwiGLU MLP
$$\text{SwiGLU}(X) = \left(\text{SiLU}(X W_{\text{gate}}) \odot (X W_{\text{up}})\right) W_{\text{down}}$$

### 4. RMSNorm
$$\text{RMSNorm}(X) = \frac{X}{\sqrt{\frac{1}{D}\sum_{i=1}^D X_i^2 + \epsilon}} \odot \gamma$$

---

## 100M Parameter Calibration Table

Both models share identical embedding size, depth (16 layers), attention heads (9 heads, 64 dim), and total parameters:

| Component | Standard MHA-100M | 5-Component Attention 100M | Difference |
| :--- | :--- | :--- | :--- |
| **Vocab Size** | 49,152 (SmolLM) | 49,152 (SmolLM) | 0 |
| **Layers ($L$)** | 16 | 16 | 0 |
| **Hidden Dimension ($D$)** | 576 | 576 | 0 |
| **Attention Heads ($H$)** | 9 | 9 | 0 |
| **Head Dimension ($D_h$)** | 64 | 64 | 0 |
| **Attention Projections** | **4** ($W_q, W_k, W_v, W_o$) | **5** ($W_q, W_k, W_v, W_g, W_o$) | +1 matrix |
| **MLP Hidden Dimension** | 1,824 | 1,632 | -192 |
| **Embedding Parameters** | 28,311,552 | 28,311,552 | 0 |
| **Attention Parameters** | 21,233,664 | 26,542,080 | +5,308,416 |
| **SwiGLU MLP Parameters** | 50,429,952 | 45,121,536 | -5,308,416 |
| **RMSNorm Parameters** | 19,008 | 19,008 | 0 |
| **Total Parameters** | **99,994,176** | **99,994,176** | **0 (EXACT MATCH)** |

---

## Installation

```bash
# Clone repository
git clone https://github.com/your-org/fiveattn.git
cd fiveattn

# Install in editable mode
pip install -e .

# Or install with dev dependencies
pip install -e ".[dev]"
```

---

## Remote Workflow (Recommended)

To run your training on a remote GPU machine (e.g. Lambda, RunPod, AWS, GCP):

### Step 1: Pre-tokenize TinyStories on Remote
Run `prepare-data` on your remote server to download and convert TinyStories into high-throughput memory-mapped binary files:

```bash
# Tokenize train split (~2.1M stories, uint16 memmap)
fiveattn prepare-data --output-dir data/tinystories --split train

# Tokenize validation split
fiveattn prepare-data --output-dir data/tinystories --split validation
```

### Step 2: Train 5-Component Attention (100M)
```bash
fiveattn train \
  --model fiveattn \
  --data-path data/tinystories/train.bin \
  --val-data-path data/tinystories/validation.bin \
  --batch-size 32 \
  --seq-len 512 \
  --lr 5e-4 \
  --steps 5000 \
  --device cuda \
  --compile \
  --autotune
```

### Step 3: Train Standard MHA Baseline (100M)
```bash
fiveattn train \
  --model mha \
  --data-path data/tinystories/train.bin \
  --val-data-path data/tinystories/validation.bin \
  --batch-size 32 \
  --seq-len 512 \
  --lr 5e-4 \
  --steps 5000 \
  --device cuda \
  --compile \
  --autotune
```

---

## CLI Command Reference

### `fiveattn count-params`
Prints detailed parameter breakdown for both models:
```bash
fiveattn count-params
```

### `fiveattn verify-compile`
Formally checks with `torch._dynamo.explain` that models have 0 graph breaks:
```bash
fiveattn verify-compile --seq-len 64
```

### `fiveattn benchmark`
Runs latency, peak VRAM, and tokens/sec throughput benchmarks comparing MHA vs 5-Component Attention:
```bash
fiveattn benchmark --batch-size 8 --seq-len 256 --steps 50 --device cuda --compile
```

### `fiveattn train`
Trains a model on TinyStories (online, offline memmap, or synthetic):
```bash
# Quick offline smoke-test on CPU with synthetic data
fiveattn train --model fiveattn --steps 10 --batch-size 2 --seq-len 64 --device cpu --synthetic

# Full training on CUDA with compiled kernels and autotuning
fiveattn train --model fiveattn --steps 2000 --batch-size 16 --seq-len 512 --device cuda --compile --autotune
```

### `fiveattn generate`
Generates text continuation from a checkpoint or initialized model:
```bash
fiveattn generate \
  --checkpoint checkpoints/fiveattn_step_2000.pt \
  --prompt "Once upon a time, there was a little robot" \
  --max-tokens 100 \
  --temperature 0.8
```

---

## Python API Usage

```python
import torch
from fiveattn import ModelConfig, TransformerLM, CutCrossEntropyLoss

# 1. Initialize 100M 5-component attention model
config = ModelConfig.get_100m_fiveattn_config()
model = TransformerLM(config).cuda()

# 2. Forward pass with cut-cross-entropy loss
loss_fn = CutCrossEntropyLoss(shift=True, impl="cce")

input_ids = torch.randint(0, config.vocab_size, (4, 512), device="cuda")
targets = input_ids.clone()

# Directly compute hidden states and linear cross-entropy without allocating (B, T, V) logits
hidden_states = model.forward_hidden_states(input_ids)
loss = loss_fn(hidden_states, model.get_output_embeddings(), targets)
loss.backward()

# 3. Autoregressive generation
from fiveattn import generate_text
output = generate_text(model, prompt="Once upon a time")
print(output)
```

---

## Running Unit Tests

```bash
pytest -v
```
All 15 tests run fully offline without network access or dataset downloads.

---

## License
Apache-2.0
