Text Generation
Transformers
PyTorch
tinymixtral
conversational
custom_code
tinymixtral / README.md
mikecovlee's picture
Update README.md
ffd87a1 verified
|
Raw
History Blame Contribute Delete
2.84 kB
metadata
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 (李登淳).