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