README.md

---
license: mit
datasets:
- allenai/c4
- HuggingFaceFW/fineweb-edu
- HuggingFaceTB/smollm-corpus
pipeline_tag: text-generation
library_name: transformers
---

# TinyMixtral

A small Mixtral-style Mixture-of-Experts causal language model (~432M total, ~176M active parameters) for pretraining research on a single consumer GPU.

## Model Architecture

| Parameter | Value |
|-----------|-------|
| hidden_size | 896 |
| num_layers | 10 |
| Attention | Grouped Query Attention (14 heads / 2 KV heads) |
| Head dim | 64 |
| RoPE theta | 1,000,000 |
| Norm | RMSNorm |
| Experts | 6 (top-2 routing) |
| Expert FFN | SwiGLU, intermediate = 2389 (8/3 × hidden_size) |
| Vocab size | 32,000 |
| Max position | 2,048 |
| **Total params** | **~432M** |
| **Active params** | **~176M** |

## Hardware & Environment

- GPU: NVIDIA RTX A5000 24GB
- CPU: AMD Ryzen 7 5800X
- RAM: 32GB

## Training Details

- **Precision**: bf16 (model, AdamW states, autocast forward/backward)
- **Optimizer**: AdamW (β=0.9,0.95, wd=0.1), weight decay only on ≥2D parameters
- **LR schedule**: Cosine decay with linear warmup (warmup_steps=2000)
- **Gradient clipping**: 1.0
- **Batch**: 22 × 1024 = 22,528 tokens/step
- **Activation checkpointing**: enabled (required for 24GB VRAM)
- **Data**: C4-en, pre-tokenized to `.pt` shards (100M tokens each), cycled round-robin

## Results

### Training Summary

| Phase | Data | Tokens | Steps | Time | Start Loss | End Loss |
|-------|------|:------:|:-----:|:----:|:----------:|:--------:|
| Pretrain | C4-en | 4B | 177,557 | 77.1 h | 10.5 | 3.0 |
| Post-train | FineWeb-Edu + Cosmopedia v2 (50:50) | 1B | 44,390 | 20.8 h | 3.05 | 2.0 |

Post-training used learning rate 5e-5 with 300-step re-warmup, continuing from the pretrain checkpoint with AdamW momentum preserved.

### GLUE (zero-shot)

| Task | Metric | Pretrain (4B C4) | Post-train (5B total) |
|------|--------|:---:|:---:|
| SST2 | accuracy | 0.470 | **0.554** |
| MRPC | accuracy / f1 | 0.338 / 0.069 | **0.706 / 0.815** |
| QQP | accuracy / f1 | 0.470 / 0.412 | **0.530** / 0.342 |
| QNLI | accuracy | 0.494 | 0.452 |
| RTE | accuracy | 0.520 | 0.484 |
| CoLA | MCC | 0.089 | 0.006 |
| MNLI | accuracy | 0.348 | 0.348 |
| MNLI-mm | accuracy | 0.368 | 0.368 |
| **Mean** | — | **0.403** | **0.483** |

### ARC

| Task | Pretrain (4B C4) | Post-train (5B total) |
|------|:---:|:---:|
| ARC-C 0-shot | 0.220 | **0.233** |
| ARC-C 5-shot | 0.223 | **0.246** |
| ARC-E 0-shot | 0.311 | **0.342** |
| ARC-E 5-shot | 0.320 | **0.348** |

Zero-shot evaluation uses conditional log-likelihood scoring over answer spans. All evals run on a single GPU with `--limit 500 --batch-size 16 --max-length 512`. Pretrain and post-train evaluated under identical settings for fair comparison.

## License

MIT License. Copyright (C) 2026 Michael Lee (李登淳).