Skip to content

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

3 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

How Much Do Common Training Optimizations Cost You?

A Systematic Ablation Study on the BabyLM Strict-Small Track

Trains a 33.6M-parameter GPT-2-style decoder on the BabyLM 2026 Strict-Small dataset (10M words, ≤10 epochs) across 15 combinations of 5 training optimizations (FP16, Flash Attention, Curriculum Learning, Dynamic Batching, INT8 post-training Quantization), each validated across n=5 seeds with GPU execution verified deterministic — to test whether practitioners' assumption that these optimizations "compose for free" actually holds.

Venue: BabyLM Challenge 2026 Workshop @ EMNLP, Budapest (submission deadline: 20 July 2026, AoE) Model (public, HuggingFace Hub): AbdulRahmanIqbal/babylm-baseline-10epmain branch is the final checkpoint; branches chck_1Mchck_100M are the 19 official word-count-milestone checkpoints (used for AoA and fast-eval-across-checkpoints) Leaderboard entry: optimization-ablation-baseline (BabyLM 2026 Strict-Small track)


Headline Results

Rank Config Mean PPL (n=5, 10 epochs) Quant $\Delta$ (n=5)
🥇 1 Baseline (plain FP32) 136.786 ± 0.838 +1.074 ± 0.130
2 FP16 + Flash Attention 136.822 ± 0.582 +0.959 ± 0.110
3–6 Flash+Quant / Flash / FP16 / FP16+Quant 137.06–137.31 0.96–1.12
7–10 Curriculum-learning variants 139.9–140.2 0.76–0.92
11–15 Dynamic-batching variants (worst: All Five) 143.1–147.1 0.64–0.89

Zero-shot downstream (n=5 seeds, final): BLiMP 62.44%±0.47, Supplement 58.13%±0.81, COMPS (base) 54.11%±0.35, WUGs 50.16%±0.22, EWoK 50.89%±0.71, Entity Tracking 18.54%±0.25, GlobalPIQA (non-parallel/parallel) 46.80%±3.90 / 20.00%±2.13, Reading eye-tracking 0.53±0.12, self-paced 0.00±0.00 (null), AoA correlation −0.031±0.070 (weak, seed-unstable). GLUE fine-tuning (n=5 seeds): average accuracy 62.57% across BoolQ/MultiRC/RTE/WSC/MRPC/QQP/MNLI.

Three headline findings:

  1. Baseline (no optimizations) wins outright at 10 epochs, statistically tied only with FP16+Flash (Welch's t-test, $p$-values and Cohen's $d$ computed for every configuration).
  2. Curriculum learning's advantage reverses with more training. It wins by -3.99 PPL at a 6-epoch budget but loses to plain Baseline at 10 epochs — traced to its easy-data phase being a fixed number of epochs, not a fixed fraction of the budget.
  3. INT8 quantization costs 50–90× more than an earlier, buggy measurement suggested, and configurations using curriculum and/or dynamic batching are consistently more quantization-robust than plain FP32/FP16/Flash configurations.

What's in this repo (and what isn't)

This repo has the training/analysis code and the 75 seed-run logs (15 configs × 5 seeds) produced by the full ablation run. It does not include:


Requirements

  • Python 3.11+ (tested on 3.13)
  • NVIDIA GPU with CUDA (tested on RTX 4070 Ti, 12GB VRAM) — training is feasible on 6GB+ GPUs but will be slower
  • ~15GB free disk for checkpoints during a full 15×5 run (checkpoints are not kept after computing results; peak usage during the run is higher)
python -m venv venv
venv\Scripts\activate          # Windows
# source venv/bin/activate     # Linux/Mac

pip install torch transformers datasets tiktoken numpy scipy matplotlib tqdm

Tested with: torch==2.12.1+cu126, transformers==5.12.1, datasets==5.0.0, tiktoken==0.13.0, numpy==2.4.6, scipy==1.17.1, matplotlib==3.11.0. No requirements.txt is pinned beyond this — the codebase has no unusual dependency requirements.


Full Reproduction Steps

1. Data (automatic on first run)

dataset.py's prepare_data() auto-downloads and tokenizes BabyLM-2026-Strict-Small (10M words, GPT-2 BPE) into data/train.bin + data/val.bin the first time any training script runs. Requires internet access to Hugging Face on first run only.

2. Train all 15 configurations × 5 seeds (deterministic)

python seed_validator.py --folder .

Takes ~24–31 hours on an RTX 4070 Ti. Produces seed_results_babyllm.csv (mean/std/per-seed PPL for all 15 configs). Safe to interrupt and resume — it skips configs/seeds already present in the CSV and in each experiments/<exp>/seed_val/seed_<N>/ folder.

Determinism (torch.backends.cudnn.deterministic, torch.use_deterministic_algorithms, CUBLAS_WORKSPACE_CONFIG) is already patched into all 15 training scripts (patch_determinism.py, applied once — re-running it is safe/idempotent and skips already-patched files). Without this, GPU execution is not bit-reproducible for a fixed seed and the top-tier configs' ranking can reshuffle between reruns — see Known Issues.

3. Measure real INT8 quantization cost

python run_quantize_validation.py --folder .     # the 4 configs originally paired with a "+Quant" variant (3C, 3F, 3H, 4)
python run_quantize_extra11.py --folder .        # post-hoc quantization of the other 11 configs' checkpoints

Both are CPU-only and fast (~5–10 min each). Produce seed_results_babyllm_quant.csv and seed_results_babyllm_quant_extra11.csv respectively — together, real quantization deltas for all 15 configurations.

4. Pick the winning model + convert to HuggingFace format

python convert_to_hf.py

Converts seed_val/seed_789/best_model.pt (Baseline, the best individual seed) to HF GPT2LMHeadModel format in a local hf_model/ folder (not checked into this repo — see above). To convert a different seed/config, set env vars before running:

HF_CKPT_OVERRIDE=path/to/best_model.pt HF_OUT_OVERRIDE=path/to/output_dir python convert_to_hf.py

5. Run the official eval pipeline

Clone the harness first — it's a separate repo, not included here:

git clone https://github.com/babylm-org/babylm-eval
cd babylm-eval/strict
python -m scripts.download_evals            # downloads BLiMP/COMPS/GLUE/etc (needs `hf auth login` first)
python -m evaluation_pipeline.ewok.dl_and_filter   # EWoK is gated — accept terms at huggingface.co/datasets/ewok-core/ewok-core-1.0 first, then run this

# Zero-shot (BLiMP, supplement, COMPS, EWoK, entity tracking, reading):
bash scripts/eval_zero_shot.sh AbdulRahmanIqbal/babylm-baseline-10ep causal

# GLUE fine-tuning (7 tasks):
touch ../.env    # empty is fine unless you pass --wandb
bash scripts/eval_finetuning.sh --model_path AbdulRahmanIqbal/babylm-baseline-10ep

You can point these scripts at our public HF model directly (as above), or at your own locally-converted hf_model/ from step 4. See Known Issues for real bugs (ours and upstream) you may hit running this.

6. Generate figures

python figures/generate_babylm_figures.py

Produces figures/fig_ablation.pdf, fig_training_curves.pdf, fig_quant_delta.pdf, fig_architecture.pdf. Colors validated against CVD/contrast checks (see the dataviz methodology referenced in the figure-generation script).


Project Structure

.
├── config.py, model.py, dataset.py, train.py    # Baseline (FP32) training
├── patch_seeds.py, patch_determinism.py          # One-time patchers (already applied)
├── seed_validator.py                             # Runs all 15 configs × 5 seeds
├── run_quantize_validation.py                    # Quant deltas for the original 4 "+Quant" configs
├── run_quantize_extra11.py                       # Quant deltas for the other 11 configs (post-hoc)
├── convert_to_hf.py                              # Our format → HuggingFace GPT2LMHeadModel
├── train_aoa_checkpoints.py, run_aoa_eval.py     # AoA: 19-checkpoint retrain + surprisal/scoring
├── extract_efficiency_metrics.py, measure_gpu_memory.py   # Wall-clock/tok-s/GPU-memory measurement
├── run_anova_analysis.py                         # Two-way interaction OLS/ANOVA
├── extract_rerun_results.py                      # Parses zero-shot report files into summary tables
├── prepare_leaderboard_checkpoints.py            # Converts the 19 AoA checkpoints to HF format
├── upload_leaderboard_repo.py                    # Pushes them to the public HF repo as branches
│
├── experiments/
│   ├── 2a_fp16/, 2b_flash/, 2c_curriculum/, 2d_dynamic_batch/
│   ├── 3a_fp16_flash/ … 3i_curriculum_dynbatch/  # all pairwise combinations
│   └── 4_all_five/                                # all 5 combined
│       each with: train_*.py, quantize_*.py (if applicable), seed_val/seed_<N>/
│
├── seed_val/seed_<42,123,456,789,2024>/           # Baseline's own seed runs (root-level config)
│                                                   # (training_log.txt / steps_log.txt only — no checkpoints)
├── seed_results_babyllm.csv                       # Main results: mean/std/per-seed PPL, all 15 configs
├── seed_results_babyllm_quant.csv                 # Quant deltas: original 4 configs
├── seed_results_babyllm_quant_extra11.csv         # Quant deltas: other 11 configs
├── efficiency_wallclock_tokspeed.csv, efficiency_gpu_memory.csv
├── anova_results.csv, anova_summary.txt
│
└── figures/generate_babylm_figures.py             # Figure generation (+ output .pdf's)

Each experiment folder follows the same layout:

  • train_*.py — standalone training script (already patched for seeds + determinism)
  • quantize_*.py — post-training INT8 quantization script (only for 3C, 3F, 3H, 4; parametrized via QUANT_CKPT_OVERRIDE/QUANT_OUT_OVERRIDE env vars so run_quantize_extra11.py can reuse them generically)
  • seed_val/seed_<N>/training_log.txt (per-eval-checkpoint), steps_log.txt (per-step) — the 75 seed-run logs from the full ablation

Known Issues We Hit and Fixed

If you're reproducing this from scratch, these will save you time:

  1. GPU nondeterminism reshuffles close rankings. Fixed by patch_determinism.py (already applied to all scripts in this repo). Without it, identical seeds do not reproduce bit-identical PPL run to run — verified 8/9 configs bit-exact after the fix; Flash-Attention-only configs retain small (<0.15 PPL) residual noise from the fused kernel's backward pass.
  2. A quantization step was silently never invoked. The original seed_validator.py design only ran each "+Quant" config's Phase-1 training script, never the Phase-2 quantize_*.py script — so early quantization deltas (~0.01–0.016 PPL) were actually just two runs of the identical pre-quantization script. Real deltas (measured via run_quantize_validation.py / run_quantize_extra11.py) are 50–90× larger.
  3. GPT-2 tokenizer has no pad_token by default, which crashes GLUE fine-tuning (ValueError: Asking to pad but the tokenizer does not have a padding token). Fixed in convert_to_hf.py (tokenizer.pad_token = tokenizer.eos_token).
  4. eval_finetuning.sh's default --sequence_length 512 exceeds our model's 256-token context window, causing a CUDA device-side assert (position-embedding index out of bounds) on any GLUE example longer than 256 tokens (hit on BoolQ/MultiRC/RTE). Fix: change all 5 --sequence_length 512 occurrences in babylm-eval/strict/scripts/eval_finetuning.sh to 256.
  5. EWoK's vocab.txt isn't UTF-8-safe under Windows' default cp1252 codec. Fix: add encoding='utf-8' to the two open() calls in evaluation_pipeline/ewok/dl_and_filter.py.
  6. On Windows, plain bash script.sh can silently resolve to the WSL launcher stub (C:\Windows\System32\bash.exe) instead of Git Bash, if WSL has no installed distro it fails instantly with no useful error in a redirected log. Call Git Bash by its full path instead: & "C:\Program Files\Git\bin\bash.exe" script.sh … (PowerShell) or ensure Git's bin precedes System32 in PATH.
  7. collate_preds.py's Entity Tracking size-check used pre-filter item counts (decode_entity_tracking() in read_files.py intentionally drops any item whose answer options contain "nothing", but the size constants weren't updated to match), and GlobalPIQA had no collation logic at all despite its eval scripts producing predictions — both silently zeroed those two benchmarks on the leaderboard. We reported both to the organizers with the exact root cause; both are now fixed upstream (as of the version referenced above).
  8. venv\Scripts\activate breaks non-interactive bash scripts on Windows if the venv was copied from another machine/path — its pyvenv.cfg/activate script can hardcode a stale absolute path, and sourcing it replaces bash's PATH with the raw Windows PATH, breaking bash/date/tee/etc. The venv's own python.exe already resolves correctly on PATH without activation; prefer export PATH="$VENV/Scripts:$PATH" in scripts that need to be explicit.

References

About

No description, website, or topics provided.

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages