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