Views
No views yet
[!IMPORTANT] 🚀 Pre-training Base Checkpoint Status Notice: This repository contains the 100% pre-trained base model checkpoint (10.0 Billion tokens). It represents the raw foundational model before instruction tuning. Supervised Fine-Tuning (SFT), DPO alignment, and specialized domain instruction-tuning pipelines are coming next! 🍿🤖⚡
1Ultron-114M Layout:
2├── Parameters : 114,053,376 (114M)
3├── Layers : 12 Transformer blocks
4├── Embedding (C) : 768 hidden dimension
5├── Attention Heads : 12 Query heads, 4 Key/Value heads (GQA 3:1 ratio)
6├── Head Dimension : 64
7├── Context Window : 1,024 tokens (RoPE frequency base 10,000)
8├── FFN Activation : SwiGLU (Tensor Core aligned to multiples of 64)
9├── Normalization : RMSNorm (with QK-head normalization, eps=1e-5)
10├── Logit Regularizer : Soft-Capping (cap=15.0 via tanh)
11├── Linear Projections: 100% Bias-Free (bias=False across all layers)
12├── Optimizer : Dual-Optimizer (Muon for 2D body, Fused AdamW for 1D/embeddings)
13└── Dataset & Tokens : FineWeb-Edu (10.0B tokens across 152,587 steps)1 Input Token IDs
2 │
3 ▼
4 Token Embedding (SmolLM Vocab: 49,152)
5 │
6 ▼
7 ┌───────────────────────────────────┐
8 │ 12 × Decoder Layer Stack │
9 │ │
10 │ ┌───────────────────────────┐ │
11 │ │ RMSNorm │ │
12 │ └─────────────┬─────────────┘ │
13 │ │ │
14 │ ▼ │
15 │ ┌───────────────────────────┐ │
16 │ │ GQA (12 Q / 4 KV) + RoPE │ │
17 │ │ └─ QK-Head RMSNorm │ │
18 │ └─────────────┬─────────────┘ │
19 │ │ │
20 │ ▼ │
21 │ Residual ───(+) │
22 │ │ │
23 │ ▼ │
24 │ ┌───────────────────────────┐ │
25 │ │ RMSNorm │ │
26 │ └─────────────┬─────────────┘ │
27 │ │ │
28 │ ▼ │
29 │ ┌───────────────────────────┐ │
30 │ │ SwiGLU FFN │ │
31 │ └─────────────┬─────────────┘ │
32 │ │ │
33 │ ▼ │
34 │ Residual ───(+) │
35 └─────────────────┬─────────────────┘
36 │
37 ▼
38 Final RMSNorm
39 │
40 ▼
41 LM Head Linear Projection
42 │
43 ▼
44 Logit Soft-Capping (cap=15.0)
45 │
46 ▼
47 Output Logits| Feature | GPT-2 (124M) | Ultron-114M | Why it Matters (Engineering Justification) |
|---|---|---|---|
| Positional Encoding | Absolute Learned (wpe) | RoPE (Rotary) | Enables zero-shot context length extension and better relative distance modeling. |
| Attention Mechanism | Multi-Head (MHA) | Grouped-Query (GQA) | 12 Q heads : 4 KV heads (3:1 ratio), reducing KV-cache memory usage during inference by 3×. |
| Attention Stability | Standard Unnormalized | QK-Head RMSNorm | Prevents logit explosion / attention entropy collapse during long pre-training runs. |
| FFN Activation | Standard GELU | SwiGLU | Gated non-linearity yielding higher model capacity per FLOP; aligned to multiples of 64 for Tensor Core throughput. |
| Layer Normalization | LayerNorm (with bias) | RMSNorm (Bias-Free) | Eliminates mean-centering overhead; 100% bias-free projections (bias=False) for cleaner gradient dynamics. |
| Logit Regularization | None | Logit Soft-Capping | Applies tanh capping (cap=15.0) to prevent overconfidence and extreme logit growth. |
| Optimizer Engine | AdamW | Muon + Fused AdamW | Uses Keller Jordan's Muon (Momentum Orthogonalized by 5th-order Newton-Schulz iterations) for 2D body weights. |
| Learning Rate Schedule | Cosine Decay | WSD Schedule | Warmup-Stable-Decay schedule (80% stable phase, 20% cosine decay), allowing flexible checkpoint annealing. |
| Mixed Precision | FP32 | Native BFloat16 (bf16) | Dynamic range stability without loss scalers on RTX 30xx/40xx/50xx GPUs. |
| Graph Compiler | None | PyTorch 2.0 (torch.compile) | Fuses element-wise operations and kernel launches via Inductor. |
15.0 * tanh(logits / 15.0) applied prior to loss calculation to prevent logit explosion.AdamW for 1D vectors and embeddings.backend_tokenizer.encode_batch streaming at ~4.34 Million tokens/sec into compact uint16 binary shards.np.memmap disk slicing streams 10.0B tokens with <500MB host RAM overhead.| Parameter | Value | Description |
|---|---|---|
| Model Name / Tag | Ultron-114M | Official parameter tag (114,053,376 total parameters) |
| Layers / Query Heads / KV Heads | 12 layers / 12 Q-heads / 4 KV-heads | GQA Transformer layout ($C=768, n_{head}=12, n_{kv_head}=4$) |
| Context Window ($T$) | 1,024 tokens | Sequence length per pass |
| Micro-Batch Size ($B$) | 16 | Per-GPU micro-batch size |
| Gradient Accumulation | 4 steps | Effective batch size = 64 sequences (65,536 tokens/step) |
| Tokenizer | SmolLM Vocab (49,152) | Efficient BPE tokenizer (HuggingFaceTB/SmolLM2-135M) |
| Precision | BFloat16 (bf16) | Native mixed precision |
| LR Schedule | WSD | Warmup-Stable-Linear-Decay (80% stable, 20% linear decay) |
| Optimizer | Muon + Fused AdamW | Newton-Schulz matrix optimizer ($LR=0.04$) + fused AdamW ($LR=1.2\times 10^{-3}$) |
| Throughput | ~186,310 tok/sec (~2.80 step/sec) | Benchmarked on single NVIDIA RTX 5090 GPU (32GB) |
| GPU VRAM Allocation | ~16.2 GB / 32 GB | Measured via nvidia-smi during active pre-training |
| Total Pre-training Time | 15 Hours 1 Minute (54,063s) | 10.0 Billion Tokens / 152,587 total steps (100% Complete) |
1ultron/
2├── model.py # PyTorch Ultron-114M (RoPE + GQA + SwiGLU + RMSNorm + QKNorm + Logit SoftCap)
3├── config.py # Model & Hyperparameter Configuration Dataclass
4├── dataset.py # Zero-Copy Memmap Sharded Dataset Loader
5├── train.py # Main Accelerated Distributed Training Runner
6├── trainer.py # Trainer Class with Keller Jordan Muon + Fused AdamW
7├── telemetry.py # Telemetry & Experiment Tracking Manager (W&B + ETA + Checkpoint state)
8├── requirements.txt # Virtual environment dependencies
9├── accelerate_checkpoint/ # Saved Accelerate model weights, optimizer state & RNG seeds
10├── shards_edu/ # Binary FineWeb-Edu tokenized data shards (.bin)
11├── logs/ # Dedicated logs directory (loss_curve.svg plot & benchmark JSON evaluations)
12├── wandb/ # Local step telemetry logs & experiment tracking runs
13├── .agents/ # Project AGENTS.md rules & workspace customization
14├── tests/ # Unit & Integration Tests (Accelerate + torch.testing)
15│ └── test_model.py # Core model architecture & generation unit tests
16└── scripts/ # Helper Scripts
17 ├── generate.py # Text generation from local Accelerate checkpoint
18 ├── tokenize_dataset.py # FineWeb-Edu dataset tokenization into binary shards
19 ├── eval_lm_harness.py # EleutherAI lm-evaluation-harness benchmark script
20 ├── upload_checkpoint.py# Hugging Face Hub model checkpoint uploader script
21 └── upload_dataset_shards.py # Hugging Face Hub dataset shards uploader scriptjaipkapoor99/ultron-124mjaipkapoor99/ultron-fineweb-edu-shards1git clone https://github.com/jaipkapoor99/ultron.git
2cd ultron
3
4# Fast environment setup using uv
5uv venv --python 3.13 venv
6source venv/bin/activate
7uv pip install -r requirements.txt nvidia-cuda-nvccpython3 scripts/tokenize_dataset.pyaccelerate config| Setting | Value | Why |
|---|---|---|
| Compute environment | Local machine | Single-node training |
| Distributed type | NO | Single GPU |
| Mixed precision | bf16 | Required for peak throughput on RTX 30xx/40xx/50xx |
| TorchDynamo backend | INDUCTOR | Enables torch.compile graph compilation |
accelerate launch train.py| Metric | Recorded Value | Description |
|---|---|---|
| Total Steps Completed | 152,587 / 152,587 (100%) | Full pre-training run on FineWeb-Edu |
| Total Tokens Processed | ~10.0 Billion Tokens | 65,536 tokens per step (batch size 64 $\times$ seq len 1,024) |
| Step Throughput | ~2.80 iterations/sec | 2.79–2.82 it/s continuous speed |
| Token Throughput | ~186,310 tokens/sec | SOTA Muon + PyTorch 2.0 compile throughput |
| Compute Hardware | NVIDIA RTX 5090 (32GB) | Native BFloat16 (bf16) mixed precision |
| Total Wall-Clock Time | 15 Hours 1 Minute (54,063s) | Completed full 10B token pre-training |
Final Validation (dev_loss) | 2.9683 | Evaluated on validation set at step 152,587 |
Final Train Loss (train_loss) | 2.9038 | 100-step moving average at step 152,587 |
lm-evaluation-harness Baseline Benchmark Reportscripts/eval_lm_harness.py (Results stored in logs/pre_training_checkpoint_eval.json):accelerate launch scripts/eval_lm_harness.py --limit=0| Benchmark Task | Benchmark Domain | Un-truncated Test Size | Pre-SFT Accuracy | Random Guess Baseline |
|---|---|---|---|---|
piqa | Physical Commonsense Reasoning | 1,838 samples | 63.33% | 50.00% |
arc_easy | Elementary Science QA | 2,376 samples | 54.42% | 25.00% |
winogrande | Pronoun Resolution & Commonsense | 1,267 samples | 51.07% | 50.00% |
hellaswag | Sentence Completion & Reasoning | 10,042 samples | 30.39% | 25.00% |
arc_challenge | Advanced Science Reasoning | 1,172 samples | 24.06% | 25.00% |
openbookqa | Open Book Science QA | 500 samples | 18.80% | 25.00% |
[!NOTE] Loss Trajectory & WSD Decay Analysis: During the final WSD cosine decay phase (steps 150,000–152,587), the moving averagetrain_loss(~2.85) dropped slightly below the validationdev_loss(2.9179). This ~0.06 delta is the expected mathematical outcome of learning rate annealing as step sizes approach zero, allowing the optimizer to settle efficiently into local minima while validation loss continuously improves.
ultron-pretraining project.telemetry.py and parse_plot_telemetry.py read binary .wandb logs directly from wandb/ on disk without network dependency.resume="allow" in setup_accelerator_trackers(). W&B runs now resume seamlessly across checkpoint restarts without step monotonicity conflicts.train/*, eval/*, and perf/* metrics are linked to the global step index with dev_loss set to summary="min".[!TIP] Engineering Takeaway — Master W&B & Telemetry Pipeline: "There is no data science without data." Resolving metric step alignment and offline binary.wandblog parsing reinforced the importance of mastering telemetry pipelines, structured metric registration (define_metric), and experiment tracking early in large-scale pre-training projects.
scripts/generate.py at step 152,587 (~70 tokens each, temperature 0.85, top-k 50):accelerate launch scripts/generate.py --prompt "..." --max-tokens 70Prompt: "The laws of thermodynamics state that"
Output: The laws of thermodynamics state that the volume of an isolated gas undergoing a reaction is always greater than the volume of the gas initiating the reaction. This is the theory behind nuclear reactions. There is also the theory of quantum mechanical systems, that properties of an object can only be represented through atomic interaction…
Prompt: "To learn mathematics effectively, one must"
Output: To learn mathematics effectively, one must pay careful attention to the steps that follow in a given order. The most important ones are presented as a series. Here we consider the square root of 7 as a series of steps to solve algebraic equations…
Prompt: "As the quantum portal activated, the crew stared into the alien"
Output: As the quantum portal activated, the crew stared into the alien space through the lens and the camera could see the entire universe as if it was inside the universe. To accomplish this, a quantum bit was stored in a device called a crystal…
python3 failed with RuntimeError due to uninitialized process groups. Standardizing accelerate launch across all entry points solved device allocation cleanly.bf16 + torch.compile provided superior stability and peak throughput (~186.3k tok/s) without framework friction.venv) & C-Header ManagementPython.h): torch.compile() failed on Python 3.14 due to missing C headers (Python.h: No such file or directory). Switching virtual environments to Python 3.13 via uv provided standalone C-headers natively, eliminating compiler breakage.muon via pip pulled down an unrelated single-cell bioinformatics library instead of Keller Jordan's neural network optimizer. Resolved by importing muon-optimizer.np.memmap)for-loop tokenization with Rust backend_tokenizer.encode_batch (num_threads=1 per worker process) increased dataset streaming speed by >100x from 40k tok/s to ~4.34 Million tokens/sec!uint16 binary shards (shards_edu/*.bin) enabled zero-copy memory mapping (np.memmap), allowing 10.0B token streaming with <500MB host RAM usage.nanoGPT project.1@misc{jordan2024muon,
2 author = {Jordan, Keller and Jin, Yuchen and Boza, Vlado and You, Jiacheng and Cesista, Franz and Newhouse, Laker and Bernstein, Jeremy},
3 title = {Muon: An optimizer for hidden layers in neural networks},
4 year = {2024},
5 url = {https://kellerjordan.github.io/posts/muon/}
6}