Blame for README.md

---
language:
- en
license: apache-2.0
tags:
- pretraining
- causal-lm
- fsdp2
- pytorch
- h100
---

# AR (Autoregressive 4B) — Full Pretraining Codebase & Infrastructure

Pretraining codebase, cluster orchestration scripts, data pipeline, and evaluation harness for **AR-4B**, a 4.02B-parameter decoder-only transformer pretrained from scratch on ~325B tokens across up to 64x NVIDIA H100 GPUs.

- **Pretrained Weights & Model Card:** [AjayRangoji/AR-4B-base](https://huggingface.co/AjayRangoji/AR-4B-base)
- **Architecture:** 4.02B parameters, 36 layers, hidden size 2,560, GQA (32/8 heads, head dim 128), SwiGLU MLP (inner size 9,728), QK-norm RMSNorm, RoPE ($\theta = 1,000,000$), tied input/output embeddings, 151,936 vocab (Qwen3 BPE).
- **Scale:** ~325 Billion tokens (Stage 1 ~300B general + Stage 2 ~25B reasoning & code anneal).
- **Framework:** PyTorch FSDP2, bf16 compute / fp32 reductions, activation checkpointing, FlexAttention with document masking.

---

## Repository Structure

```
AR/
├── ar/                         # Core Python package
│   ├── model/                  # Model architecture (Qwen3 / Llama-style transformer)
│   │   ├── attention.py        # Grouped-Query Attention with QK-norm & RoPE
│   │   ├── block.py            # Transformer decoder block with pre-RMSNorm & residual
│   │   ├── config.py           # ModelConfig dataclass
│   │   ├── mask.py             # Intra-document causal attention masking
│   │   ├── mlp.py              # SwiGLU feed-forward network
│   │   ├── model.py            # Full AR transformer causal LM implementation
│   │   ├── rmsnorm.py          # fp32 RMSNorm implementation
│   │   └── rope.py             # Rotary Position Embedding
│   ├── train/                  # Distributed training routines
│   │   ├── train_fsdp.py       # Distributed FSDP2 training engine (multi-node, 64x H100)
│   │   ├── train.py            # Single-GPU training loop (development & testing)
│   │   ├── train_config.py     # Training hyperparameters & config
│   │   └── stage2.py           # Stage 2 data loader & reasoning anneal curriculum
│   ├── data/                   # Data preparation & dataset iterators
│   │   ├── prepare_data.py     # Tokenization pipeline for FineWeb, FineMath, The Stack
│   │   ├── dataset.py          # Chunked binary dataset reader
│   │   ├── readonce.py         # Memory-efficient streaming data loader
│   │   └── tokenizer.py        # Tokenizer wrapper
│   └── eval/                   # Benchmarking & checkpoint export
│       ├── val_loss.py         # Held-out validation loss evaluation
│       ├── gsm8k.py            # 8-shot GSM8K evaluation
│       ├── humaneval.py        # HumanEval pass@1 coding benchmark
│       ├── generate.py         # Greedy & top-p text generation
│       ├── batch_gen.py        # Batched inference
│       └── export_ckpt.py      # Weights-only bf16 checkpoint exporter
│
├── Slurm Job Scripts:
│   ├── train_real.sbatch       # Production multi-node training job launcher (FSDP2)
│   ├── tokenize.sbatch         # Parallel tokenization batch job
│   ├── dataeval.sbatch         # Dataset validation job
│   ├── sweep_lr.sbatch         # Learning rate sweep job
│   ├── test_ckpt.sbatch        # Checkpoint sanity testing
│   └── sanity_4node.sbatch     # Multi-node interconnect & cluster sanity check
│
├── Cluster Operations & Automation:
│   ├── watchdog.sh             # Autonomous monitoring daemon (detects hangs/crashes, auto-resumes)
│   ├── save_stage1.sh          # Waits for Stage 1 checkpoint & exports bf16 model
│   ├── save_final.sh           # Waits for final checkpoint (ckpt_619699) & exports bf16 model
│   ├── run_args.sh             # Training hyperparameters, data mix, learning rate schedule
│   ├── status.sh               # Live cluster & GPU status summary
│   ├── autoreload.sh           # Hot-reloads training code upon checkpoint completion
│   ├── cleanup.sh              # Checkpoint garbage collection & disk space management
│   ├── diskwatch.sh            # Disk utilization watcher
│   ├── eval_here.sh / eval_latest.sh # Quick local and remote evaluation runners
│   ├── find_nodes.sh / nodecheck.sh / testnodes.sh # Node diagnostics & bad-node isolation
│   └── swap.sh / swap2.sh      # Node swapping utilities
│
├── upload/                     # Hugging Face export & conversion utilities
│   ├── convert_and_check.py    # Converts PyTorch FSDP weights to safetensors & validates parity
│   ├── modeling_ar.py          # Hugging Face Transformers custom model definition
│   └── README.md               # Model card published with weights
│
└── PROGRESS.md                 # Training log notes, milestones, and decisions
```

---

## Training Methodology

### 1. Model Architecture
| Specification | Value |
|---|---|
| **Parameters** | 4,022,238,720 (4.02B) |
| **Layers** | 36 decoder blocks |
| **Hidden dimension** | 2,560 |
| **Heads (Q / KV)** | 32 query heads / 8 key-value heads (GQA) |
| **Head dimension** | 128 |
| **MLP inner dimension** | 9,728 (SwiGLU) |
| **Position encoding** | RoPE (base $\theta = 1,000,000$) |
| **Normalization** | RMSNorm ($\epsilon = 10^{-6}$), computed in fp32; QK-norm per head |
| **Embeddings** | Tied input embedding and output projection |
| **Vocabulary** | 151,936 BPE tokens (`Qwen/Qwen3-4B`) |
| **Context Length** | 4,096 tokens per sequence |

### 2. Multi-Stage Pretraining Curriculum
- **Stage 1 (General Pretraining, Steps 0 - 572,000, ~300B tokens):**
  - 65% FineWeb-Edu (English web)
  - 20% The Stack (dedup) across 9 programming languages
  - 15% FineMath 4+ / 3+
  - LR schedule: Warmup to $3 \times 10^{-4} \to$ Cosine decay to $3 \times 10^{-5}$
- **Stage 2 (Reasoning Anneal, Steps 572,000 - 619,699, ~25B tokens):**
  - 40% FineMath & InfiWebMath (high quality math)
  - 30% The Stack (code)
  - 20% FineWeb-Edu (educational quality $\ge 4$)
  - 10% Cosmopedia v2 (synthetic textbooks & tutorials)
  - LR schedule: Warmup to $1 \times 10^{-4} \to$ Cosine decay to $1 \times 10^{-5}$

### 3. Benchmarks & Results
| Benchmark | Stage 1 Final | Stage 2 Final |
|---|---|---|
| **Python held-out loss** | 0.966 | **0.929** |
| **English held-out loss** | 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)** |

---

## Quickstart

### Environment Setup
```bash
python3 -m venv .venv
source .venv/bin/activate
pip install torch transformers datasets safetensors huggingface_hub
```

### Single GPU Smoke Test
```bash
python -m ar.train.train --model_size small --max_steps 100 --batch_size 2 --seq_len 1024
```

### Distributed Multi-GPU Training (FSDP2)
```bash
# Launch across 8 GPUs on a single node:
torchrun --nproc_per_node=8 -m ar.train.train_fsdp

# Launch on Slurm cluster (64x H100):
sbatch train_real.sbatch
```

### Checkpoint Export & Hugging Face Conversion
```bash
# Export trained checkpoint to weights-only bf16:
python -m ar.eval.export_ckpt --ckpt /path/to/ckpt --out /path/to/model.pt

# Convert to Hugging Face safetensors format and verify numerical parity:
python upload/convert_and_check.py \
    --ckpt /path/to/model.pt \
    --out /path/to/hf_export \
    --tokenizer_src /path/to/tokenizer
```

---

## License
Apache-2.0