File size: 5,389 Bytes
2ae6691 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | ---
license: apache-2.0
library_name: transformers
pipeline_tag: text-generation
tags:
- text-generation
- causal-lm
- pretrained
- en
language:
- en
datasets:
- HuggingFaceFW/fineweb-edu
- HuggingFaceTB/finemath
- HuggingFaceTB/smollm-corpus
- bigcode/the-stack-dedup
---
# AR-4B (base)
AR-4B is a 4.02B-parameter decoder-only transformer, pretrained from scratch on about **325B tokens** of English, code and math. It's a **base** model: no instruction tuning or RLHF, so it continues text rather than following chat instructions.
- **Parameters:** 4.02B (36 layers, hidden size 2,560, grouped-query attention, SwiGLU MLP)
- **Context length:** 4,096 during pretraining (position encoding supports longer)
- **Pretraining tokens:** ~325B (Stage 1 ~300B general + Stage 2 ~25B math/code anneal)
- **Precision:** bf16
## Load the model
```python
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
name = "AjayRangoji/AR-4B-base"
tok = AutoTokenizer.from_pretrained(name)
model = AutoModelForCausalLM.from_pretrained(
name, torch_dtype=torch.bfloat16, device_map="auto",
trust_remote_code=True, # loads modeling_ar.py from this repo
)
prompt = "The quickest way to sort a list in Python is"
inputs = tok(prompt, return_tensors="pt").to(model.device)
out = model.generate(**inputs, max_new_tokens=64, do_sample=False)
print(tok.decode(out[0], skip_special_tokens=True))
```
`trust_remote_code=True` is required because the repo ships its own small `modeling_ar.py`.
## Architecture
| | |
|---|---|
| Layers | 36 decoder blocks, pre-norm, residual |
| Hidden size | 2,560 |
| Attention | Grouped-query: 32 query heads, 8 key/value heads, head dim 128 |
| Attention stability | RMSNorm on each head's Q and K (QK-norm), before RoPE |
| Positions | RoPE, base theta = 1,000,000 |
| MLP | SwiGLU, inner size 9,728 |
| Normalization | RMSNorm (eps = 1e-6), computed in fp32 |
| Embeddings | Input embedding tied to output head (saves 389M parameters) |
| Vocabulary | 151,936 BPE tokens |
| Document separator | `<|im_end|>` (id 151645) |
Forward pass: token embedding -> 36 x [RMSNorm -> GQA (QK-norm, RoPE) -> add residual -> RMSNorm -> SwiGLU -> add residual] -> final RMSNorm -> tied output head -> 151,936 next-token scores per position.
## Training
Pure data parallelism with PyTorch FSDP2, bf16 compute / fp32 reductions, activation checkpointing on every block, FlexAttention with a causal + same-document mask.
| | Stage 1: general | Stage 2: reasoning anneal |
|---|---|---|
| Steps | 0 - 572,000 | 572,000 - 619,699 |
| Tokens | ~300B | ~25B |
| English | 65% FineWeb-Edu (191B unique) | 20% FineWeb-Edu, int_score >= 4 only |
| Textbooks | - | 10% Cosmopedia v2 |
| Code | 20% The Stack (dedup), 9 languages | 30% same code |
| Math | 15% FineMath 4+ / 3+ | 40% FineMath 4+, InfiWebMath 4+, FineMath 3+ |
| LR | warmup 2,000 -> 3e-4 -> cosine -> 3e-5 | warmup 1,000 -> 1e-4 -> cosine -> 1e-5 |
- **Batch:** 524,288 tokens per optimizer step (64 GPUs x 2 sequences x 4,096 tokens).
- **Optimizer:** AdamW, beta = (0.9, 0.95), weight decay 0.1, gradient clip 1.0.
- **Hardware:** Up to 8 nodes x 8 H100 80GB (64 GPUs) on an AWS cluster, EFA network, Lustre storage.
- **Throughput:** about 500K tokens/s on 64 GPUs (~22% MFU).
## Evaluation (base model, no instruction tuning)
Held-out loss on data the model never trained on; HumanEval and GSM8K are standard benchmarks for 4B base models.
| Metric | Stage 1 final | **Stage 2 final** |
|---|---|---|
| Python held-out loss (lower is better) | 0.966 | **0.929** |
| English held-out loss (lower is better) | 2.220 | 2.300 |
| HumanEval pass@1 (greedy) | 18.9% (31/164) | **25.0% (41/164)** |
| GSM8K 8-shot exact match | 31.6% (79/250) | **43.6% (109/250)** |
As expected from the Stage 2 mix (40% math, 30% code), reasoning benchmarks improved substantially while general English loss moved slightly up.
## Tokenizer
Uses the Qwen3 BPE tokenizer (`Qwen/Qwen3-4B`) - 151,936 vocab, strong on English, code and math. Credit to the tokenizer's authors. Document separator in training: `<|im_end|>` (id 151645).
## Intended use and limitations
- **Base model** for continued pretraining, SFT, RLHF, or distillation. Not a chat model.
- **No safety training.** Outputs may be factually wrong, biased or unsafe. Review and filter before any downstream use.
- English-dominant data. Behavior in other languages is likely weaker, even where the tokenizer handles them.
- 4,096-token context in pretraining; longer contexts work mechanically through RoPE but have not been trained.
- Not evaluated on standard safety benchmarks.
## Data sources
- **FineWeb-Edu** - [HuggingFaceFW/fineweb-edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu)
- **FineMath 4+ / 3+** and **InfiWebMath 4+** - [HuggingFaceTB/finemath](https://huggingface.co/datasets/HuggingFaceTB/finemath)
- **Cosmopedia v2** - [HuggingFaceTB/smollm-corpus](https://huggingface.co/datasets/HuggingFaceTB/smollm-corpus)
- **The Stack (dedup)** - [bigcode/the-stack-dedup](https://huggingface.co/datasets/bigcode/the-stack-dedup) (Python, JavaScript, Java, C++, TypeScript, Go, Rust, Shell, SQL)
Each dataset keeps its original license; see the links above.
## License
Model weights and code: **Apache 2.0**. The tokenizer's license is the one it ships with.
|