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