Skip to content

Releases: AntonioRoye/TinyRecursiveModels

v1.0.1-arxiv — TRM Reproduction Bundle

Choose a tag to compare

@AntonioRoye AntonioRoye released this 23 Nov 19:29

Provenance

  • Upstream repository: SamsungSAILMontreal/TinyRecursiveModels
  • Upstream commit pinned: e7b6871
  • Additions in this fork: experiments/* (evaluation utilities), modal_trm.py (TRM via Modal), modal_llama.py (Llama baseline via Modal), cleaned READMEs with exact commands.

Summary

  • Exp1: Paper-mode (≤1000× voting) vs single‑augmentation evaluation with per‑example logs and bootstrap CIs.
  • Exp3: Naive Llama‑3‑8B QLoRA baseline (Modal).
  • Exp6: Efficiency benchmark (Modal) reporting throughput, latency, and peak VRAM.
  • CUDA/Torch standardized to CUDA 12.1; local and Modal paths both supported.

Requirements (summary)

  • Python 3.10, CUDA 12.1 (or similar)
pip install --upgrade pip wheel setuptools
pip install --upgrade torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
pip install -r requirements.txt

Data preparation

python -m experiments.prepare_arc_datasets
# writes data/arc-aug-0, data/arc-aug-1000 and data/arc_exp1_manifest.json

Fetch verification checkpoint

python -m experiments.fetch_hf_checkpoints \
  --repo_id arcprize/trm_arc_prize_verification \
  --dest checkpoints/hf_trm

Exp1 — Evaluate TRM (local, paper vs single‑aug)

python -m experiments.run_exp1 \
  data_paths='[data/arc-aug-1000]' \
  data_paths_test='[data/arc-aug-0]' \
  load_checkpoint=/ABS/PATH/TO/step_XXXXX \
  checkpoint_path=checkpoints/exp1_arc \
  arch=trm arch.L_cycles=4 arch.H_cycles=3 arch.L_layers=2 \
  global_batch_size=256
# If you encounter OOM, try: global_batch_size=128

Exp1 — Evaluate TRM (Modal)

# Build datasets in Modal volume
modal run modal_trm.py::main --action prepare_data
# or smaller pair matching HF v1
modal run modal_trm.py::main --action build_arc1

# Optional: fetch checkpoint into Modal volume
modal run modal_trm.py::main --action fetch_ckpt \
  --repo arcprize/trm_arc_prize_verification \
  --dest checkpoints/hf_trm

# Run both modes (paper + single)
modal run modal_trm.py::main --action exp1 \
  --checkpoint /workspace/TinyRecursiveModels/checkpoints/hf_trm/<subdir>/step_XXXXX

# Paper-only / Single-only
modal run modal_trm.py::main --action exp1_paper_only \
  --checkpoint /workspace/TinyRecursiveModels/checkpoints/hf_trm/<subdir>/step_XXXXX
modal run modal_trm.py::main --action exp1_single_only \
  --checkpoint /workspace/TinyRecursiveModels/checkpoints/hf_trm/<subdir>/step_XXXXX

# Quick sanity (few batches)
modal run modal_trm.py::main --action quick_exp1 \
  --checkpoint /workspace/TinyRecursiveModels/checkpoints/hf_trm/<subdir>/step_XXXXX

Optional — Accuracy vs augmentation curve (local)

python -m experiments.acc_vs_aug \
  load_checkpoint=/ABS/PATH/TO/step_XXXXX \
  checkpoint_path=checkpoints/exp1_arc_curve \
  arch=trm

Exp6 — Efficiency benchmark (Modal)

modal run modal_trm.py::main --action exp6 \
  --checkpoint /workspace/TinyRecursiveModels/checkpoints/hf_trm/<subdir>/step_XXXXX

Exp3 — Llama‑3‑8B QLoRA baseline (Modal)

  • If the base model is gated, provide an HF token via a Modal secret named pipeline-secrets (key: HF_TOKEN), or set HUGGING_FACE_HUB_TOKEN/HF_TOKEN in the environment.
modal run modal_llama.py::main --action prepare
modal run modal_llama.py::main --action train
modal run modal_llama.py::main --action eval

Optional quick dataset check

python -m experiments.quick_checks --data_dir data/arc-aug-1000

Outputs

  • checkpoints/exp1_arc/exp1_report.json (environment, timing, memory, dataset/checkpoint hashes)
  • checkpoints/exp1_arc/exp1_table.json (means, 95% CIs, vote histograms)
  • checkpoints/exp1_arc/evaluator_ARC_step_0/per_example.jsonl (per-example results, num_votes)
  • data/arc_exp1_manifest.json (dataset locations)

Environment capture (optional)

pip freeze > pip-freeze-v1.0.1-arxiv.txt
# For driver snapshot (optional): nvidia-smi -q > nvidia-smi-q.txt

Data and licensing

  • Do not redistribute ARC data. Generate processed arrays locally via:
python -m experiments.prepare_arc_datasets

Citation

  • See CITATION.bib for BibTeX entries (includes this report and the TRM paper).