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

python3 -m venv .venv
source .venv/bin/activate
pip install torch transformers datasets safetensors huggingface_hub

Single GPU Smoke Test

python -m ar.train.train --model_size small --max_steps 100 --batch_size 2 --seq_len 1024

Distributed Multi-GPU Training (FSDP2)

# 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

# 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