From ced865deba6a0da1fdb90a662b996f635d35be76 Mon Sep 17 00:00:00 2001 From: Sijia Liu Date: Mon, 20 Jul 2026 21:09:13 -0700 Subject: [PATCH 1/5] [2.0] Add nanoslm_hybrid_arch_design problem + reproduction rig Hybrid LM architecture-design task (Olmo Hybrid, ~190M): agents submit a model.py scored on held-out val_bpb vs a locked olmo3_190M baseline, trained under a fixed wall-clock budget on a single H100. - Embedding tying + param_cap reconciled: 400M cap in config.yaml and settings.py; baseline_model.py and reference.py both tie embeddings so the arms differ only in the sequence mixer; reference.py analytic/self-check updated to the tied count. - repro/: paper-faithful rig for Olmo Hybrid Table 5 (190M Base-Easy BPB) -- paper-spec pre-norm transformer + GDN-3:1 hybrid, WSD-S schedule, Modal training on the dolma3 mix. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../nanoslm_hybrid_arch_design/DESIGN.md | 264 ++++++++ .../nanoslm_hybrid_arch_design/PR_SUMMARY.md | 119 ++++ .../nanoslm_hybrid_arch_design/config.yaml | 67 ++ .../docker/agent/Dockerfile | 59 ++ .../docker/build_images.sh | 150 +++++ .../docker/judge/Dockerfile | 73 +++ .../docker/prep_assets.py | 294 +++++++++ .../nanoslm_hybrid_arch_design/evaluate.sh | 19 + .../nanoslm_hybrid_arch_design/evaluator.py | 584 ++++++++++++++++++ .../harbor/app/README.md | 97 +++ .../harbor/app/model.py | 434 +++++++++++++ .../harbor/app/public_test.py | 64 ++ .../harbor/app/public_test.sh | 4 + .../harness/__init__.py | 9 + .../harness/baseline_model.py | 314 ++++++++++ .../harness/data.py | 197 ++++++ .../harness/eval_ppl.py | 120 ++++ .../harness/modal_app.py | 277 +++++++++ .../harness/model_config.py | 37 ++ .../harness/policy.py | 124 ++++ .../harness/runner.py | 407 ++++++++++++ .../harness/scoring.py | 113 ++++ .../harness/settings.py | 281 +++++++++ .../harness/train.py | 148 +++++ .../nanoslm_hybrid_arch_design/readme | 169 +++++ .../nanoslm_hybrid_arch_design/reference.py | 434 +++++++++++++ .../repro/__init__.py | 1 + .../repro/modal_repro.py | 483 +++++++++++++++ .../repro/paper_models.py | 344 +++++++++++ .../nanoslm_hybrid_arch_design/repro/wsd.py | 91 +++ 30 files changed, 5777 insertions(+) create mode 100644 2.0/problems/nanoslm_hybrid_arch_design/DESIGN.md create mode 100644 2.0/problems/nanoslm_hybrid_arch_design/PR_SUMMARY.md create mode 100644 2.0/problems/nanoslm_hybrid_arch_design/config.yaml create mode 100644 2.0/problems/nanoslm_hybrid_arch_design/docker/agent/Dockerfile create mode 100755 2.0/problems/nanoslm_hybrid_arch_design/docker/build_images.sh create mode 100644 2.0/problems/nanoslm_hybrid_arch_design/docker/judge/Dockerfile create mode 100644 2.0/problems/nanoslm_hybrid_arch_design/docker/prep_assets.py create mode 100755 2.0/problems/nanoslm_hybrid_arch_design/evaluate.sh create mode 100644 2.0/problems/nanoslm_hybrid_arch_design/evaluator.py create mode 100644 2.0/problems/nanoslm_hybrid_arch_design/harbor/app/README.md create mode 100644 2.0/problems/nanoslm_hybrid_arch_design/harbor/app/model.py create mode 100644 2.0/problems/nanoslm_hybrid_arch_design/harbor/app/public_test.py create mode 100755 2.0/problems/nanoslm_hybrid_arch_design/harbor/app/public_test.sh create mode 100644 2.0/problems/nanoslm_hybrid_arch_design/harness/__init__.py create mode 100644 2.0/problems/nanoslm_hybrid_arch_design/harness/baseline_model.py create mode 100644 2.0/problems/nanoslm_hybrid_arch_design/harness/data.py create mode 100644 2.0/problems/nanoslm_hybrid_arch_design/harness/eval_ppl.py create mode 100644 2.0/problems/nanoslm_hybrid_arch_design/harness/modal_app.py create mode 100644 2.0/problems/nanoslm_hybrid_arch_design/harness/model_config.py create mode 100644 2.0/problems/nanoslm_hybrid_arch_design/harness/policy.py create mode 100644 2.0/problems/nanoslm_hybrid_arch_design/harness/runner.py create mode 100644 2.0/problems/nanoslm_hybrid_arch_design/harness/scoring.py create mode 100644 2.0/problems/nanoslm_hybrid_arch_design/harness/settings.py create mode 100644 2.0/problems/nanoslm_hybrid_arch_design/harness/train.py create mode 100644 2.0/problems/nanoslm_hybrid_arch_design/readme create mode 100644 2.0/problems/nanoslm_hybrid_arch_design/reference.py create mode 100644 2.0/problems/nanoslm_hybrid_arch_design/repro/__init__.py create mode 100644 2.0/problems/nanoslm_hybrid_arch_design/repro/modal_repro.py create mode 100644 2.0/problems/nanoslm_hybrid_arch_design/repro/paper_models.py create mode 100644 2.0/problems/nanoslm_hybrid_arch_design/repro/wsd.py diff --git a/2.0/problems/nanoslm_hybrid_arch_design/DESIGN.md b/2.0/problems/nanoslm_hybrid_arch_design/DESIGN.md new file mode 100644 index 000000000..da77fb9b7 --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/DESIGN.md @@ -0,0 +1,264 @@ +# Design — nanoslm_hybrid_arch_design + +## 1. Task + +Design a **hybrid language-model architecture** — mixing attention with +linear-recurrent sequence mixers — that reaches the lowest held-out +**bits-per-byte (`val_bpb`)** when trained from scratch under a **fixed +wall-clock budget on a single H100**. The agent submits one file, `model.py` +(full freedom over the model definition); the hidden judge drops it into a +**locked** dolma2 BPE training + evaluation harness, trains it for a fixed +wall-clock `T`, and scores `val_bpb` against a **locked baseline architecture** +trained under the identical budget. + +## 1.1 Where the starting setup comes from + +The initial configuration is taken from **Olmo Hybrid: From Theory to Practice +and Back** (Merrill, Li, Romero et al., arXiv:2604.03444). That work shows +theoretically that hybrids express capabilities beyond both transformers and +linear RNNs, then validates it at 7B by replacing sliding-window attention +layers with Gated DeltaNet layers, outperforming the comparable Olmo 3 baseline. + +Concretely inherited from it and from the OLMo-core implementation: + + * the **3:1 interleaved pattern** (three Gated DeltaNet layers per full + attention layer, i.e. 25% attention), which is what `reference.py` ships as + the agent's starting point; + * the practice of **rebalancing head count / width so the hybrid is + parameter-matched** against the pure-attention baseline + (`REMOVE_HEADS` in the upstream script) — **now applied** as + `REMOVE_HEADS = 1` (d_model 768 -> 704, n_heads 12 -> 11); + * **dolma2** tokenization and the ~190M OLMo3 shape (d=768, L=12, H=12). + +The scientific questions this task poses are the ones that paper answers at +scale but leaves open at 190M under a tight, parameter-matched, wall-clock +budget — enumerated in §1.2. The agent **builds on the baseline architecture** +(`harness/baseline_model.py`, a faithful pure-attention `olmo3_190M`) and +hybridizes it; `reference.py` ships one worked hybrid — the 3:1 GDN pattern — as +an existence proof, not a ceiling. It doubles as the repo-required reference +*solution* and the CI gate: `reference.py` must beat the hidden, score-0 +`baseline_model.py` (the pure-attention `olmo3_190M`, judge-side only). + +The optimizer, LR schedule, data, tokenizer, EVALUATION context (8192) and +budget are all fixed by the judge. The agent changes the architecture — and the +TRAINING context length, which trades optimizer steps against length +extrapolation. + +## 1.2 The three research questions + +The task is organized around the three architecture questions the Olmo Hybrid +paper settles at 7B and ablates on a 60M–1B ladder (its Table 5, `val_bpb` +averaged over Math/Code/QA; lower is better). The agent re-derives them at 190M, +**from the pure-attention baseline up**, under a fixed wall-clock budget and a +hard parameter cap — conditions under which the paper's answers are *priors, not +conclusions* (see the caveat below). + +1. **RNN architecture — GDN vs. Mamba2.** Which linear-recurrent mixer makes the + best hybrid? Paper prior @190M: Gated DeltaNet (3:1) reaches 0.891 bpb vs. + Mamba2 (3:1) at 0.921, and pure GDN (0.895) beats pure Mamba2 (0.941). GDN is + the paper's pick and is what `reference.py` ships; whether that margin + survives parameter-matching at this scale is open. + +2. **Layer placement — interleaved vs. concentrated.** Given a fixed count of + attention layers, where do they go? Paper prior @190M: uniformly interleaved + (0.891) beats concentrating them in the middle of the stack (0.899). The + paper's argument is structural — concentrating attention forces all global + information through one bottleneck, whereas interleaving maximizes the number + of alternations between mixer types, which its expressivity theory (Theorems + 1 and 3) predicts should help. Clustering attention early (context in) vs. + late (read-out) is unexplored and on the table. + +3. **Attention ratio — how much attention is enough?** What fraction of layers + should stay full attention? Paper prior @190M: 1:1 / 50% (0.896), 3:1 / 25% + (0.891), 7:1 / 12.5% (0.892) — nearly a wash at this scale (a 0.001 bpb + spread), with the paper selecting 3:1 because it wins at 600M–1B, not at + 190M. Fewer attention layers also buy more optimizer steps under the + wall-clock budget (§5), so here the ratio trades context against *both* + parameters and steps. + +**Why the paper's answers are priors, not conclusions here.** The paper's 190M +hybrids are *not* parameter-matched: its GDN (3:1) carries ~254M non-embedding +params against the transformer's ~190M — **+34%** — and the paper flags exactly +this (Table 22: "per-size comparisons should be interpreted with care, as +architectures at the same nominal scale differ in actual parameter count"). This +task instead matches non-embedding params to ≈±2% (the `REMOVE_HEADS` +compensation in `reference.py`, measured at −2.19%) and enforces a hard +`param_cap`, so the portion of the paper's measured gain that was *bought with +parameters* is unavailable. The scored question is which of these three choices +still pays once the parameters are held fixed and the clock — not the token +count — is the budget. + +Beyond the three headline axes, `reference.py` enumerates the finer knobs each +one interacts with: recurrent state size (`expand_v`, `num_v_heads` — cheap at +8192 since it does not grow with sequence length), non-uniform layers, and the +rest of the block (norm placement, gating, MLP ratio, head count, embedding +tying). + +## 2. Metric: held-out bits-per-byte (`val_bpb`) + +The tokenizer is **dolma2 BPE** (`allenai/dolma2-tokenizer`, the OLMo-3 +tokenizer, vocab 100278), locked by the judge. Token-level cross-entropy is not +comparable across models unless tokenization is fixed, so the metric normalizes +it by the **raw byte count** of the held-out text rather than by tokens: +`val_bpb = (held-out token cross-entropy in bits) / (held-out bytes)`. Dividing +by bytes makes the number **tokenizer-independent and ungameable** — the agent +submits a model, not a tokenizer, so it cannot lower `val_bpb` by retokenizing, +and the metric stays comparable across architectures the way a byte-level vocab +used to make it by construction. Lower is better; §8 scores the absolute gain +over the locked baseline. + +## 3. Why this is a real task + +Under a fixed wall-clock budget the frontier is genuinely architectural. Because +the baseline is already a tuned `olmo3_190M`, the easy block-level wins (RMSNorm, +rotary, QK-norm, SwiGLU) are **already in it**; what is left on the table is the +hybrid direction — the mixer, ratio, placement, and state size of §1.2. + +## 4. Submission surface & interface + +Submission is a **single file**, `/app/model.py` (`submission.kind: file`). Only +this file travels to the judge; edits to the harness in the agent workspace are +ignored by scoring (black-box judge uses its own locked harness copies). + +`model.py` must expose a factory the harness can call: + +```python +def build_model(config) -> torch.nn.Module: ... +# or a class usable as NanoSLM(config) +class NanoSLM(torch.nn.Module): ... +``` + +`config` is the harness-owned `ModelConfig`: `vocab_size` (100278), `block_size` +(the TRAINING context, which the submission may set via a module-level +`BLOCK_SIZE`), `eval_block_size` (the SCORING context, always 8192 and never +negotiable), plus read-only budget hints. The returned module must implement: + +```python +def forward(self, idx): # idx: LongTensor [B, T] of token ids in [0, vocab_size) + return logits # FloatTensor [B, T, vocab_size] + # returning (logits, loss) is accepted, but the judge IGNORES any returned + # loss and computes cross-entropy itself for BOTH training and val_bpb. +``` + +The judge owns the loss so a model cannot report a fake low loss. The optimizer, +LR schedule, weight-decay grouping (no WD on 1-D params), data order and +wall-clock budget are all locked in `harness/train.py`; the training context is +the one training-side quantity the submission controls, and the evaluation +context is locked in `harness/data.py::val_windows`. + +## 5. Iso-wallclock protocol (the anti-gaming axis) + +- The judge trains the baseline and the submission for the **same fixed + wall-clock `T`** on the **same H100**, from the **same seed and data order** + (common random numbers), then evaluates both on the **same hidden held-out + validation set**. The scored (final/verifier) run measures baseline + submission + back-to-back in one job/GPU/process (no cache), mirroring + `nanowm`'s `run_pair(role="final")`. A cached baseline keyed by + `settings.config_fingerprint()` is used only for the cheap iterative + (agent-role) feedback path. Role is selected by the judge via + `FRONTIER_NANOSLM_ROLE` (`agent` = train submission only + cached baseline, ~T; + `final` = fresh baseline+submission CRN pair, ~2T; default `final`). GPU is + served on Modal from a CPU judge; a directly-attached H100 is + used only in local testing mode. +- The wall-clock cutoff is enforced **by the harness timer**, not by anything the + model can influence: training stops at the first optimizer step whose start + exceeds `T`. Model `forward` cost therefore trades directly against step count + — a more efficient architecture legitimately completes more useful steps. This + is the intended lever, and it is the mechanism by which a hybrid can win at + all: at ctx 8192 attention is estimated to dominate layer FLOPs (~78–84%, an + analytic estimate not yet confirmed on the judge GPU), so replacing most of it + with a cheaper mixer converts directly into extra optimizer steps. +- Determinism: fixed seeds; `torch.use_deterministic_algorithms(True, + warn_only=True)`, `cudnn.deterministic=True`, `benchmark=False`, TF32 + **disabled** during eval, `CUBLAS_WORKSPACE_CONFIG=:4096:8`. Iso-wallclock has + irreducible timing noise (step count varies run-to-run); the CRN pairing keeps + the *comparison* stable even so. The seed-to-seed noise floor itself is an open + calibration item, not yet measured on the judge GPU. + +## 6. Submission policy (validated before training; torch-free, unit-tested) + +`model.py` only (`.py`), ≤ 256 KB, no file deletion, safe path; the submission +must define `build_model(config)` or `class NanoSLM`. + +Two deny lists in `harness/policy.py`, scanned differently (representative, not +exhaustive): + +- `POLICY_DENY_TOKENS`, matched over the **full source** — a leak signal in a + comment is still a leak: + - **Escape / leakage:** `os.environ`, `os.getenv`, `putenv`, `subprocess`, + `socket`, `requests`, `urllib`, `httpx`, `FRONTIER_`/`JUDGE_`/`HARBOR_`/`MODAL_`, + `HF_TOKEN`, judge paths (`/judge`, `/opt/`, `/tests/`). + - **Pretrained-weight loading (must train from scratch):** `from_pretrained`, + `torch.load`, `load_state_dict`, `hf_hub`, `huggingface`, `safetensors`, `AutoModel`. + - **Filesystem reads:** `open(`, `Path(`, `np.load`, `np.fromfile`, `mmap`, `pickle.load`. + - **Timer / control-flow / concurrency:** `time.time`, `time.perf_counter`, + `time.sleep`, `while True`, `threading`, `multiprocessing`, `os.fork`, `ctypes`, + `exec(`, `__import__`. +- `POLICY_DENY_TOKENS_CODE`, matched over **code only** (comments and string + literals stripped) so a docstring *may* name the metric it targets: `val_ppl`, + `perplexity`, `val_bpb`, `bits_per_byte`, `holdout`, `val_data`, `val.bin`. + +`eval(` and `compile(` are deliberately **allowed** so `model.eval()` and +`torch.compile()` work — the sandbox and judge-owned loss cover the rest. + +The policy is a **static allow/deny gate**; §7 guards the residual dynamic +risks it cannot see. + +## 7. Dynamic guards (what the static scan can't catch) + +- **Judge owns the loss.** `val_bpb` is computed by the harness from the model's + `logits` via its own cross-entropy on **hidden** held-out bytes (never mounted + in `/app`), so a model cannot report a fake loss or memorize the val set. +- **Trained-from-scratch guard.** Params are fingerprinted at init and re-checked + after training; a model whose weights did not change (untrained / frozen + constant) or that produces a (near-)constant logit distribution over inputs is + scored 0. Blocks "return a cached distribution" degenerate submissions. +- **Resource caps.** Param count ≤ `PARAM_CAP` and peak activation memory within + the H100; OOM or over-cap → score 0 with a public message (no traceback). Stops + memory-bomb / timer-dodge attempts. +- **Sandboxing.** Submitted code is imported and run as an unprivileged user with + the evaluator source chmod-protected (same pattern as `erdos_demo`); evaluator + output returns public metrics and concise errors only — never raw submission + stdout/stderr or tracebacks (2.0 black-box safety rules). + +## 8. Scoring — ABSOLUTE bits-per-byte gain + +`harness/scoring.py` (shared with the public test). Lower `val_bpb` is better. +With `base_bpb` = locked-baseline held-out bits-per-byte and `sub_bpb` = the +submission's, both trained under the identical wall-clock budget with common +random numbers and both evaluated at the fixed 8192-token window: + +``` +gain = base_bpb - sub_bpb # THE MEASUREMENT +score = clip(100 * gain / bpb_score_scale, 0, 100) +``` + +Baseline-tying and worse-than-baseline both score 0. `score_unbounded` is the +un-clipped ratio and keeps rising past 100. + +### 8.1 `bpb_score_scale` is a display convention, NOT a calibration target + +The only measurement is `gain`, in bits per byte. `bpb_score_scale` exists solely +to map that gain onto the 0–100 range Harbor's `reward = score / 100` expects — a +raw gain of 0.05 bpb would otherwise surface as reward 0.0005. It is **not** a +calibrated definition of "a full win": no such number has been measured, and the +clip discards information (a submission at 2× the scale scores the same 100 as one +exactly at it), which is why `score_unbounded` stays un-clipped for an operator to +read. + +The one real requirement is **discrimination**: too large and everything pins at +0, too small and everything pins at 100 — checkable from a single real-corpus run, +no headroom study needed. Current value: `bpb_score_scale = 0.05` bpb. (The +predecessor `r_target` is dropped, not satisfied: it asked for a measured full-win +definition that was never well-posed.) + +## 9. Compute / infra + +- `tag: systems`; H100 served on Modal (one per environment), CPU judge — + identical shape to `nanowm_rollout_speedup` and `vllm_llm_serving_optimization`. +- `runtime.timeout_seconds`: long-horizon (agent iterates for hours); + per-submission training is capped at ≤ 1 h/train (single H100). +- Hidden judge assets (judge image only): the tokenized training shard, the + **held-out** validation byte stream, the locked baseline architecture, and the + cached baseline `val_bpb` (keyed by config fingerprint). None are mounted in + `/app`. \ No newline at end of file diff --git a/2.0/problems/nanoslm_hybrid_arch_design/PR_SUMMARY.md b/2.0/problems/nanoslm_hybrid_arch_design/PR_SUMMARY.md new file mode 100644 index 000000000..0bf7ecfe2 --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/PR_SUMMARY.md @@ -0,0 +1,119 @@ +# PR: add `nanoslm_hybrid_arch_design` (Frontier-CS 2.0) + +## What + +A new open-ended 2.0 problem: design a **hybrid language-model architecture** +(attention + linear-recurrent mixers) that reaches the lowest held-out +**bits-per-byte (`val_bpb`)** when trained from scratch under a **fixed +wall-clock budget on a single H100**. The agent submits one `model.py` (full +architecture freedom); the hidden judge trains it under a locked dolma2-BPE +recipe and scores `val_bpb` against a locked baseline architecture +(pure-attention `olmo3_190M`) trained under the identical budget. + +Adapts the Genesys system ("Language Modeling by Language Models", +arXiv:2506.20249) into a bounded, ungameable, single-objective task — the +hybrid architecture discovery starting from Olmo Hybrid (arXiv:2604.03444). Structurally mirrors +`nanowm_rollout_speedup` / `vllm_llm_serving_optimization`: patch/file +submission, Modal H100, CPU judge, latency/quality-style guardrails, CRN +determinism. + +## Problem id / scoring + +- id: `nanoslm_hybrid_arch_design` (Harbor id `frontier-cs-2-0-nanoslm-hybrid-arch-design`) +- metric: held-out `val_bpb` = total NLL (bits) / held-out BYTES, over the + dolma2 BPE tokenizer (vocab 100278). Normalizing by BYTES (not tokens) is what + makes it tokenizer-independent and ungameable. Lower is better. (`val_ppl` is + reported for readability only and is NO LONGER equal to `2^val_bpb`.) +- score: `clip(100 * (base_bpb - sub_bpb) / bpb_score_scale, 0, 100)` — the + ABSOLUTE bpb gain; baseline and worse-than-baseline score 0. + `bpb_score_scale` is a display convention that maps bpb onto 0-100, NOT a + calibration target (DESIGN.md §8.1); `score_unbounded` stays un-clipped. +- context: EVALUATION is fixed at 8192 for every submission; the TRAINING + context is agent-controlled via a module-level `BLOCK_SIZE` in model.py + (power of two in [256, 8192], default 8192). Shorter buys optimizer steps and + costs length extrapolation. + +## Resource budget + +- `tag: systems`, single Modal **H100**, CPU judge. +- `runtime.timeout_seconds: 21600` (long-horizon agent iteration). +- Per-submission training capped at ≤ 1 h wall-clock (`train_seconds` fixed `T` + + `max_train_seconds` hard abort). +- Hidden judge assets: tokenized train shard, held-out val byte stream, locked + baseline architecture, cached baseline `val_bpb` (keyed by config fingerprint). + +## Anti-gaming + +- Static policy gate (`harness/policy.py`): model.py only, ≤256 KB, denies + env/network/subprocess, pretrained-weight loading, file reads, metric/data/timer + peeking; allows `model.eval()` / `torch.compile()`. +- Dynamic guards (`harness/runner.py`): judge owns the loss (computes CE from the + model's logits on hidden val); trained-from-scratch param-delta check; degenerate + constant-logit check; param cap; OOM → 0. Black-box safe (no traceback/stdout leak). +- Iso-wallclock CRN: baseline + submission trained back-to-back on one GPU, same + seed/data order, scored on the same hidden bytes. + +## Validation (this branch) + +Torch-free layers unit-tested and **full harness smoke-tested on CPU** (no GPU): + +```bash +# torch-free: policy + scoring + fingerprint (17 assertions) +bash evaluate.sh --selftest # -> SELFTEST OK + +# full pipeline on CPU (tiny model/budget), needs torch: +FRONTIER_NANOSLM_SMOKE=1 bash evaluate.sh reference.py +# base_val_bpb=; sub_val_bpb=; abs_bpb_delta=+; score=<...> (reference > baseline) +FRONTIER_NANOSLM_SMOKE=1 python3 evaluator.py # -> guard: not trained; score 0 +python3 evaluator.py # -> policy_rejected; score 0 + +# role split (GPU via Modal from CPU judge; direct H100 only in testing mode): +FRONTIER_NANOSLM_ROLE=final ... evaluator.py reference.py # fresh baseline+submission CRN pair (~2T) +FRONTIER_NANOSLM_ROLE=agent ... evaluator.py reference.py # trains submission only; reuses +# fingerprint-cached baseline (note=iterative(cached-baseline)), ~T per iteration +``` + +Standard 2.0 CLI checks to run in the repo before merge: + +```bash +uv run frontier list 2.0 +uv run frontier show 2.0 nanoslm_hybrid_arch_design +python3 -m py_compile 2.0/problems/nanoslm_hybrid_arch_design/evaluator.py +``` + +## PENDING before ship (see DESIGN.md §3) + +- **[BLOCKER] Single-H100 calibration** of the wall-clock budget `T`, the + model-scale band, and the `val_bpb` noise floor (protocol in DESIGN.md §3). + (`r_target` was on this list; it is now `bpb_score_scale` and is no longer a + calibration gate — DESIGN.md §8.1.) + All such constants are flagged `CALIBRATE` in `harness/settings.py` / `config.yaml`. +- Modal end-to-end run + deployed app name / H100 SKU confirmation (GPU path awaits + maintainer credentials, as with `nanowm`). +- Judge-image bake-asset provenance + data license (FineWeb-Edu, dolma2 BPE); + the held-out val byte ranges are the only hidden component. +- Decide: keep the optimizer locked (current: architecture-only) vs. expose an + optional `configure_optimizers` hook (more autoresearch-faithful). + +## Files + +``` +2.0/problems/nanoslm_hybrid_arch_design/ + config.yaml # systems/H100/Modal, file submission /app/model.py + readme # public task statement (agent-facing) + evaluator.py # evaluate() + prepare() + --selftest; torch-free top-level + evaluate.sh # local CLI wrapper (+ --selftest) + reference.py # reference solution model.py (3:1 GDN hybrid on the olmo3_190M baseline) + DESIGN.md # full design, calibration protocol, open items + PR_SUMMARY.md # this file + harness/ + settings.py # locked config + smoke overrides + config fingerprint (torch-free) + policy.py # static submission policy (torch-free) + scoring.py # val_bpb gain -> [0,100] (torch-free) + model_config.py # ModelConfig contract object (torch-free) + data.py # dolma2-BPE token data (synthetic fallback for smoke) + baseline_model.py # locked baseline transformer (score-0 reference point) + train.py # locked iso-wallclock training loop (harness-owned loss) + eval_ppl.py # held-out val_bpb eval + determinism knobs + runner.py # arm run + CRN pair + dynamic guards +``` diff --git a/2.0/problems/nanoslm_hybrid_arch_design/config.yaml b/2.0/problems/nanoslm_hybrid_arch_design/config.yaml new file mode 100644 index 000000000..a2d25e416 --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/config.yaml @@ -0,0 +1,67 @@ +tag: systems +runtime: + # Submission is a single Python file, the model definition (/app/model.py). + # `language: python` keeps the extension/CLI conventions standard. + language: python + timeout_seconds: 21600 # long-horizon: the agent iterates for hours + environment: >- + Full-freedom PyTorch model.py submitted against a locked + training harness over a dolma2-BPE-tokenized FineWeb-Edu corpus; a single + Modal H100 trains the model from scratch for a fixed wall-clock budget; the + judge scores held-out validation bits-per-byte vs a locked baseline + architecture trained under the same budget. + apt_packages: + - bash + - ca-certificates + - curl + - git + - python3 + - python3-pip + judge_apt_packages: + - bash + - ca-certificates + - curl + - git + - python3 + - python3-pip + judge_pip_packages: + - modal + docker: + # Experimental images; build before a local Harbor trial. The agent image + # carries the harness + a starter model.py; the judge image additionally + # vendors the tokenized training shard, the HELD-OUT validation token stream + # and its manifest (which carries val_bytes, the bpb denominator), the + # locked baseline architecture, and the cached baseline metric. + image: frontiercs/nanoslm-hybrid-arch-design-agent:experimental-v0 + judge_image: frontiercs/nanoslm-hybrid-arch-design-judge:experimental-v0 +environment: + cpus: 8 + memory_mb: 32768 + storage_mb: 32768 + build_timeout_seconds: 5400 +evaluation: + # GPU served on Modal (one per environment); judge container is CPU-only. + gpu: H100 + # dolma2 BPE (the OLMo-3 tokenizer). UNGATED -- no HF token required. + tokenizer: allenai/dolma2-tokenizer + vocab_size: 100278 # > 65535, so token streams are uint32 + dataset: HuggingFaceFW/fineweb-edu:sample-10BT + block_size: 8192 # default training context (agent may override) + eval_block_size: 8192 # fixed scoring context + train_seconds: 1800 # fixed wall-clock budget T per training run + max_train_seconds: 3300 # hard cap (< 1h) + val_tokens: 1048576 # held-out TARGET tokens scored + # The scored metric is bits per BYTE: nll_nats / (val_bytes * ln2). val_bytes + # is measured by docker/prep_assets.py and shipped in the judge image's + # manifest.json -- normalizing by the TOKEN count would make the metric + # tokenizer-dependent. + metric: val_bpb + # Scoring is on the ABSOLUTE val_bpb gain over the locked baseline: + # score = clip(100 * (base_bpb - sub_bpb) / bpb_score_scale, 0, 100) + bpb_score_scale: 0.05 # bpb gain that saturates the bounded score + param_cap: 400000000 # trainable-param cap; over-cap -> score 0 (matches harness/settings.py) + seed: 1337 +submission: + kind: file + path: /app/model.py + max_queue_size: 2 diff --git a/2.0/problems/nanoslm_hybrid_arch_design/docker/agent/Dockerfile b/2.0/problems/nanoslm_hybrid_arch_design/docker/agent/Dockerfile new file mode 100644 index 000000000..5556a36f8 --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/docker/agent/Dockerfile @@ -0,0 +1,59 @@ +# Agent workspace image for nanoslm_hybrid_arch_design. +# +# CPU-ONLY. The agent designs an architecture and submits /app/model.py; the +# JUDGE trains it (on a Modal H100). So this image ships no torch, no CUDA and +# no corpus -- an agent cannot and need not train here. +# +# What it DOES ship is the judge's own static policy gate, so +# `bash /app/public_test.sh` runs the exact accept/reject rules the judge +# applies, instead of the agent discovering them by burning submissions. +# +# CRITICALLY: /opt/nanoslm_arch/data/val.bin -- the held-out stream -- is NOT +# here and must never be. build_images.sh asserts its absence post-build. +FROM python:3.11-slim + +ENV DEBIAN_FRONTEND=noninteractive \ + PIP_NO_CACHE_DIR=1 \ + PYTHONUNBUFFERED=1 + +RUN apt-get update && apt-get install -y --no-install-recommends \ + bash ca-certificates curl git ripgrep \ + && rm -rf /var/lib/apt/lists/* + +# Parity with the adapter template's tooling. +RUN pip install --no-cache-dir "numpy<2" \ + && curl -fsSL https://claude.ai/install.sh | bash 2>/dev/null || true + +WORKDIR /app + +# Workspace scaffolding: starter model.py, README, public_test. +COPY task_ctx/harbor_app/ /app/ + +# The judge's static gate, shipped verbatim so the local check and the judge's +# check cannot drift. policy.py is settings-free and names no hidden value: +# it contains the DENYLIST (which is public by design and stated in the readme) +# but no thresholds, no paths, and no held-out identifiers. build_images.sh +# greps it before copying. +COPY task_ctx/agent_task/policy.py /app/policy.py +RUN chmod +x /app/*.sh 2>/dev/null || true + +# Fail the build if the starter submission does not pass the shipped gate -- +# an agent must never open a workspace whose default file is already rejected. +# +# Assert through public_test.sh, the SAME entrypoint the agent runs, not by +# importing check_source directly. A direct `sys.path`+import once passed here +# while `bash /app/public_test.sh` crashed, because public_test.py loads +# policy.py via spec_from_file_location and that path had its own bug. Testing +# the real entrypoint is what makes this assertion mean what it claims. +# The __pycache__ removal is load-bearing, not tidiness: running the gate +# compiles policy.py, and the resulting .pyc embeds the denylist strings +# ("val.bin", "/opt/"). build_images.sh greps every file under /app except +# policy.py itself, so leaving the byte-compiled twin behind trips that check. +RUN bash /app/public_test.sh \ + || (echo 'FATAL: shipped starter model.py fails the policy gate' && exit 1) \ + && rm -rf /app/__pycache__ + +# Assert the held-out stream is absent (defence in depth; build_images.sh also +# checks the built image). +RUN test ! -e /opt/nanoslm_arch/data/val.bin || \ + (echo 'FATAL: held-out val.bin present in the AGENT image' && exit 1) diff --git a/2.0/problems/nanoslm_hybrid_arch_design/docker/build_images.sh b/2.0/problems/nanoslm_hybrid_arch_design/docker/build_images.sh new file mode 100755 index 000000000..0c7406d10 --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/docker/build_images.sh @@ -0,0 +1,150 @@ +#!/usr/bin/env bash +# Build the agent + judge images for nanoslm_hybrid_arch_design. +# +# bash docker/build_images.sh [tag] +# +# NOTE ON THE THREE-IMAGE ARCHITECTURE. This script builds TWO images; a third +# is built remotely by Modal from harness/modal_app.py and is never produced +# here: +# +# agent (this script, CPU) - workspace: starter model.py + the static gate +# judge (this script, CPU) - dispatches to Modal, scores, holds val.bin +# Modal (modal deploy, GPU) - torch 2.6.0 / triton 3.2.0 / fla 0.5.1 +# +# The GPU stack lives ONLY in the Modal image. Keep its pins in +# harness/modal_app.py in sync: that exact combination is +# the one whose GDN backward kernel compiles (triton 3.1.0 hung for 30+ min; +# fla 0.4.1 failed to lower the backward). +set -euo pipefail + +TAG="${1:-experimental-v0}" +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PROB="$(cd "$HERE/.." && pwd)" +ASSETS="${NANOSLM_ARCH_ASSETS:-$PROB/.assets}" +AGENT_IMG="frontiercs/nanoslm-hybrid-arch-design-agent:$TAG" +JUDGE_IMG="frontiercs/nanoslm-hybrid-arch-design-judge:$TAG" + +# --- assets ------------------------------------------------------------------- +if [ ! -s "$ASSETS/train.bin" ] || [ ! -s "$ASSETS/val.bin" ] \ + || [ ! -s "$ASSETS/manifest.json" ]; then + echo "[build] staging tokenized corpus into $ASSETS (cold cache)" + NANOSLM_ARCH_ASSETS="$ASSETS" python3 "$HERE/prep_assets.py" +fi +# manifest.json is NOT optional metadata: it carries `val_bytes`, the denominator +# of the scored bits-per-byte. Without it harness/data.py raises rather than +# silently normalizing by a token count -- so a judge image built without it +# would fail every run. See docker/prep_assets.py. +for f in train.bin val.bin manifest.json; do + [ -s "$ASSETS/$f" ] || { echo "FATAL: missing $ASSETS/$f"; exit 1; } +done +MANIFEST="$ASSETS/manifest.json" python3 - <<'PY' +import json, os, sys +m = json.load(open(os.environ["MANIFEST"])) +b, t = m.get("val_bytes", 0), m.get("val_target_tokens", 0) +if not (b > 0 and t > 0): + sys.exit("FATAL: manifest has no positive val_bytes/val_target_tokens") +print("[check] val_bytes=%d over %d target tokens (%.2f B/tok)" % (b, t, b / t)) +PY + +# --- build context ------------------------------------------------------------ +CTX="$(mktemp -d)" +trap 'rm -rf "$CTX"' EXIT +mkdir -p "$CTX/task_ctx/task_pkg" "$CTX/task_ctx/agent_task" "$CTX/task_ctx/assets" + +cp -r "$PROB/harness" "$CTX/task_ctx/task_pkg/harness" +cp "$PROB/evaluator.py" "$CTX/task_ctx/evaluator.py" +cp -r "$PROB/harbor/app" "$CTX/task_ctx/harbor_app" +cp "$PROB/harness/policy.py" "$CTX/task_ctx/agent_task/policy.py" +cp "$ASSETS/train.bin" "$ASSETS/val.bin" "$ASSETS/manifest.json" \ + "$CTX/task_ctx/assets/" + +# modal_app.py declares the GPU image and is judge-side only; it is already in +# task_pkg via harness/. Nothing else needs to move. + +# --- leak grep: what ships to the AGENT --------------------------------------- +# policy.py legitimately CONTAINS the denylist, which names the metric and +# held-out identifiers -- those strings are public by design (the readme states +# them) and are what the gate matches on. So grep for hidden VALUES and real +# PATHS, not for the denylist's own words. +POLICY="$CTX/task_ctx/agent_task/policy.py" +# +# TWO greps, not one. policy.py's DENY tuples literally contain "val.bin", +# "/opt/", "holdout" etc -- that IS the denylist, it is public by design, and a +# naive single grep fails the build on the gate's own source every time. (This +# exact trap already exists in lm_arch_discovery/docker/build_images.sh and in +# align_overopt_stability's; it is easy to re-introduce.) +# +# (1) hidden VALUES over the whole file. These are strictly more specific +# than anything the denylist contains, so they cannot collide with it. +# (2) asset identifiers with the DENY tuples EXCISED, which catches a real +# path leaking in via a default or a stray constant while letting the +# denylist keep naming them. +# `bpb_score_scale` replaced the old `r_target`; both are matched so a stale +# checkout cannot slip the old name through this gate. +if grep -qiE "r_target|bpb_score_scale|train_seconds|param_cap|baseline_ppl|[0-9]{6,}" "$POLICY"; then + echo "FATAL: policy.py would leak a hidden VALUE to the agent image" + exit 1 +fi +if sed -e '/^POLICY_DENY_TOKENS[A-Z_]*: tuple/,/^)/d' -e '/^[[:space:]]*#/d' "$POLICY" \ + | grep -qiE "/opt/nanoslm_arch|val\.bin"; then + echo "FATAL: policy.py names a judge asset outside its denylist block" + exit 1 +fi +# The starter model.py must not name the held-out stream either. +if grep -qiE "/opt/nanoslm_arch|val\.bin" "$CTX/task_ctx/harbor_app/model.py"; then + echo "FATAL: starter model.py references judge assets"; exit 1 +fi + +# --- build -------------------------------------------------------------------- +echo "[build] agent -> $AGENT_IMG" +docker build -f "$HERE/agent/Dockerfile" -t "$AGENT_IMG" "$CTX" +echo "[build] judge -> $JUDGE_IMG" +docker build -f "$HERE/judge/Dockerfile" -t "$JUDGE_IMG" "$CTX" + +# --- post-build isolation ------------------------------------------------------ +echo "[check] held-out stream must be unreachable from the AGENT image" +# policy.py is EXEMPT from the content grep and only from it: its denylist +# names "val.bin" by design, which is the third place in this script where that +# distinction matters. EVERY other file under /app is grepped whole -- the loop +# enumerates with `find`, not a fixed list, so a leak in a file added to the +# workspace later still fails the build. (It used to name four files explicitly, +# which made the policy.py exemption below dead code and left anything new +# unchecked.) File PRESENCE is checked unconditionally -- no exemption there. +if docker run --rm --entrypoint bash "$AGENT_IMG" -c ' + ls /opt/nanoslm_arch/data/val.bin 2>/dev/null; + ls /opt/nanoslm_arch 2>/dev/null; + find /app -type f | while read -r f; do + [ "$f" = /app/policy.py ] && continue; + grep -il "val\.bin\|/opt/nanoslm_arch" "$f" 2>/dev/null; + done + ' | grep -q .; then + echo "FATAL: held-out data or its identity is reachable from the AGENT image" + exit 1 +fi + +echo "[check] judge must NOT carry a GPU stack (it dispatches to Modal)" +if docker run --rm --entrypoint bash "$JUDGE_IMG" -c \ + 'python3 -c "import torch" 2>/dev/null && echo present' | grep -q present; then + echo "FATAL: torch present in the CPU-only judge image -- a silent LOCAL" + echo " fallback would try to train ~190M at ctx 8192 on CPU" + exit 1 +fi + +echo "[check] judge must hold both token streams AND the byte-count manifest" +docker run --rm --entrypoint bash "$JUDGE_IMG" -c \ + 'test -s /opt/nanoslm_arch/data/train.bin \ + && test -s /opt/nanoslm_arch/data/val.bin \ + && test -s /opt/nanoslm_arch/data/manifest.json' \ + || { echo "FATAL: judge image missing a token stream or manifest.json"; exit 1; } + +cat <=1.0" "numpy<2" + +WORKDIR /opt/nanoslm_arch + +# The task package. Preserve the DIRECTORY layout: evaluator.py does +# `from harness import policy, scoring, settings` and the harness modules use +# relative imports, so flattening breaks both. +COPY task_ctx/task_pkg/harness/ /opt/nanoslm_arch/task/harness/ +COPY task_ctx/evaluator.py /opt/nanoslm_arch/task/evaluator.py + +# Hidden assets. val.bin is the HELD-OUT stream and exists ONLY here -- it is +# never copied into the agent image, and build_images.sh asserts that. Both +# streams are flat uint32 dolma2 token ids (vocab 100278), not raw bytes. +# +# manifest.json is REQUIRED, not documentation: it carries `val_bytes`, the +# denominator of the scored bits-per-byte. harness/data.py raises rather than +# normalizing by a token count, so a judge without it fails every run. +COPY task_ctx/assets/train.bin /opt/nanoslm_arch/data/train.bin +COPY task_ctx/assets/val.bin /opt/nanoslm_arch/data/val.bin +COPY task_ctx/assets/manifest.json /opt/nanoslm_arch/data/manifest.json +RUN mkdir -p /opt/nanoslm_arch/baseline + +# Paths must match harness/settings.py TaskConfig defaults. +ENV PYTHONPATH=/opt/nanoslm_arch/task \ + NANOSLM_ARCH_ASSETS=/opt/nanoslm_arch/data + +# Fail the build if the judge cannot import its own entrypoint, or if it can +# import torch (which would mean the CPU-only contract was broken and a silent +# LOCAL fallback is possible). +RUN python3 -c "\ +import sys; sys.path.insert(0, '/opt/nanoslm_arch/task');\ +import evaluator, harness.policy, harness.scoring, harness.settings;\ +print('judge imports OK');\ +" +# MUST stay ONE physical line. Docker splices a trailing `\` by deleting the +# newline, so a multi-line try/except collapses into `...)except ImportError:` +# -- a SyntaxError, which exits non-zero and thus "passed" for the wrong reason +# while never testing torch at all. find_spec also beats `import torch` here: it +# flags a present-but-broken torch, which a bare ImportError would swallow. +RUN python3 -c "import importlib.util, sys; sys.exit('FATAL: torch present in the CPU-only judge image') if importlib.util.find_spec('torch') else print('torch absent as intended')" diff --git a/2.0/problems/nanoslm_hybrid_arch_design/docker/prep_assets.py b/2.0/problems/nanoslm_hybrid_arch_design/docker/prep_assets.py new file mode 100644 index 000000000..e7bfbc564 --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/docker/prep_assets.py @@ -0,0 +1,294 @@ +"""Stage the tokenized corpus for nanoslm_hybrid_arch_design. Judge/operator side only. + +Idempotent. Writes two flat uint32 TOKEN-ID streams plus a manifest: + + $NANOSLM_ARCH_ASSETS/train.bin ~512M tokens training ids + $NANOSLM_ARCH_ASSETS/val.bin 1M+1 tokens HELD OUT, judge image only + $NANOSLM_ARCH_ASSETS/manifest.json provenance + THE HELD-OUT BYTE COUNT + +TOKENIZER +--------- +`allenai/dolma2-tokenizer` (vocab 100278) -- the OLMo-3 tokenizer. UNGATED: it +downloads with no HF token, unlike lm_arch_discovery's license-gated Llama-2 +tokenizer. Pinned in `harness/settings.py` as `TaskConfig.tokenizer_name` and +fingerprinted, so changing it invalidates any cached baseline. + +DTYPE: uint32, NOT uint16 +------------------------- +100278 ids do not fit in uint16 (max 65535). A uint16 stream would wrap the +upper ~35k of the vocabulary into low ids -- a corpus that still loads, still +trains, and shows up only as an unexplained floor on val_bpb. `harness/data.py` +reads `np.fromfile(path, dtype=np.uint32)` and asserts `max(id) < vocab_size`. + +THE HELD-OUT BYTE COUNT -- WHY THIS FILE COMPUTES IT +---------------------------------------------------- +The scored metric is bits per BYTE: + + val_bpb = total_nll_nats / (val_bytes * ln 2) + +While the tokenizer was byte-level, 1 token == 1 byte and normalizing by the +token count was accidentally identical. Under BPE it is not: per-token CE/ln2 is +a tokenizer-dependent quantity that is no longer comparable across setups, which +would destroy the reason bpb was chosen (DESIGN.md §2). Nothing at +eval time can recover the byte count from a token stream, so it is measured HERE +-- by decoding the exact span of target tokens the harness scores and taking its +UTF-8 length -- and recorded in the manifest as `val_bytes` alongside +`val_target_tokens`. `harness/settings.resolve_val_bytes_per_token()` reads it +and `harness/data.TokenData` raises rather than falling back to a token count. + +DISJOINTNESS +------------ +`val.bin` is filled FIRST from the head of the document stream; the document +that straddles the boundary is then DISCARDED; training fills from what +follows. So no document contributes to both, structurally rather than by +sampling. Asserted at write time. + +SIZING +------ +`val.bin` holds `TaskConfig.val_tokens + 1` ids: the harness scores +`val_tokens` TARGET tokens over non-overlapping windows, and the last target +needs one more input token behind it. + +`train.bin` defaults to 512M tokens (~2 GB on disk at uint32). A 30-minute H100 +run at batch 8 x ctx 8192 touches on the order of 10^7-10^8 tokens, so this +leaves ample headroom without the sampler ever exhausting the stream. + +Usage: + NANOSLM_ARCH_ASSETS=./assets python3 docker/prep_assets.py + NANOSLM_ARCH_ASSETS=./assets python3 docker/prep_assets.py --train-tokens 2000000 +""" + +from __future__ import annotations + +import argparse +import json +import os +import pathlib +import sys + +import numpy as np + +# HuggingFaceFW/fineweb-edu, sample-10BT: public, ungated, English, and the +# `dataset_name` recorded in TaskConfig is derived from it. +CORPUS = "HuggingFaceFW/fineweb-edu" +CORPUS_CONFIG = "sample-10BT" + +TOKENIZER = "allenai/dolma2-tokenizer" # matches TaskConfig.tokenizer_name +VOCAB_SIZE = 100278 # matches TaskConfig.vocab_size +TOKEN_DTYPE = np.uint32 # matches harness/data.TOKEN_DTYPE + +VAL_TOKENS = 1 << 20 # matches TaskConfig.val_tokens +TRAIN_TOKENS = 512_000_000 +DOC_BATCH = 256 # docs per tokenizer call + + +def _stream_docs(): + from datasets import load_dataset + + ds = load_dataset(CORPUS, CORPUS_CONFIG, split="train", streaming=True) + for rec in ds: + text = rec.get("text") or rec.get("content") + if text: + yield text + + +def _load_tokenizer(): + import logging + + from transformers import AutoTokenizer + + tok = AutoTokenizer.from_pretrained(TOKENIZER) + # "Token indices sequence length is longer than the specified maximum + # sequence length for this model (15040 > 8192)" fires on nearly every batch + # and is IRRELEVANT here: we are building a flat corpus stream, not feeding + # documents to a model. The harness slices its own block_size windows out of + # it later. Silenced so a real error is not buried in thousands of lines. + tok.model_max_length = int(1e12) + logging.getLogger("transformers.tokenization_utils_base").setLevel(logging.ERROR) + if len(tok) > VOCAB_SIZE: + raise SystemExit( + f"FATAL: tokenizer has {len(tok)} ids but VOCAB_SIZE is {VOCAB_SIZE}; " + "harness/settings.py TaskConfig.vocab_size must be raised to match" + ) + return tok + + +def _encode_batched(tok, docs, quota: int): + """Yield uint32 id chunks from `docs` until `quota` ids are produced. + + Documents are concatenated with NO separator token, exactly as the + byte-level version concatenated raw text. A separator would carry zero bytes + of text while still consuming a token, biasing the byte accounting. + """ + produced, used, buf = 0, 0, [] + + def _flush(): + nonlocal produced, buf + if not buf: + return None + enc = tok(buf, add_special_tokens=False)["input_ids"] + buf = [] + flat = np.fromiter( + (i for seq in enc for i in seq), dtype=np.int64, + ) + if produced + flat.size > quota: + flat = flat[: quota - produced] + produced += flat.size + return flat.astype(TOKEN_DTYPE) + + for text in docs: + buf.append(text) + used += 1 + if len(buf) >= DOC_BATCH: + chunk = _flush() + if chunk is not None and chunk.size: + yield chunk, used + if produced >= quota: + return + chunk = _flush() + if chunk is not None and chunk.size: + yield chunk, used + + +def _fill(path: pathlib.Path, quota: int, docs, tok, *, keep: bool): + """Write `quota` token ids to `path`; return (written, docs_used, ids|None).""" + written, used, kept = 0, 0, [] + with open(path, "wb") as fh: + for chunk, used in _encode_batched(tok, docs, quota): + fh.write(chunk.tobytes()) + written += chunk.size + if keep: + kept.append(chunk) + if written >= quota: + break + ids = np.concatenate(kept) if (keep and kept) else None + return written, used, ids + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--assets", default=os.environ.get("NANOSLM_ARCH_ASSETS", "./assets")) + ap.add_argument("--train-tokens", type=int, default=TRAIN_TOKENS) + ap.add_argument("--val-tokens", type=int, default=VAL_TOKENS, + help="TARGET tokens scored; val.bin holds this many + 1") + args = ap.parse_args() + + out = pathlib.Path(args.assets) + out.mkdir(parents=True, exist_ok=True) + train_p, val_p = out / "train.bin", out / "val.bin" + man_p = out / "manifest.json" + + itemsize = np.dtype(TOKEN_DTYPE).itemsize + val_ids_needed = args.val_tokens + 1 + if (train_p.exists() and val_p.exists() and man_p.exists() + and train_p.stat().st_size >= args.train_tokens * itemsize + and val_p.stat().st_size >= val_ids_needed * itemsize): + print(f"[prep] cached: train={train_p.stat().st_size // itemsize:,} tok " + f"val={val_p.stat().st_size // itemsize:,} tok") + return 0 + + tok = _load_tokenizer() + docs = _stream_docs() + + # HELD-OUT FIRST, then discard the straddling document, then train. + v_written, v_docs, v_ids = _fill(val_p, val_ids_needed, docs, tok, keep=True) + straddler = next(docs, None) + t_written, t_docs, _ = _fill(train_p, args.train_tokens, docs, tok, keep=False) + + # CLOSE THE STREAM EXPLICITLY, or this script hangs after finishing its work. + # `_stream_docs` wraps a `datasets` STREAMING iterator, and every quota above + # is satisfied by breaking out early -- so the generator is left suspended + # mid-iteration with its HTTP/fsspec resources open. Measured: the script + # wrote train.bin, val.bin and manifest.json correctly and then sat for 14+ + # minutes at interpreter shutdown waiting on those non-daemon threads. + # docker/build_images.sh calls this on a cold cache and would hang with it. + docs.close() + + if v_written < val_ids_needed: + print(f"FATAL: val short ({v_written} < {val_ids_needed})", file=sys.stderr) + return 1 + if t_written < args.train_tokens: + print(f"FATAL: train short ({t_written} < {args.train_tokens})", file=sys.stderr) + return 1 + # STRUCTURAL disjointness assertion, unchanged in intent from the byte-level + # version: the boundary document is consumed and thrown away, so the last + # document in val and the first in train are different documents. + if straddler is None: + print("FATAL: corpus exhausted at the val/train boundary", file=sys.stderr) + return 1 + + hi = int(v_ids.max()) + if hi >= VOCAB_SIZE: + print(f"FATAL: token id {hi} >= vocab {VOCAB_SIZE}", file=sys.stderr) + return 1 + + # THE BYTE COUNT. Decode exactly the span the harness scores -- the TARGET + # tokens, ids[1 : val_tokens+1], since the model predicts token i+1 from + # token i -- and measure its UTF-8 length. + target_ids = v_ids[1 : args.val_tokens + 1] + val_text = tok.decode(target_ids.tolist()) + val_bytes = len(val_text.encode("utf-8")) + if val_bytes <= 0: + print("FATAL: decoded held-out span is empty", file=sys.stderr) + return 1 + # CORRECTNESS CHECK: the decode must be byte-ADDITIVE across the input/target + # split, i.e. decode(ids[:1]) + decode(ids[1:]) reproduces decode(ids) + # exactly. That is what makes `val_bytes` the true byte length of the span + # the model is scored on rather than an approximation. + # + # NOT a re-encode check. `tok(decode(span)) == span` is the obvious test and + # it is WRONG here: BPE re-merges greedily, so a span that begins or ends + # mid-document re-tokenizes differently at its boundaries while decoding to + # byte-identical text. Measured on a real FineWeb-Edu val span, the re-encode + # test fails and the additivity test passes -- the byte count is fine, the + # test was not. + full_bytes = len(tok.decode(v_ids.tolist()).encode("utf-8")) + head_bytes = len(tok.decode(v_ids[:1].tolist()).encode("utf-8")) + byte_additive = (head_bytes + val_bytes == full_bytes) + if not byte_additive: + print(f"FATAL: decode is not byte-additive across the target split " + f"({head_bytes} + {val_bytes} != {full_bytes}); val_bytes would " + f"not describe the scored span", file=sys.stderr) + return 1 + + man_p.write_text(json.dumps({ + "corpus": f"{CORPUS}:{CORPUS_CONFIG}", + "tokenizer": TOKENIZER, + "vocab_size": VOCAB_SIZE, + "format": "flat uint32 token ids, no header (np.fromfile dtype=uint32)", + "val_ids": int(v_written), "val_docs": int(v_docs), + # THE FIELD THE HARNESS READS. bpb = nll / (val_bytes * ln2), scaled by + # val_target_tokens when a run scores fewer tokens than were staged. + "val_target_tokens": int(target_ids.size), + "val_bytes": int(val_bytes), + "val_bytes_per_token": round(val_bytes / target_ids.size, 4), + "val_decode_byte_additive": byte_additive, + "train_tokens": int(t_written), "train_docs": int(t_docs), + "disjoint": "val filled first; straddling document discarded; train follows", + }, indent=2)) + + print(f"[prep] val {v_written:,} ids from {v_docs:,} docs -> {val_p}") + print(f"[prep] val {val_bytes:,} BYTES over {target_ids.size:,} target tokens " + f"({val_bytes / target_ids.size:.2f} B/tok, byte-additive={byte_additive})") + print(f"[prep] train {t_written:,} ids from {t_docs:,} docs -> {train_p}") + print("[prep] disjoint by construction (straddling document discarded)") + return 0 + + +if __name__ == "__main__": + # os._exit, NOT SystemExit, and this is load-bearing rather than a style + # choice. `datasets` streaming leaves non-daemon fsspec/aiohttp threads + # alive; even after `docs.close()` above, normal interpreter shutdown blocks + # joining them. MEASURED: this script wrote train.bin, val.bin and + # manifest.json correctly within a minute and then hung for 14+ minutes + # doing nothing. docker/build_images.sh runs it on a cold asset cache and + # would hang with it, with no output to explain why. + # + # Everything this script produces is already durable at this point -- the + # .bin files are written and closed by _fill(), and manifest.json by + # pathlib.write_text -- so skipping teardown costs nothing. The explicit + # flush is required because os._exit does NOT flush Python's stdio buffers. + _rc = main() + sys.stdout.flush() + sys.stderr.flush() + os._exit(_rc) diff --git a/2.0/problems/nanoslm_hybrid_arch_design/evaluate.sh b/2.0/problems/nanoslm_hybrid_arch_design/evaluate.sh new file mode 100755 index 000000000..3e8e5a0ff --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/evaluate.sh @@ -0,0 +1,19 @@ +#!/usr/bin/env bash +# Local CLI wrapper for non-Harbor evaluation. +# +# Usage: +# bash evaluate.sh /path/to/model.py # score a submission (needs torch+GPU) +# bash evaluate.sh --selftest # torch-free policy/scoring/fingerprint tests +# +# Tip: FRONTIER_NANOSLM_SMOKE=1 shrinks the model/budget for a fast CPU wiring +# check (never used for real scoring). +set -euo pipefail + +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ARG="${1:-}" + +if [[ -z "${ARG}" ]]; then + ARG="${HERE}/reference.py" +fi + +exec python3 "${HERE}/evaluator.py" "${ARG}" diff --git a/2.0/problems/nanoslm_hybrid_arch_design/evaluator.py b/2.0/problems/nanoslm_hybrid_arch_design/evaluator.py new file mode 100644 index 000000000..0e9fa25f6 --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/evaluator.py @@ -0,0 +1,584 @@ +"""Evaluator for nanoslm_hybrid_arch_design (Frontier-CS 2.0). + +Contract: ``evaluate(solution_path) -> (score, score_unbounded, message, metrics)``. +The submission is a single ``model.py`` (full architecture freedom). The judge +trains it from scratch for a fixed wall-clock budget on a single H100 and scores +held-out validation perplexity against a locked baseline architecture trained +under the identical budget. See DESIGN.md. + +Top-level imports are torch-free so this module loads (and self-tests) without a +GPU; the training/eval path is lazy-imported inside :func:`evaluate`. +""" + +from __future__ import annotations + +import importlib.util +import json +import os +import sys +import traceback +from pathlib import Path + +# Torch-free harness layers (safe to import anywhere). +_HERE = Path(__file__).resolve().parent +if str(_HERE) not in sys.path: + sys.path.insert(0, str(_HERE)) + +from harness import policy, scoring, settings # noqa: E402 + +# cuBLAS determinism MUST be configured before the first CUDA call, which means +# before torch is imported anywhere. This module's top-level imports are +# deliberately torch-free (see the docstring), so setting it here is early +# enough on every path -- Modal, judge container, or a local GPU box. +# +# TaskConfig has carried `cublas_workspace_config` since the problem was +# written, and DESIGN.md 5 lists CUBLAS_WORKSPACE_CONFIG=:4096:8 as part of the +# determinism story, but NOTHING READ IT. So every run so far had +# torch.use_deterministic_algorithms(True) set while cuBLAS remained free to be +# non-deterministic -- which showed up as a UserWarning in the first GPU run and +# quietly weakened the CRN pairing the score depends on. +# `setdefault`, so an operator can still override from outside. +os.environ.setdefault( + "CUBLAS_WORKSPACE_CONFIG", settings.DEFAULT.cublas_workspace_config +) + + +def _protect_evaluator_source() -> None: + """Hide evaluator source from unprivileged submitted code in containers.""" + try: + p = Path(__file__).resolve() + if str(p).startswith(("/judge/", "/tests/")) and os.geteuid() == 0: + p.chmod(0o600) + except Exception: + pass + + +_protect_evaluator_source() + + +def _backend() -> str: + """MODAL (CPU judge -> Modal GPU) or LOCAL (directly-attached GPU). + + The design specifies a CPU judge with the H100 served on Modal, the same + shape as nanowm_rollout_* and vllm_llm_serving_optimization -- but until + now nothing implemented it: `evaluate()` only called `_select_device()`, + which in a CPU-only judge container returns "cpu" and would try to train a + ~190M model at ctx 8192 on CPU. Harbor invokes THIS module, not the manual + `modal_app.evaluate_remote` wrapper, so without this the problem is not + runnable under Harbor at all. + + Explicit env wins so either path can be forced in testing; otherwise a + Modal token implies Modal, and a local CUDA device implies local. + """ + b = os.environ.get("FRONTIER_NANOSLM_BACKEND", "").strip().lower() + if b in ("modal", "local"): + return b + if os.environ.get("MODAL_TOKEN_ID") or os.environ.get("MODAL_TOKEN_SECRET"): + return "modal" + return "local" + + +def _run_pair_modal(solution_path: str, role: str): + """Run both arms on a Modal GPU; score judge-side. + + Only ARM METRICS cross the boundary -- scoring and the hidden constants + (bpb_score_scale) stay in the judge, so a GPU worker never sees them and cannot + hand back a score the judge did not compute. + """ + from harness.modal_app import app, run_pair_remote + + source = Path(solution_path).read_text(encoding="utf-8", errors="replace") + with app.run(): + res = run_pair_remote.remote( + solution_source=source, + smoke=settings.smoke_enabled(), + role=role, + ) + return res + + +def _select_device() -> str: + try: + import torch + except Exception as exc: # torch not installed (dev box) + raise RuntimeError(f"torch unavailable: {type(exc).__name__}") + return "cuda" if torch.cuda.is_available() else "cpu" + + +def _role() -> str: + """'agent' (cheap iterative feedback, cached baseline) or 'final' (fresh CRN).""" + r = os.environ.get("FRONTIER_NANOSLM_ROLE", "final").strip().lower() + return "agent" if r == "agent" else "final" + + +def _baseline_cache_file(cfg) -> str: + return os.environ.get("FRONTIER_NANOSLM_BASELINE_CACHE", cfg.baseline_cache_path) + + +def _baseline_cache_get(cfg, fingerprint: str): + """Return cached baseline perplexity for this fingerprint, or None.""" + try: + with open(_baseline_cache_file(cfg), "r", encoding="utf-8") as fh: + return float(json.load(fh).get(fingerprint)) + except Exception: + return None + + +def _baseline_cache_put(cfg, fingerprint: str, ppl: float) -> None: + """Best-effort write of the baseline perplexity (path may be read-only).""" + path = _baseline_cache_file(cfg) + try: + os.makedirs(os.path.dirname(path) or ".", exist_ok=True) + try: + with open(path, "r", encoding="utf-8") as fh: + data = json.load(fh) + except Exception: + data = {} + data[fingerprint] = ppl + with open(path, "w", encoding="utf-8") as fh: + json.dump(data, fh) + except Exception: + pass + + +def _load_submission_module(path: str): + """Import the submitted model.py AFTER the static policy gate has passed.""" + spec = importlib.util.spec_from_file_location("submission_model", path) + if spec is None or spec.loader is None: + raise RuntimeError("could not load submission module") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def evaluate(solution_path: str): + cfg = settings.active_config() + fp = settings.config_fingerprint(cfg) + + # 1) Static policy gate (torch-free, adversarial-safe). + pol = policy.check_file(solution_path) + if not pol.ok: + return 0.0, 0.0, f"policy_rejected: {pol.reason}", {"config_fingerprint": fp} + + # 2) Backend. On MODAL the judge never imports torch -- the arms run in a + # Modal GPU container and only plain metrics come back. + role = _role() + if _backend() == "modal": + try: + res = _run_pair_modal(solution_path, role) + except Exception as exc: # noqa: BLE001 - classify, never leak + return (0.0, 0.0, f"environment_error: modal dispatch failed " + f"({type(exc).__name__})", {"config_fingerprint": fp}) + b, c = res["baseline"], res["submission"] + sr = scoring.score_from_bpb(b["val_bpb"], c["val_bpb"], cfg.bpb_score_scale) + note = ["smoke"] if settings.smoke_enabled() else [] + message = scoring.format_message( + b["val_bpb"], c["val_bpb"], sr, + steps=c["steps"], wall_seconds=c["wall_seconds"], + extra=",".join(note), + ) + return sr.score, sr.score_unbounded, message, { + "base_val_bpb": b["val_bpb"], "sub_val_bpb": c["val_bpb"], + "abs_bpb_delta": sr.abs_bpb_delta, "sub_val_ppl": c["val_ppl"], + # abs_bpb_delta is the SCORED quantity; rel_improvement is kept + # alongside it so an operator never has to recompute it. + "rel_improvement": sr.rel_improvement, + "bpb_score_scale": cfg.bpb_score_scale, + "sub_steps": c["steps"], "base_steps": b["steps"], + "sub_params": c["n_params"], "sub_wall_seconds": c["wall_seconds"], + "sub_train_block_size": c.get("train_block_size"), + "base_train_block_size": b.get("train_block_size"), + "eval_block_size": cfg.eval_block_size, + "role": role, "device": "modal:H100", + "config_fingerprint": fp, "stack": res.get("stack"), + } + + # LOCAL backend: torch + a directly-attached GPU. + try: + device = _select_device() + except RuntimeError as exc: + return ( + 0.0, + 0.0, + f"environment_error: {exc} (training path requires torch + GPU)", + {"config_fingerprint": fp}, + ) + + # 3) Train + eval; score perplexity vs the locked baseline. + # agent role: train submission only (~T), reuse fingerprint-cached baseline. + # final role: fresh baseline+submission CRN pair (~2T) — authoritative. + try: + from harness import runner # lazy: imports torch + from harness.data import TokenData + + module = _load_submission_module(solution_path) + + if role == "agent": + data = TokenData(cfg) + # Range-guarded before anything expensive runs. + sub_block = runner.resolve_train_block_size(module, cfg) + base_bpb = _baseline_cache_get(cfg, fp) + base_steps = None + if base_bpb is None: + base_arm = runner.run_arm( + runner.baseline_factory(), data, cfg, device, + runner.baseline_block_size(cfg), + ) + base_bpb = base_arm.val_bpb + base_steps = base_arm.steps + _baseline_cache_put(cfg, fp, base_bpb) + sub = runner.run_arm( + runner.load_factory(module), data, cfg, device, sub_block + ) + else: + base_arm, sub = runner.run_pair(module, cfg, device) + base_bpb = base_arm.val_bpb + base_steps = base_arm.steps + except Exception as exc: + from harness.runner import GuardError + + if isinstance(exc, GuardError): + return 0.0, 0.0, f"guard: {exc}", {"config_fingerprint": fp, "role": role} + # Never leak submission tracebacks/stdout (2.0 black-box safety). + return ( + 0.0, + 0.0, + f"submission_error: {type(exc).__name__}", + {"config_fingerprint": fp, "role": role}, + ) + + sr = scoring.score_from_bpb(base_bpb, sub.val_bpb, cfg.bpb_score_scale) + note = ["smoke"] if settings.smoke_enabled() else [] + if role == "agent": + note.append("iterative(cached-baseline)" if base_steps is None else "iterative") + message = scoring.format_message( + base_bpb, + sub.val_bpb, + sr, + steps=sub.steps, + wall_seconds=sub.wall_seconds, + extra=",".join(note), + ) + metrics = { + "base_val_bpb": base_bpb, + "sub_val_bpb": sub.val_bpb, + "abs_bpb_delta": sr.abs_bpb_delta, + "sub_val_ppl": sub.val_ppl, # derived, readability only + "rel_improvement": sr.rel_improvement, + "bpb_score_scale": cfg.bpb_score_scale, + "sub_steps": sub.steps, + "base_steps": base_steps, + "sub_params": sub.n_params, + "sub_wall_seconds": sub.wall_seconds, + # Agent-controlled TRAINING context vs the FIXED scoring window, both + # reported so a submission can see the steps-vs-extrapolation trade it + # actually made rather than inferring it. + "sub_train_block_size": sub.train_block_size, + "eval_block_size": cfg.eval_block_size, + "role": role, + "device": device, + "config_fingerprint": fp, + } + return sr.score, sr.score_unbounded, message, metrics + + +def prepare() -> dict: + """Warm assets + cache the baseline perplexity for the iterative feedback path. + + Best-effort: if torch/GPU is unavailable (dev box) we still return the + fingerprint so the judge readiness log is meaningful. + """ + cfg = settings.active_config() + info = {"config_fingerprint": settings.config_fingerprint(cfg)} + try: + device = _select_device() + from harness import runner + from harness.data import TokenData + + data = TokenData(cfg) + base = runner.run_arm( + runner.baseline_factory(), data, cfg, device, + runner.baseline_block_size(cfg), + ) + info["base_val_bpb"] = base.val_bpb + info["base_steps"] = base.steps + info["device"] = device + _baseline_cache_put(cfg, info["config_fingerprint"], base.val_bpb) + info["baseline_cached"] = True + except Exception as exc: + info["warm_skipped"] = f"{type(exc).__name__}" + return info + + +# --------------------------------------------------------------------------- # +# Torch-free self-test: exercises policy + scoring + fingerprint without a GPU. +# --------------------------------------------------------------------------- # +def _selftest() -> int: + ok = True + + def check(name, cond): + nonlocal ok + ok = ok and bool(cond) + print(f"[{'PASS' if cond else 'FAIL'}] {name}", file=sys.stderr) + + # --- policy: accept reference + baseline, reject malicious variants --- + ref_src = (_HERE / "reference.py").read_text() + base_src = (_HERE / "harness" / "baseline_model.py").read_text() + check("policy accepts reference.py", policy.check_source(ref_src).ok) + check("policy accepts baseline_model.py", policy.check_source(base_src).ok) + + malicious = { + "env leak": "import os\ndef build_model(c):\n x=os.environ['HF_TOKEN']\n class NanoSLM: pass\n", + "torch.load": "def build_model(c):\n import torch; torch.load('/opt/x'); return None\n", + "timer peek": "import time\ndef build_model(c):\n time.time()\n class NanoSLM: pass\n", + "metric peek": "def build_model(c):\n val_ppl=1.0\n class NanoSLM: pass\n", + "no factory": "x = 1\n", + "empty": " ", + } + for name, src in malicious.items(): + check(f"policy rejects {name}", not policy.check_source(src).ok) + + # legitimate idioms must survive + check( + "policy allows model.eval()/torch.compile()", + policy.check_source( + "class NanoSLM:\n def forward(self,x):\n self.eval(); return x\n" + ).ok, + ) + + # --- scoring: ABSOLUTE bpb gain. baseline->0, worse->0, scale->100 --- + # Cases are constructed by SUBTRACTING bits per byte, not by scaling + # base_bpb: the scored quantity is now `base_bpb - sub_bpb`, so a relative + # construction would no longer test what it claims to. + scale = 0.05 # bpb gain that saturates the bounded score + base_bpb = 4.0 + s_tie = scoring.score_from_bpb(base_bpb, base_bpb, scale) + s_worse = scoring.score_from_bpb(base_bpb, base_bpb + 0.10, scale) + s_sat = scoring.score_from_bpb(base_bpb, base_bpb - scale, scale) + s_half = scoring.score_from_bpb(base_bpb, base_bpb - scale / 2, scale) + s_over = scoring.score_from_bpb(base_bpb, base_bpb - 2 * scale, scale) + check("baseline scores 0", abs(s_tie.score) < 1e-9) + check("worse-than-baseline scores 0", abs(s_worse.score) < 1e-9) + check("a gain of BPB_SCORE_SCALE saturates to 100", abs(s_sat.score - 100.0) < 1e-6) + check("half the scale ~= 50", abs(s_half.score - 50.0) < 1e-6) + check("bounded score caps at 100", abs(s_over.score - 100.0) < 1e-9) + check("unbounded keeps rewarding past 100", s_over.score_unbounded > 100.0) + check("unbounded is exactly the un-clipped ratio", + abs(s_over.score_unbounded - 200.0) < 1e-9) + # The score must depend on the ABSOLUTE gain only -- the same 0.05 bpb gain + # scores the same at any operating point. (This is the property the old + # relative form did NOT have, and the reason the operating point now + # matters: see DESIGN.md 8.) + s_lowbase = scoring.score_from_bpb(1.5, 1.5 - scale, scale) + check("score depends on absolute gain, not on base_bpb", + abs(s_lowbase.score - s_sat.score) < 1e-9) + # ... while BOTH figures are still reported. + check("relative improvement still reported", + abs(s_sat.rel_improvement - (scale / base_bpb)) < 1e-12 + and abs(s_sat.abs_bpb_delta - scale) < 1e-12) + # A SCORING constant, not a training knob: changing it must not invalidate + # a cached baseline, so it must stay out of the fingerprint. + check("bpb_score_scale is NOT fingerprinted", + "bpb_score_scale" not in settings._FINGERPRINT_KEYS + and "r_target" not in settings._FINGERPRINT_KEYS) + + # --- fingerprint: stable, and changes iff a locked knob changes --- + from dataclasses import replace + + fp0 = settings.config_fingerprint(settings.DEFAULT) + fp_same = settings.config_fingerprint(settings.DEFAULT) + fp_diff = settings.config_fingerprint(replace(settings.DEFAULT, learning_rate=1e-3)) + check("fingerprint stable", fp0 == fp_same) + check("fingerprint changes on locked-knob change", fp0 != fp_diff) + + # The EVAL window is fingerprinted (changing it changes what val_bpb means, + # so a cached baseline must not be reused) ... + check( + "fingerprint changes on eval_block_size change", + fp0 != settings.config_fingerprint( + replace(settings.DEFAULT, eval_block_size=4096) + ), + ) + # ... while the TRAINING context is NOT. It varies per submission now, so + # fingerprinting it would give every submission a unique key and the cached + # baseline would never hit -- doubling the cost of the agent-role path. + check( + "fingerprint IGNORES the default training block_size", + fp0 == settings.config_fingerprint( + replace(settings.DEFAULT, block_size=2048) + ), + ) + + # --- agent-controlled training context: resolution + bounds --- + # Imported lazily: harness.runner pulls in torch, which the CPU dev box and + # the CPU judge container may not have. The selftest must still run there, + # so an absent torch skips these rather than failing them. + try: + from harness import runner as _runner + except Exception as exc: # pragma: no cover - torch-free host + print(f"[SKIP] BLOCK_SIZE checks ({type(exc).__name__})", file=sys.stderr) + _runner = None + + if _runner is not None: + import types + + cfg = settings.DEFAULT + + def _mod(**attrs): + m = types.ModuleType("sub") + for k, v in attrs.items(): + setattr(m, k, v) + return m + + check( + "absent BLOCK_SIZE -> cfg default", + _runner.resolve_train_block_size(_mod(), cfg) == cfg.block_size, + ) + check( + "BLOCK_SIZE=2048 accepted", + _runner.resolve_train_block_size(_mod(BLOCK_SIZE=2048), cfg) == 2048, + ) + check( + "baseline declares 8192 explicitly", + _runner.baseline_block_size(cfg) == cfg.eval_block_size == 8192, + ) + for bad in (128, 16384, 3000, -1, 0, "2048", 2048.0, True): + try: + _runner.resolve_train_block_size(_mod(BLOCK_SIZE=bad), cfg) + rejected = False + except _runner.GuardError: + rejected = True + check(f"BLOCK_SIZE={bad!r} rejected", rejected) + + # A model that sizes a positional table off the TRAINING context trains + # fine and then fails at the 8192 eval -- the characteristic new failure + # mode of the split contexts. It must come back as a GuardError naming + # BOTH contexts, not as a bare RuntimeError/IndexError. (torch raises + # IndexError here, which an earlier `except RuntimeError` missed.) + import torch + import torch.nn as nn + + from harness.data import TokenData as _TD + + class _BadPos(nn.Module): + def __init__(self, config): + super().__init__() + self.emb = nn.Embedding(config.vocab_size, 8) + # WRONG on purpose: should be config.eval_block_size. + self.pos = nn.Embedding(config.block_size, 8) + self.lin = nn.Linear(8, config.vocab_size) + + def forward(self, idx): + p = self.pos(torch.arange(idx.shape[1], device=idx.device)) + return self.lin(self.emb(idx) + p) + + tiny = replace( + settings.DEFAULT, vocab_size=256, eval_block_size=512, batch_size=2, + train_seconds=0.5, max_train_seconds=10.0, val_tokens=1024, + train_tokens_path="", val_tokens_path="", manifest_path="", + ) + try: + _runner.run_arm( + _runner.load_factory(_mod(build_model=_BadPos)), + _TD(tiny), tiny, "cpu", 256, + ) + got = "no error" + except _runner.GuardError as exc: + got = str(exc) + check( + "eval-context failure -> GuardError naming both contexts", + "evaluation failed at context 512" in got and "trained at 256" in got, + ) + + # --- THE CRUX: eval windows are 8192 whatever the training context is --- + # Numpy-only, no torch: exercise the window arithmetic val_windows uses. + try: + import numpy as np + + from harness.data import TokenData + except Exception as exc: # pragma: no cover + print(f"[SKIP] val_windows checks ({type(exc).__name__})", file=sys.stderr) + TokenData = None + + if TokenData is not None: + smoke = replace( + settings.DEFAULT, val_tokens=32_768, train_tokens_path="", + val_tokens_path="", manifest_path="", val_bytes=0, + ) + widths, denoms = set(), set() + for train_ctx in (8192, 2048, 256): + d = TokenData(replace(smoke, block_size=train_ctx)) + # val_windows yields torch tensors, so replicate its slicing here + # rather than importing torch: same expression, no device. + eb = d.cfg.eval_block_size + n = min(d.val.size - 1, d.cfg.val_tokens) + n_windows = n // eb + widths.add((eb, n_windows)) + denoms.add(round(d.val_bytes_for(n_windows * eb), 6)) + # ... and the TRAINING batch really does follow the training ctx. + x = np.stack([d.train[i: i + train_ctx] for i in (0, 1)]) + check(f"train batch width == {train_ctx}", x.shape[1] == train_ctx) + check("eval window width/count independent of training ctx", len(widths) == 1) + check("bpb denominator independent of training ctx", len(denoms) == 1) + check("eval window width is 8192", widths.pop() == (8192, 4)) + + # --- baseline cache HITS across submissions with different train contexts --- + # This is the reason block_size left the fingerprint. The key is derived + # from the CONFIG only; a submission's BLOCK_SIZE is resolved per-arm inside + # run_arm and never reaches the key, so one cached baseline serves both. + import tempfile + + with tempfile.TemporaryDirectory() as td: + cache = os.path.join(td, "baseline_ppl.json") + prev = os.environ.get("FRONTIER_NANOSLM_BASELINE_CACHE") + os.environ["FRONTIER_NANOSLM_BASELINE_CACHE"] = cache + try: + cfg_c = settings.DEFAULT + key = settings.config_fingerprint(cfg_c) + _baseline_cache_put(cfg_c, key, 1.2345) + # Two submissions, training at 8192 and at 2048, both look the + # baseline up under the SAME config -> the same key -> a hit. + hits = [ + _baseline_cache_get(cfg_c, settings.config_fingerprint(cfg_c)) + for _ in (8192, 2048) + ] + check("baseline cache hits for both training contexts", + hits == [1.2345, 1.2345]) + check( + "cache MISSES when the eval window changes", + _baseline_cache_get( + cfg_c, + settings.config_fingerprint( + replace(cfg_c, eval_block_size=4096)), + ) is None, + ) + finally: + if prev is None: + os.environ.pop("FRONTIER_NANOSLM_BASELINE_CACHE", None) + else: + os.environ["FRONTIER_NANOSLM_BASELINE_CACHE"] = prev + + print(("SELFTEST OK" if ok else "SELFTEST FAILED"), file=sys.stderr) + return 0 if ok else 1 + + +def main(argv) -> int: + if len(argv) == 2 and argv[1] == "--selftest": + return _selftest() + if len(argv) != 2: + print("usage: evaluator.py /path/to/model.py | --selftest", file=sys.stderr) + return 1 + try: + result = evaluate(argv[1]) + score, score_unbounded, message = result[0], result[1], result[2] + print(message, file=sys.stderr) + print(f"{score:.12f} {score_unbounded:.12f}") + return 0 + except Exception: + print(traceback.format_exc(), file=sys.stderr) + print("0.0 0.0") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv)) diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harbor/app/README.md b/2.0/problems/nanoslm_hybrid_arch_design/harbor/app/README.md new file mode 100644 index 000000000..3062eee82 --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/harbor/app/README.md @@ -0,0 +1,97 @@ +# NanoSLM Hybrid Architecture Design — workspace + +Design an architecture that reaches the lowest held-out **bits-per-byte** when +trained from scratch under a **fixed wall-clock budget** on one H100. + +The metric is bits per **byte** of held-out text, not per token: the judge +divides your total cross-entropy by the byte length of the validation text, so +it does not depend on how many bytes a token happens to cover. + +## Loop + +```bash +vim /app/model.py # 1. edit your architecture +bash /app/public_test.sh # 2. static gate (seconds, no GPU) +bash /app/submit.sh # 3. enqueue for the judge +``` + +You do **not** train here — this container has no GPU and no torch. The judge +trains both your architecture and a locked baseline under the identical budget +and scores your **absolute bpb gain** over it, `base_bpb - sub_bpb`. Bits per +byte is the unit the literature quotes, so the gain is directly comparable to +published numbers; the constant that maps it onto 0–100 is a scaling convention, +not a target, so simply maximize the gain. + +## What you submit + +`/app/model.py`, defining `build_model(config)` or `class NanoSLM`, whose +`forward(idx)` returns logits `[B, T, 100278]` for `idx` of BPE token ids. + +`config` gives you `vocab_size` (always **100278** — the dolma2 BPE tokenizer, +`allenai/dolma2-tokenizer`), `block_size` (the context you are **trained** at), +`eval_block_size` (the context you are **scored** at — always 8192), and +read-only hints. You control the architecture **and the training context +length**; the optimizer, data, tokenizer, evaluation context and the wall-clock +budget are fixed by the judge. + +## Training context length — a real lever, with a real catch + +Declare it with a module-level integer in `model.py`: + +```python +BLOCK_SIZE = 2048 # power of two in [256, 8192]; omit it to train at 8192 +``` + +Out-of-range or non-power-of-two values are rejected before training (score 0). + +**Evaluation is always at 8192**, whatever you train at. So: + +- shorter training context → cheaper steps → **more optimizer steps** in the + same fixed wall-clock budget (at 8192, attention is ~78–84% of layer FLOPs, so + this is a big effect); +- but you are still scored on 8192-token windows, so a model trained at 1024 has + to produce sensible logits **8x past** any position it ever saw. + +That second half depends heavily on your **position encoding**, which is a +different thing from architectural quality under a compute budget: plain RoPE +degrades sharply beyond its training length, while NTK-aware/YaRN scaling, +position interpolation and ALiBi extrapolate far better. Choosing a scheme that +survives the gap is now part of the design problem — do not shorten the context +without addressing it. + +Also: anything you size off `config.block_size` (RoPE tables, learned positional +embeddings, mask buffers) must still run at `config.eval_block_size`. Size them +against `eval_block_size`, or build them lazily from the actual `T`. + +The baseline always trains at 8192, so you are trading against a fixed point. + +Note the vocabulary is large: the embedding table alone is ~70M parameters at +d=704 (and the model ships it **tied**, one shared table), so how you handle it +(tying, factorizing, resizing) is part of the design problem rather than a +detail. + +## The starting point + +`/app/model.py` ships as the **3:1 Gated DeltaNet hybrid** from *Olmo Hybrid*: +three linear-recurrent layers per full-attention layer (25% attention), which +already beats the attention-only baseline. Your job is to push further. + +The 3:1 ratio is inherited, not tuned — it is the OLMo-3 sliding-window pattern +`[4096, 4096, 4096, -1]`, and the recipe replaces exactly the sliding-window +layers with GDN. It was chosen for a sliding window, not for a linear RNN, so +there is no reason to believe it is optimal here. + +The shipped model is also **parameter-matched** against the baseline the way +upstream does it (`REMOVE_HEADS`: d_model 768 -> 704, heads 12 -> 11, head_dim +64 preserved), which is why it is narrower than the shape its name suggests. + +Open questions it does not answer: the ratio (3:1 vs 5:1 vs 7:1), where the +attention layers belong, recurrent state size, whether every layer should be +identical, and the rest of the block (norms, gating, MLP ratio, head count). + +## Why wall-clock matters + +A cheaper mixer completes more optimizer steps in the same budget, and that is +the intended lever — at ctx 8192 attention is ~78–84% of layer FLOPs. But a +slower architecture is genuinely penalized, so throughput is part of the design +problem, not an afterthought. diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harbor/app/model.py b/2.0/problems/nanoslm_hybrid_arch_design/harbor/app/model.py new file mode 100644 index 000000000..6170d9453 --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/harbor/app/model.py @@ -0,0 +1,434 @@ +"""Reference solution for nanoslm_hybrid_arch_design -- the Olmo-Hybrid recipe at 190M. + +This is the STARTING POINT, not the ceiling. The locked baseline is a faithful +``olmo3_190M``; this file applies OLMo-core's own hybrid recipe +(``src/scripts/train/OLMo_hybrid/OLMo-hybrid-7B.py``) to it. + +WHERE THE 3:1 RATIO ACTUALLY COMES FROM +--------------------------------------- +It is NOT an arbitrary choice. ``olmo3_190M`` already carries:: + + SlidingWindowAttentionConfig(pattern=[4096, 4096, 4096, -1], + force_full_attention_on_first_layer=False, + force_full_attention_on_last_layer=True) + +so 9 of its 12 layers are sliding-window and 3 (layers 3, 7, 11) are full +attention. The paper "replaces sliding window layers with Gated DeltaNet +layers", and upstream's recipe encodes exactly that:: + + config.block_pattern = ["gdn", "gdn", "gdn", "attn"] + +The three sliding-window layers of each 4-layer group become GDN; the +full-attention layer of each group survives untouched. The 3:1 ratio IS olmo3's +own sliding-window pattern. Note the consequence: the hybrid has NO sliding +window left anywhere -- every surviving attention layer sat at a ``-1`` position. + +PARAMETER MATCHING (upstream's REMOVE_HEADS practice) +----------------------------------------------------- +GDN's mixer is larger than an attention mixer, so upstream shrinks the model to +compensate rather than letting the hybrid quietly buy capacity:: + + REMOVE_HEADS = 2 + config.d_model -= REMOVE_HEADS * 128 # 4096 -> 3840 + num_heads -= REMOVE_HEADS # 32 -> 30 + assert config.d_model / num_heads == 128 # head_dim preserved + +At 190M with head_dim 64 the equivalent is ``REMOVE_HEADS = 1``:: + + d_model 768 -> 704, n_heads 12 -> 11, 704 / 11 == 64 (head_dim preserved) + +Matching is done on NON-EMBEDDING parameters, which is the quantity upstream +itself feeds to ``Duration.chinchilla_tokens(model_params=...)``. Both arms TIE +their single vocab table, which at this scale is ~40% of all parameters -- +including it would drown the mixer difference being matched. Measured +analytically: + + baseline non-embedding 113,283,840 + hybrid non-embedding 110,805,734 -2.19% + +which is the same tolerance upstream's own 7B pair achieves (+2.84%). One caveat +worth stating: because ``d_model`` shrinks 768 -> 704, the tied table shrinks too +(77.0M -> 70.6M), so the hybrid's TOTAL is 4.7% below the baseline's +(181,401,446 vs 190,297,344) even though its non-embedding count is within 2.2%. + +Note also that ``hidden_size`` stays 3072. Upstream mutates ``d_model`` in place +and never recomputes the feed-forward width, so the FFN keeps the width derived +from the ORIGINAL d_model. Reproduced here deliberately. + +Everything else is identical to the baseline -- RMSNorm eps 1e-6, SwiGLU 3072, +reordered_norm blocks, qk_norm, RoPE theta 500_000, TIED embeddings (the baseline +ties, so this arm ties too for comparability) -- so the ONLY architectural +difference is the sequence mixer in those nine layers. + +YOUR TASK: PUSH val_bpb FURTHER +------------------------------- +Everything above is upstream's recipe, transposed to this scale. The open +questions it does NOT answer, each worth real bpb: + + * THE RATIO. 3:1 is inherited from olmo3's sliding-window pattern, which was + chosen for a sliding window, not for a linear RNN. 5:1 and 7:1 buy more + steps but give up more global context. + * PLACEMENT. Every 4th, or clustered (early layer for global context in, late + layer for read-out)? Same cost, different models. + * STATE SIZE. ``expand_v`` and ``num_v_heads`` set the recurrent state, the + main capacity knob of a linear RNN -- unlike a KV cache it does not grow + with sequence length, so capacity is cheap at 8192. + * THE MIXER ITSELF. GDN is one choice; GLA, RetNet and Mamba2-style mixers are + all expressible in the same chunkwise matmul form. + * NON-UNIFORMITY. Nothing requires every GDN layer to be identical, or the + attention layers to be full-width. + * THE REST OF THE BLOCK. Norm placement, gating, MLP ratio, head count, + embedding tying -- all still on the table, and all interact with the above. + * THE TRAINING CONTEXT. Declare a module-level ``BLOCK_SIZE`` int (a power of + two in [256, 8192]) to train at a shorter context than the default 8192. + Shorter steps are cheaper, so you complete more of them in the fixed + wall-clock budget -- but EVALUATION IS ALWAYS AT 8192, so the model must + extrapolate to positions well beyond anything it trained on. How well it + does that is mostly a property of the POSITION ENCODING (plain RoPE degrades + sharply; NTK-aware/YaRN scaling and ALiBi hold up far better), which is a + different question from mixer efficiency. This file declares no BLOCK_SIZE + and therefore trains at 8192. Anything sized off ``config.block_size`` must + still run at ``config.eval_block_size``. + +Locked and not yours to change: optimizer, data, tokenizer, the EVALUATION +context (8192), and the wall-clock budget. Interface: ``build_model(config)`` / +``NanoSLM(config)`` returning ``forward(idx) -> logits [B, T, vocab_size]``. +""" + +from __future__ import annotations + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +# --------------------------------------------------------------------------- # +# Shape: olmo3_190M with upstream's REMOVE_HEADS compensation applied. +# --------------------------------------------------------------------------- # +REMOVE_HEADS = 1 + +N_LAYER = 12 +N_HEAD = 12 - REMOVE_HEADS # 11 +HEAD_DIM = 64 # preserved, exactly as upstream asserts +N_EMBD = 768 - REMOVE_HEADS * HEAD_DIM # 704 +assert N_EMBD // N_HEAD == HEAD_DIM, "REMOVE_HEADS must preserve head_dim" + +# NOT recomputed from the reduced d_model -- upstream mutates d_model in place +# and leaves the feed-forward width at the value llama_like derived from the +# ORIGINAL 768. Reproduced deliberately. +HIDDEN_SIZE = 3072 + +ROPE_THETA = 500_000.0 +NORM_EPS = 1e-6 + +# GatedDeltaNetConfig(n_heads=num_heads, head_dim=int(0.75*d_model/num_heads), +# allow_neg_eigval=True) with expand_v at its 2.0 default. +GDN_HEAD_DIM = int(0.75 * N_EMBD / N_HEAD) # 48 +GDN_EXPAND_V = 2.0 +GDN_HEAD_V_DIM = int(GDN_HEAD_DIM * GDN_EXPAND_V) # 96 +GDN_KEY_DIM = N_HEAD * GDN_HEAD_DIM # 528 +GDN_VALUE_DIM = N_HEAD * GDN_HEAD_V_DIM # 1056 +GDN_CONV_SIZE = 4 + +# config.block_pattern = ["gdn", "gdn", "gdn", "attn"] -- the three +# sliding-window layers of each group become GDN, the full-attention layer +# survives. Attention therefore lands at layers 3, 7, 11. +PATTERN = ("gdn", "gdn", "gdn", "attn") +assert N_LAYER % len(PATTERN) == 0 + + +class _RMSNorm(nn.Module): + def __init__(self, d: int, eps: float = NORM_EPS): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(d)) + + def forward(self, x): + dt = x.dtype + xf = x.float() + n = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps) + return n.to(dt) * self.weight + + +def _rope(x, theta: float = ROPE_THETA): + _, _, T, D = x.shape + half = D // 2 + freq = theta ** (-torch.arange(0, half, device=x.device, dtype=torch.float32) / half) + ang = torch.arange(T, device=x.device, dtype=torch.float32)[:, None] * freq[None, :] + cos, sin = ang.cos()[None, None], ang.sin()[None, None] + x1, x2 = x[..., :half].float(), x[..., half:].float() + return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1).to(x.dtype) + + +class _Attention(nn.Module): + """Identical to the baseline's attention, always FULL causal. + + Full, not windowed, and that is not a simplification: the attention layers + the hybrid keeps are exactly the ones sitting at the ``-1`` positions of + olmo3's [4096, 4096, 4096, -1] pattern, so they were already full attention + before the GDN substitution. + + qk_norm spans the full n_heads*head_dim projection and is applied BEFORE the + head reshape -- upstream's behaviour when ``use_head_qk_norm`` is False, + which is the olmo3_190M default. + """ + + def __init__(self, n_embd: int, n_head: int): + super().__init__() + assert n_embd % n_head == 0 + self.n_head, self.n_embd = n_head, n_embd + self.head_dim = n_embd // n_head + self.w_q = nn.Linear(n_embd, n_embd, bias=False) + self.w_k = nn.Linear(n_embd, n_embd, bias=False) + self.w_v = nn.Linear(n_embd, n_embd, bias=False) + self.w_out = nn.Linear(n_embd, n_embd, bias=False) + self.q_norm = _RMSNorm(n_embd) + self.k_norm = _RMSNorm(n_embd) + + def forward(self, x): + B, T, C = x.shape + q, k, v = self.w_q(x), self.w_k(x), self.w_v(x) + q, k = self.q_norm(q), self.k_norm(k) + q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2) + k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2) + v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2) + q, k = _rope(q), _rope(k) + y = F.scaled_dot_product_attention(q, k, v, is_causal=True) + return self.w_out(y.transpose(1, 2).contiguous().view(B, T, C)) + + +class _ChunkedDeltaNet(nn.Module): + """CPU-only chunkwise fallback, PARAMETER-IDENTICAL to fla's GatedDeltaNet. + + Exists so the CPU smoke path can build and report the same parameter count + the CUDA path will. Its module inventory is fla 0.5.1's exactly -- + q/k/v/a/b/g/o projections, three depthwise short convolutions, A_log, + dt_bias and the gated output norm -- so ``n_params`` is device-independent. + + The recurrence is a chunkwise GATED LINEAR ATTENTION, not the full delta + rule: it carries the decay and the gates but skips the delta-rule inverse. + That is a deliberate approximation. It is never scored -- CUDA raises rather + than falling back here (see ``_make_gdn``) -- and it exists to prove wiring, + not to reproduce GDN's numerics. + """ + + def __init__(self, chunk: int = 512): + super().__init__() + self.chunk = chunk + d, H = N_EMBD, N_HEAD + self.q_proj = nn.Linear(d, GDN_KEY_DIM, bias=False) + self.k_proj = nn.Linear(d, GDN_KEY_DIM, bias=False) + self.v_proj = nn.Linear(d, GDN_VALUE_DIM, bias=False) + self.a_proj = nn.Linear(d, H, bias=False) + self.b_proj = nn.Linear(d, H, bias=False) + self.g_proj = nn.Linear(d, GDN_VALUE_DIM, bias=False) + self.o_proj = nn.Linear(GDN_VALUE_DIM, d, bias=False) + self.A_log = nn.Parameter(torch.zeros(H)) + self.dt_bias = nn.Parameter(torch.zeros(H)) + # ShortConvolution weights: depthwise, (channels, 1, kernel), no bias. + self.q_conv1d = nn.Parameter(torch.zeros(GDN_KEY_DIM, 1, GDN_CONV_SIZE)) + self.k_conv1d = nn.Parameter(torch.zeros(GDN_KEY_DIM, 1, GDN_CONV_SIZE)) + self.v_conv1d = nn.Parameter(torch.zeros(GDN_VALUE_DIM, 1, GDN_CONV_SIZE)) + self.o_norm = nn.Parameter(torch.ones(GDN_HEAD_V_DIM)) + + @staticmethod + def _short_conv(x, w): + """Causal depthwise conv over time. x [B, T, C], w [C, 1, K].""" + C, K = w.shape[0], w.shape[2] + xt = F.pad(x.transpose(1, 2), (K - 1, 0)) + return F.conv1d(xt, w, groups=C).transpose(1, 2) + + def forward(self, x): + B, T, _ = x.shape + H, DK, DV, S = N_HEAD, GDN_HEAD_DIM, GDN_HEAD_V_DIM, self.chunk + + q = F.silu(self._short_conv(self.q_proj(x), self.q_conv1d)) + k = F.silu(self._short_conv(self.k_proj(x), self.k_conv1d)) + v = self._short_conv(self.v_proj(x), self.v_conv1d) + + q = q.view(B, T, H, DK).transpose(1, 2) + k = k.view(B, T, H, DK).transpose(1, 2) + v = v.view(B, T, H, DV).transpose(1, 2) + + # Per-head decay in (0, 1): the gate that makes this "gated". + dt = F.softplus(self.a_proj(x) + self.dt_bias) # [B, T, H] + g = torch.exp(-torch.exp(self.A_log) * dt) # [B, T, H] + g = g.permute(0, 2, 1).unsqueeze(-1) # [B, H, T, 1] + beta = torch.sigmoid(self.b_proj(x)).permute(0, 2, 1).unsqueeze(-1) + v = v * beta + + pad = (S - T % S) % S + if pad: + q, k, v = (F.pad(t, (0, 0, 0, pad)) for t in (q, k, v)) + g = F.pad(g, (0, 0, 0, pad), value=1.0) + nC = (T + pad) // S + rs = lambda t: t.view(B, H, nC, S, -1) # noqa: E731 + qc, kc, vc, gc = rs(q), rs(k), rs(v), rs(g) + + causal = torch.tril(torch.ones(S, S, device=x.device, dtype=torch.bool)) + intra = (qc @ kc.transpose(-1, -2)).masked_fill(~causal, 0.0) @ vc + + state = torch.zeros(B, H, DK, DV, device=x.device, dtype=x.dtype) + outs = [] + for c in range(nC): + outs.append(intra[:, :, c] + qc[:, :, c] @ state) + decay = gc[:, :, c].prod(dim=2, keepdim=True) + state = state * decay + kc[:, :, c].transpose(-1, -2) @ (vc[:, :, c] * gc[:, :, c]) + y = torch.stack(outs, dim=2).view(B, H, T + pad, DV)[:, :, :T] + + # Gated RMSNorm over head_v_dim, then merge heads. + yf = y.float() + y = (yf * torch.rsqrt(yf.pow(2).mean(-1, keepdim=True) + 1e-5)).to(x.dtype) + y = y * self.o_norm + y = y.transpose(1, 2).contiguous().view(B, T, GDN_VALUE_DIM) + y = y * F.silu(self.g_proj(x)) + return self.o_proj(y) + + +def _make_gdn(): + """fla's fused GatedDeltaNet on CUDA; chunkwise only on CPU (smoke). + + WHY fla IS REQUIRED HERE, not merely preferred: PyTorch ships a fused kernel + for softmax attention (SDPA) but none for a linear/recurrent mixer, so a + hand-written GDN is unfused eager code. Scoring is fixed WALL-CLOCK, so an + unfused mixer loses on throughput regardless of architectural merit -- + measured at ctx 8192, the chunkwise fallback did 6% of attention's FLOPs and + still ran 3.5x slower. Both arms must be kernel-matched or the score + measures kernel quality, not architecture. + + HARD FAILURE ON CUDA, deliberately: a silent fallback here once produced a + scored-looking run where the reference "lost" 18 steps to 65 purely because + fla was missing. A missing kernel must raise, not quietly change what is + being measured. + """ + if not torch.cuda.is_available(): + return _ChunkedDeltaNet() + + from fla.layers import GatedDeltaNet # ImportError here is intentional + + # head_dim MUST be passed: fla defaults it to 256, so omitting it builds + # num_heads*256-wide projections instead of the intended width -- a silent + # param blowup that trains a much larger model than the baseline. + return GatedDeltaNet( + hidden_size=N_EMBD, + num_heads=N_HEAD, + head_dim=GDN_HEAD_DIM, + expand_v=GDN_EXPAND_V, + use_gate=True, + use_short_conv=True, + ) + + +class _GDNLayer(nn.Module): + def __init__(self): + super().__init__() + self.impl = _make_gdn() + + def forward(self, x): + out = self.impl(x) + # fla layers return (hidden_states, attentions, past_kv); the chunkwise + # fallback returns a tensor. Normalize so _Block need not care. + return out[0] if isinstance(out, tuple) else out + + +class _FeedForward(nn.Module): + """SwiGLU, hidden 3072, bias=False -- unchanged from the baseline.""" + + def __init__(self, n_embd: int, hidden: int = HIDDEN_SIZE): + super().__init__() + self.w1 = nn.Linear(n_embd, hidden, bias=False) + self.w3 = nn.Linear(n_embd, hidden, bias=False) + self.w2 = nn.Linear(hidden, n_embd, bias=False) + + def forward(self, x): + return self.w2(F.silu(self.w1(x)) * self.w3(x)) + + +class _ReorderedNormBlock(nn.Module): + """Upstream builds the GDN block as ``attn_block.replace(sequence_mixer=...)`` + -- same reordered_norm structure, same norms, only the mixer swapped.""" + + def __init__(self, n_embd: int, n_head: int, kind: str): + super().__init__() + self.attention = _Attention(n_embd, n_head) if kind == "attn" else _GDNLayer() + self.attention_norm = _RMSNorm(n_embd) + self.feed_forward = _FeedForward(n_embd) + self.feed_forward_norm = _RMSNorm(n_embd) + + def forward(self, x): + h = x + self.attention_norm(self.attention(x)) + return h + self.feed_forward_norm(self.feed_forward(h)) + + +class NanoSLM(nn.Module): + def __init__(self, config): + super().__init__() + self.block_size = config.block_size + self.wte = nn.Embedding(config.vocab_size, N_EMBD) + kinds = [PATTERN[i % len(PATTERN)] for i in range(N_LAYER)] + self.blocks = nn.ModuleList([_ReorderedNormBlock(N_EMBD, N_HEAD, k) for k in kinds]) + self.norm_f = _RMSNorm(N_EMBD) + # TIED -- the baseline ties its embeddings (a deliberate deviation from + # upstream's untied default), so the reference ties too and the two arms + # stay comparable: the ONLY architectural difference between them is the + # sequence mixer in the GDN layers, not the output-head param budget. The + # tie is set AFTER init (below) so the shared table carries wte's init. + self.lm_head = nn.Linear(N_EMBD, config.vocab_size, bias=False) + self.apply(self._init) + for name, p in self.named_parameters(): + if name.endswith("w_out.weight") or name.endswith("w2.weight"): + nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * N_LAYER)) + self.lm_head.weight = self.wte.weight # weight tying + + @staticmethod + def _init(m): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, mean=0.0, std=0.02) + if m.bias is not None: + nn.init.zeros_(m.bias) + elif isinstance(m, nn.Embedding): + nn.init.normal_(m.weight, mean=0.0, std=0.02) + + def forward(self, idx): + x = self.wte(idx) + for block in self.blocks: + x = block(x) + x = self.norm_f(x) + return self.lm_head(x) # logits [B, T, vocab_size] + + +def build_model(config) -> NanoSLM: + return NanoSLM(config) + + +def analytic_n_params(vocab_size: int) -> int: + """Analytic total, derived from the spec rather than from a measurement.""" + d, H, hid = N_EMBD, N_HEAD, HIDDEN_SIZE + block_norms = 2 * d + ffn = 3 * d * hid + attn_mix = 4 * d * d + 2 * d # projections + q/k norms + gdn_mix = ( + 2 * d * GDN_KEY_DIM # q_proj, k_proj + + 2 * d * GDN_VALUE_DIM # v_proj, g_proj + + 2 * d * H # a_proj, b_proj + + GDN_VALUE_DIM * d # o_proj + + 2 * H # A_log, dt_bias + + GDN_CONV_SIZE * (2 * GDN_KEY_DIM + GDN_VALUE_DIM) # short convs + + GDN_HEAD_V_DIM # o_norm + ) + n_attn = sum(1 for i in range(N_LAYER) if PATTERN[i % len(PATTERN)] == "attn") + n_gdn = N_LAYER - n_attn + non_emb = n_attn * (attn_mix + block_norms + ffn) + n_gdn * (gdn_mix + block_norms + ffn) + d + # ONE vocab-sized table: lm_head is TIED to wte (see NanoSLM.__init__). + return non_emb + vocab_size * d + + +def _self_check(vocab_size: int = 100278) -> int: + from harness.model_config import ModelConfig + + m = NanoSLM(ModelConfig(vocab_size=vocab_size, block_size=8192)) + got = sum(p.numel() for p in m.parameters() if p.requires_grad) + want = analytic_n_params(vocab_size) + assert got == want, f"param mismatch: built {got} != analytic {want}" + return got diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harbor/app/public_test.py b/2.0/problems/nanoslm_hybrid_arch_design/harbor/app/public_test.py new file mode 100644 index 000000000..e8c59675b --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/harbor/app/public_test.py @@ -0,0 +1,64 @@ +"""Run the judge's exact static gate on a submission, locally. + +Imports /app/policy.py -- the same module the judge uses -- rather than +restating the rules, so the two cannot drift apart. + +This checks only what is checkable without a GPU: size, the required factory, +and the forbidden-token rules. It cannot tell you whether your architecture is +GOOD; only a scored submission does that. +""" +from __future__ import annotations + +import importlib.util +import pathlib +import sys + + +def _load(name: str, path: pathlib.Path): + spec = importlib.util.spec_from_file_location(name, path) + mod = importlib.util.module_from_spec(spec) + # Register BEFORE exec_module. policy.py defines a @dataclass, and + # dataclasses resolves string annotations via sys.modules[cls.__module__]; + # if the module is absent from sys.modules that lookup returns None and the + # decorator dies with "NoneType has no attribute __dict__". This is the + # documented spec_from_file_location idiom, not a workaround. + sys.modules[name] = mod + spec.loader.exec_module(mod) + return mod + + +def main() -> int: + target = pathlib.Path(sys.argv[1] if len(sys.argv) > 1 else "/app/model.py") + if not target.exists(): + print(f"FAIL {target} does not exist") + return 1 + + policy = _load("policy", pathlib.Path("/app/policy.py")) + src = target.read_text(encoding="utf-8", errors="replace") + res = policy.check_source(src) + + print(f"file : {target} ({len(src.encode()):,} bytes)") + print(f"policy : {'ok' if res.ok else 'REJECTED'}") + if not res.ok: + print(f"reason : {res.reason}") + return 1 + + # Cheap structural check the judge would otherwise charge a training run for. + import ast + try: + tree = ast.parse(src) + except SyntaxError as exc: + print(f"FAIL does not parse: line {exc.lineno}") + return 1 + names = {n.name for n in ast.walk(tree) + if isinstance(n, (ast.FunctionDef, ast.ClassDef))} + if "build_model" not in names and "NanoSLM" not in names: + print("FAIL must define build_model(config) or class NanoSLM") + return 1 + + print("RESULT : PASS (static only -- architecture quality is not checked here)") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harbor/app/public_test.sh b/2.0/problems/nanoslm_hybrid_arch_design/harbor/app/public_test.sh new file mode 100755 index 000000000..0dd4ed07f --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/harbor/app/public_test.sh @@ -0,0 +1,4 @@ +#!/usr/bin/env bash +# Static policy gate -- the judge's exact rules. Fast, no GPU, no training. +set -euo pipefail +exec python3 /app/public_test.py "${1:-/app/model.py}" diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harness/__init__.py b/2.0/problems/nanoslm_hybrid_arch_design/harness/__init__.py new file mode 100644 index 000000000..75043f4dc --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/__init__.py @@ -0,0 +1,9 @@ +"""Locked training/eval harness for nanoslm_hybrid_arch_design. + +Only ``model.py`` is submitted by the agent; everything in this package is +judge-owned and not part of the submission. The torch-free modules +(:mod:`settings`, :mod:`policy`, :mod:`scoring`) are unit-testable on CPU with +no GPU or PyTorch (as is :mod:`model_config`); the torch-dependent modules +(:mod:`data`, :mod:`train`, +:mod:`eval_ppl`, :mod:`baseline_model`, :mod:`runner`) run on the H100. +""" diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harness/baseline_model.py b/2.0/problems/nanoslm_hybrid_arch_design/harness/baseline_model.py new file mode 100644 index 000000000..eb2e1e2ad --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/baseline_model.py @@ -0,0 +1,314 @@ +"""Locked baseline architecture (hidden; the score-0 reference point). + +A plain-PyTorch reimplementation of OLMo-core's ``TransformerConfig.olmo3_190M``, +faithful in every respect EXCEPT that the word embeddings are TIED (see +"ON THE PARAMETER COUNT" below). Upstream unties them; this benchmark ties BOTH +arms (this baseline and ``reference.py``) instead, as a deliberate choice so the +only architectural difference between them is the sequence mixer, not the +output-head param budget. Tying is NOT forced by the ``param_cap`` (400M -- the +untied 267.3M model would fit); it is a comparability choice, and the single +deliberate deviation from upstream. + +Implements the exact interface the agent must also honor: ``build_model(config)`` +/ ``NanoSLM(config)`` returning a module whose ``forward(idx)`` yields logits +``[B, T, vocab_size]``. + +WHAT olmo3_190M ACTUALLY IS (verified against upstream, not reconstructed) +------------------------------------------------------------------------- +``olmo3_190M`` == ``olmo2_190M`` + sliding window + flash_2, and ``olmo2_190M`` +is ``llama_like(...)`` with:: + + d_model=768, n_layers=12, n_heads=12 -> head_dim 64 + hidden_size_multiplier=1.5 + block_name=TransformerBlockType.reordered_norm + qk_norm=True + rope_theta=500_000 + layer_norm_eps=1e-6 + +``llama_like`` then fixes the rest:: + + hidden_size = int(8 * d_model / 3) # 2048 + hidden_size = int(1.5 * hidden_size) # 3072 + hidden_size = ensure_multiple_of(., 256) # 3072 + bias = False everywhere + layer_norm = LayerNormType.rms -> RMSNorm (NOT LayerNorm) + +and ``olmo3_190M`` adds:: + + SlidingWindowAttentionConfig(pattern=[4096, 4096, 4096, -1], + force_full_attention_on_first_layer=False, + force_full_attention_on_last_layer=True) + +THE reordered_norm BLOCK IS NOT PRE-NORM. From upstream ``block.py``:: + + h = x + attention_norm(attention(x)) + out = h + feed_forward_norm(feed_forward(h)) + +The norm is applied to the residual branch's OUTPUT, not to its input. An +earlier revision of this file was a nanoGPT-style pre-norm block with only the +dimensions swapped -- LayerNorm, GELU 4x MLP, biases everywhere, no qk_norm, no +sliding window -- which measured 239,083,008 params and was not olmo3_190M in +any respect other than d/L/H. That is retracted. + +ON THE PARAMETER COUNT -- READ THIS BEFORE "FIXING" IT +------------------------------------------------------ +Upstream leaves ``tie_word_embeddings`` at its dataclass default of False, so a +faithful olmo3_190M carries TWO vocab-sized matrices and totals **267,310,848** +trainable params at vocab 100278. That fits comfortably under this task's +``param_cap`` of 400,000,000, so the tie is NOT a legality workaround. This +baseline (and ``reference.py``) TIE the embeddings so ``lm_head`` shares +``wte``'s weight, one 77,013,504-param table is removed, and the total drops to:: + + 12 blocks (113,283,072) + ONE shared embedding table (77,013,504) + + final norm (768) = 190,297,344 + +which is the "190M" the model's name refers to. Tying BOTH arms keeps them +differing only in the sequence mixer; the shared table is still ~40% of all +params -- at the 190M shape with a 100278-id vocabulary the embedding dominates; +at 7B the same table is ~11% and near-negligible. + +WHY THE BASELINE IS NOT A HYBRID +-------------------------------- +``reference.py`` is the Olmo-Hybrid recipe applied to this exact model, and it +is the floor the agent starts from. For that reference to be a meaningful floor +it has to beat something, so the score-0 point is the pure-attention model it +improves on. The scored question is therefore "how much further past a +competent hybrid can you push val_bpb". +""" + +from __future__ import annotations + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +# --------------------------------------------------------------------------- # +# olmo3_190M's shape and hyper-parameters. Every constant below is upstream's. +# --------------------------------------------------------------------------- # +N_LAYER = 12 +N_HEAD = 12 +N_EMBD = 768 # head_dim = 768 / 12 = 64 +HEAD_DIM = N_EMBD // N_HEAD + +# int(8*768/3) = 2048 -> int(1.5*2048) = 3072 -> ensure_multiple_of(3072, 256) +HIDDEN_SIZE = 3072 + +ROPE_THETA = 500_000.0 +NORM_EPS = 1e-6 + +# SlidingWindowAttentionConfig(pattern=[4096,4096,4096,-1], +# force_full_attention_on_first_layer=False, +# force_full_attention_on_last_layer=True) +SWA_PATTERN = (4096, 4096, 4096, -1) +FORCE_FULL_FIRST = False +FORCE_FULL_LAST = True + +# TRAINING CONTEXT OF THE LOCKED ARM, declared through the same module-level +# ``BLOCK_SIZE`` protocol a submission uses (runner.resolve_train_block_size). +# +# Stated EXPLICITLY rather than inheriting the task's ``block_size``: that field +# (8192 in the task config) is now only the submission-facing DEFAULT, and the +# score-0 reference point must not silently move if it is ever retuned. It also +# has to equal ``eval_block_size`` (also 8192) for the cached baseline to remain +# a valid comparison partner for submissions that train shorter -- the baseline +# never trades context for steps, so the trade a submission makes is measured +# against a fixed point. +BLOCK_SIZE = 8192 + + +def window_size_for_layer(layer_idx: int, n_layers: int) -> int: + """Port of ``SlidingWindowAttentionConfig._get_window_size``. -1 == full. + + Note the ``force_full_attention_on_first_layer`` branch also SHIFTS the + pattern index (upstream applies the pattern starting from the second layer + in that case). olmo3 sets it False, so there is no shift here -- but the + shift is reproduced anyway so the function is a faithful port. + """ + if FORCE_FULL_FIRST and layer_idx == 0: + return -1 + if FORCE_FULL_LAST and layer_idx == (n_layers - 1): + return -1 + eff = layer_idx - 1 if FORCE_FULL_FIRST else layer_idx + return SWA_PATTERN[eff % len(SWA_PATTERN)] + + +# For N_LAYER=12 this yields full attention at layers 3, 7, 11 and a 4096-wide +# window everywhere else -- the 9:3 split the reference's 3:1 GDN:attention +# ratio comes from. It is olmo3's own pattern, not an arbitrary choice. +LAYER_WINDOWS = tuple(window_size_for_layer(i, N_LAYER) for i in range(N_LAYER)) + + +class _RMSNorm(nn.Module): + """LayerNormType.rms with bias=False, eps=1e-6. Computed in fp32.""" + + def __init__(self, d: int, eps: float = NORM_EPS): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(d)) + + def forward(self, x): + dt = x.dtype + xf = x.float() + n = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps) + return n.to(dt) * self.weight + + +# --------------------------------------------------------------------------- # +# Sliding-window mask cache. +# +# torch's SDPA has no native window flag, so a banded causal mask is built and +# handed to `attn_mask`. It is cached per (T, window, device): one bool mask at +# T=8192 is 64 MiB, so rebuilding it 9x per forward would be ruinous, and all +# nine windowed layers share the same window (4096) hence the same mask. +# +# KERNEL CAVEAT, stated plainly: passing an explicit `attn_mask` disqualifies +# SDPA's fused flash kernel and falls back to the memory-efficient backend, +# whereas upstream uses flash_2's native `window_size=(4095, 0)`. The windowed +# layers are therefore SLOWER here than upstream would be, despite doing fewer +# FLOPs. At ctx 8192 a 4096 window removes only ~25% of attention FLOPs anyway +# (full causal already averages T/2 = 4096 keys per query), so the window is a +# faithfulness feature here, not a speed feature. +# --------------------------------------------------------------------------- # +_MASK_CACHE: dict = {} + + +def _sliding_causal_mask(T: int, window: int, device) -> torch.Tensor: + """Bool mask [T, T], True == attend. Keys in [i-window+1, i] inclusive. + + Matches upstream's flash window_size=(window-1, 0), documented there as + "window is [i - window_size[0], i + window_size[1]] inclusive". + """ + key = (T, window, str(device)) + m = _MASK_CACHE.get(key) + if m is None: + i = torch.arange(T, device=device) + q, k = i[:, None], i[None, :] + m = (k <= q) & (k > q - window) + _MASK_CACHE[key] = m + return m + + +def _rope(x, theta: float = ROPE_THETA): + """Rotary position embedding over the head dim, computed in fp32. + + theta=500_000 (upstream ``rope_theta``), and ``rope_full_precision=True`` + upstream, hence the fp32 angle computation before casting back. + """ + _, _, T, D = x.shape + half = D // 2 + freq = theta ** (-torch.arange(0, half, device=x.device, dtype=torch.float32) / half) + ang = torch.arange(T, device=x.device, dtype=torch.float32)[:, None] * freq[None, :] + cos, sin = ang.cos()[None, None], ang.sin()[None, None] + x1, x2 = x[..., :half].float(), x[..., half:].float() + return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1).to(x.dtype) + + +class _Attention(nn.Module): + """olmo_core AttentionConfig: bias=False, qk_norm=RMSNorm, RoPE, optional SWA. + + qk_norm is applied over the FULL n_heads*head_dim projection, BEFORE the + head reshape. That is upstream's behaviour when ``use_head_qk_norm`` is + False, which is the olmo3_190M default (``llama_like`` only forwards + ``use_head_qk_norm`` when explicitly asked, and olmo2_190M never asks). It + is NOT a per-head norm -- q_norm/k_norm are 768-wide, not 64-wide. + """ + + def __init__(self, n_embd: int, n_head: int, window: int): + super().__init__() + assert n_embd % n_head == 0 + self.n_head, self.n_embd = n_head, n_embd + self.head_dim = n_embd // n_head + self.window = window # -1 == full causal attention + self.w_q = nn.Linear(n_embd, n_embd, bias=False) + self.w_k = nn.Linear(n_embd, n_embd, bias=False) + self.w_v = nn.Linear(n_embd, n_embd, bias=False) + self.w_out = nn.Linear(n_embd, n_embd, bias=False) + self.q_norm = _RMSNorm(n_embd) + self.k_norm = _RMSNorm(n_embd) + + def forward(self, x): + B, T, C = x.shape + q, k, v = self.w_q(x), self.w_k(x), self.w_v(x) + q, k = self.q_norm(q), self.k_norm(k) + q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2) + k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2) + v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2) + q, k = _rope(q), _rope(k) + + if self.window == -1: + y = F.scaled_dot_product_attention(q, k, v, is_causal=True) + else: + mask = _sliding_causal_mask(T, self.window, x.device) + y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask) + return self.w_out(y.transpose(1, 2).contiguous().view(B, T, C)) + + +class _FeedForward(nn.Module): + """SwiGLU, hidden_size=3072, bias=False (olmo_core FeedForwardConfig).""" + + def __init__(self, n_embd: int, hidden: int = HIDDEN_SIZE): + super().__init__() + self.w1 = nn.Linear(n_embd, hidden, bias=False) # gate + self.w3 = nn.Linear(n_embd, hidden, bias=False) # up + self.w2 = nn.Linear(hidden, n_embd, bias=False) # down + + def forward(self, x): + return self.w2(F.silu(self.w1(x)) * self.w3(x)) + + +class _ReorderedNormBlock(nn.Module): + """TransformerBlockType.reordered_norm -- norm AFTER the residual branch.""" + + def __init__(self, n_embd: int, n_head: int, window: int): + super().__init__() + self.attention = _Attention(n_embd, n_head, window) + self.attention_norm = _RMSNorm(n_embd) + self.feed_forward = _FeedForward(n_embd) + self.feed_forward_norm = _RMSNorm(n_embd) + + def forward(self, x): + h = x + self.attention_norm(self.attention(x)) + return h + self.feed_forward_norm(self.feed_forward(h)) + + +class NanoSLM(nn.Module): + def __init__(self, config): + super().__init__() + self.block_size = config.block_size + self.wte = nn.Embedding(config.vocab_size, N_EMBD) + self.blocks = nn.ModuleList( + [_ReorderedNormBlock(N_EMBD, N_HEAD, w) for w in LAYER_WINDOWS] + ) + # LMHeadConfig(layer_norm=rms, bias=False). Embeddings are TIED: the + # lm_head reuses wte's weight (a comparability choice so both arms differ + # only in the mixer -- see the module docstring; NOT forced by param_cap). + # The tie is set AFTER init so the shared tensor carries wte's init. + self.norm_f = _RMSNorm(N_EMBD) + self.lm_head = nn.Linear(N_EMBD, config.vocab_size, bias=False) + self.apply(self._init_weights) + for name, p in self.named_parameters(): + if name.endswith("w_out.weight") or name.endswith("w2.weight"): + nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * N_LAYER)) + self.lm_head.weight = self.wte.weight # weight tying + + @staticmethod + def _init_weights(module): + if isinstance(module, nn.Linear): + nn.init.normal_(module.weight, mean=0.0, std=0.02) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=0.02) + + def forward(self, idx): + x = self.wte(idx) + for block in self.blocks: + x = block(x) + x = self.norm_f(x) + return self.lm_head(x) # logits [B, T, vocab_size] + + +def build_model(config) -> NanoSLM: + return NanoSLM(config) \ No newline at end of file diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harness/data.py b/2.0/problems/nanoslm_hybrid_arch_design/harness/data.py new file mode 100644 index 000000000..761ccd1ce --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/data.py @@ -0,0 +1,197 @@ +"""Token-level data loading (torch path). + +Train/val data are flat ``uint32`` TOKEN-ID streams (``.bin``) baked into the +judge image by ``docker/prep_assets.py`` using the dolma2 BPE tokenizer +(``allenai/dolma2-tokenizer``, vocab 100278). The validation stream is HELD OUT +and never mounted in ``/app``. When the files are absent (CPU smoke / local dev) +a deterministic synthetic token stream is generated so the harness wiring runs +end-to-end without baked assets. + +WHY uint32 AND NOT uint16 +------------------------- +The dolma2 vocabulary is 100278 ids. ``uint16`` tops out at 65535, so the upper +~35k of the vocabulary would wrap silently into low ids -- a corrupted corpus +that still loads, still trains, and only shows up as an unexplained bpb floor. +The stream is therefore uint32 and ``_load_bin`` asserts every id is in range. + +BYTE ACCOUNTING +--------------- +``val_bytes_for(n)`` reports how many BYTES OF TEXT ``n`` scored target tokens +cover. ``eval_ppl`` divides by that, not by the token count -- see +``settings.resolve_val_bytes_per_token`` for why that distinction is the whole +point of the metric. + +Data order is fixed by the harness (common random numbers across the baseline +and the submission), so the two arms see identical batches. +""" + +from __future__ import annotations + +import numpy as np + +from .settings import ( + SYNTHETIC_BYTES_PER_TOKEN, + TaskConfig, + resolve_val_bytes_per_token, +) + +TOKEN_DTYPE = np.uint32 + + +class DataError(RuntimeError): + """Corpus is present but unusable (bad dtype, out-of-range ids, no byte count).""" + + +def _synthetic_tokens(n: int, seed: int, vocab_size: int) -> np.ndarray: + """Deterministic pseudo-text token stream for smoke runs (no baked data). + + Emits plausible BPE ids: a repeating "vocabulary" of a few thousand common + ids (real text is heavily Zipfian, so a smoke model can actually reduce + perplexity below uniform) plus occasional rare ids from the full range, so + the embedding table is exercised beyond its first rows. + """ + rng = np.random.default_rng(seed) + common = rng.integers(0, min(4096, vocab_size), size=max(1024, n // 8)) + stream = np.tile(common, (n // common.size) + 1)[:n].astype(np.int64) + rare_mask = rng.random(n) < 0.02 + stream[rare_mask] = rng.integers(0, vocab_size, size=int(rare_mask.sum())) + return stream.astype(TOKEN_DTYPE) + + +def _load_bin(path: str, *, fallback_n: int, seed: int, vocab_size: int) -> tuple[np.ndarray, bool]: + """Return ``(tokens, is_real)``; synthesize when the file is absent.""" + try: + arr = np.fromfile(path, dtype=TOKEN_DTYPE) + except OSError: + return _synthetic_tokens(fallback_n, seed, vocab_size), False + if arr.size == 0: + return _synthetic_tokens(fallback_n, seed, vocab_size), False + # A uint8 stream left over from the byte-level era reads as uint32 without + # error -- it just produces garbage ids. Catch it here rather than as a + # mystery bpb. (An out-of-range id would also index past the embedding + # table and raise a device-side assert deep inside the model.) + hi = int(arr.max()) + if hi >= vocab_size: + raise DataError( + f"{path}: token id {hi} >= vocab_size {vocab_size} " + f"(stale corpus, or written with the wrong dtype)" + ) + return arr, True + + +class TokenData: + """Holds train/val token arrays and yields fixed-order training batches. + + TWO CONTEXT LENGTHS -- READ BEFORE EDITING + --------------------------------------------------------- + ``val_windows`` is cut at ``cfg.eval_block_size`` (8192, judge-owned) and + NOTHING here may make it depend on the training context: the scored + bits-per-byte is only comparable across submissions because every arm is + evaluated on the identical windows of the identical held-out stream. The + TRAINING path (``batch``) uses the per-arm training context instead, which a + submission may set below 8192 to buy optimizer steps. + """ + + def __init__(self, cfg: TaskConfig): + self.cfg = cfg + # Longest window this data object will ever have to serve: the eval + # window is fixed at 8192 and the training context is capped there too + # (runner.MAX_TRAIN_BLOCK), so sizing off the max keeps both paths safe + # no matter what a submission asks for. + self.max_block = max(cfg.block_size, cfg.eval_block_size) + self.train, train_real = _load_bin( + cfg.train_tokens_path, fallback_n=1 << 20, + seed=cfg.seed, vocab_size=cfg.vocab_size, + ) + self.val, self.val_is_real = _load_bin( + cfg.val_tokens_path, fallback_n=cfg.val_tokens + cfg.eval_block_size + 1, + seed=cfg.seed + 999, vocab_size=cfg.vocab_size, + ) + if self.train.size < self.max_block + 1: + self.train = _synthetic_tokens(1 << 20, cfg.seed, cfg.vocab_size) + train_real = False + self.train_is_real = train_real + + # Bytes of text per scored target token. ASSERTED, never defaulted to 1: + # a ratio of 1 would silently turn val_bpb back into per-token CE/ln2. + ratio = resolve_val_bytes_per_token(cfg) + if ratio is None: + if self.val_is_real: + raise DataError( + "held-out byte count unavailable: no val_bytes in " + f"{cfg.manifest_path!r}, no FRONTIER_NANOSLM_VAL_BYTES, and " + "TaskConfig.val_bytes is 0. bits-per-byte cannot be " + "computed from a token count -- re-run docker/prep_assets.py." + ) + # Synthetic stream: there is no underlying text, so use the measured + # dolma2 compression ratio to keep smoke bpb in a plausible range. + ratio = SYNTHETIC_BYTES_PER_TOKEN + self.val_bytes_estimated = True + else: + self.val_bytes_estimated = not self.val_is_real + self.bytes_per_token = float(ratio) + + def val_bytes_for(self, n_target_tokens: int) -> float: + """Bytes of held-out text covered by ``n_target_tokens`` target tokens.""" + return self.bytes_per_token * float(n_target_tokens) + + def batch(self, step: int, device: str, block_size: int | None = None): + """Deterministic (step-keyed) training batch of token windows. + + ``block_size`` is the TRAINING context of the arm being trained, which + may differ per submission; it defaults to the config's. The window START + INDICES are drawn against ``self.max_block`` rather than against the + requested width, so two arms with different training contexts still + begin their windows at the SAME offsets -- a short-context arm reads a + prefix of what a long-context arm reads, instead of an unrelated slice. + That keeps common random numbers as close to intact as differing window + widths permit. + """ + import torch + + cfg = self.cfg + bs = int(block_size or cfg.block_size) + # Common random numbers: batch content depends only on (seed, step), so + # baseline and submission arms train on identical data. + g = np.random.default_rng(cfg.seed * 1_000_003 + step) + hi = max(1, self.train.size - self.max_block - 1) + ix = g.integers(0, hi, size=cfg.batch_size) + x = np.stack([self.train[i : i + bs] for i in ix]) + y = np.stack([self.train[i + 1 : i + 1 + bs] for i in ix]) + xt = torch.from_numpy(x.astype(np.int64)) + yt = torch.from_numpy(y.astype(np.int64)) + if device.startswith("cuda"): + xt = xt.pin_memory().to(device, non_blocking=True) + yt = yt.pin_memory().to(device, non_blocking=True) + else: + xt, yt = xt.to(device), yt.to(device) + return xt, yt + + def val_windows(self, device: str): + """Yield non-overlapping (x, y) windows over the held-out token stream. + + ALWAYS ``cfg.eval_block_size`` wide -- never the training context. A + submission that trained at 2048 is still scored on 8192-wide windows and + must extrapolate to those positions; that is the deliberate trade + described above, and the reason the bpb denominator (bytes per + scored target token x number of target tokens) is unaffected by the + training context. + """ + import torch + + cfg = self.cfg + eb = cfg.eval_block_size + n = min(self.val.size - 1, cfg.val_tokens) + n_windows = n // eb + for w in range(n_windows): + s = w * eb + x = self.val[s : s + eb][None, :] + y = self.val[s + 1 : s + 1 + eb][None, :] + xt = torch.from_numpy(x.astype(np.int64)).to(device) + yt = torch.from_numpy(y.astype(np.int64)).to(device) + yield xt, yt + + +# The class was ``ByteData`` while the tokenizer was byte-level. Kept as an +# alias so any out-of-tree caller keeps working; new code uses TokenData. +ByteData = TokenData diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harness/eval_ppl.py b/2.0/problems/nanoslm_hybrid_arch_design/harness/eval_ppl.py new file mode 100644 index 000000000..3a00dbbea --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/eval_ppl.py @@ -0,0 +1,120 @@ +"""Held-out bits-per-byte evaluation (torch path). + +The harness computes cross-entropy from the model's ``logits`` on the HELD-OUT +token stream and returns + + val_bpb = total_nll_nats / (total_val_BYTES * ln 2) + +Any loss returned by the model is ignored. + +WHY bpb IS THE PRIMARY METRIC (not perplexity): + * It is tokenizer-independent, which is what makes it comparable across + architectures and ungameable -- see DESIGN.md §2. + * It needs no ``exp``. ``val_ppl = exp(mean_ce)`` overflows to ``inf`` for a + sufficiently bad submission (mean_ce > ~709), turning a merely-poor model + into a non-finite number the guards then have to special-case. bpb is linear + in the loss and cannot overflow. + +WHY THE DENOMINATOR IS BYTES AND NOT TOKENS -- READ BEFORE EDITING +------------------------------------------------------------------ +While the tokenizer was byte-level, 1 token == 1 byte and dividing the summed +NLL by the TOKEN count was accidentally identical to dividing by the byte count. +Under dolma2 BPE the two differ by the compression ratio (~4.4x), and per-token +CE/ln2 is a tokenizer-dependent number: a tokenizer that packs more text per +token raises it for the same model quality. Normalizing that way would destroy +exactly the property bpb was adopted for. So the denominator comes from +``data.val_bytes_for(...)``, which is backed by a byte count measured at +corpus-prep time and asserted present (``data.DataError`` otherwise) rather than +defaulted. + +``val_ppl`` is still derived and reported, for human readability only -- it is +no longer the scored quantity, and it remains a PER-TOKEN perplexity (so +``val_ppl != 2**val_bpb`` any more; that identity held only at byte level). It +is computed defensively so a diverged model reports ``inf`` rather than raising. +Determinism knobs (no TF32, deterministic algorithms) are set on the GPU path so +a cached/separate baseline is a valid common-random-numbers partner. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from .data import DataError, TokenData +from .settings import TaskConfig +from .train import _logits_only + + +@dataclass +class EvalOutput: + val_bpb: float # SCORED quantity: total_nll_nats / (bytes * ln 2) + val_ppl: float # derived, per-TOKEN, readability only + val_ce_nats: float # mean per-token CE + n_tokens: int + mean_abs_logit_std: float # degeneracy signal (near-constant logits -> ~0) + n_bytes: float = 0.0 # denominator actually used for val_bpb + + +def set_determinism(cfg: TaskConfig) -> None: + import torch + + try: + torch.use_deterministic_algorithms(True, warn_only=True) + except Exception: + pass + if torch.cuda.is_available(): + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + torch.manual_seed(cfg.seed) + + +def evaluate_perplexity(model, data: TokenData, cfg: TaskConfig, device: str) -> EvalOutput: + import math + + import torch + import torch.nn.functional as F + + model.eval() + total_ce = 0.0 + total_tok = 0 + logit_std_accum = 0.0 + n_batches = 0 + + with torch.no_grad(): + for x, y in data.val_windows(device): + logits = _logits_only(model(x)).float() + B, T, V = logits.shape + ce = F.cross_entropy( + logits.reshape(B * T, V), y.reshape(B * T), reduction="sum" + ) + total_ce += float(ce.item()) + total_tok += B * T + # degeneracy probe: spread of the logits across vocab + logit_std_accum += float(logits.std(dim=-1).mean().item()) + n_batches += 1 + + if total_tok == 0: + return EvalOutput(float("inf"), float("inf"), float("inf"), 0, 0.0, 0.0) + + # BYTES, not tokens -- see the module docstring. `total_ce` is the SUM of + # NLL in nats over every scored target token, so this is exactly + # total_nll / (bytes * ln 2). + total_bytes = data.val_bytes_for(total_tok) + if not (total_bytes > 0.0): + raise DataError( + "held-out byte count is zero/unavailable; refusing to normalize " + "bits-per-byte by a token count (see harness/eval_ppl.py)" + ) + val_bpb = total_ce / (total_bytes * math.log(2.0)) + + mean_ce = total_ce / total_tok + # Derived for readability only, and PER TOKEN. Guarded: exp overflows above + # mean_ce ~709, and a diverged submission can get there. + try: + val_ppl = math.exp(mean_ce) + except OverflowError: + val_ppl = float("inf") + mean_logit_std = logit_std_accum / max(1, n_batches) + return EvalOutput(val_bpb, val_ppl, mean_ce, total_tok, mean_logit_std, + float(total_bytes)) diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harness/modal_app.py b/2.0/problems/nanoslm_hybrid_arch_design/harness/modal_app.py new file mode 100644 index 000000000..8d19ffe80 --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/modal_app.py @@ -0,0 +1,277 @@ +"""Modal app: run the nanoslm_hybrid_arch_design judge on one GPU. + +Mirrors the GPU-on-Modal pattern of vllm_llm_serving_optimization and nanowm: +the judge container stays CPU-only and calls into a Modal GPU function. + +Two things make this app much lighter than lm_arch_discovery's: + + * NO EXTERNAL REPO AND NO CUDA SOURCE BUILD. The harness is self-contained + PyTorch; the only third-party imports across harness/*.py are torch and + numpy. torch's own wheels ship the CUDA runtime, so a slim base works and + the image builds in ~1 min rather than needing nvcc. + * NO DATASET DOWNLOAD AND NO GATED TOKENIZER. The corpus is PRE-TOKENIZED by + docker/prep_assets.py with dolma2 (allenai/dolma2-tokenizer, vocab 100278 -- + ungated, no HF token), so this image never tokenizes anything; and + harness/data.py falls back to a deterministic synthetic TOKEN stream when + the .bin paths are absent. So the wiring is provable before any corpus is + staged; a real shard is mounted later for calibration. + + CAVEAT while no shard is mounted here: with the synthetic stream there is no + underlying text, so the bits-per-byte denominator comes from + settings.SYNTHETIC_BYTES_PER_TOKEN (the measured dolma2 compression ratio) + rather than from a manifest. Smoke bpb is therefore in the right RANGE but + is not a measurement of anything. Staging the real train.bin/val.bin/ + manifest.json on this image is required before calibration. + +Rather than reimplement the training loop here, this calls evaluator.evaluate() +-- the SAME entrypoint Harbor calls -- so there is no second code path to drift. + +Parametrized via env vars: + LMARCH_MODAL_GPU Modal GPU string (default "H100") + LMARCH_MODAL_APP Modal app name + +Deploy: modal deploy harness/modal_app.py +""" +from __future__ import annotations + +import os +import pathlib + +GPU = os.environ.get("LMARCH_MODAL_GPU", "H100") +APP_NAME = os.environ.get("LMARCH_MODAL_APP", "nanoslm-hybrid-arch-design") +REMOTE_TASK = "/opt/nanoslm_arch/task" + +_PROBLEM_DIR = pathlib.Path(__file__).resolve().parent.parent + +def _ver(mod): + """Version string as a PLAIN str. + + str() is required, not cosmetic: torch.__version__ is a + torch.torch_version.TorchVersion -- a str SUBCLASS defined inside torch -- + so returning it raw makes the whole result unpicklable on a client without + torch: DeserializationError: 'torch' module is not available. + The judge is deliberately CPU-only and has no torch. + """ + try: + return str(__import__(mod).__version__) + except Exception as exc: + return f"<{type(exc).__name__}>" + + +try: + import modal + + image = ( + modal.Image.debian_slim(python_version="3.11") + # torch wheels bundle the CUDA runtime, so no nvidia/cuda base and no + # nvcc are needed -- the whole reason this image builds in ~1 min. + .pip_install("torch==2.6.0", "numpy<2") + + # flash-linear-attention: REQUIRED, not an optimization. PyTorch ships a + # fused kernel for softmax attention (SDPA/flash) but none for linear or + # recurrent mixers, so a hand-written GDN is unfused eager code. Measured + # at ctx 8192: chunkwise did 6% of attention's FLOPs and still ran 3.5x + # slower (11 optimizer steps vs 38). Under a fixed WALL-CLOCK budget that + # makes the comparison "fused kernel vs unfused kernel" rather than + # "architecture vs architecture", and a hybrid can never win. + # + # torch 2.6.0 (not 2.5.1) because fla warns it needs Triton >= 3.2.0 and + # 2.5.1 ships 3.1.0 -- on 3.1.0 every GatedDeltaNet call hung 30+ min. + # + # fla 0.5.1 (not 0.4.1): on 0.4.1 the FORWARD compiled fine but the + # BACKWARD did not -- + # fla/ops/gated_delta_rule/wy_fast.py:303 prepare_wy_repr_bwd + # -> triton make_ttgir -> RuntimeError: PassManager::run failed + # i.e. 0.4.1's backward kernels predate Triton 3.2's IR. This cost three + # wrong guesses (head count twice) because an early probe timed FORWARD + # ONLY and looked healthy; always exercise fwd+bwd when validating a + # fused kernel. + .pip_install("flash-linear-attention==0.5.1", "transformers==4.46.3") + # Ship the problem directory itself: harness/, evaluator.py, + # reference.py. The GPU function calls the real evaluator, so what runs + # remotely is byte-identical to what Harbor runs. + .add_local_dir( + _PROBLEM_DIR, REMOTE_TASK, copy=True, + ignore=["__pycache__", "*.pyc", ".git", "docker"], + ) + .env({ + "PYTHONPATH": REMOTE_TASK, + # MUST be set before the first CUDA call, so it belongs in the image + # env -- setting it inside the function is too late. settings.py + # declares this value but nothing was exporting it, so the Modal path + # ran with torch.use_deterministic_algorithms(True) while cuBLAS was + # still free to be non-deterministic (visible as a UserWarning in the + # container log). CRN pairing depends on this. + "CUBLAS_WORKSPACE_CONFIG": ":4096:8", + # fla's Triton kernels JIT-compile and autotune on first use. Give + # them a writable cache so a warm container does not pay it twice. + "TRITON_CACHE_DIR": "/tmp/triton-cache", + }) + ) + + app = modal.App(APP_NAME) + + @app.function(image=image, gpu=GPU, timeout=6 * 60 * 60) + def evaluate_remote(solution_source: str, smoke: bool = True, + role: str = "final", overrides: dict | None = None) -> dict: + """Run the judge on `solution_source` and return its full result. + + `solution_source` travels as TEXT, not a path: the submission is written + into a scratch file remotely, exactly as the judge receives it. + """ + import json + import sys + import tempfile + import time + + sys.path.insert(0, REMOTE_TASK) + + os.environ["FRONTIER_NANOSLM_SMOKE"] = "1" if smoke else "0" + os.environ["FRONTIER_NANOSLM_ROLE"] = role + for k, v in (overrides or {}).items(): + os.environ[k] = str(v) + + import torch + + gpu_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else None + + # Record the actual stack in every result. Version drift between what a + # Dockerfile pins and what a warm container runs has already produced + # two misleading results in this task's history. + def _ver(mod): # noqa: F811 (module-level _ver is the shared one) + # str() is REQUIRED, not cosmetic. torch.__version__ is a + # torch.torch_version.TorchVersion (a str SUBCLASS defined inside + # torch), so returning it raw makes the whole result unpicklable on + # any client without torch installed: + # DeserializationError: 'torch' module is not available + # The judge/caller is deliberately CPU-only and has no torch, so + # every value crossing this boundary must be a plain builtin. + try: + return str(__import__(mod).__version__) + except Exception as exc: + return f"<{type(exc).__name__}>" + + stack = {"torch": _ver("torch"), "triton": _ver("triton"), "fla": _ver("fla")} + + import evaluator # noqa: E402 (must follow the env setup above) + + with tempfile.TemporaryDirectory() as td: + sub = pathlib.Path(td) / "model.py" + sub.write_text(solution_source) + + # DEBUG PATH (operator-only, never reachable from a submission): + # rerun the arm OUTSIDE the evaluator's guard so the real traceback + # survives. evaluator.py deliberately collapses every training + # exception to "training runtime error: RuntimeError" for black-box + # safety -- correct for agents, useless for debugging the harness. + if overrides and overrides.get("DEBUG_TRACEBACK"): + import traceback as _tb + + from harness import runner, settings as _s + from harness.data import TokenData + + cfg = _s.active_config() + try: + mod = evaluator._load_submission_module(str(sub)) + runner.run_arm(runner.load_factory(mod), TokenData(cfg), cfg, "cuda") + # `stack` on the SUCCESS branch too. It was only on the + # failure branch, so a passing run could not be attributed + # to a version -- exactly the warm-container trap that + # already produced one stale, misread result here. + return {"debug": "submission arm completed without error", + "stack": stack} + except Exception: + return {"debug_traceback": _tb.format_exc()[-4000:], "stack": stack} + + t0 = time.perf_counter() + score, unbounded, message, metrics = evaluator.evaluate(str(sub)) + wall = time.perf_counter() - t0 + + return { + "score": score, + "score_unbounded": unbounded, + "message": message, + "metrics": json.loads(json.dumps(metrics, default=str)), + "wall_seconds": round(wall, 2), + "gpu": gpu_name, + "cuda_available": gpu_name is not None, + "stack": stack, + } + + @app.function(image=image, gpu=GPU, timeout=6 * 60 * 60) + def run_pair_remote(solution_source: str, smoke: bool = False, + role: str = "final") -> dict: + """Run BOTH arms on one GPU and return their metrics as plain dicts. + + This is the entrypoint the JUDGE uses. It deliberately returns ARM + METRICS, not a score: scoring stays judge-side so the scoring code and + the hidden constants (bpb_score_scale) never ship to a GPU worker, and so the + judge cannot be handed a score it did not compute. + + Contrast with evaluate_remote(), which runs the whole evaluator + remotely -- useful for manual testing, but circular now that + evaluator.evaluate() dispatches here. + + Everything crossing this boundary must be a plain builtin: the judge + container is CPU-only and has no torch, so e.g. torch.__version__ + (a str SUBCLASS defined in torch) would make the result unpicklable. + """ + import sys + import tempfile + + sys.path.insert(0, REMOTE_TASK) + os.environ["FRONTIER_NANOSLM_SMOKE"] = "1" if smoke else "0" + os.environ["FRONTIER_NANOSLM_ROLE"] = role + + import torch # noqa: F401 (ensures CUDA init before harness import) + + import evaluator # noqa: F401 (sets CUBLAS_WORKSPACE_CONFIG early) + from harness import runner, settings as _s + + cfg = _s.active_config() + with tempfile.TemporaryDirectory() as td: + sub = pathlib.Path(td) / "model.py" + sub.write_text(solution_source) + mod = evaluator._load_submission_module(str(sub)) + base, cand = runner.run_pair(mod, cfg, "cuda") + + def _arm(a): + # Every value here must be a PLAIN builtin -- the judge is CPU-only + # and has no torch, so anything torch-defined is unpicklable there. + return { + "val_bpb": float(a.val_bpb), + "val_ppl": float(a.val_ppl), + "steps": int(a.steps), + "wall_seconds": float(a.wall_seconds), + "n_params": int(a.n_params), + # Context this arm TRAINED at. The baseline is always 8192; + # a submission may declare a shorter one via BLOCK_SIZE. Both + # arms are always EVALUATED at cfg.eval_block_size (8192). + "train_block_size": int(a.train_block_size), + # Warmup is where Triton JIT is supposed to land. If + # warmup_error is non-empty, warmup silently no-opped and its + # cost was deferred INTO the timed loop -- the exact failure + # that produced sub_steps=1 at 103.2s. + "warmup_seconds": float(a.warmup_seconds), + "warmup_error": str(a.warmup_error), + } + + return { + "baseline": _arm(base), + "submission": _arm(cand), + # FIXED scoring window, reported so a result is self-describing: + # both arms' val_bpb come from windows of exactly this width. + "eval_block_size": int(cfg.eval_block_size), + # Bump on every behavioral change to this file. A deploy does NOT + # evict warm containers, so the first call after `modal deploy` can + # still execute the OLD code -- this has produced misleading results + # twice in this task's history. Check this value before trusting a + # result. + "code_version": "agent-train-ctx-v5", + "stack": {"torch": str(torch.__version__), + "triton": _ver("triton"), "fla": _ver("fla")}, + } + + +except ImportError: # modal absent (e.g. CPU-only unit-test host) + app = None diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harness/model_config.py b/2.0/problems/nanoslm_hybrid_arch_design/harness/model_config.py new file mode 100644 index 000000000..0cbd946ad --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/model_config.py @@ -0,0 +1,37 @@ +"""ModelConfig: the read-only contract object passed to the agent's factory. + +Torch-free. The agent's ``build_model(config)`` / ``NanoSLM(config)`` receives one +of these. ``vocab_size`` is fixed by the judge; the agent chooses internal +width/depth/etc. freely inside the model. The ``*_hint`` fields are +informational (so a model can size itself to the budget) and carry no +guarantees. + +TWO CONTEXT LENGTHS +------------------- +``block_size`` the context this model will be TRAINED at. It is the judge's + default (8192) unless model.py declares a module-level + ``BLOCK_SIZE``, in which case it is that value. +``eval_block_size`` the context it will be SCORED at. ALWAYS 8192, whatever the + training context is. + +When those differ the model is asked for logits at positions it never saw during +training, so anything sized or cached off ``block_size`` (RoPE tables, +positional embeddings, mask buffers) must still work at ``eval_block_size`` -- +build such buffers against ``eval_block_size``, or lazily against the actual +``T``. +""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class ModelConfig: + vocab_size: int + block_size: int # TRAINING context (per-submission) + eval_block_size: int = 8192 # SCORING context (fixed by the judge) + # Read-only budget hints (informational; not part of the scored contract). + train_seconds_hint: float = 0.0 + param_cap_hint: int = 0 + device_hint: str = "cuda" diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harness/policy.py b/2.0/problems/nanoslm_hybrid_arch_design/harness/policy.py new file mode 100644 index 000000000..454c247ae --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/policy.py @@ -0,0 +1,124 @@ +"""Submission policy: static allow/deny gate for the submitted model.py. + +Torch-free and unit-tested on CPU. Treat submission content as adversarial +(2.0 black-box safety). This is the *static* layer; runtime guards that the +scan cannot see (trained-from-scratch check, judge-owned loss, resource caps) +live in the runner/evaluator. See DESIGN.md §6–§7. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +MAX_BYTES = 256 * 1024 # 256 KB + +# Substrings that must not appear in the submitted source. Deliberately broad; +# a legitimate architecture never needs any of these. +POLICY_DENY_TOKENS: tuple[str, ...] = ( + # --- escape / environment / network leakage --- + "os.environ", "os.getenv", "getenv", "putenv", + "subprocess", "socket", "requests", "urllib", "httpx", "http.client", + "FRONTIER_", "JUDGE_", "HARBOR_", "MODAL_", "HF_TOKEN", + "/judge", "/opt/", "/tests/", + # --- pretrained-weight loading (must train from scratch) --- + "from_pretrained", "torch.load", "load_state_dict", "safetensors", + "hf_hub", "huggingface", "timm.create_model", "AutoModel", + # --- filesystem reads (no data/weight/val access) --- + "open(", "Path(", "np.load", "numpy.load", "np.fromfile", "mmap", + "pickle.load", "joblib.load", + # NOTE: metric/data names are NOT here -- see POLICY_DENY_TOKENS_CODE. + # --- timer / control-flow short-circuits & concurrency tricks --- + "time.time", "time.perf_counter", "perf_counter", "time.sleep", + "while True", "threading", "multiprocessing", "os.fork", "ctypes", + "exec(", "__import__", + # NOTE: "eval(" and "compile(" are intentionally NOT banned so that the + # legitimate `model.eval()` and `torch.compile()` idioms are allowed + # (mirrors nanowm's allowlist). Sandboxing + judge-owned loss cover the rest. +) + + +# Scanned over CODE ONLY (comments and string literals stripped). +# +# These names exist to stop a submission from READING the metric or the held-out +# data, which takes executable code. Scanning them over prose rejected any +# submission that merely *documented its intent* -- "# should improve val_bpb" +# -- with a confusing "forbidden token" error. Both the reference and the locked +# baseline tripped it once their docstrings explained the task, which is how it +# was found. A submission should be able to say what it optimizes. +# +# Everything in POLICY_DENY_TOKENS above is still scanned over the FULL source: +# a judge path or an env var appearing in a comment is a leak signal in itself, +# whereas the word "perplexity" in a docstring reads nothing. +POLICY_DENY_TOKENS_CODE: tuple[str, ...] = ( + "val_ppl", "perplexity", "val_bpb", "bits_per_byte", + "holdout", "hold_out", "val_data", "val.bin", "validation_bytes", +) + + +def _code_only(source: str) -> str: + """Return `source` with comments and string literals removed. + + Fails CLOSED: if the source will not tokenize, return it unchanged so the + scan still sees everything rather than waving a malformed file through. + """ + import io + import tokenize + + try: + out, last = [], (1, 0) + for tok in tokenize.generate_tokens(io.StringIO(source).readline): + if tok.type in (tokenize.COMMENT, tokenize.STRING): + continue + (srow, scol), (erow, ecol) = tok.start, tok.end + if srow != last[0]: + out.append("\n") + last = (srow, 0) + out.append(" " * max(0, scol - last[1])) + out.append(tok.string) + last = (erow, ecol) + return "".join(out) + except Exception: + return source + + +@dataclass +class PolicyResult: + ok: bool + reason: str = "" + + +def check_source(source: str, *, size_bytes: int | None = None) -> PolicyResult: + """Validate submitted model.py source text against the static policy.""" + if size_bytes is None: + size_bytes = len(source.encode("utf-8", errors="replace")) + if size_bytes > MAX_BYTES: + return PolicyResult(False, f"submission exceeds {MAX_BYTES} bytes") + if not source.strip(): + return PolicyResult(False, "submission is empty") + + for tok in POLICY_DENY_TOKENS: + if tok in source: + return PolicyResult(False, f"forbidden token in submission: {tok!r}") + + # Metric/data names: code only, so docstrings and comments may name the + # metric the submission is trying to improve. + code = _code_only(source) + for tok in POLICY_DENY_TOKENS_CODE: + if tok in code: + return PolicyResult(False, f"forbidden token in submission code: {tok!r}") + + # Must expose a model factory the harness can call. + if ("def build_model" not in source) and ("class NanoSLM" not in source): + return PolicyResult( + False, "submission must define build_model(config) or class NanoSLM" + ) + return PolicyResult(True, "policy ok") + + +def check_file(path: str) -> PolicyResult: + try: + with open(path, "r", encoding="utf-8", errors="replace") as fh: + source = fh.read() + except OSError as exc: # pragma: no cover - defensive + return PolicyResult(False, f"cannot read submission: {exc.__class__.__name__}") + return check_source(source) diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harness/runner.py b/2.0/problems/nanoslm_hybrid_arch_design/harness/runner.py new file mode 100644 index 000000000..b53a9d199 --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/runner.py @@ -0,0 +1,407 @@ +"""Orchestration (torch path): train + eval one arm, and CRN-pair two arms. + +Isolation model mirrors ``nanowm``: the judge/Modal boundary separates this code +from the agent workspace ``/app``; the submitted ``model.py`` is imported here +only after the static policy gate (evaluator) passes. All dynamic guards that the +static scan cannot see live in this module. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field, replace +from typing import Callable + +from . import baseline_model +from .data import DataError, TokenData +from .eval_ppl import EvalOutput, evaluate_perplexity, set_determinism +from .model_config import ModelConfig +from .settings import TaskConfig +from .train import TrainOutput, train_model + + +class GuardError(Exception): + """Raised when a submission trips a dynamic guard -> score 0, public reason.""" + + +@dataclass +class ArmMetrics: + val_bpb: float # SCORED quantity + val_ppl: float # derived, readability only + steps: int + wall_seconds: float + n_params: int + train_block_size: int = 0 # context this arm actually TRAINED at + warmup_seconds: float = 0.0 # time absorbed OUTSIDE the timed window + warmup_error: str = "" # "" when warmup completed cleanly + eval: EvalOutput = field(repr=False, default=None) + + +def load_factory(module) -> Callable: + fn = getattr(module, "build_model", None) + if callable(fn): + return fn + cls = getattr(module, "NanoSLM", None) + if cls is not None: + return lambda config: cls(config) + raise GuardError("model.py must define build_model(config) or class NanoSLM") + + +# --------------------------------------------------------------------------- # +# Agent-controlled TRAINING context. +# +# A submission declares it with a module-level ``BLOCK_SIZE`` int in model.py. +# That is deliberately the whole protocol: no new required entrypoint, and the +# existing build_model(config) / class NanoSLM contract is untouched, so a +# submission that says nothing keeps training at the judge's 8192 default. +# +# The EVALUATION context is never negotiable -- see settings.eval_block_size. +# --------------------------------------------------------------------------- # +MIN_TRAIN_BLOCK = 256 + + +def _is_power_of_two(n: int) -> bool: + return n > 0 and (n & (n - 1)) == 0 + + +def resolve_train_block_size(module, cfg: TaskConfig) -> int: + """Training context for `module`: its ``BLOCK_SIZE``, else the cfg default. + + Enforced as a GuardError (public message), exactly like the param cap: + + * power of two -- kernels, chunkwise mixers and mask caches all assume it, + and it keeps the space small enough to reason about; + * >= 256 -- below that the context degenerates and a "language model" + stops modelling anything the 8192-wide evaluation rewards; + * <= eval_block_size (8192) -- training LONGER than the scoring window + buys nothing measurable and only costs memory and steps. + """ + raw = getattr(module, "BLOCK_SIZE", None) + if raw is None: + return int(cfg.block_size) + if isinstance(raw, bool) or not isinstance(raw, int): + raise GuardError( + f"model.py BLOCK_SIZE must be an int, got {type(raw).__name__}" + ) + hi = int(cfg.eval_block_size) + if not (MIN_TRAIN_BLOCK <= raw <= hi) or not _is_power_of_two(raw): + raise GuardError( + f"model.py BLOCK_SIZE={raw} is out of range: the training context " + f"must be a power of two in [{MIN_TRAIN_BLOCK}, {hi}] " + f"(evaluation is always at {hi})" + ) + return int(raw) + + +def _model_config(cfg: TaskConfig, device: str) -> ModelConfig: + return ModelConfig( + vocab_size=cfg.vocab_size, + block_size=cfg.block_size, # per-arm TRAINING context + eval_block_size=cfg.eval_block_size, # fixed SCORING context + train_seconds_hint=cfg.train_seconds, + param_cap_hint=cfg.param_cap, + device_hint="cuda" if device.startswith("cuda") else device, + ) + + +def _count_params(model) -> int: + return sum(p.numel() for p in model.parameters() if p.requires_grad) + + +def _param_snapshot(model): + import torch + + with torch.no_grad(): + return torch.cat([p.detach().reshape(-1).float().cpu() for p in model.parameters()]) + + +def _warmup(model, data: TokenData, cfg: TaskConfig, device: str): + """Compile/autotune kernels BEFORE the wall-clock timer starts. + + This is required for a fair fixed-wall-clock comparison, not an + optimization. Triton JIT-compiles per shape on first use: measured on an + H100, fla's GatedDeltaNet takes 15.5s for its first forward at T=8192 and + 74.5s at T=64, then 3ms warm. torch's SDPA path pays essentially none. + + Left inside the timed window, an fla-based hybrid would spend a chunk of its + budget T compiling while a pure-attention arm spent none -- so the score + would partly measure COMPILATION rather than throughput, and the smaller T + is the worse the distortion. + + THREE THINGS THIS MUST MATCH ABOUT ``train_model``, or it warms kernels the + training loop never calls and the JIT cost lands inside the timer anyway: + + 1. DEVICE. The factories (``baseline_model.build_model``, the submission's) + construct on CPU; only ``train_model`` calls ``model.to(device)``. This + function used to run BEFORE that move, so ``model(x)`` with a CUDA batch + raised "Expected all tensors to be on the same device", which the old + bare ``except Exception: pass`` swallowed. Warmup was a silent no-op for + BOTH arms -- harmless for SDPA, fatal for fla, whose entire Triton + compile then landed in step 0 (measured: one 103.2s step against a 45s + budget). Hence ``model.to(device)`` here, and hence the error is now + REPORTED rather than swallowed. + 2. DTYPE. ``train_model`` runs under ``torch.autocast(bfloat16)``. Triton + autotunes per dtype, so an fp32 warmup warms the wrong kernels. + 3. THE OPTIMIZER STEP. ``clip_grad_norm_`` and AdamW's foreach/fused + kernels are first-call-compiled too, so warmup takes a real step and + then RESTORES the parameters, leaving training bit-identical to a run + that never warmed up. + + Two iterations, not one: Triton autotuning picks a config on the first call + and can still recompile on the second. The second is ~ms once warm. + + FOURTH THING, added with the agent-controlled training context (DESIGN.md + 15): the training and evaluation SHAPES can now differ. Warming only the + training shape leaves every eval-shape kernel cold for a submission that + trained at, say, 2048 -- and Triton autotunes per shape, so that is a fresh + compile of tens of seconds. It does NOT currently land inside the scored + window (``train_model`` stops its timer before ``evaluate_perplexity`` is + called), so this is not a scoring bug today; it is warmed here anyway + because (a) it is a real cost against the container/`max_train_seconds` + budget, (b) it makes the property robust to any future reordering rather + than accidental, and (c) it surfaces an eval-shape failure (a model whose + buffers were sized at the training context and cannot run at 8192) HERE, in + ``warmup_error``, instead of as a mystery exception after training. The + eval warm mirrors ``eval_ppl.evaluate_perplexity`` exactly: ``model.eval()``, + ``no_grad``, and NO autocast -- evaluation runs in fp32, so warming it under + bfloat16 would warm the wrong kernels again. + + Errors are captured and returned, never raised: a genuinely broken model + must fail in the real training loop where the guards can classify it. But it + is no longer INVISIBLE -- the message rides out in ArmMetrics.warmup_error. + """ + import time + + import torch + + t0 = time.monotonic() + err = "" + try: + model.to(device) + model.train() + autocast = ( + torch.autocast(device_type="cuda", dtype=torch.bfloat16) + if device.startswith("cuda") + else _nullctx() + ) + # Snapshot so the warmup optimizer step leaves no trace on training. + saved = [p.detach().clone() for p in model.parameters()] + opt = torch.optim.AdamW(model.parameters(), lr=0.0) + + for i in range(2): + # cfg.block_size here is the ARM's training context (run_arm passes + # a per-arm cfg), so this warms the shape train_model will use. + x, y = data.batch(i, device, cfg.block_size) # same API train.py uses + with autocast: + out = model(x) + logits = out[0] if isinstance(out, tuple) else out + loss = torch.nn.functional.cross_entropy( + logits.reshape(-1, logits.size(-1)).float(), y.reshape(-1) + ) + opt.zero_grad(set_to_none=True) + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip) + opt.step() + + with torch.no_grad(): + for p, s in zip(model.parameters(), saved): + p.copy_(s) + model.zero_grad(set_to_none=True) + del saved, opt + + # Eval-shape warm: [1, eval_block_size] in eval mode, no_grad, fp32 -- + # exactly how evaluate_perplexity will call the model. Uses the first + # real held-out window so the shape and dtype match to the element; it + # runs under no_grad and takes no optimizer step, so it cannot influence + # any parameter or leak the held-out stream into training. + model.eval() + with torch.no_grad(): + first = next(iter(data.val_windows(device)), None) + if first is not None: + for _ in range(2): + out = model(first[0]) + _ = out[0] if isinstance(out, tuple) else out + del out + del first + model.train() + + if device.startswith("cuda"): + torch.cuda.synchronize() + torch.cuda.empty_cache() + except Exception as exc: + err = f"{type(exc).__name__}: {exc}"[:300] + + return round(time.monotonic() - t0, 2), err + + +class _nullctx: + def __enter__(self): + return None + + def __exit__(self, *a): + return False + + +def run_arm(factory: Callable, data: TokenData, cfg: TaskConfig, device: str, + train_block_size: int | None = None) -> ArmMetrics: + """Instantiate -> guard params -> train -> eval -> guard training/degeneracy. + + ``train_block_size`` is this arm's TRAINING context (default: cfg's). It is + threaded through as a per-arm ``cfg`` so that every consumer -- the model + config handed to the factory, ``_warmup``, and ``train_model`` -- sees one + consistent value. ``cfg.eval_block_size`` is untouched by it, so evaluation + stays at 8192 for every arm. + """ + import torch + + if train_block_size is not None and int(train_block_size) != int(cfg.block_size): + cfg = replace(cfg, block_size=int(train_block_size)) + + set_determinism(cfg) + try: + model = factory(_model_config(cfg, device)) + except GuardError: + raise + except Exception as exc: + raise GuardError(f"model construction failed: {type(exc).__name__}") + + if not isinstance(model, torch.nn.Module): + raise GuardError("factory did not return a torch.nn.Module") + + n_params = _count_params(model) + if n_params <= 0: + raise GuardError("model has no trainable parameters") + if n_params > cfg.param_cap: + raise GuardError(f"model exceeds param cap ({n_params} > {cfg.param_cap})") + + # COMPILE BOTH ARMS, UNIFORMLY, IN THE HARNESS. + # + # Two reasons, and the second is a correctness issue: + # + # 1. Throughput. The chunkwise linear mixer is unfused eager PyTorch and is + # dominated by kernel-launch overhead, so it loses to SDPA's single fused + # flash-attention kernel despite doing ~6% of the FLOPs (measured: 11 + # optimizer steps vs 38 at ctx 8192). Compilation fuses the elementwise + # work and is what makes a linear mixer competitive at all here. + # + # 2. Fairness. `torch.compile()` is explicitly allowed in submissions + # (policy.py). If the harness did not compile, a submission could take + # the BASELINE ARCHITECTURE UNCHANGED, wrap it in torch.compile, gain + # wall-clock and score well with zero architectural insight -- the same + # kernel-mismatch confound that fla created, by another route. Compiling + # both arms here removes compilation as a lever and puts the comparison + # back on architecture. + # + # Failures fall back to eager: a submitted architecture may contain + # something inductor cannot trace, and that should cost throughput, not be + # a hard rejection. + # DEFAULT OFF. The fairness argument above is real and unresolved, but + # compilation measured WORSE than eager on the chunkwise mixer (1 step vs + # 11 at ctx 8192 -- inductor almost certainly graph-breaks on the chunk + # loop's data-dependent state), and its interaction with fla's Triton + # kernels is unvalidated. Re-enable only after measuring both arms with it. + if getattr(cfg, "compile_model", False): + try: + model = torch.compile(model) + except Exception: + pass + + # Warm kernels OUTSIDE the timed window -- see _warmup. This now also + # absorbs inductor's compilation, which is the whole reason it must precede + # the timer. + warmup_seconds, warmup_error = _warmup(model, data, cfg, device) + + before = _param_snapshot(model) + try: + tout: TrainOutput = train_model(model, data, cfg, device) + except RuntimeError as exc: + # e.g. CUDA OOM or a shape error from a malformed architecture. + msg = str(exc).lower() + if "out of memory" in msg: + raise GuardError("training ran out of GPU memory") + raise GuardError(f"training runtime error: {type(exc).__name__}") + + if tout.steps <= 0: + raise GuardError("model completed zero optimizer steps in the budget") + + after = _param_snapshot(model) + if before.numel() == after.numel(): + mean_abs_delta = float((after - before).abs().mean().item()) + if mean_abs_delta < cfg.min_param_delta: + raise GuardError("parameters did not change during training (not trained)") + + # A NEW FAILURE MODE, since the training and evaluation contexts can differ: + # a model whose buffers were sized off the TRAINING context + # trains happily and then blows up at eval, where T is always + # eval_block_size. Left unclassified that surfaced as a bare + # "submission_error: RuntimeError" with no hint about the cause, so it is + # turned into a GuardError naming both contexts. + # `Exception`, not `RuntimeError`: the commonest form of this bug is a + # positional embedding table indexed past its end, which torch raises as an + # IndexError, not a RuntimeError. Catching only RuntimeError let exactly the + # failure this guard exists for escape unclassified. DataError is re-raised + # untouched -- a missing held-out byte count is a JUDGE-side asset problem + # and must not be reported as a submission fault. + try: + eout = evaluate_perplexity(model, data, cfg, device) + except (GuardError, DataError): + raise + except Exception as exc: + msg = str(exc).lower() + if "out of memory" in msg: + raise GuardError( + f"evaluation ran out of GPU memory at context " + f"{cfg.eval_block_size}" + ) + raise GuardError( + f"evaluation failed at context {cfg.eval_block_size} " + f"(trained at {cfg.block_size}): {type(exc).__name__}. Buffers sized " + f"off config.block_size must still work at config.eval_block_size" + ) + # Guard on bpb, the scored quantity. bpb cannot overflow the way exp(ce) + # can, so this catches a diverged model as a large-but-finite bpb rather + # than relying on an inf that only appears past mean_ce ~709. + if not (eout.val_bpb > 0.0) or eout.val_bpb == float("inf"): + raise GuardError("evaluation produced a non-finite bits-per-byte") + if eout.mean_abs_logit_std < 1e-4: + raise GuardError("model produced near-constant logits (degenerate)") + + return ArmMetrics( + val_bpb=eout.val_bpb, + val_ppl=eout.val_ppl, + steps=tout.steps, + wall_seconds=tout.wall_seconds, + n_params=n_params, + train_block_size=int(cfg.block_size), + warmup_seconds=warmup_seconds, + warmup_error=warmup_error, + eval=eout, + ) + + +def baseline_factory() -> Callable: + return baseline_model.build_model + + +def baseline_block_size(cfg: TaskConfig) -> int: + """Training context of the LOCKED arm. + + Read through the same ``BLOCK_SIZE`` protocol a submission uses, from + ``baseline_model``, which declares 8192 EXPLICITLY rather than inheriting a + default. The locked arm's context must be unambiguous and must not move if + ``TaskConfig.block_size`` (now only the submission-facing default) is ever + retuned -- and it is what makes the fingerprint-keyed baseline cache valid + across submissions with different training contexts. + """ + return resolve_train_block_size(baseline_model, cfg) + + +def run_pair(submission_module, cfg: TaskConfig, device: str): + """Scored path: train baseline + submission back-to-back on one GPU (CRN).""" + # Resolve (and range-guard) the submission's training context BEFORE + # spending a baseline run on it: an out-of-range BLOCK_SIZE should cost the + # judge nothing. + sub_block = resolve_train_block_size(submission_module, cfg) + data = TokenData(cfg) + base = run_arm(baseline_factory(), data, cfg, device, baseline_block_size(cfg)) + sub = run_arm(load_factory(submission_module), data, cfg, device, sub_block) + return base, sub diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harness/scoring.py b/2.0/problems/nanoslm_hybrid_arch_design/harness/scoring.py new file mode 100644 index 000000000..ec4b25d1b --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/scoring.py @@ -0,0 +1,113 @@ +"""Scoring: ABSOLUTE held-out bits-per-byte gain over the locked baseline. + +Torch-free and unit-tested on CPU. Shared by the evaluator and the public test +so the agent sees the same math the judge uses. See DESIGN.md §8. +""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass +class ScoreResult: + score: float # bounded [0, 100] reported to Harbor + score_unbounded: float # keeps rewarding past 100 + rel_improvement: float # (base_bpb - sub_bpb) / base_bpb — reported only + abs_bpb_delta: float # base_bpb - sub_bpb — THE SCORED QUANTITY + + +def _clip(x: float, lo: float, hi: float) -> float: + return max(lo, min(hi, x)) + + +def score_from_bpb( + base_bpb: float, + sub_bpb: float, + bpb_score_scale: float, +) -> ScoreResult: + """Map a submission's held-out bits-per-byte to a [0, 100] score. + + Lower ``sub_bpb`` is better. The MEASUREMENT is the absolute gain:: + + gain = base_bpb - sub_bpb # bits per byte + score = clip(100 * gain / bpb_score_scale, 0, 100) + + The baseline architecture (``sub_bpb == base_bpb``) scores 0, and a + submission worse than the baseline scores 0. + + WHAT ``bpb_score_scale`` IS, AND WHAT IT IS NOT + ----------------------------------------------- + It is a DISPLAY CONVENTION and it is arbitrary by design. Its one necessary + job is to map bits-per-byte into the 0-100 range Harbor expects: Harbor + computes ``reward = score / 100``, so a raw ``gain`` of 0.05 bpb would + surface as reward 0.0005 -- indistinguishable from zero for every + submission. That is the whole reason the constant exists. + + It is NOT a calibrated definition of "what counts as a full win". No such + number has been measured, and presenting an arbitrary constant as a target + would (a) attach the score's meaning to a figure nobody determined and + (b) throw away information at the clip, where a submission at 2x the scale + scores the same 100 as one exactly at it. ``score_unbounded`` is therefore + the un-clipped value and an operator should read it whenever the bounded + score pins at 100. + + The one weak requirement that does remain is DISCRIMINATION: set far too + large and every submission pins at 0, far too small and every submission + pins at 100. That is a much weaker condition than calibration and can be set + from a single real-corpus run. + + WHY ABSOLUTE AND NOT RELATIVE + ----------------------------- + Absolute bpb is the unit the literature quotes (e.g. Olmo Hybrid), so a + result is directly comparable to published numbers; it is linear in + cross-entropy, so equal absolute gains are equal information gains, and + gains compose roughly additively. The cost is that it is NOT scale-free: the + same absolute gain is roughly twice as hard at ``base_bpb`` 1.5 as at 2.9, + so the operating point matters -- see DESIGN.md §8. ``rel_improvement`` is + still computed and reported so an operator can see the relative figure + without recomputing it. + """ + # Invalid / degenerate values -> zero (defensive; runner also guards). + if not (base_bpb > 0.0) or not (sub_bpb > 0.0): + return ScoreResult(0.0, 0.0, 0.0, 0.0) + if not (base_bpb < float("inf")) or not (sub_bpb < float("inf")): + return ScoreResult(0.0, 0.0, 0.0, 0.0) + if bpb_score_scale <= 0.0: + raise ValueError("bpb_score_scale must be positive") + + abs_delta = base_bpb - sub_bpb # the scored measurement + rel = abs_delta / base_bpb # reported for readability only + unbounded = 100.0 * abs_delta / bpb_score_scale + bounded = _clip(unbounded, 0.0, 100.0) + # No credit for tying or losing to the baseline. + if abs_delta <= 0.0: + bounded = 0.0 + return ScoreResult(bounded, unbounded, rel, abs_delta) + + +def format_message( + base_bpb: float, + sub_bpb: float, + result: ScoreResult, + *, + steps: int, + wall_seconds: float, + extra: str = "", +) -> str: + """Public feedback string — metrics only, no submission stdout/tracebacks. + + Reports BOTH the scored absolute gain and the relative figure: only the + SCORED quantity changed to absolute, and an operator should not have to + recompute the other one. + """ + msg = ( + f"base_val_bpb={base_bpb:.5f}; sub_val_bpb={sub_bpb:.5f}; " + f"abs_bpb_delta={result.abs_bpb_delta:+.5f}; " # SCORED + f"rel_improvement={result.rel_improvement:+.4%}; " + f"steps={steps}; train_wall_s={wall_seconds:.1f}; " + f"score={result.score:.4f}; score_unbounded={result.score_unbounded:.4f}" + ) + if extra: + msg += f"; note={extra}" + return msg diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harness/settings.py b/2.0/problems/nanoslm_hybrid_arch_design/harness/settings.py new file mode 100644 index 000000000..fe9af3deb --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/settings.py @@ -0,0 +1,281 @@ +"""Locked task settings + config fingerprint (torch-free). + +All numbers here are read by the judge only. The agent never sees this file at +scoring time. Values tagged ``CALIBRATE`` are placeholders pending the +single-H100 calibration described in DESIGN.md §3 and must be frozen before the +task ships. +""" + +from __future__ import annotations + +import hashlib +import json +import os +from dataclasses import asdict, dataclass + + +# --------------------------------------------------------------------------- # +# Locked training / evaluation configuration. +# --------------------------------------------------------------------------- # +@dataclass(frozen=True) +class TaskConfig: + # --- fixed model I/O contract (agent must honor these) --- + # dolma2 BPE (allenai/dolma2-tokenizer), the OLMo-3 tokenizer. UNGATED -- + # no HF token is needed to download it, unlike lm_arch_discovery's Llama-2 + # tokenizer. 100278 ids > 65535, so token streams are uint32 on disk; uint16 + # would silently wrap the upper ~35k of the vocabulary. + vocab_size: int = 100278 + # HYBRID-DISCOVERY SETTING. 8192, not 1024: hybrids exist to escape + # attention's O(L^2), and at ctx 1024 attention is only ~31-40% of layer + # FLOPs, so a cheaper sequence mixer has little to win and the fixed + # wall-clock budget barely rewards it. At 8192 attention is ~78-84% of + # layer FLOPs, so the attention/recurrence trade is the dominant design + # question -- which is the point of this task. + # + # TWO CONTEXT LENGTHS, AND THE DISTINCTION IS LOAD-BEARING. + # + # eval_block_size FIXED, judge-owned. Every submission is scored on + # non-overlapping windows of exactly this width, so + # val_bpb stays comparable across submissions. NEVER + # derive this from anything a submission controls. + # block_size DEFAULT training context. A submission may override it + # with a module-level ``BLOCK_SIZE`` int in model.py + # (runner.resolve_train_block_size), because trading + # context length for optimizer steps inside the fixed + # wall-clock budget is now part of the design space -- + # at the cost of having to extrapolate to 8192 at eval. + block_size: int = 8192 # DEFAULT training context (agent may override) + eval_block_size: int = 8192 # FIXED scoring window -- judge-owned + + # --- fixed optimization recipe (locked; agent designs architecture only) --- + # 8x4 rather than 32x1: at ctx 8192 a 32-sequence micro-batch is 262k tokens + # of activations per forward. Accumulating keeps the EFFECTIVE batch at 32 + # sequences (unchanged from the ctx-1024 recipe, so the optimization recipe + # is still locked and comparable) while bounding peak memory. + batch_size: int = 8 # sequences per micro-batch (per device) + grad_accum: int = 4 # -> effective batch 32 sequences + learning_rate: float = 3.0e-3 # AdamW peak LR (nanoGPT-speedrun class) + min_lr: float = 3.0e-4 + weight_decay: float = 0.1 + beta1: float = 0.9 + beta2: float = 0.95 + grad_clip: float = 1.0 + warmup_steps: int = 100 + + # --- iso-wallclock budget (CALIBRATE: pick largest T leaving >=15min headroom) --- + train_seconds: float = 1800.0 # CALIBRATE wall-clock T per training run (30 min) + max_train_seconds: float = 3300.0 # hard cap (< 1h) — abort if exceeded + + # --- data / tokenizer (locked; hidden tokens live in the judge image) --- + dataset_name: str = "HuggingFaceFW/fineweb-edu:sample-10BT" # provenance; matches config.yaml + tokenizer_name: str = "allenai/dolma2-tokenizer" + token_dtype: str = "uint32" # vocab 100278 > 65535 -> uint16 is WRONG + train_tokens_path: str = "/opt/nanoslm_arch/data/train.bin" + val_tokens_path: str = "/opt/nanoslm_arch/data/val.bin" # HELD OUT — never in /app + # Written by docker/prep_assets.py next to the streams; carries val_bytes. + manifest_path: str = "/opt/nanoslm_arch/data/manifest.json" + val_tokens: int = 1_048_576 # held-out TOKENS scored (non-overlapping windows) + # BYTE length of the held-out text that `val_tokens` target tokens cover. + # 0 means "resolve from the manifest / FRONTIER_NANOSLM_VAL_BYTES at load + # time" -- see resolve_val_bytes(). This must NEVER silently become a token + # count: bpb is normalized by BYTES so that it stays comparable across + # architectures the way the byte-level tokenizer used to make it by + # construction. + val_bytes: int = 0 + + # --- iterative-role cached baseline (GPU via Modal; judge container CPU-only) --- + # Mirrors nanowm's `baseline_cache`. The cheap agent-role feedback path trains + # only the submission (~T on a Modal H100) and reuses a baseline perplexity + # cached by config fingerprint; the final/verifier path recomputes a fresh + # baseline+submission CRN pair (~2T) on one Modal GPU/process. In local testing + # mode the judge can use a directly-attached H100 instead of Modal. + baseline_cache_path: str = "/opt/nanoslm_arch/baseline/baseline_ppl.json" + + # --- scoring --- + # ABSOLUTE bits-per-byte gain is the scored measurement: + # gain = base_bpb - sub_bpb + # score = clip(100 * gain / bpb_score_scale, 0, 100) + # + # NOT A CALIBRATION TARGET, AND DELIBERATELY NOT LISTED AS ONE. + # `bpb_score_scale` is a DISPLAY CONVENTION whose only necessary job is to + # map bpb into the 0-100 range Harbor wants: Harbor computes + # `reward = score / 100`, so a raw gain of 0.05 bpb would arrive as reward + # 0.0005 -- zero, for every submission. It does NOT define "what counts as a + # full win"; no such number has been measured, and dressing an arbitrary + # constant as a measured target would put the score's meaning on a figure + # nobody determined. The only real requirement is DISCRIMINATION (too coarse + # and everyone pins at 0, too fine and everyone pins at 100), which is far + # weaker than calibration and settable from one real-corpus run. + # `score_unbounded` stays un-clipped so strong submissions remain visible. + # + # NOTE this replaced a RELATIVE `r_target` of the same numeric value, and + # the two are NOT equivalent: 0.05 relative was ~0.145 bpb absolute at the + # synthetic operating point (base_bpb ~2.9) and would be ~0.075-0.10 bpb + # once a real corpus puts the baseline nearer 1.5-2.0. So 0.05 ABSOLUTE is a + # STRICTER bar on real data than the relative form it replaced. + bpb_score_scale: float = 0.05 # bpb gain that saturates the bounded score + + # --- guards / resource caps --- + param_cap: int = 400_000_000 # max trainable params (over-cap -> score 0) + min_param_delta: float = 1e-6 # trained-from-scratch guard: mean|Δparam| threshold + seed: int = 1337 + # OFF pending validation -- see runner.run_arm. Closing the "compile the + # baseline architecture and win on wall-clock" loophole still needs doing. + compile_model: bool = False + + # --- determinism knobs applied on the GPU path --- + cublas_workspace_config: str = ":4096:8" + + +DEFAULT = TaskConfig() + + +# --------------------------------------------------------------------------- # +# Byte accounting for the held-out stream. +# +# THIS IS THE SUBTLE PART OF THE BPE MIGRATION. At byte level 1 token == 1 byte, +# so `mean per-token CE / ln2` WAS bits-per-byte and dividing the total NLL by +# the token count happened to be right. With a BPE tokenizer the two diverge by +# the compression ratio (~4.4 bytes/token for dolma2 on English), and per-token +# CE/ln2 is a TOKENIZER-DEPENDENT quantity -- it is not comparable across setups +# and it is not bits per byte. The whole reason bpb was chosen as the scored +# metric (DESIGN.md §2) is that it is tokenizer-independent, so the +# denominator must be BYTES. +# +# The byte count is produced at corpus-prep time (docker/prep_assets.py decodes +# the scored token span and measures its UTF-8 length) and travels in +# manifest.json. eval_ppl.py ASSERTS it is available rather than falling back. +# --------------------------------------------------------------------------- # + +# Only used when no real corpus is staged (synthetic smoke stream). Roughly the +# measured dolma2 compression ratio on English web text, so a smoke bpb lands in +# a plausible range instead of being off by ~4.4x. +SYNTHETIC_BYTES_PER_TOKEN = 4.4 + + +def resolve_val_bytes_per_token(cfg: TaskConfig | None = None) -> float | None: + """Bytes of held-out TEXT per scored target token, or ``None`` if unknown. + + A RATIO rather than a bare total, because the number of tokens actually + scored is ``min(len(val)-1, cfg.val_tokens)`` rounded down to whole windows + -- the smoke config scores far fewer than the corpus contains. The ratio is + exact for the production config (which scores the whole staged stream) and + proportional otherwise, and it can never be confused for a token count. + + Resolution order (first hit wins): + 1. ``FRONTIER_NANOSLM_VAL_BYTES`` (+ optional ``..._VAL_BYTES_TOKENS``) -- + operator / Modal-image override. + 2. ``cfg.val_bytes`` when frozen to a non-zero literal (paired with + ``cfg.val_tokens``). + 3. ``val_bytes`` / ``val_target_tokens`` from the manifest that + ``docker/prep_assets.py`` writes next to the token streams. + """ + cfg = cfg or active_config() + + def _ratio(nbytes, ntok) -> float | None: + try: + nbytes, ntok = int(nbytes), int(ntok) + except (TypeError, ValueError): + return None + return nbytes / ntok if nbytes > 0 and ntok > 0 else None + + env = os.environ.get("FRONTIER_NANOSLM_VAL_BYTES", "").strip() + if env: + r = _ratio(env, os.environ.get("FRONTIER_NANOSLM_VAL_BYTES_TOKENS") + or cfg.val_tokens) + if r: + return r + + if cfg.val_bytes > 0: + r = _ratio(cfg.val_bytes, cfg.val_tokens) + if r: + return r + + for path in (os.environ.get("FRONTIER_NANOSLM_MANIFEST", ""), cfg.manifest_path): + if not path: + continue + try: + with open(path, "r", encoding="utf-8") as fh: + man = json.load(fh) + except Exception: + continue + r = _ratio(man.get("val_bytes"), man.get("val_target_tokens")) + if r: + return r + return None + + +# --------------------------------------------------------------------------- # +# Smoke overrides: tiny, fast, CPU-friendly. Enabled by FRONTIER_NANOSLM_SMOKE=1. +# Used only to prove the harness wiring end-to-end without a GPU; never scored. +# --------------------------------------------------------------------------- # +def smoke_enabled() -> bool: + return os.environ.get("FRONTIER_NANOSLM_SMOKE", "") == "1" + + +def active_config() -> TaskConfig: + if not smoke_enabled(): + return DEFAULT + # SMOKE SHRINKS THE BUDGET, NOT THE SHAPE. + # + # The obvious smoke config (block_size=64) is actively harmful here: fla's + # Triton kernels autotune per shape, and a 64-token sequence is far outside + # the regime they are tuned for. Measured on an H100, GatedDeltaNet's first + # forward costs 74.5s at T=64 but only 15.5s at T=8192 -- the tiny sequence + # is ~5x SLOWER to compile. A block_size=64 smoke therefore appears to hang + # for tens of minutes while compiling kernels no scored run will ever use. + # + # So the smoke keeps the REAL sequence length and cuts batch size, steps and + # eval tokens instead. It is slower than a toy config but it exercises the + # kernels the scored run actually uses, which is the point of a smoke test. + return TaskConfig( + block_size=8192, # REAL shape -- see above + eval_block_size=8192, # scoring window is NEVER shrunk by the smoke + batch_size=1, + grad_accum=1, + warmup_steps=1, + # 45s, not 3s. The score is a function of THROUGHPUT, so a smoke that + # yields one optimizer step per arm cannot validate the thing being + # measured -- at train_seconds=3.0 both arms did exactly 1 step and the + # reported +4.2% was noise. 45s gives tens of steps and a real + # steps-per-arm comparison, at the cost of a slower smoke. + train_seconds=45.0, + max_train_seconds=900.0, + val_tokens=32_768, + param_cap=400_000_000, + ) + + +# --------------------------------------------------------------------------- # +# Config fingerprint: the cache key for the iterative-role cached baseline. +# A change to ANY locked knob invalidates a cached baseline rather than +# mispairing it (mirrors nanowm settings.config_fingerprint()). +# --------------------------------------------------------------------------- # +# +# NOTE `block_size` IS DELIBERATELY ABSENT, and `eval_block_size` replaces it. +# `block_size` is now only the DEFAULT training context: each submission may +# pick its own (runner.resolve_train_block_size). Fingerprinting a per-arm value +# would give every submission a distinct key, so the cached baseline would never +# hit and the agent-role feedback path would silently cost ~2T instead of ~T -- +# the exact thing the cache exists to avoid. The BASELINE always trains at 8192 +# (baseline_model.BLOCK_SIZE) and is always scored at `eval_block_size`, so +# neither of those depends on the submission and the cached number stays valid. +# `eval_block_size` DOES belong here: changing the scoring window changes what +# val_bpb means, so a baseline cached under the old one must not be reused. +_FINGERPRINT_KEYS = ( + "vocab_size", "eval_block_size", "batch_size", "grad_accum", + "learning_rate", "min_lr", "weight_decay", "beta1", "beta2", + "grad_clip", "warmup_steps", "train_seconds", "dataset_name", + # tokenizer_name is fingerprinted alongside vocab_size: changing the + # tokenizer changes what val_bpb MEANS, so a baseline cached under the old + # one must not be reused. + "tokenizer_name", "val_tokens", "seed", +) + + +def config_fingerprint(cfg: TaskConfig | None = None) -> str: + cfg = cfg or active_config() + d = asdict(cfg) + payload = {k: d[k] for k in _FINGERPRINT_KEYS} + blob = json.dumps(payload, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(blob.encode("utf-8")).hexdigest()[:16] diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harness/train.py b/2.0/problems/nanoslm_hybrid_arch_design/harness/train.py new file mode 100644 index 000000000..260bcc14b --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/train.py @@ -0,0 +1,148 @@ +"""Locked training loop (torch path). + +The optimizer, LR schedule, weight-decay grouping, gradient accumulation, data +order, gradient clipping, and the wall-clock cutoff are all fixed here so the +task measures *architecture*, not training tricks. The harness computes the +cross-entropy loss from the model's logits and IGNORES any loss the model +returns. + +Each optimizer step accumulates ``cfg.grad_accum`` micro-batches of +``cfg.batch_size`` sequences, for an effective batch of +``batch_size * grad_accum`` at the peak activation memory of a single +micro-batch. + +Timing note: wall-clock is measured by ``time.monotonic`` inside the harness +(never by the model). Training stops at the first optimizer step whose start +exceeds ``cfg.train_seconds``. ``max_train_seconds`` is a hard safety abort. +""" + +from __future__ import annotations + +import math +import time +from dataclasses import dataclass + +from .data import TokenData +from .settings import TaskConfig + + +@dataclass +class TrainOutput: + steps: int + wall_seconds: float + final_train_loss: float + + +def _lr_at(step: int, cfg: TaskConfig, total_steps_guess: int) -> float: + if step < cfg.warmup_steps: + return cfg.learning_rate * (step + 1) / max(1, cfg.warmup_steps) + # cosine decay toward min_lr over an estimated horizon + denom = max(1, total_steps_guess - cfg.warmup_steps) + frac = min(1.0, (step - cfg.warmup_steps) / denom) + coeff = 0.5 * (1.0 + math.cos(math.pi * frac)) + return cfg.min_lr + coeff * (cfg.learning_rate - cfg.min_lr) + + +def _param_groups(model, cfg: TaskConfig): + """No weight decay on 1-D params (norms, biases, embeddings).""" + decay, no_decay = [], [] + for _, p in model.named_parameters(): + if not p.requires_grad: + continue + (decay if p.dim() >= 2 else no_decay).append(p) + return [ + {"params": decay, "weight_decay": cfg.weight_decay}, + {"params": no_decay, "weight_decay": 0.0}, + ] + + +def _cross_entropy(logits, targets): + import torch.nn.functional as F + + B, T, V = logits.shape + return F.cross_entropy( + logits.reshape(B * T, V).float(), targets.reshape(B * T) + ) + + +def train_model(model, data: TokenData, cfg: TaskConfig, device: str) -> TrainOutput: + import torch + + model.to(device) + model.train() + opt = torch.optim.AdamW( + _param_groups(model, cfg), + lr=cfg.learning_rate, + betas=(cfg.beta1, cfg.beta2), + ) + # Rough horizon estimate for the cosine schedule; the true stop is wall-clock. + total_steps_guess = 10_000 + + use_amp = device.startswith("cuda") + autocast = ( + torch.autocast(device_type="cuda", dtype=torch.bfloat16) + if use_amp + else _nullcontext() + ) + + step = 0 + last_loss = float("nan") + t0 = time.monotonic() + while True: + elapsed = time.monotonic() - t0 + if elapsed >= cfg.train_seconds: + break + if elapsed >= cfg.max_train_seconds: # hard safety abort + break + + lr = _lr_at(step, cfg, total_steps_guess) + for pg in opt.param_groups: + pg["lr"] = lr + + # Gradient accumulation: grad_accum micro-batches of cfg.batch_size make + # one optimizer step (effective batch = batch_size * grad_accum) while + # peak activation memory stays at a single micro-batch. cfg.block_size is + # this arm's TRAINING context (run_arm passes a per-arm cfg); evaluation + # is at cfg.eval_block_size regardless -- see data.val_windows. + accum = max(1, int(cfg.grad_accum)) + opt.zero_grad(set_to_none=True) + step_ce = 0.0 + for micro in range(accum): + # Distinct deterministic micro-batch, identical across arms (CRN): + # keyed by (step * accum + micro) so an optimizer step's micro-batches + # never repeat and both arms see the same data. accum == 1 reproduces + # the pre-accumulation key exactly. + x, y = data.batch(step * accum + micro, device, cfg.block_size) + with autocast: + logits = _logits_only(model(x)) + loss = _cross_entropy(logits, y) / accum # harness-owned loss + loss.backward() + step_ce += float(loss.detach().float().item()) + torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip) + opt.step() + + last_loss = step_ce # sum of per-micro (already /accum) losses = mean CE + step += 1 + + if device.startswith("cuda"): + import torch as _t + + _t.cuda.synchronize() + return TrainOutput( + steps=step, wall_seconds=time.monotonic() - t0, final_train_loss=last_loss + ) + + +def _logits_only(out): + """Accept either ``logits`` or ``(logits, loss)``; keep only logits.""" + if isinstance(out, (tuple, list)): + return out[0] + return out + + +class _nullcontext: + def __enter__(self): + return None + + def __exit__(self, *a): + return False diff --git a/2.0/problems/nanoslm_hybrid_arch_design/readme b/2.0/problems/nanoslm_hybrid_arch_design/readme new file mode 100644 index 000000000..49b6fac2f --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/readme @@ -0,0 +1,169 @@ +# NanoSLM Hybrid Architecture Design (fixed-wall-clock, held-out bits-per-byte) + +## Problem + +Design a language-model **architecture** that reaches the lowest possible +held-out **bits-per-byte** (`val_bpb`) when trained **from scratch** under a +**fixed wall-clock budget on a single H100**. You submit one file — `model.py`, +the model definition — and nothing else. The hidden judge plugs it into a locked +training + evaluation harness, trains it for a fixed wall-clock +`T`, and scores its `val_bpb` against a **locked baseline architecture** trained +under the identical budget. + +The starting point is the 3:1 Gated DeltaNet hybrid from *Olmo Hybrid* +(arXiv:2604.03444); the question is how much further it can be pushed. The optimizer, learning-rate schedule, +weight-decay grouping, data, tokenizer, and the wall-clock budget are all +**fixed by the judge**. You control the architecture — and, as of this revision, +the **training context length** (see below). A more compute-efficient +architecture legitimately completes more useful training steps within `T` — that +is the intended lever. + +## Metric + +The tokenizer is fixed by the judge: **dolma2 BPE** +(`allenai/dolma2-tokenizer`, vocabulary 100278) over a FineWeb-Edu corpus. The +scored quantity is held-out **bits per byte**: + +``` +val_bpb = total_cross_entropy_nats / (val_bytes * ln 2) +``` + +where `val_bytes` is the byte length of the hidden held-out **text** — not the +number of tokens it was split into. Normalizing by bytes is what makes the +number independent of the tokenizer, so it is comparable across architectures +and cannot be gamed by changing tokenization. **Lower is better.** + +A per-token validation perplexity is also reported, for readability only. It is +not the scored quantity, and because the tokenizer is not byte-level it is *not* +equal to `2**val_bpb`. + +## Program interface + +Submit a single Python file `model.py` that defines a factory the harness calls: + +```python +def build_model(config) -> torch.nn.Module: ... +# or, equivalently, a class usable as NanoSLM(config): +class NanoSLM(torch.nn.Module): ... +``` + +`config` is provided by the harness and has: + +- `config.vocab_size` — always `100278` (dolma2 BPE; do not change). +- `config.block_size` — the context length you will be **trained** at (8192 by + default; see "Training context length" below). +- `config.eval_block_size` — the context length you will be **scored** at. + Always `8192`, whatever you train at. +- `config.train_seconds_hint`, `config.param_cap_hint`, `config.device_hint` — + read-only informational hints (no guarantees). + +The returned module must implement: + +```python +def forward(self, idx): + # idx: LongTensor [B, T] of BPE token ids in [0, 100278) + return logits # FloatTensor [B, T, vocab_size] +``` + +Returning `(logits, loss)` is accepted, but the judge **ignores any returned +loss** and computes cross-entropy itself, from your `logits`, for both training +and validation. You choose everything inside the model: width, depth, attention +mechanism, normalization, positional scheme, embedding sharing, initialization, +and so on. You do **not** control the optimizer or training loop. + +The vocabulary is large enough that the embedding table is a first-class design +concern — at `d=768` it is ~77M parameters on its own, so tying, factorizing or +otherwise reshaping it is a real lever rather than a detail. + +The 3:1 Gated DeltaNet hybrid `model.py` (the *Olmo Hybrid* recipe applied to the +`olmo3_190M` baseline) is provided as a starting point. + +## Training context length — yours to choose, with a catch + +You may declare the context length you are **trained** at by defining a +module-level integer in `model.py`: + +```python +BLOCK_SIZE = 2048 # power of two in [256, 8192]; omit it to train at 8192 +``` + +Omit it and you train at the default **8192**. A value outside `[256, 8192]`, or +one that is not a power of two, is rejected before training (score 0). + +**You are always evaluated at 8192.** That is fixed for every submission — it is +what makes `val_bpb` comparable across submissions — and it is the whole trade: + +- A shorter training context makes each optimizer step cheaper, so you complete + **more steps** inside the same fixed wall-clock budget `T`. At 8192 attention + is ~78–84% of layer FLOPs, so this is a large effect. +- But your model is still scored on 8192-token windows. Train at 1024 and you + are asked for logits at positions **8x beyond anything you saw in training**. + +How well a model survives that depends heavily on its **position encoding**, +which is not the same thing as architectural quality under a compute budget. +Plain RoPE degrades sharply past its training length; NTK-aware / YaRN-style +scaling, position interpolation and ALiBi extrapolate considerably better. So if +you shorten the training context, treat the position encoding as part of the +decision rather than an afterthought — and be aware that some of what you would +then be measuring is extrapolation behaviour, not mixer efficiency. + +Practical warning: anything you size or cache off `config.block_size` — RoPE +tables, learned positional embeddings, causal or sliding-window mask buffers — +must still work at `config.eval_block_size`. Build such buffers against +`config.eval_block_size`, or lazily against the actual `T` you are handed. A +model that crashes or silently truncates at 8192 scores 0. + +The baseline always trains at 8192, so this trade is measured against a fixed +point. + +## Validity constraints + +- Submission is `model.py` only, `.py`, at most 256 KB. +- The file must define `build_model(config)` or `class NanoSLM`. +- Train from scratch: no loading pretrained weights, no reading files, no + network, no environment access, no reading the clock, no timing + short-circuits. Submissions containing such calls are rejected before running. +- Trainable parameters must be within the judge's param cap, and the model must + fit and train within GPU memory. Over-cap or out-of-memory scores 0. +- The model must actually train (its parameters must change) and must produce a + non-degenerate output distribution. Untrained or constant-output models score 0. + +## Scoring + +Let `base_bpb` be the locked baseline architecture's held-out bits-per-byte and +`sub_bpb` be your submission's, both trained for the same wall-clock `T` on the +same H100 with the same data order (common random numbers), and both scored on +the same hidden held-out text at the same fixed 8192-token windows. The scored +measurement is your **absolute** bits-per-byte gain: + +``` +gain = base_bpb - sub_bpb # bits per byte +score = clip(100 * gain / BPB_SCORE_SCALE, 0, 100) +``` + +- The baseline architecture (tying its `val_bpb`) scores **0**. +- A submission worse than the baseline scores **0**. +- `BPB_SCORE_SCALE` is only a **scaling convention** — it exists to map bits per + byte onto the 0–100 range the platform expects, not to declare what counts as + a full result. Maximize `gain`; `score_unbounded` keeps rewarding improvements + past the cap. + +Bits per byte is the unit the language-modelling literature quotes, so a gain +here is directly comparable to published figures. A relative improvement is also +reported, for readability. + +Public feedback reports `base_val_bpb`, `sub_val_bpb`, the absolute gain and the +relative improvement, the training context you used, the number of optimizer +steps completed, and the training wall-clock. Hidden data, evaluator internals, +and the baseline definition are not exposed. + +## Iterating + +During a trial you can package and score the current `model.py`: + +```bash +bash /app/submit.sh +``` + +The best successful iterative submission is kept if a later artifact is worse or +you time out. diff --git a/2.0/problems/nanoslm_hybrid_arch_design/reference.py b/2.0/problems/nanoslm_hybrid_arch_design/reference.py new file mode 100644 index 000000000..6170d9453 --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/reference.py @@ -0,0 +1,434 @@ +"""Reference solution for nanoslm_hybrid_arch_design -- the Olmo-Hybrid recipe at 190M. + +This is the STARTING POINT, not the ceiling. The locked baseline is a faithful +``olmo3_190M``; this file applies OLMo-core's own hybrid recipe +(``src/scripts/train/OLMo_hybrid/OLMo-hybrid-7B.py``) to it. + +WHERE THE 3:1 RATIO ACTUALLY COMES FROM +--------------------------------------- +It is NOT an arbitrary choice. ``olmo3_190M`` already carries:: + + SlidingWindowAttentionConfig(pattern=[4096, 4096, 4096, -1], + force_full_attention_on_first_layer=False, + force_full_attention_on_last_layer=True) + +so 9 of its 12 layers are sliding-window and 3 (layers 3, 7, 11) are full +attention. The paper "replaces sliding window layers with Gated DeltaNet +layers", and upstream's recipe encodes exactly that:: + + config.block_pattern = ["gdn", "gdn", "gdn", "attn"] + +The three sliding-window layers of each 4-layer group become GDN; the +full-attention layer of each group survives untouched. The 3:1 ratio IS olmo3's +own sliding-window pattern. Note the consequence: the hybrid has NO sliding +window left anywhere -- every surviving attention layer sat at a ``-1`` position. + +PARAMETER MATCHING (upstream's REMOVE_HEADS practice) +----------------------------------------------------- +GDN's mixer is larger than an attention mixer, so upstream shrinks the model to +compensate rather than letting the hybrid quietly buy capacity:: + + REMOVE_HEADS = 2 + config.d_model -= REMOVE_HEADS * 128 # 4096 -> 3840 + num_heads -= REMOVE_HEADS # 32 -> 30 + assert config.d_model / num_heads == 128 # head_dim preserved + +At 190M with head_dim 64 the equivalent is ``REMOVE_HEADS = 1``:: + + d_model 768 -> 704, n_heads 12 -> 11, 704 / 11 == 64 (head_dim preserved) + +Matching is done on NON-EMBEDDING parameters, which is the quantity upstream +itself feeds to ``Duration.chinchilla_tokens(model_params=...)``. Both arms TIE +their single vocab table, which at this scale is ~40% of all parameters -- +including it would drown the mixer difference being matched. Measured +analytically: + + baseline non-embedding 113,283,840 + hybrid non-embedding 110,805,734 -2.19% + +which is the same tolerance upstream's own 7B pair achieves (+2.84%). One caveat +worth stating: because ``d_model`` shrinks 768 -> 704, the tied table shrinks too +(77.0M -> 70.6M), so the hybrid's TOTAL is 4.7% below the baseline's +(181,401,446 vs 190,297,344) even though its non-embedding count is within 2.2%. + +Note also that ``hidden_size`` stays 3072. Upstream mutates ``d_model`` in place +and never recomputes the feed-forward width, so the FFN keeps the width derived +from the ORIGINAL d_model. Reproduced here deliberately. + +Everything else is identical to the baseline -- RMSNorm eps 1e-6, SwiGLU 3072, +reordered_norm blocks, qk_norm, RoPE theta 500_000, TIED embeddings (the baseline +ties, so this arm ties too for comparability) -- so the ONLY architectural +difference is the sequence mixer in those nine layers. + +YOUR TASK: PUSH val_bpb FURTHER +------------------------------- +Everything above is upstream's recipe, transposed to this scale. The open +questions it does NOT answer, each worth real bpb: + + * THE RATIO. 3:1 is inherited from olmo3's sliding-window pattern, which was + chosen for a sliding window, not for a linear RNN. 5:1 and 7:1 buy more + steps but give up more global context. + * PLACEMENT. Every 4th, or clustered (early layer for global context in, late + layer for read-out)? Same cost, different models. + * STATE SIZE. ``expand_v`` and ``num_v_heads`` set the recurrent state, the + main capacity knob of a linear RNN -- unlike a KV cache it does not grow + with sequence length, so capacity is cheap at 8192. + * THE MIXER ITSELF. GDN is one choice; GLA, RetNet and Mamba2-style mixers are + all expressible in the same chunkwise matmul form. + * NON-UNIFORMITY. Nothing requires every GDN layer to be identical, or the + attention layers to be full-width. + * THE REST OF THE BLOCK. Norm placement, gating, MLP ratio, head count, + embedding tying -- all still on the table, and all interact with the above. + * THE TRAINING CONTEXT. Declare a module-level ``BLOCK_SIZE`` int (a power of + two in [256, 8192]) to train at a shorter context than the default 8192. + Shorter steps are cheaper, so you complete more of them in the fixed + wall-clock budget -- but EVALUATION IS ALWAYS AT 8192, so the model must + extrapolate to positions well beyond anything it trained on. How well it + does that is mostly a property of the POSITION ENCODING (plain RoPE degrades + sharply; NTK-aware/YaRN scaling and ALiBi hold up far better), which is a + different question from mixer efficiency. This file declares no BLOCK_SIZE + and therefore trains at 8192. Anything sized off ``config.block_size`` must + still run at ``config.eval_block_size``. + +Locked and not yours to change: optimizer, data, tokenizer, the EVALUATION +context (8192), and the wall-clock budget. Interface: ``build_model(config)`` / +``NanoSLM(config)`` returning ``forward(idx) -> logits [B, T, vocab_size]``. +""" + +from __future__ import annotations + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +# --------------------------------------------------------------------------- # +# Shape: olmo3_190M with upstream's REMOVE_HEADS compensation applied. +# --------------------------------------------------------------------------- # +REMOVE_HEADS = 1 + +N_LAYER = 12 +N_HEAD = 12 - REMOVE_HEADS # 11 +HEAD_DIM = 64 # preserved, exactly as upstream asserts +N_EMBD = 768 - REMOVE_HEADS * HEAD_DIM # 704 +assert N_EMBD // N_HEAD == HEAD_DIM, "REMOVE_HEADS must preserve head_dim" + +# NOT recomputed from the reduced d_model -- upstream mutates d_model in place +# and leaves the feed-forward width at the value llama_like derived from the +# ORIGINAL 768. Reproduced deliberately. +HIDDEN_SIZE = 3072 + +ROPE_THETA = 500_000.0 +NORM_EPS = 1e-6 + +# GatedDeltaNetConfig(n_heads=num_heads, head_dim=int(0.75*d_model/num_heads), +# allow_neg_eigval=True) with expand_v at its 2.0 default. +GDN_HEAD_DIM = int(0.75 * N_EMBD / N_HEAD) # 48 +GDN_EXPAND_V = 2.0 +GDN_HEAD_V_DIM = int(GDN_HEAD_DIM * GDN_EXPAND_V) # 96 +GDN_KEY_DIM = N_HEAD * GDN_HEAD_DIM # 528 +GDN_VALUE_DIM = N_HEAD * GDN_HEAD_V_DIM # 1056 +GDN_CONV_SIZE = 4 + +# config.block_pattern = ["gdn", "gdn", "gdn", "attn"] -- the three +# sliding-window layers of each group become GDN, the full-attention layer +# survives. Attention therefore lands at layers 3, 7, 11. +PATTERN = ("gdn", "gdn", "gdn", "attn") +assert N_LAYER % len(PATTERN) == 0 + + +class _RMSNorm(nn.Module): + def __init__(self, d: int, eps: float = NORM_EPS): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(d)) + + def forward(self, x): + dt = x.dtype + xf = x.float() + n = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps) + return n.to(dt) * self.weight + + +def _rope(x, theta: float = ROPE_THETA): + _, _, T, D = x.shape + half = D // 2 + freq = theta ** (-torch.arange(0, half, device=x.device, dtype=torch.float32) / half) + ang = torch.arange(T, device=x.device, dtype=torch.float32)[:, None] * freq[None, :] + cos, sin = ang.cos()[None, None], ang.sin()[None, None] + x1, x2 = x[..., :half].float(), x[..., half:].float() + return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1).to(x.dtype) + + +class _Attention(nn.Module): + """Identical to the baseline's attention, always FULL causal. + + Full, not windowed, and that is not a simplification: the attention layers + the hybrid keeps are exactly the ones sitting at the ``-1`` positions of + olmo3's [4096, 4096, 4096, -1] pattern, so they were already full attention + before the GDN substitution. + + qk_norm spans the full n_heads*head_dim projection and is applied BEFORE the + head reshape -- upstream's behaviour when ``use_head_qk_norm`` is False, + which is the olmo3_190M default. + """ + + def __init__(self, n_embd: int, n_head: int): + super().__init__() + assert n_embd % n_head == 0 + self.n_head, self.n_embd = n_head, n_embd + self.head_dim = n_embd // n_head + self.w_q = nn.Linear(n_embd, n_embd, bias=False) + self.w_k = nn.Linear(n_embd, n_embd, bias=False) + self.w_v = nn.Linear(n_embd, n_embd, bias=False) + self.w_out = nn.Linear(n_embd, n_embd, bias=False) + self.q_norm = _RMSNorm(n_embd) + self.k_norm = _RMSNorm(n_embd) + + def forward(self, x): + B, T, C = x.shape + q, k, v = self.w_q(x), self.w_k(x), self.w_v(x) + q, k = self.q_norm(q), self.k_norm(k) + q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2) + k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2) + v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2) + q, k = _rope(q), _rope(k) + y = F.scaled_dot_product_attention(q, k, v, is_causal=True) + return self.w_out(y.transpose(1, 2).contiguous().view(B, T, C)) + + +class _ChunkedDeltaNet(nn.Module): + """CPU-only chunkwise fallback, PARAMETER-IDENTICAL to fla's GatedDeltaNet. + + Exists so the CPU smoke path can build and report the same parameter count + the CUDA path will. Its module inventory is fla 0.5.1's exactly -- + q/k/v/a/b/g/o projections, three depthwise short convolutions, A_log, + dt_bias and the gated output norm -- so ``n_params`` is device-independent. + + The recurrence is a chunkwise GATED LINEAR ATTENTION, not the full delta + rule: it carries the decay and the gates but skips the delta-rule inverse. + That is a deliberate approximation. It is never scored -- CUDA raises rather + than falling back here (see ``_make_gdn``) -- and it exists to prove wiring, + not to reproduce GDN's numerics. + """ + + def __init__(self, chunk: int = 512): + super().__init__() + self.chunk = chunk + d, H = N_EMBD, N_HEAD + self.q_proj = nn.Linear(d, GDN_KEY_DIM, bias=False) + self.k_proj = nn.Linear(d, GDN_KEY_DIM, bias=False) + self.v_proj = nn.Linear(d, GDN_VALUE_DIM, bias=False) + self.a_proj = nn.Linear(d, H, bias=False) + self.b_proj = nn.Linear(d, H, bias=False) + self.g_proj = nn.Linear(d, GDN_VALUE_DIM, bias=False) + self.o_proj = nn.Linear(GDN_VALUE_DIM, d, bias=False) + self.A_log = nn.Parameter(torch.zeros(H)) + self.dt_bias = nn.Parameter(torch.zeros(H)) + # ShortConvolution weights: depthwise, (channels, 1, kernel), no bias. + self.q_conv1d = nn.Parameter(torch.zeros(GDN_KEY_DIM, 1, GDN_CONV_SIZE)) + self.k_conv1d = nn.Parameter(torch.zeros(GDN_KEY_DIM, 1, GDN_CONV_SIZE)) + self.v_conv1d = nn.Parameter(torch.zeros(GDN_VALUE_DIM, 1, GDN_CONV_SIZE)) + self.o_norm = nn.Parameter(torch.ones(GDN_HEAD_V_DIM)) + + @staticmethod + def _short_conv(x, w): + """Causal depthwise conv over time. x [B, T, C], w [C, 1, K].""" + C, K = w.shape[0], w.shape[2] + xt = F.pad(x.transpose(1, 2), (K - 1, 0)) + return F.conv1d(xt, w, groups=C).transpose(1, 2) + + def forward(self, x): + B, T, _ = x.shape + H, DK, DV, S = N_HEAD, GDN_HEAD_DIM, GDN_HEAD_V_DIM, self.chunk + + q = F.silu(self._short_conv(self.q_proj(x), self.q_conv1d)) + k = F.silu(self._short_conv(self.k_proj(x), self.k_conv1d)) + v = self._short_conv(self.v_proj(x), self.v_conv1d) + + q = q.view(B, T, H, DK).transpose(1, 2) + k = k.view(B, T, H, DK).transpose(1, 2) + v = v.view(B, T, H, DV).transpose(1, 2) + + # Per-head decay in (0, 1): the gate that makes this "gated". + dt = F.softplus(self.a_proj(x) + self.dt_bias) # [B, T, H] + g = torch.exp(-torch.exp(self.A_log) * dt) # [B, T, H] + g = g.permute(0, 2, 1).unsqueeze(-1) # [B, H, T, 1] + beta = torch.sigmoid(self.b_proj(x)).permute(0, 2, 1).unsqueeze(-1) + v = v * beta + + pad = (S - T % S) % S + if pad: + q, k, v = (F.pad(t, (0, 0, 0, pad)) for t in (q, k, v)) + g = F.pad(g, (0, 0, 0, pad), value=1.0) + nC = (T + pad) // S + rs = lambda t: t.view(B, H, nC, S, -1) # noqa: E731 + qc, kc, vc, gc = rs(q), rs(k), rs(v), rs(g) + + causal = torch.tril(torch.ones(S, S, device=x.device, dtype=torch.bool)) + intra = (qc @ kc.transpose(-1, -2)).masked_fill(~causal, 0.0) @ vc + + state = torch.zeros(B, H, DK, DV, device=x.device, dtype=x.dtype) + outs = [] + for c in range(nC): + outs.append(intra[:, :, c] + qc[:, :, c] @ state) + decay = gc[:, :, c].prod(dim=2, keepdim=True) + state = state * decay + kc[:, :, c].transpose(-1, -2) @ (vc[:, :, c] * gc[:, :, c]) + y = torch.stack(outs, dim=2).view(B, H, T + pad, DV)[:, :, :T] + + # Gated RMSNorm over head_v_dim, then merge heads. + yf = y.float() + y = (yf * torch.rsqrt(yf.pow(2).mean(-1, keepdim=True) + 1e-5)).to(x.dtype) + y = y * self.o_norm + y = y.transpose(1, 2).contiguous().view(B, T, GDN_VALUE_DIM) + y = y * F.silu(self.g_proj(x)) + return self.o_proj(y) + + +def _make_gdn(): + """fla's fused GatedDeltaNet on CUDA; chunkwise only on CPU (smoke). + + WHY fla IS REQUIRED HERE, not merely preferred: PyTorch ships a fused kernel + for softmax attention (SDPA) but none for a linear/recurrent mixer, so a + hand-written GDN is unfused eager code. Scoring is fixed WALL-CLOCK, so an + unfused mixer loses on throughput regardless of architectural merit -- + measured at ctx 8192, the chunkwise fallback did 6% of attention's FLOPs and + still ran 3.5x slower. Both arms must be kernel-matched or the score + measures kernel quality, not architecture. + + HARD FAILURE ON CUDA, deliberately: a silent fallback here once produced a + scored-looking run where the reference "lost" 18 steps to 65 purely because + fla was missing. A missing kernel must raise, not quietly change what is + being measured. + """ + if not torch.cuda.is_available(): + return _ChunkedDeltaNet() + + from fla.layers import GatedDeltaNet # ImportError here is intentional + + # head_dim MUST be passed: fla defaults it to 256, so omitting it builds + # num_heads*256-wide projections instead of the intended width -- a silent + # param blowup that trains a much larger model than the baseline. + return GatedDeltaNet( + hidden_size=N_EMBD, + num_heads=N_HEAD, + head_dim=GDN_HEAD_DIM, + expand_v=GDN_EXPAND_V, + use_gate=True, + use_short_conv=True, + ) + + +class _GDNLayer(nn.Module): + def __init__(self): + super().__init__() + self.impl = _make_gdn() + + def forward(self, x): + out = self.impl(x) + # fla layers return (hidden_states, attentions, past_kv); the chunkwise + # fallback returns a tensor. Normalize so _Block need not care. + return out[0] if isinstance(out, tuple) else out + + +class _FeedForward(nn.Module): + """SwiGLU, hidden 3072, bias=False -- unchanged from the baseline.""" + + def __init__(self, n_embd: int, hidden: int = HIDDEN_SIZE): + super().__init__() + self.w1 = nn.Linear(n_embd, hidden, bias=False) + self.w3 = nn.Linear(n_embd, hidden, bias=False) + self.w2 = nn.Linear(hidden, n_embd, bias=False) + + def forward(self, x): + return self.w2(F.silu(self.w1(x)) * self.w3(x)) + + +class _ReorderedNormBlock(nn.Module): + """Upstream builds the GDN block as ``attn_block.replace(sequence_mixer=...)`` + -- same reordered_norm structure, same norms, only the mixer swapped.""" + + def __init__(self, n_embd: int, n_head: int, kind: str): + super().__init__() + self.attention = _Attention(n_embd, n_head) if kind == "attn" else _GDNLayer() + self.attention_norm = _RMSNorm(n_embd) + self.feed_forward = _FeedForward(n_embd) + self.feed_forward_norm = _RMSNorm(n_embd) + + def forward(self, x): + h = x + self.attention_norm(self.attention(x)) + return h + self.feed_forward_norm(self.feed_forward(h)) + + +class NanoSLM(nn.Module): + def __init__(self, config): + super().__init__() + self.block_size = config.block_size + self.wte = nn.Embedding(config.vocab_size, N_EMBD) + kinds = [PATTERN[i % len(PATTERN)] for i in range(N_LAYER)] + self.blocks = nn.ModuleList([_ReorderedNormBlock(N_EMBD, N_HEAD, k) for k in kinds]) + self.norm_f = _RMSNorm(N_EMBD) + # TIED -- the baseline ties its embeddings (a deliberate deviation from + # upstream's untied default), so the reference ties too and the two arms + # stay comparable: the ONLY architectural difference between them is the + # sequence mixer in the GDN layers, not the output-head param budget. The + # tie is set AFTER init (below) so the shared table carries wte's init. + self.lm_head = nn.Linear(N_EMBD, config.vocab_size, bias=False) + self.apply(self._init) + for name, p in self.named_parameters(): + if name.endswith("w_out.weight") or name.endswith("w2.weight"): + nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * N_LAYER)) + self.lm_head.weight = self.wte.weight # weight tying + + @staticmethod + def _init(m): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, mean=0.0, std=0.02) + if m.bias is not None: + nn.init.zeros_(m.bias) + elif isinstance(m, nn.Embedding): + nn.init.normal_(m.weight, mean=0.0, std=0.02) + + def forward(self, idx): + x = self.wte(idx) + for block in self.blocks: + x = block(x) + x = self.norm_f(x) + return self.lm_head(x) # logits [B, T, vocab_size] + + +def build_model(config) -> NanoSLM: + return NanoSLM(config) + + +def analytic_n_params(vocab_size: int) -> int: + """Analytic total, derived from the spec rather than from a measurement.""" + d, H, hid = N_EMBD, N_HEAD, HIDDEN_SIZE + block_norms = 2 * d + ffn = 3 * d * hid + attn_mix = 4 * d * d + 2 * d # projections + q/k norms + gdn_mix = ( + 2 * d * GDN_KEY_DIM # q_proj, k_proj + + 2 * d * GDN_VALUE_DIM # v_proj, g_proj + + 2 * d * H # a_proj, b_proj + + GDN_VALUE_DIM * d # o_proj + + 2 * H # A_log, dt_bias + + GDN_CONV_SIZE * (2 * GDN_KEY_DIM + GDN_VALUE_DIM) # short convs + + GDN_HEAD_V_DIM # o_norm + ) + n_attn = sum(1 for i in range(N_LAYER) if PATTERN[i % len(PATTERN)] == "attn") + n_gdn = N_LAYER - n_attn + non_emb = n_attn * (attn_mix + block_norms + ffn) + n_gdn * (gdn_mix + block_norms + ffn) + d + # ONE vocab-sized table: lm_head is TIED to wte (see NanoSLM.__init__). + return non_emb + vocab_size * d + + +def _self_check(vocab_size: int = 100278) -> int: + from harness.model_config import ModelConfig + + m = NanoSLM(ModelConfig(vocab_size=vocab_size, block_size=8192)) + got = sum(p.numel() for p in m.parameters() if p.requires_grad) + want = analytic_n_params(vocab_size) + assert got == want, f"param mismatch: built {got} != analytic {want}" + return got diff --git a/2.0/problems/nanoslm_hybrid_arch_design/repro/__init__.py b/2.0/problems/nanoslm_hybrid_arch_design/repro/__init__.py new file mode 100644 index 000000000..ff7c20caa --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/repro/__init__.py @@ -0,0 +1 @@ +"""Repro rig for Olmo Hybrid Table 5 (190M Base-Easy BPB).""" diff --git a/2.0/problems/nanoslm_hybrid_arch_design/repro/modal_repro.py b/2.0/problems/nanoslm_hybrid_arch_design/repro/modal_repro.py new file mode 100644 index 000000000..9c8279e19 --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/repro/modal_repro.py @@ -0,0 +1,483 @@ +"""Modal app to reproduce Olmo Hybrid Table 5 (190M Base-Easy BPB). + +Two remote functions on one shared Volume: + + prep_data (CPU) stream the Olmo 3 mix as raw text, dolma2-tokenize it, and + write uint32 train.bin / val.bin (+ manifest with val_bytes) + to the Volume. Parameterized by target_tokens so a tiny + wiring test and a full 1x/8x-Chinchilla stage share code. + + train (H100) load the tokenized bins, build the paper-faithful 190M + transformer or GDN-3:1 hybrid (repro.paper_models), train + with a WSD-S schedule (repro.wsd), log train + held-out CE, + and checkpoint to the Volume. + +The DOWNSTREAM OlmoBaseEval-Easy bpb eval (the actual 0.950/0.891 metric) is a +separate stage (oe-eval) run on the produced checkpoints; this app produces the +checkpoints and a held-out-CE trajectory to sanity-check the training recipe. + +Run (from the driver venv, with ~/.modal.toml authed): + modal run repro/modal_repro.py::prep_data --target-tokens 20000000 # wiring + modal run repro/modal_repro.py::train --arm transformer --target-tokens 20000000 +""" + +from __future__ import annotations + +import os +import pathlib + +import modal + +APP_NAME = "nanoslm-repro" +DATA = "/data" +REPRO_REMOTE = "/root/repro" +_HERE = pathlib.Path(__file__).resolve().parent + +VOL = modal.Volume.from_name("nanoslm-repro", create_if_missing=True) + +# Same torch/fla/triton stack the benchmark's modal_app pins -- fla 0.5.1 needs +# Triton >= 3.2 (torch 2.6 ships it); 0.4.1's backward kernels predate it and +# hang/err. transformers carries the dolma2 tokenizer. +image = ( + modal.Image.debian_slim(python_version="3.11") + .pip_install( + "torch==2.6.0", "numpy<2", + "flash-linear-attention==0.5.1", "transformers==4.46.3", + "huggingface_hub>=0.25", "tokenizers>=0.20", "zstandard>=0.22", + "hf_transfer>=0.1.6", # HF_HUB_ENABLE_HF_TRANSFER=1 needs this present + ) + .add_local_dir(str(_HERE), REPRO_REMOTE, copy=True, + ignore=["__pycache__", "*.pyc"]) + .env({ + "CUBLAS_WORKSPACE_CONFIG": ":4096:8", + "TRITON_CACHE_DIR": "/tmp/triton-cache", + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", + "HF_HUB_ENABLE_HF_TRANSFER": "1", + "PYTHONPATH": "/root", + }) +) + +app = modal.App(APP_NAME) + +TOKENIZER = "allenai/dolma2-tokenizer" # vocab 100278 (padded to 100352) +VOCAB_PADDED = 100352 +DEFAULT_DATASET = "allenai/dolma3_mix-150B-1025" # 150B sample of the Olmo 3 mix + + +# --------------------------------------------------------------------------- # +# Data prep: stream raw text -> dolma2 token ids -> uint32 bins on the Volume. +# --------------------------------------------------------------------------- # +@app.function(image=image, cpu=16.0, memory=32768, timeout=6 * 60 * 60, + volumes={DATA: VOL}) +def prep_data(target_tokens: int, val_tokens: int = 8_000_000, + dataset: str = DEFAULT_DATASET, text_field: str = "text", + tag: str = "smoke", seed: int = 0) -> dict: + """Tokenize `target_tokens` (+ val) tokens of `dataset` into the Volume. + + The dataset is 6k+ ``.jsonl.zst`` shards grouped by domain; the per-domain + shard COUNT already encodes the Olmo 3 mixture proportions. We list every + shard, DETERMINISTICALLY SHUFFLE it (so a small subset stays mixture- + representative rather than all-one-domain), then read shards via zstandard + until the token target is met. Train and val come from DISJOINT shards. + + Writes {DATA}/{tag}/train.bin, val.bin, manifest.json. manifest carries + val_bytes (UTF-8 length of the decoded val span) so bpb is normalized by + BYTES, not tokens. + """ + import io + import json + import time + + import numpy as np + import zstandard + from huggingface_hub import HfApi, hf_hub_download + from transformers import AutoTokenizer + + outdir = pathlib.Path(DATA) / tag + outdir.mkdir(parents=True, exist_ok=True) + tok = AutoTokenizer.from_pretrained(TOKENIZER) + eos = tok.eos_token_id if tok.eos_token_id is not None else 0 + + api = HfApi() + shards = sorted(f.rfilename for f in api.dataset_info(dataset).siblings + if f.rfilename.endswith(".jsonl.zst")) + order = np.random.default_rng(seed).permutation(len(shards)) + shards = [shards[i] for i in order] + print(f"{len(shards)} shards; shuffled seed={seed}") + + need = int(target_tokens) + int(val_tokens) + buf = np.empty(need + 4_000_000, dtype=np.uint32) + n = 0 + dctx = zstandard.ZstdDecompressor() + t0 = time.time() + + def append(texts): + nonlocal n + enc = tok(texts, add_special_tokens=False)["input_ids"] + for ids in enc: + m = len(ids) + 1 + if n + m > buf.shape[0]: + m = buf.shape[0] - n + if m <= 0: + return + buf[n:n + m - 1] = np.asarray(ids[:m - 1], dtype=np.uint32) + buf[n + m - 1] = eos + n += m + + shards_used = 0 + for path in shards: + if n >= need: + break + local = hf_hub_download(dataset, path, repo_type="dataset") + batch = [] + with open(local, "rb") as fh: + reader = io.TextIOWrapper(dctx.stream_reader(fh), encoding="utf-8") + for line in reader: + if not line.strip(): + continue + try: + t = json.loads(line).get(text_field) + except Exception: + continue + if not t: + continue + batch.append(t) + if len(batch) >= 1000: + append(batch); batch = [] + if n >= need: + break + if batch and n < need: + append(batch) + os.remove(local) # keep container disk bounded + shards_used += 1 + if shards_used % 10 == 0: + print(f" {shards_used} shards, {n/1e6:.1f}M/{need/1e6:.1f}M tok, " + f"{n/max(1e-9,time.time()-t0)/1e3:.0f}k tok/s") + + n = min(n, need) + ids = buf[:n] + hi = int(ids.max()) if n else 0 + assert hi < 100278, f"token id {hi} out of dolma2 vocab" + + val_ids = ids[-val_tokens:] + train_ids = ids[:-val_tokens] + train_ids.tofile(outdir / "train.bin") + val_ids.tofile(outdir / "val.bin") + + val_text = tok.decode(val_ids.tolist(), skip_special_tokens=True) + val_bytes = len(val_text.encode("utf-8")) + manifest = { + "tag": tag, "dataset": dataset, "tokenizer": TOKENIZER, + "shards_used": shards_used, "seed": seed, + "train_tokens": int(train_ids.size), "val_tokens": int(val_ids.size), + "val_bytes": int(val_bytes), + "val_bytes_per_token": val_bytes / max(1, val_ids.size), + "seconds": round(time.time() - t0, 1), + } + (outdir / "manifest.json").write_text(json.dumps(manifest, indent=2)) + VOL.commit() + print("prep_data done:", manifest) + return manifest + + +# --------------------------------------------------------------------------- # +# Training: paper-faithful 190M model, WSD-S schedule, checkpoint to Volume. +# --------------------------------------------------------------------------- # +@app.function(image=image, cpu=8.0, memory=16384, timeout=6 * 60 * 60, + volumes={DATA: VOL}) +def _tokenize_shards(args) -> dict: + """Worker: tokenize a list of shards -> one part_{idx}.bin on the Volume.""" + import io + import json + import time + + import numpy as np + import zstandard + from huggingface_hub import hf_hub_download + from transformers import AutoTokenizer + + idx, paths, dataset, text_field, tag = args + tok = AutoTokenizer.from_pretrained(TOKENIZER) + eos = tok.eos_token_id if tok.eos_token_id is not None else 0 + dctx = zstandard.ZstdDecompressor() + parts = pathlib.Path(DATA) / tag / "parts" + parts.mkdir(parents=True, exist_ok=True) + + chunks, ntok, t0 = [], 0, time.time() + for path in paths: + local = hf_hub_download(dataset, path, repo_type="dataset") + texts = [] + with open(local, "rb") as fh: + reader = io.TextIOWrapper(dctx.stream_reader(fh), encoding="utf-8") + for line in reader: + if not line.strip(): + continue + try: + t = json.loads(line).get(text_field) + except Exception: + continue + if t: + texts.append(t) + if len(texts) >= 1000: + enc = tok(texts, add_special_tokens=False)["input_ids"] + for ids in enc: + a = np.asarray(ids + [eos], dtype=np.uint32); chunks.append(a); ntok += a.size + texts = [] + if texts: + enc = tok(texts, add_special_tokens=False)["input_ids"] + for ids in enc: + a = np.asarray(ids + [eos], dtype=np.uint32); chunks.append(a); ntok += a.size + os.remove(local) + arr = np.concatenate(chunks) if chunks else np.empty(0, dtype=np.uint32) + hi = int(arr.max()) if arr.size else 0 + assert hi < 100278, f"token id {hi} out of dolma2 vocab" + outp = parts / f"part_{idx:04d}.bin" + arr.tofile(outp) + VOL.commit() + return {"idx": idx, "tokens": int(arr.size), "shards": len(paths), + "seconds": round(time.time() - t0, 1)} + + +@app.function(image=image, cpu=32.0, memory=65536, timeout=6 * 60 * 60, + volumes={DATA: VOL}) +def prep_data_parallel(target_tokens: int, val_tokens: int = 8_000_000, + dataset: str = DEFAULT_DATASET, tag: str = "run", + n_workers: int = 32, tokens_per_shard_est: int = 22_000_000, + headroom: float = 1.35, seed: int = 0) -> dict: + """Fan-out tokenization: split a shuffled shard subset across workers, then + concatenate the parts into train.bin / val.bin (+ manifest).""" + import json + import time + + import numpy as np + from huggingface_hub import HfApi + from transformers import AutoTokenizer + + outdir = pathlib.Path(DATA) / tag + outdir.mkdir(parents=True, exist_ok=True) + api = HfApi() + shards = sorted(f.rfilename for f in api.dataset_info(dataset).siblings + if f.rfilename.endswith(".jsonl.zst")) + order = np.random.default_rng(seed).permutation(len(shards)) + shards = [shards[i] for i in order] + + need = int(target_tokens) + int(val_tokens) + n_shards = min(len(shards), int(np.ceil(need / tokens_per_shard_est * headroom))) + picked = shards[:n_shards] + # Round-robin assign so each worker gets a mixture-representative slice. + buckets = [[] for _ in range(n_workers)] + for i, p in enumerate(picked): + buckets[i % n_workers].append(p) + jobs = [(i, b, dataset, "text", tag) for i, b in enumerate(buckets) if b] + print(f"{len(picked)} shards over {len(jobs)} workers (need {need/1e6:.0f}M tok)") + + t0 = time.time() + results = sorted(_tokenize_shards.map(jobs), key=lambda r: r["idx"]) + total = sum(r["tokens"] for r in results) + print(f"tokenized {total/1e6:.0f}M tok in {(time.time()-t0)/60:.1f}m across workers") + if total < need: + print(f"WARNING: only {total/1e6:.0f}M < need {need/1e6:.0f}M; increase headroom") + + # Concatenate parts in worker order into train.bin, then val.bin. + VOL.reload() + parts = sorted((outdir / "parts").glob("part_*.bin")) + train_target = int(target_tokens) + written_train = 0 + ftrain = open(outdir / "train.bin", "wb") + val_arrs, val_have = [], 0 + for pf in parts: + a = np.fromfile(pf, dtype=np.uint32) + if written_train < train_target: + take = min(a.size, train_target - written_train) + a[:take].tofile(ftrain); written_train += take + rest = a[take:] + else: + rest = a + if rest.size and val_have < val_tokens: + need_v = val_tokens - val_have + val_arrs.append(rest[:need_v]); val_have += min(rest.size, need_v) + ftrain.close() + val_ids = np.concatenate(val_arrs) if val_arrs else np.empty(0, dtype=np.uint32) + val_ids.tofile(outdir / "val.bin") + + tok = AutoTokenizer.from_pretrained(TOKENIZER) + val_text = tok.decode(val_ids.tolist(), skip_special_tokens=True) + val_bytes = len(val_text.encode("utf-8")) + for pf in parts: # free part files + pf.unlink() + manifest = { + "tag": tag, "dataset": dataset, "tokenizer": TOKENIZER, + "shards_used": len(picked), "n_workers": len(jobs), "seed": seed, + "train_tokens": int(written_train), "val_tokens": int(val_ids.size), + "val_bytes": int(val_bytes), + "val_bytes_per_token": val_bytes / max(1, val_ids.size), + "prep_minutes": round((time.time() - t0) / 60, 1), + } + (outdir / "manifest.json").write_text(json.dumps(manifest, indent=2)) + VOL.commit() + print("prep_data_parallel done:", manifest) + return manifest + + +@app.function(image=image, gpu="H100", memory=32768, timeout=12 * 60 * 60, + volumes={DATA: VOL}) +def train(arm: str = "transformer", target_tokens: int = 20_000_000, + tag: str = "smoke", seq_len: int = 4096, + # micro_batch defaults are ARM-AWARE (0 = auto): the transformer fits + # mb=8 (fp32 logits ~13GB), but the hybrid's fla GDN autotuner needs + # more headroom and OOMs at mb=8, so it uses mb=4. Effective batch is + # held at 128 sequences via grad_accum = 128 // micro_batch. + micro_batch: int = 0, grad_accum: int = 0, + peak_lr: float = 2.0e-3, warmup_frac: float = 0.02, + weight_decay: float = 0.1, beta1: float = 0.9, beta2: float = 0.95, + grad_clip: float = 1.0, log_every: int = 25, seed: int = 1337) -> dict: + """Train one arm for ~target_tokens tokens with WSD-S; log held-out CE.""" + import json + import sys + import time + + import numpy as np + import torch + import torch.nn.functional as F + + sys.path.insert(0, "/root") + from repro import paper_models, wsd + + torch.manual_seed(seed) + dev = "cuda" + ddir = pathlib.Path(DATA) / tag + # memmap the 15GB train stream (read-only) so it isn't loaded into RAM. + train_ids = np.memmap(ddir / "train.bin", dtype=np.uint32, mode="r") + val_ids = np.fromfile(ddir / "val.bin", dtype=np.uint32) # small (~32MB) + manifest = json.loads((ddir / "manifest.json").read_text()) + bytes_per_tok = float(manifest["val_bytes_per_token"]) + + class Cfg: + vocab_size = VOCAB_PADDED + block_size = seq_len + eval_block_size = seq_len + + build = paper_models.build_transformer if arm == "transformer" else paper_models.build_hybrid + model = build(Cfg()).to(dev) + n_params = sum(p.numel() for p in model.parameters() if p.requires_grad) + n_non_emb = n_params - model.wte.weight.numel() + + if not micro_batch: + micro_batch = 4 if arm == "hybrid" else 8 + if not grad_accum: + grad_accum = max(1, 128 // micro_batch) + tokens_per_step = micro_batch * grad_accum * seq_len + total_steps = max(1, int(target_tokens) // tokens_per_step) + warmup_steps = max(1, int(warmup_frac * total_steps)) + + decay, no_decay = [], [] + for p in model.parameters(): + (decay if p.dim() >= 2 else no_decay).append(p) + opt = torch.optim.AdamW( + [{"params": decay, "weight_decay": weight_decay}, + {"params": no_decay, "weight_decay": 0.0}], + lr=peak_lr, betas=(beta1, beta2), + ) + + rng = np.random.default_rng(seed) + hi = train_ids.size - seq_len - 1 + + def get_batch(nseq): + ix = rng.integers(0, hi, size=nseq) + x = np.stack([train_ids[i:i + seq_len] for i in ix]).astype(np.int64) + y = np.stack([train_ids[i + 1:i + 1 + seq_len] for i in ix]).astype(np.int64) + return (torch.from_numpy(x).pin_memory().to(dev, non_blocking=True), + torch.from_numpy(y).pin_memory().to(dev, non_blocking=True)) + + @torch.no_grad() + def eval_ce(max_windows=24): + model.eval() + n = min(val_ids.size - 1, max_windows * seq_len) + nw = n // seq_len + tot_ce, tot_tok = 0.0, 0 + for w in range(nw): + s = w * seq_len + x = torch.from_numpy(val_ids[s:s + seq_len][None].astype(np.int64)).to(dev) + y = torch.from_numpy(val_ids[s + 1:s + 1 + seq_len][None].astype(np.int64)).to(dev) + logits = model(x).float() + tot_ce += float(F.cross_entropy(logits.reshape(-1, logits.size(-1)), + y.reshape(-1), reduction="sum")) + tot_tok += y.numel() + model.train() + mean_ce = tot_ce / max(1, tot_tok) + val_bpb = tot_ce / (tot_tok * bytes_per_tok * np.log(2)) + return mean_ce, val_bpb + + logpath = ddir / f"train_{arm}.log" + log_lines = [] + + def log(msg): + print(msg, flush=True) + log_lines.append(msg) + + log(f"[{arm}] non_emb={n_non_emb/1e6:.1f}M total={n_params/1e6:.1f}M " + f"tokens/step={tokens_per_step} total_steps={total_steps} warmup={warmup_steps} " + f"target_tokens={target_tokens}") + + t0 = time.time() + for step in range(total_steps): + lr = wsd.lr_wsd(step, peak_lr=peak_lr, total_steps=total_steps, + warmup_steps=warmup_steps) + for pg in opt.param_groups: + pg["lr"] = lr + opt.zero_grad(set_to_none=True) + loss_acc = 0.0 + for _ in range(grad_accum): + x, y = get_batch(micro_batch) + with torch.autocast("cuda", dtype=torch.bfloat16): + logits = model(x) + loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)).float(), + y.reshape(-1)) / grad_accum + loss.backward() + loss_acc += float(loss) * grad_accum + torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip) + opt.step() + + if step % log_every == 0 or step == total_steps - 1: + toks = (step + 1) * tokens_per_step + dt = time.time() - t0 + tps = toks / max(1e-9, dt) + log(f"[{arm}] step {step:5d}/{total_steps} lr={lr:.2e} " + f"train_ce={loss_acc/grad_accum:.4f} tok={toks/1e6:.1f}M " + f"{tps/1e3:.0f}k tok/s wall={dt/60:.1f}m") + # Flush progress to the Volume so DETACHED runs are observable + # without a live client (see modal run --detach). + logpath.write_text("\n".join(log_lines)) + VOL.commit() + + mean_ce, val_bpb = eval_ce() + wall = time.time() - t0 + ckpt = ddir / f"ckpt_{arm}.pt" + torch.save({"model": model.state_dict(), "arm": arm, "cfg_vocab": VOCAB_PADDED, + "seq_len": seq_len, "target_tokens": target_tokens}, ckpt) + logpath.write_text("\n".join(log_lines)) + VOL.commit() + result = { + "arm": arm, "n_non_emb_M": round(n_non_emb / 1e6, 1), + "total_steps": total_steps, "tokens_trained_M": round(total_steps * tokens_per_step / 1e6, 1), + "final_train_ce": round(loss_acc / grad_accum, 4), + "val_ce_nats": round(mean_ce, 4), "val_bpb_lm": round(val_bpb, 4), + "wall_min": round(wall / 60, 1), + "tok_per_s_k": round(total_steps * tokens_per_step / wall / 1e3, 1), + "gpu": torch.cuda.get_device_name(0), + } + log(f"[{arm}] DONE {result}") + return result + + +@app.local_entrypoint() +def main(action: str = "prep", arm: str = "transformer", + target_tokens: int = 20_000_000, tag: str = "smoke"): + if action == "prep": + print(prep_data.remote(target_tokens=target_tokens, tag=tag)) + elif action == "prep_par": + print(prep_data_parallel.remote(target_tokens=target_tokens, tag=tag)) + elif action == "train": + print(train.remote(arm=arm, target_tokens=target_tokens, tag=tag)) + else: + raise SystemExit(f"unknown action {action!r}") diff --git a/2.0/problems/nanoslm_hybrid_arch_design/repro/paper_models.py b/2.0/problems/nanoslm_hybrid_arch_design/repro/paper_models.py new file mode 100644 index 000000000..5ac59d4b4 --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/repro/paper_models.py @@ -0,0 +1,344 @@ +"""Paper-faithful 190M ablation models for reproducing Olmo Hybrid Table 5. + +These are DELIBERATELY NOT the benchmark's ``baseline_model.py`` / ``reference.py``. +The Table 5 numbers (Transformer 0.950, GDN-3:1 0.891 Base-Easy BPB @ 190M) were +produced by the scaling-ladder ablation models described in arXiv:2604.03444 +Section D.3 + Table 22, which differ from the benchmark variants in three ways: + + * NORM PLACEMENT: pre-norm (RMSNorm BEFORE each sub-layer), per D.3 + ("RMSNorm applied before each sub-layer (pre-norm)"). The benchmark's + baseline_model.py uses reordered_norm (norm AFTER the residual branch). + * BASELINE ATTENTION: plain full multi-head causal attention. D.3 describes + no sliding window for the ablation ladder; the benchmark baseline uses a + 9:3 SWA (4096-window) pattern. + * HYBRID SIZING: the ablation hybrid keeps the SAME d=768, h=12, l=12 as the + transformer (non-embedding params grow 190M -> 254M). The benchmark's + reference.py instead param-matches via REMOVE_HEADS=1 (d704, h11). + +Shape (Table 22, 190M column): d=768, h=12, l=12, head_dim=64. +MLP: SwiGLU, hidden = round_up_256(1.5 * 8d/3) = 3072 (D.3). +Untied embeddings; RoPE theta 500_000; QK-norm; RMSNorm eps 1e-6. +Vocab: dolma2, 100352 padded (D.3). BPB is vocab-padding-invariant, but we use +100352 to match the paper's embedding shape exactly. + +GDN (D.3 + Appendix A.1): head sizing proportional to an attention head -- +d_k = 3/4 * head_dim = 48, d_v = 2 * d_k = 96 (expand_v = 2.0), allow_neg_eigval += True, use_gate = True, use_short_conv = True. fla's GatedDeltaNet is REQUIRED +on CUDA (a hand-written mixer is unfused and would confound any timing); a +chunkwise CPU fallback exists only so param counts / wiring are checkable off-GPU. + +Placement (D.3): "every r-th layer is a full transformer block" with r=4 (3:1), +"and we additionally enforce the final layer be an attention layer". For l=12 +that is attention at layers 3, 7, 11 (0-indexed), i.e. pattern (gdn,gdn,gdn,attn). +""" + +from __future__ import annotations + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +# --------------------------------------------------------------------------- # +# Shared shape (Table 22, 190M). +# --------------------------------------------------------------------------- # +N_LAYER = 12 +N_HEAD = 12 +N_EMBD = 768 +HEAD_DIM = N_EMBD // N_HEAD # 64 +HIDDEN_SIZE = 3072 # round_up_256(1.5 * 8*768/3) = 3072 +ROPE_THETA = 500_000.0 +NORM_EPS = 1e-6 + +# GDN head sizing (D.3): hGDN = ceil_128(0.75 * d/h), where ceil_128 rounds UP +# to the nearest multiple of 128. For 190M: 0.75 * 768/12 = 48 -> 128. Key dim +# = h * hGDN, value dim = h * 2*hGDN (expand_v = 2). This ceil-to-128 is what +# reference.py omitted (it used 48), and it is why the ablation hybrid is 254M +# non-embed, not ~201M -- the GDN mixer is ~9.5M/layer, not ~3.6M. +def _ceil_mult(x: int, m: int) -> int: + return ((x + m - 1) // m) * m + +GDN_HEAD_DIM = _ceil_mult(int(0.75 * HEAD_DIM), 128) # ceil_128(48) = 128 +GDN_EXPAND_V = 2.0 +GDN_HEAD_V_DIM = int(GDN_HEAD_DIM * GDN_EXPAND_V) # 256 +GDN_KEY_DIM = N_HEAD * GDN_HEAD_DIM # 1536 +GDN_VALUE_DIM = N_HEAD * GDN_HEAD_V_DIM # 3072 +GDN_CONV_SIZE = 4 + +# r=4 (3:1) interleave, final layer forced to attention -> attn at 3, 7, 11. +PATTERN = ("gdn", "gdn", "gdn", "attn") +assert N_LAYER % len(PATTERN) == 0 + + +class RMSNorm(nn.Module): + def __init__(self, d: int, eps: float = NORM_EPS): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(d)) + + def forward(self, x): + dt = x.dtype + xf = x.float() + n = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps) + return n.to(dt) * self.weight + + +def _rope(x, theta: float = ROPE_THETA): + _, _, T, D = x.shape + half = D // 2 + freq = theta ** (-torch.arange(0, half, device=x.device, dtype=torch.float32) / half) + ang = torch.arange(T, device=x.device, dtype=torch.float32)[:, None] * freq[None, :] + cos, sin = ang.cos()[None, None], ang.sin()[None, None] + x1, x2 = x[..., :half].float(), x[..., half:].float() + return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1).to(x.dtype) + + +class Attention(nn.Module): + """Full causal MHA with QK-norm (full-width, pre-reshape) + RoPE. No SWA.""" + + def __init__(self, n_embd: int = N_EMBD, n_head: int = N_HEAD): + super().__init__() + assert n_embd % n_head == 0 + self.n_head, self.n_embd = n_head, n_embd + self.head_dim = n_embd // n_head + self.w_q = nn.Linear(n_embd, n_embd, bias=False) + self.w_k = nn.Linear(n_embd, n_embd, bias=False) + self.w_v = nn.Linear(n_embd, n_embd, bias=False) + self.w_out = nn.Linear(n_embd, n_embd, bias=False) + self.q_norm = RMSNorm(n_embd) + self.k_norm = RMSNorm(n_embd) + + def forward(self, x): + B, T, C = x.shape + q, k, v = self.w_q(x), self.w_k(x), self.w_v(x) + q, k = self.q_norm(q), self.k_norm(k) + q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2) + k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2) + v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2) + q, k = _rope(q), _rope(k) + y = F.scaled_dot_product_attention(q, k, v, is_causal=True) + return self.w_out(y.transpose(1, 2).contiguous().view(B, T, C)) + + +class _ChunkedDeltaNet(nn.Module): + """CPU-only chunkwise fallback, param-identical to fla's GatedDeltaNet. + + Numerically approximate (gated linear attention, no delta-rule inverse); + exists ONLY so param counts / wiring are checkable off-GPU. Never used on + CUDA -- _make_gdn raises there if fla is missing. + """ + + def __init__(self, chunk: int = 512): + super().__init__() + self.chunk = chunk + d, H = N_EMBD, N_HEAD + self.q_proj = nn.Linear(d, GDN_KEY_DIM, bias=False) + self.k_proj = nn.Linear(d, GDN_KEY_DIM, bias=False) + self.v_proj = nn.Linear(d, GDN_VALUE_DIM, bias=False) + self.a_proj = nn.Linear(d, H, bias=False) + self.b_proj = nn.Linear(d, H, bias=False) + self.g_proj = nn.Linear(d, GDN_VALUE_DIM, bias=False) + self.o_proj = nn.Linear(GDN_VALUE_DIM, d, bias=False) + self.A_log = nn.Parameter(torch.zeros(H)) + self.dt_bias = nn.Parameter(torch.zeros(H)) + self.q_conv1d = nn.Parameter(torch.zeros(GDN_KEY_DIM, 1, GDN_CONV_SIZE)) + self.k_conv1d = nn.Parameter(torch.zeros(GDN_KEY_DIM, 1, GDN_CONV_SIZE)) + self.v_conv1d = nn.Parameter(torch.zeros(GDN_VALUE_DIM, 1, GDN_CONV_SIZE)) + self.o_norm = nn.Parameter(torch.ones(GDN_HEAD_V_DIM)) + + @staticmethod + def _short_conv(x, w): + C, K = w.shape[0], w.shape[2] + xt = F.pad(x.transpose(1, 2), (K - 1, 0)) + return F.conv1d(xt, w, groups=C).transpose(1, 2) + + def forward(self, x): + B, T, _ = x.shape + H, DK, DV, S = N_HEAD, GDN_HEAD_DIM, GDN_HEAD_V_DIM, self.chunk + q = F.silu(self._short_conv(self.q_proj(x), self.q_conv1d)) + k = F.silu(self._short_conv(self.k_proj(x), self.k_conv1d)) + v = self._short_conv(self.v_proj(x), self.v_conv1d) + q = q.view(B, T, H, DK).transpose(1, 2) + k = k.view(B, T, H, DK).transpose(1, 2) + v = v.view(B, T, H, DV).transpose(1, 2) + dt = F.softplus(self.a_proj(x) + self.dt_bias) + g = torch.exp(-torch.exp(self.A_log) * dt).permute(0, 2, 1).unsqueeze(-1) + beta = torch.sigmoid(self.b_proj(x)).permute(0, 2, 1).unsqueeze(-1) + v = v * beta + pad = (S - T % S) % S + if pad: + q, k, v = (F.pad(t, (0, 0, 0, pad)) for t in (q, k, v)) + g = F.pad(g, (0, 0, 0, pad), value=1.0) + nC = (T + pad) // S + rs = lambda t: t.view(B, H, nC, S, -1) # noqa: E731 + qc, kc, vc, gc = rs(q), rs(k), rs(v), rs(g) + causal = torch.tril(torch.ones(S, S, device=x.device, dtype=torch.bool)) + intra = (qc @ kc.transpose(-1, -2)).masked_fill(~causal, 0.0) @ vc + state = torch.zeros(B, H, DK, DV, device=x.device, dtype=x.dtype) + outs = [] + for c in range(nC): + outs.append(intra[:, :, c] + qc[:, :, c] @ state) + decay = gc[:, :, c].prod(dim=2, keepdim=True) + state = state * decay + kc[:, :, c].transpose(-1, -2) @ (vc[:, :, c] * gc[:, :, c]) + y = torch.stack(outs, dim=2).view(B, H, T + pad, DV)[:, :, :T] + yf = y.float() + y = (yf * torch.rsqrt(yf.pow(2).mean(-1, keepdim=True) + 1e-5)).to(x.dtype) + y = y * self.o_norm + y = y.transpose(1, 2).contiguous().view(B, T, GDN_VALUE_DIM) + y = y * F.silu(self.g_proj(x)) + return self.o_proj(y) + + +def _make_gdn(): + if not torch.cuda.is_available(): + return _ChunkedDeltaNet() + from fla.layers import GatedDeltaNet # ImportError here is intentional + return GatedDeltaNet( + hidden_size=N_EMBD, + num_heads=N_HEAD, + head_dim=GDN_HEAD_DIM, + expand_v=GDN_EXPAND_V, + use_gate=True, + use_short_conv=True, + allow_neg_eigval=True, + ) + + +class GDNLayer(nn.Module): + def __init__(self): + super().__init__() + self.impl = _make_gdn() + + def forward(self, x): + out = self.impl(x) + return out[0] if isinstance(out, tuple) else out + + +class FeedForward(nn.Module): + """SwiGLU, hidden 3072, bias=False.""" + + def __init__(self, n_embd: int = N_EMBD, hidden: int = HIDDEN_SIZE): + super().__init__() + self.w1 = nn.Linear(n_embd, hidden, bias=False) + self.w3 = nn.Linear(n_embd, hidden, bias=False) + self.w2 = nn.Linear(hidden, n_embd, bias=False) + + def forward(self, x): + return self.w2(F.silu(self.w1(x)) * self.w3(x)) + + +class PreNormBlock(nn.Module): + """Pre-norm (D.3): norm BEFORE each sub-layer, residual add after.""" + + def __init__(self, kind: str): + super().__init__() + self.mixer = Attention() if kind == "attn" else GDNLayer() + self.mixer_norm = RMSNorm(N_EMBD) + self.ffn = FeedForward() + self.ffn_norm = RMSNorm(N_EMBD) + + def forward(self, x): + x = x + self.mixer(self.mixer_norm(x)) + return x + self.ffn(self.ffn_norm(x)) + + +class PaperModel(nn.Module): + """Paper-faithful ablation model. ``hybrid=False`` -> pure transformer.""" + + def __init__(self, config, hybrid: bool): + super().__init__() + self.block_size = config.block_size + self.wte = nn.Embedding(config.vocab_size, N_EMBD) + if hybrid: + kinds = [PATTERN[i % len(PATTERN)] for i in range(N_LAYER)] + else: + kinds = ["attn"] * N_LAYER + self.blocks = nn.ModuleList([PreNormBlock(k) for k in kinds]) + self.norm_f = RMSNorm(N_EMBD) + self.lm_head = nn.Linear(N_EMBD, config.vocab_size, bias=False) # untied + self.apply(self._init) + for name, p in self.named_parameters(): + if name.endswith("w_out.weight") or name.endswith("w2.weight"): + nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * N_LAYER)) + + @staticmethod + def _init(m): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, mean=0.0, std=0.02) + if m.bias is not None: + nn.init.zeros_(m.bias) + elif isinstance(m, nn.Embedding): + nn.init.normal_(m.weight, mean=0.0, std=0.02) + + def forward(self, idx): + x = self.wte(idx) + for block in self.blocks: + x = block(x) + x = self.norm_f(x) + return self.lm_head(x) + + +def build_transformer(config): + return PaperModel(config, hybrid=False) + + +def build_hybrid(config): + return PaperModel(config, hybrid=True) + + +# --------------------------------------------------------------------------- # +# Analytic param counts (Table 22 convention: "non-embedding" = blocks + LM +# head, i.e. everything except the INPUT embedding table). At 190M this should +# reproduce Transformer 190 and GDN-3:1 254 (millions). +# --------------------------------------------------------------------------- # +def _attn_params(): + d = N_EMBD + return 4 * d * d + 2 * d # q/k/v/out + q/k norm + + +def _gdn_params(): + d, H = N_EMBD, N_HEAD + return ( + 2 * d * GDN_KEY_DIM # q_proj, k_proj + + 2 * d * GDN_VALUE_DIM # v_proj, g_proj + + 2 * d * H # a_proj, b_proj + + GDN_VALUE_DIM * d # o_proj + + 2 * H # A_log, dt_bias + + GDN_CONV_SIZE * (2 * GDN_KEY_DIM + GDN_VALUE_DIM) # short convs + + GDN_HEAD_V_DIM # o_norm + ) + + +def analytic_non_embed(hybrid: bool, vocab_size: int) -> int: + d = N_EMBD + block_norms = 2 * d + ffn = 3 * d * HIDDEN_SIZE + n_attn = sum(1 for i in range(N_LAYER) if PATTERN[i % len(PATTERN)] == "attn") if hybrid else N_LAYER + n_gdn = N_LAYER - n_attn + blocks = n_attn * (_attn_params() + block_norms + ffn) + n_gdn * (_gdn_params() + block_norms + ffn) + lm_head = vocab_size * d # counted as non-embedding (Table 22 convention) + return blocks + d + lm_head # + final norm + + +if __name__ == "__main__": + # Local self-check (CPU): param counts vs Table 22 (190 / 254 million). + class _Cfg: + vocab_size = 100352 + block_size = 256 + eval_block_size = 256 + + for hybrid, want in ((False, 190), (True, 254)): + m = PaperModel(_Cfg(), hybrid=hybrid) + total = sum(p.numel() for p in m.parameters() if p.requires_grad) + non_emb_built = total - m.wte.weight.numel() + non_emb_analytic = analytic_non_embed(hybrid, _Cfg.vocab_size) + assert non_emb_built == non_emb_analytic, (hybrid, non_emb_built, non_emb_analytic) + name = "hybrid GDN-3:1" if hybrid else "transformer" + print(f"{name:16s} non-embed(Table22 conv)={non_emb_built/1e6:6.1f}M " + f"(want ~{want}M) total={total/1e6:6.1f}M") + # tiny forward + idx = torch.randint(0, _Cfg.vocab_size, (2, 16)) + out = m(idx) + assert out.shape == (2, 16, _Cfg.vocab_size), out.shape + print("forward + param-count self-check OK") diff --git a/2.0/problems/nanoslm_hybrid_arch_design/repro/wsd.py b/2.0/problems/nanoslm_hybrid_arch_design/repro/wsd.py new file mode 100644 index 000000000..f09539c02 --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/repro/wsd.py @@ -0,0 +1,91 @@ +"""WSD-S learning-rate schedule (warmup - stable - decay with periodic resets). + +Paper (arXiv:2604.03444, Section 4.1): a WSD-S schedule that is token-agnostic +-- the LR at step t does not depend on the total token budget T. "We decay the +learning rate for 5% of the training tokens to 0 at each of the five Chinchilla +factors to obtain the trained checkpoint for that factor." Between factors the +LR resets back to the stable peak (the '-S' = periodic reset), so a single long +run yields checkpoints at 0.5x/1x/2x/4x/8x Chinchilla. + +Refs: Hu et al. 2024 (MiniCPM WSD); Wen et al. 2025b (river-valley view of WSD). + +Two schedule shapes here: + * ``lr_wsd``: single warmup -> stable -> terminal 5% decay-to-0. Used for a + fixed-budget run (e.g. the 1x-Chinchilla smoke). + * ``lr_wsd_s``: the periodic-reset variant. Given the ordered list of + decay-endpoint steps (the Chinchilla factors), it produces a checkpoint at + each: warmup once, hold at peak, and for the last ``decay_frac`` of the span + BEFORE each endpoint, decay to 0; immediately after an endpoint, jump back + to peak. Evaluate/snapshot AT each endpoint step. + +Decay shape: the WSD "river-valley" analysis (Wen et al.) favours a 1-sqrt decay +over linear; we use 1 - sqrt(progress), which spends more steps near the low LR. +Set ``decay_shape='linear'`` for a plain linear ramp to 0. +""" + +from __future__ import annotations + +import math + + +def _decay_mult(progress: float, shape: str) -> float: + """LR multiplier in [0,1] as decay progresses 0->1 (1 at start, 0 at end).""" + progress = min(1.0, max(0.0, progress)) + if shape == "linear": + return 1.0 - progress + # '1-sqrt' (river-valley): stays high then drops fast near the end. + return 1.0 - math.sqrt(progress) + + +def lr_wsd(step: int, *, peak_lr: float, total_steps: int, warmup_steps: int, + decay_frac: float = 0.05, decay_shape: str = "1-sqrt") -> float: + """Single warmup -> stable -> terminal decay-to-0 over the last ``decay_frac``.""" + if step < warmup_steps: + return peak_lr * (step + 1) / max(1, warmup_steps) + decay_steps = max(1, int(round(decay_frac * total_steps))) + decay_start = total_steps - decay_steps + if step < decay_start: + return peak_lr + progress = (step - decay_start) / decay_steps + return peak_lr * _decay_mult(progress, decay_shape) + + +def lr_wsd_s(step: int, *, peak_lr: float, endpoints: list[int], warmup_steps: int, + decay_frac: float = 0.05, decay_shape: str = "1-sqrt") -> float: + """Periodic-reset WSD-S. ``endpoints`` = sorted decay-endpoint steps. + + For each consecutive span (prev_endpoint, endpoint], the last ``decay_frac`` + of the span decays to 0; the rest holds at peak (after the one-time warmup). + Snapshot the model exactly AT each endpoint step to get that factor's ckpt. + """ + if step < warmup_steps: + return peak_lr * (step + 1) / max(1, warmup_steps) + prev = 0 + for end in endpoints: + if step <= end: + span = end - prev + decay_steps = max(1, int(round(decay_frac * span))) + decay_start = end - decay_steps + if step < decay_start: + return peak_lr + progress = (step - decay_start) / decay_steps + return peak_lr * _decay_mult(progress, decay_shape) + prev = end + return 0.0 # past the final endpoint + + +if __name__ == "__main__": + # Sanity checks. + P, T, W = 2e-3, 1000, 50 + assert abs(lr_wsd(0, peak_lr=P, total_steps=T, warmup_steps=W) - P / W) < 1e-12 + assert abs(lr_wsd(W, peak_lr=P, total_steps=T, warmup_steps=W) - P) < 1e-12 # stable + assert abs(lr_wsd(940, peak_lr=P, total_steps=T, warmup_steps=W) - P) < 1e-9 # still stable (decay=last 50) + assert lr_wsd(999, peak_lr=P, total_steps=T, warmup_steps=W) < P # decaying + assert abs(lr_wsd(T, peak_lr=P, total_steps=T, warmup_steps=W)) < 1e-9 # ~0 at end + # WSD-S: peak restored right after an endpoint. + eps = [500, 1000] + assert abs(lr_wsd_s(600, peak_lr=P, endpoints=eps, warmup_steps=W) - P) < 1e-9 # reset to peak after 500 + assert lr_wsd_s(500, peak_lr=P, endpoints=eps, warmup_steps=W) < P # decayed at endpoint 500 + print("WSD / WSD-S schedule self-check OK") + for s in (0, 25, 50, 500, 900, 950, 975, 1000): + print(f" step {s:4d} lr={lr_wsd(s, peak_lr=P, total_steps=T, warmup_steps=W):.3e}") From a85e1866af8cf78fcd7ceb5fa330766534f6245a Mon Sep 17 00:00:00 2001 From: Sijia Liu Date: Mon, 20 Jul 2026 21:32:48 -0700 Subject: [PATCH 2/5] updated DESIGN.md and README.md --- 2.0/README.md | 24 ++++ .../nanoslm_hybrid_arch_design/DESIGN.md | 66 ++-------- .../nanoslm_hybrid_arch_design/PR_SUMMARY.md | 119 ------------------ 3 files changed, 31 insertions(+), 178 deletions(-) delete mode 100644 2.0/problems/nanoslm_hybrid_arch_design/PR_SUMMARY.md diff --git a/2.0/README.md b/2.0/README.md index 628e3f314..516d04f8d 100644 --- a/2.0/README.md +++ b/2.0/README.md @@ -118,3 +118,27 @@ tail-frame (≥60) LPIPS-vs-GT over the unpatched baseline, gated by a wall-cloc guardrail (so drift can't be bought with more compute). A history-stabilization reference reliably beats baseline (validated t≈2.5/22 clips); beating it substantially is the open challenge. + +## NanoSLM Hybrid Architecture Design + +Architecture design as a scored task, framed on *Olmo Hybrid: From Theory to +Practice and Back* (arXiv:2604.03444). Its problem ID is +`nanoslm_hybrid_arch_design`. +Agents submit a single `/app/model.py` defining `build_model(config)` (or a +`class NanoSLM`), with full freedom over the model definition. The judge trains +it from scratch under a fixed wall-clock budget `T` on one H100 — dolma2-BPE +FineWeb-Edu, gradient accumulation to an effective batch of 32 — and scores the +**absolute** reduction in held-out **bits-per-byte** (`val_bpb`, normalized by +bytes so it is tokenizer-independent) against a CRN-paired, locked pure-attention +`olmo3_190M` baseline. Both arms tie their embeddings and sit under a hard +parameter cap; there is no iso-parameter guardrail, so the lever is efficiency +under the clock — a cheaper mixer completes more optimizer steps within `T`. +Evaluation is always at an 8192-token context, but the agent may lower its +*training* context via a module-level `BLOCK_SIZE`, trading steps against length +extrapolation. The 3:1 Gated-DeltaNet/attention hybrid (the Olmo Hybrid recipe at +190M) is the reference floor; beating it is the open challenge. + +Note: the open questions are the ones the paper settles at scale but leaves open +at 190M under a tight, parameter-matched budget — the recurrent mixer (GDN vs. +Mamba2), the placement of attention layers, and the attention-to-recurrence +ratio. diff --git a/2.0/problems/nanoslm_hybrid_arch_design/DESIGN.md b/2.0/problems/nanoslm_hybrid_arch_design/DESIGN.md index da77fb9b7..d6dad5b2f 100644 --- a/2.0/problems/nanoslm_hybrid_arch_design/DESIGN.md +++ b/2.0/problems/nanoslm_hybrid_arch_design/DESIGN.md @@ -11,13 +11,11 @@ wall-clock budget on a single H100**. The agent submits one file, `model.py` wall-clock `T`, and scores `val_bpb` against a **locked baseline architecture** trained under the identical budget. -## 1.1 Where the starting setup comes from - The initial configuration is taken from **Olmo Hybrid: From Theory to Practice and Back** (Merrill, Li, Romero et al., arXiv:2604.03444). That work shows theoretically that hybrids express capabilities beyond both transformers and linear RNNs, then validates it at 7B by replacing sliding-window attention -layers with Gated DeltaNet layers, outperforming the comparable Olmo 3 baseline. +layers with Gated DeltaNet (GDN) layers, outperforming the comparable Olmo 3 baseline. Concretely inherited from it and from the OLMo-core implementation: @@ -26,9 +24,8 @@ Concretely inherited from it and from the OLMo-core implementation: the agent's starting point; * the practice of **rebalancing head count / width so the hybrid is parameter-matched** against the pure-attention baseline - (`REMOVE_HEADS` in the upstream script) — **now applied** as - `REMOVE_HEADS = 1` (d_model 768 -> 704, n_heads 12 -> 11); - * **dolma2** tokenization and the ~190M OLMo3 shape (d=768, L=12, H=12). + (`REMOVE_HEADS` in the upstream script); + * **dolma2** tokenization and the ~190M OLMo3 shape. The scientific questions this task poses are the ones that paper answers at scale but leaves open at 190M under a tight, parameter-matched, wall-clock @@ -39,60 +36,11 @@ an existence proof, not a ceiling. It doubles as the repo-required reference *solution* and the CI gate: `reference.py` must beat the hidden, score-0 `baseline_model.py` (the pure-attention `olmo3_190M`, judge-side only). -The optimizer, LR schedule, data, tokenizer, EVALUATION context (8192) and +The optimizer, LR schedule, data, tokenizer, evaluation context and budget are all fixed by the judge. The agent changes the architecture — and the -TRAINING context length, which trades optimizer steps against length +training context length, which trades optimizer steps against length extrapolation. -## 1.2 The three research questions - -The task is organized around the three architecture questions the Olmo Hybrid -paper settles at 7B and ablates on a 60M–1B ladder (its Table 5, `val_bpb` -averaged over Math/Code/QA; lower is better). The agent re-derives them at 190M, -**from the pure-attention baseline up**, under a fixed wall-clock budget and a -hard parameter cap — conditions under which the paper's answers are *priors, not -conclusions* (see the caveat below). - -1. **RNN architecture — GDN vs. Mamba2.** Which linear-recurrent mixer makes the - best hybrid? Paper prior @190M: Gated DeltaNet (3:1) reaches 0.891 bpb vs. - Mamba2 (3:1) at 0.921, and pure GDN (0.895) beats pure Mamba2 (0.941). GDN is - the paper's pick and is what `reference.py` ships; whether that margin - survives parameter-matching at this scale is open. - -2. **Layer placement — interleaved vs. concentrated.** Given a fixed count of - attention layers, where do they go? Paper prior @190M: uniformly interleaved - (0.891) beats concentrating them in the middle of the stack (0.899). The - paper's argument is structural — concentrating attention forces all global - information through one bottleneck, whereas interleaving maximizes the number - of alternations between mixer types, which its expressivity theory (Theorems - 1 and 3) predicts should help. Clustering attention early (context in) vs. - late (read-out) is unexplored and on the table. - -3. **Attention ratio — how much attention is enough?** What fraction of layers - should stay full attention? Paper prior @190M: 1:1 / 50% (0.896), 3:1 / 25% - (0.891), 7:1 / 12.5% (0.892) — nearly a wash at this scale (a 0.001 bpb - spread), with the paper selecting 3:1 because it wins at 600M–1B, not at - 190M. Fewer attention layers also buy more optimizer steps under the - wall-clock budget (§5), so here the ratio trades context against *both* - parameters and steps. - -**Why the paper's answers are priors, not conclusions here.** The paper's 190M -hybrids are *not* parameter-matched: its GDN (3:1) carries ~254M non-embedding -params against the transformer's ~190M — **+34%** — and the paper flags exactly -this (Table 22: "per-size comparisons should be interpreted with care, as -architectures at the same nominal scale differ in actual parameter count"). This -task instead matches non-embedding params to ≈±2% (the `REMOVE_HEADS` -compensation in `reference.py`, measured at −2.19%) and enforces a hard -`param_cap`, so the portion of the paper's measured gain that was *bought with -parameters* is unavailable. The scored question is which of these three choices -still pays once the parameters are held fixed and the clock — not the token -count — is the budget. - -Beyond the three headline axes, `reference.py` enumerates the finer knobs each -one interacts with: recurrent state size (`expand_v`, `num_v_heads` — cheap at -8192 since it does not grow with sequence length), non-uniform layers, and the -rest of the block (norm placement, gating, MLP ratio, head count, embedding -tying). ## 2. Metric: held-out bits-per-byte (`val_bpb`) @@ -112,7 +60,7 @@ over the locked baseline. Under a fixed wall-clock budget the frontier is genuinely architectural. Because the baseline is already a tuned `olmo3_190M`, the easy block-level wins (RMSNorm, rotary, QK-norm, SwiGLU) are **already in it**; what is left on the table is the -hybrid direction — the mixer, ratio, placement, and state size of §1.2. +hybrid direction — the mixer, RNN architecture choices, layer placement, attention ratio, and training context length, etc. ## 4. Submission surface & interface @@ -149,7 +97,7 @@ context is locked in `harness/data.py::val_windows`. ## 5. Iso-wallclock protocol (the anti-gaming axis) - The judge trains the baseline and the submission for the **same fixed - wall-clock `T`** on the **same H100**, from the **same seed and data order** + wall-clock `T`** on the same H100, from the **same seed and data order** (common random numbers), then evaluates both on the **same hidden held-out validation set**. The scored (final/verifier) run measures baseline + submission back-to-back in one job/GPU/process (no cache), mirroring diff --git a/2.0/problems/nanoslm_hybrid_arch_design/PR_SUMMARY.md b/2.0/problems/nanoslm_hybrid_arch_design/PR_SUMMARY.md deleted file mode 100644 index 0bf7ecfe2..000000000 --- a/2.0/problems/nanoslm_hybrid_arch_design/PR_SUMMARY.md +++ /dev/null @@ -1,119 +0,0 @@ -# PR: add `nanoslm_hybrid_arch_design` (Frontier-CS 2.0) - -## What - -A new open-ended 2.0 problem: design a **hybrid language-model architecture** -(attention + linear-recurrent mixers) that reaches the lowest held-out -**bits-per-byte (`val_bpb`)** when trained from scratch under a **fixed -wall-clock budget on a single H100**. The agent submits one `model.py` (full -architecture freedom); the hidden judge trains it under a locked dolma2-BPE -recipe and scores `val_bpb` against a locked baseline architecture -(pure-attention `olmo3_190M`) trained under the identical budget. - -Adapts the Genesys system ("Language Modeling by Language Models", -arXiv:2506.20249) into a bounded, ungameable, single-objective task — the -hybrid architecture discovery starting from Olmo Hybrid (arXiv:2604.03444). Structurally mirrors -`nanowm_rollout_speedup` / `vllm_llm_serving_optimization`: patch/file -submission, Modal H100, CPU judge, latency/quality-style guardrails, CRN -determinism. - -## Problem id / scoring - -- id: `nanoslm_hybrid_arch_design` (Harbor id `frontier-cs-2-0-nanoslm-hybrid-arch-design`) -- metric: held-out `val_bpb` = total NLL (bits) / held-out BYTES, over the - dolma2 BPE tokenizer (vocab 100278). Normalizing by BYTES (not tokens) is what - makes it tokenizer-independent and ungameable. Lower is better. (`val_ppl` is - reported for readability only and is NO LONGER equal to `2^val_bpb`.) -- score: `clip(100 * (base_bpb - sub_bpb) / bpb_score_scale, 0, 100)` — the - ABSOLUTE bpb gain; baseline and worse-than-baseline score 0. - `bpb_score_scale` is a display convention that maps bpb onto 0-100, NOT a - calibration target (DESIGN.md §8.1); `score_unbounded` stays un-clipped. -- context: EVALUATION is fixed at 8192 for every submission; the TRAINING - context is agent-controlled via a module-level `BLOCK_SIZE` in model.py - (power of two in [256, 8192], default 8192). Shorter buys optimizer steps and - costs length extrapolation. - -## Resource budget - -- `tag: systems`, single Modal **H100**, CPU judge. -- `runtime.timeout_seconds: 21600` (long-horizon agent iteration). -- Per-submission training capped at ≤ 1 h wall-clock (`train_seconds` fixed `T` + - `max_train_seconds` hard abort). -- Hidden judge assets: tokenized train shard, held-out val byte stream, locked - baseline architecture, cached baseline `val_bpb` (keyed by config fingerprint). - -## Anti-gaming - -- Static policy gate (`harness/policy.py`): model.py only, ≤256 KB, denies - env/network/subprocess, pretrained-weight loading, file reads, metric/data/timer - peeking; allows `model.eval()` / `torch.compile()`. -- Dynamic guards (`harness/runner.py`): judge owns the loss (computes CE from the - model's logits on hidden val); trained-from-scratch param-delta check; degenerate - constant-logit check; param cap; OOM → 0. Black-box safe (no traceback/stdout leak). -- Iso-wallclock CRN: baseline + submission trained back-to-back on one GPU, same - seed/data order, scored on the same hidden bytes. - -## Validation (this branch) - -Torch-free layers unit-tested and **full harness smoke-tested on CPU** (no GPU): - -```bash -# torch-free: policy + scoring + fingerprint (17 assertions) -bash evaluate.sh --selftest # -> SELFTEST OK - -# full pipeline on CPU (tiny model/budget), needs torch: -FRONTIER_NANOSLM_SMOKE=1 bash evaluate.sh reference.py -# base_val_bpb=; sub_val_bpb=; abs_bpb_delta=+; score=<...> (reference > baseline) -FRONTIER_NANOSLM_SMOKE=1 python3 evaluator.py # -> guard: not trained; score 0 -python3 evaluator.py # -> policy_rejected; score 0 - -# role split (GPU via Modal from CPU judge; direct H100 only in testing mode): -FRONTIER_NANOSLM_ROLE=final ... evaluator.py reference.py # fresh baseline+submission CRN pair (~2T) -FRONTIER_NANOSLM_ROLE=agent ... evaluator.py reference.py # trains submission only; reuses -# fingerprint-cached baseline (note=iterative(cached-baseline)), ~T per iteration -``` - -Standard 2.0 CLI checks to run in the repo before merge: - -```bash -uv run frontier list 2.0 -uv run frontier show 2.0 nanoslm_hybrid_arch_design -python3 -m py_compile 2.0/problems/nanoslm_hybrid_arch_design/evaluator.py -``` - -## PENDING before ship (see DESIGN.md §3) - -- **[BLOCKER] Single-H100 calibration** of the wall-clock budget `T`, the - model-scale band, and the `val_bpb` noise floor (protocol in DESIGN.md §3). - (`r_target` was on this list; it is now `bpb_score_scale` and is no longer a - calibration gate — DESIGN.md §8.1.) - All such constants are flagged `CALIBRATE` in `harness/settings.py` / `config.yaml`. -- Modal end-to-end run + deployed app name / H100 SKU confirmation (GPU path awaits - maintainer credentials, as with `nanowm`). -- Judge-image bake-asset provenance + data license (FineWeb-Edu, dolma2 BPE); - the held-out val byte ranges are the only hidden component. -- Decide: keep the optimizer locked (current: architecture-only) vs. expose an - optional `configure_optimizers` hook (more autoresearch-faithful). - -## Files - -``` -2.0/problems/nanoslm_hybrid_arch_design/ - config.yaml # systems/H100/Modal, file submission /app/model.py - readme # public task statement (agent-facing) - evaluator.py # evaluate() + prepare() + --selftest; torch-free top-level - evaluate.sh # local CLI wrapper (+ --selftest) - reference.py # reference solution model.py (3:1 GDN hybrid on the olmo3_190M baseline) - DESIGN.md # full design, calibration protocol, open items - PR_SUMMARY.md # this file - harness/ - settings.py # locked config + smoke overrides + config fingerprint (torch-free) - policy.py # static submission policy (torch-free) - scoring.py # val_bpb gain -> [0,100] (torch-free) - model_config.py # ModelConfig contract object (torch-free) - data.py # dolma2-BPE token data (synthetic fallback for smoke) - baseline_model.py # locked baseline transformer (score-0 reference point) - train.py # locked iso-wallclock training loop (harness-owned loss) - eval_ppl.py # held-out val_bpb eval + determinism knobs - runner.py # arm run + CRN pair + dynamic guards -``` From 0f91bc680a443315cbcb44852773dfbc8397318c Mon Sep 17 00:00:00 2001 From: Sijia Liu Date: Tue, 21 Jul 2026 00:11:51 -0700 Subject: [PATCH 3/5] [2.0] nanoslm: fix ctx-8192 OOM (batch_size) + train on real corpus via Modal - settings.py: batch_size 8 -> 2, grad_accum 4 -> 16 (effective batch 32 unchanged). MEASURED: micro-batch 8 at ctx 8192 OOMs an H100 -- the fp32 CE logits are ~26GB and the baseline arm alone tried to allocate 24.48GiB, so the reference could not run under production settings. micro-batch 2 fits both arms (incl. the GDN hybrid) with margin. A fused/chunked CE would let it grow again (noted as a TODO in-file). - modal_app.py: mount the real FineWeb-Edu corpus (Modal Volume nanoslm-corpus) at the data path so the GPU arms train on real tokens instead of data.py's synthetic fallback; add PYTORCH_CUDA_ALLOC_CONF=expandable_segments to reduce allocator fragmentation for the large CE logits. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../harness/modal_app.py | 17 +++++++++++++++-- .../harness/settings.py | 15 +++++++++------ 2 files changed, 24 insertions(+), 8 deletions(-) diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harness/modal_app.py b/2.0/problems/nanoslm_hybrid_arch_design/harness/modal_app.py index 8d19ffe80..61100f952 100644 --- a/2.0/problems/nanoslm_hybrid_arch_design/harness/modal_app.py +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/modal_app.py @@ -106,12 +106,24 @@ def _ver(mod): # fla's Triton kernels JIT-compile and autotune on first use. Give # them a writable cache so a warm container does not pay it twice. "TRITON_CACHE_DIR": "/tmp/triton-cache", + # Reduce allocator fragmentation for the large fp32-logits CE at + # ctx 8192 (paired with the batch_size=2 fix in settings.py). + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", }) ) app = modal.App(APP_NAME) - @app.function(image=image, gpu=GPU, timeout=6 * 60 * 60) + # Real corpus (FineWeb-Edu, dolma2-tokenized) staged on a Modal Volume and + # mounted at the data path harness/settings.py expects, so the GPU arms train + # on REAL tokens instead of data.py's synthetic fallback. Populated once via + # `modal volume put nanoslm-corpus {train,val}.bin manifest.json /`. This is + # the "real shard mounted for calibration" the design leaves as a TODO. + CORPUS_VOL = modal.Volume.from_name("nanoslm-corpus", create_if_missing=True) + REMOTE_DATA = "/opt/nanoslm_arch/data" + + @app.function(image=image, gpu=GPU, timeout=6 * 60 * 60, + volumes={REMOTE_DATA: CORPUS_VOL}) def evaluate_remote(solution_source: str, smoke: bool = True, role: str = "final", overrides: dict | None = None) -> dict: """Run the judge on `solution_source` and return its full result. @@ -198,7 +210,8 @@ def _ver(mod): # noqa: F811 (module-level _ver is the shared one) "stack": stack, } - @app.function(image=image, gpu=GPU, timeout=6 * 60 * 60) + @app.function(image=image, gpu=GPU, timeout=6 * 60 * 60, + volumes={REMOTE_DATA: CORPUS_VOL}) def run_pair_remote(solution_source: str, smoke: bool = False, role: str = "final") -> dict: """Run BOTH arms on one GPU and return their metrics as plain dicts. diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harness/settings.py b/2.0/problems/nanoslm_hybrid_arch_design/harness/settings.py index fe9af3deb..5dae66924 100644 --- a/2.0/problems/nanoslm_hybrid_arch_design/harness/settings.py +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/settings.py @@ -48,12 +48,15 @@ class TaskConfig: eval_block_size: int = 8192 # FIXED scoring window -- judge-owned # --- fixed optimization recipe (locked; agent designs architecture only) --- - # 8x4 rather than 32x1: at ctx 8192 a 32-sequence micro-batch is 262k tokens - # of activations per forward. Accumulating keeps the EFFECTIVE batch at 32 - # sequences (unchanged from the ctx-1024 recipe, so the optimization recipe - # is still locked and comparable) while bounding peak memory. - batch_size: int = 8 # sequences per micro-batch (per device) - grad_accum: int = 4 # -> effective batch 32 sequences + # 2x16 (effective batch 32). MEASURED: micro-batch 8 at ctx 8192 OOMs on an + # H100 -- the [8, 8192, 100352] fp32 logits the harness materializes for CE + # are ~26GB, and the baseline transformer alone tried to allocate 24.48GiB + # with only 20.77GiB free. micro-batch 2 keeps the fp32-logits copy ~6.6GB + # and fits both arms (incl. the GDN hybrid) with margin; grad_accum 16 holds + # the effective batch at 32 so the optimization recipe stays comparable. + # (A fused/chunked cross-entropy would let micro-batch grow again -- TODO.) + batch_size: int = 2 # sequences per micro-batch (per device) + grad_accum: int = 16 # -> effective batch 32 sequences learning_rate: float = 3.0e-3 # AdamW peak LR (nanoGPT-speedrun class) min_lr: float = 3.0e-4 weight_decay: float = 0.1 From 6a6dc9f92310fbf0a4cf1263c21e9445eba7fc27 Mon Sep 17 00:00:00 2001 From: Sijia Liu Date: Wed, 22 Jul 2026 01:02:59 -0700 Subject: [PATCH 4/5] [2.0] nanoslm: harness overhaul, paper-faithful reference, 2B corpus, hardened gate Training/eval harness: - Cosine LR now decays over wall-clock fraction (elapsed/train_seconds), not a hardcoded step horizon -- identical schedule for every architecture under the fixed budget; dead max_train_seconds loop check removed. - Modal run_pair_remote honors the agent/final role split: agent role reuses a baseline cached on the corpus Volume (keyed by config fingerprint + train.bin size + CODE_VERSION v8) and trains only the submission (~T); guard rejections return their public reason instead of environment_error. - Baseline SWA layers run FlexAttention's block-sparse kernel, with autocast dtype unification at the flex boundary (SDPA silently unified mixed q/k fp32 vs v bf16; flex refuses -- the old silent mask fallback had been hiding both the bug and a 2x baseline slowdown). - CUDA-only, no fallbacks: reference requires fla or raises; baseline requires flex or raises; FRONTIER_NANOSLM_SMOKE mode removed entirely. - Warmup fast-fails a model that cannot run at the 8192 eval context before spending the 6h budget; training exceptions of any type now classify as guard errors instead of escaping as environment errors. - compile_model added to the config fingerprint. Reference / baseline fidelity (arXiv 2604.03444 + OLMo-core @ fa6c5014): - GDN head_dim 128 kept per the paper's ladder convention ceil_128(0.75*d/h), with the 7B script's divergent int() rule documented; head_dim 48 probed on H100 (works, but slower via Triton autotune). - allow_neg_eigval=True passed explicitly (fla defaults False; paper requires it); GDN o_proj included in the depth-scaled output init. - vocab padded to 100352 (Olmo 3 convention; ids on disk stay < 100278). - reordered_norm documented against the paper's "pre-norm" wording. Data & scoring: - Corpus staged at 2B train tokens for the 6h budget (prep sizing rationale updated); val re-staged with byte-additivity verified. - AST policy gate: strict import allowlist (torch/numpy/fla/einops/triton + pure-stdlib), banned dynamic-access primitives, banned dunder attributes, banned raw-file-reader names; plausibility floor (bpb < 0.4 -> guard) as the runtime backstop. - Docs: problem-structure section in readme, agent README updated (starter is a byte-identical copy of reference.py). Co-Authored-By: Claude Fable 5 --- .gitignore | 2 + .../nanoslm_hybrid_arch_design/DESIGN.md | 212 -------- .../nanoslm_hybrid_arch_design/config.yaml | 35 +- .../docker/prep_assets.py | 20 +- .../nanoslm_hybrid_arch_design/evaluate.sh | 3 - .../nanoslm_hybrid_arch_design/evaluator.py | 260 +++++++--- .../harbor/app/README.md | 36 +- .../harbor/app/model.py | 297 ++++------- .../harness/baseline_model.py | 197 ++----- .../harness/data.py | 125 ++--- .../harness/eval_ppl.py | 33 +- .../harness/modal_app.py | 235 ++++++--- .../harness/model_config.py | 10 +- .../harness/policy.py | 166 +++++- .../harness/runner.py | 182 ++++--- .../harness/scoring.py | 30 +- .../harness/settings.py | 213 ++++---- .../harness/train.py | 43 +- .../nanoslm_hybrid_arch_design/readme | 103 +++- .../nanoslm_hybrid_arch_design/reference.py | 297 ++++------- .../repro/__init__.py | 1 - .../repro/modal_repro.py | 483 ------------------ .../repro/paper_models.py | 344 ------------- .../nanoslm_hybrid_arch_design/repro/wsd.py | 91 ---- .../src/frontier_cs_2_0/adapter.py | 6 + .../frontier_cs_2_0/task-template/task.toml | 2 +- 26 files changed, 1209 insertions(+), 2217 deletions(-) delete mode 100644 2.0/problems/nanoslm_hybrid_arch_design/DESIGN.md delete mode 100644 2.0/problems/nanoslm_hybrid_arch_design/repro/__init__.py delete mode 100644 2.0/problems/nanoslm_hybrid_arch_design/repro/modal_repro.py delete mode 100644 2.0/problems/nanoslm_hybrid_arch_design/repro/paper_models.py delete mode 100644 2.0/problems/nanoslm_hybrid_arch_design/repro/wsd.py diff --git a/.gitignore b/.gitignore index 54db130d3..b86f34608 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,8 @@ build dist *.egg-info **/__pycache__/ +# Baked corpus / large local assets (staged on Modal volumes, never committed) +**/.assets/ *.log **/*.lic .vscode diff --git a/2.0/problems/nanoslm_hybrid_arch_design/DESIGN.md b/2.0/problems/nanoslm_hybrid_arch_design/DESIGN.md deleted file mode 100644 index d6dad5b2f..000000000 --- a/2.0/problems/nanoslm_hybrid_arch_design/DESIGN.md +++ /dev/null @@ -1,212 +0,0 @@ -# Design — nanoslm_hybrid_arch_design - -## 1. Task - -Design a **hybrid language-model architecture** — mixing attention with -linear-recurrent sequence mixers — that reaches the lowest held-out -**bits-per-byte (`val_bpb`)** when trained from scratch under a **fixed -wall-clock budget on a single H100**. The agent submits one file, `model.py` -(full freedom over the model definition); the hidden judge drops it into a -**locked** dolma2 BPE training + evaluation harness, trains it for a fixed -wall-clock `T`, and scores `val_bpb` against a **locked baseline architecture** -trained under the identical budget. - -The initial configuration is taken from **Olmo Hybrid: From Theory to Practice -and Back** (Merrill, Li, Romero et al., arXiv:2604.03444). That work shows -theoretically that hybrids express capabilities beyond both transformers and -linear RNNs, then validates it at 7B by replacing sliding-window attention -layers with Gated DeltaNet (GDN) layers, outperforming the comparable Olmo 3 baseline. - -Concretely inherited from it and from the OLMo-core implementation: - - * the **3:1 interleaved pattern** (three Gated DeltaNet layers per full - attention layer, i.e. 25% attention), which is what `reference.py` ships as - the agent's starting point; - * the practice of **rebalancing head count / width so the hybrid is - parameter-matched** against the pure-attention baseline - (`REMOVE_HEADS` in the upstream script); - * **dolma2** tokenization and the ~190M OLMo3 shape. - -The scientific questions this task poses are the ones that paper answers at -scale but leaves open at 190M under a tight, parameter-matched, wall-clock -budget — enumerated in §1.2. The agent **builds on the baseline architecture** -(`harness/baseline_model.py`, a faithful pure-attention `olmo3_190M`) and -hybridizes it; `reference.py` ships one worked hybrid — the 3:1 GDN pattern — as -an existence proof, not a ceiling. It doubles as the repo-required reference -*solution* and the CI gate: `reference.py` must beat the hidden, score-0 -`baseline_model.py` (the pure-attention `olmo3_190M`, judge-side only). - -The optimizer, LR schedule, data, tokenizer, evaluation context and -budget are all fixed by the judge. The agent changes the architecture — and the -training context length, which trades optimizer steps against length -extrapolation. - - -## 2. Metric: held-out bits-per-byte (`val_bpb`) - -The tokenizer is **dolma2 BPE** (`allenai/dolma2-tokenizer`, the OLMo-3 -tokenizer, vocab 100278), locked by the judge. Token-level cross-entropy is not -comparable across models unless tokenization is fixed, so the metric normalizes -it by the **raw byte count** of the held-out text rather than by tokens: -`val_bpb = (held-out token cross-entropy in bits) / (held-out bytes)`. Dividing -by bytes makes the number **tokenizer-independent and ungameable** — the agent -submits a model, not a tokenizer, so it cannot lower `val_bpb` by retokenizing, -and the metric stays comparable across architectures the way a byte-level vocab -used to make it by construction. Lower is better; §8 scores the absolute gain -over the locked baseline. - -## 3. Why this is a real task - -Under a fixed wall-clock budget the frontier is genuinely architectural. Because -the baseline is already a tuned `olmo3_190M`, the easy block-level wins (RMSNorm, -rotary, QK-norm, SwiGLU) are **already in it**; what is left on the table is the -hybrid direction — the mixer, RNN architecture choices, layer placement, attention ratio, and training context length, etc. - -## 4. Submission surface & interface - -Submission is a **single file**, `/app/model.py` (`submission.kind: file`). Only -this file travels to the judge; edits to the harness in the agent workspace are -ignored by scoring (black-box judge uses its own locked harness copies). - -`model.py` must expose a factory the harness can call: - -```python -def build_model(config) -> torch.nn.Module: ... -# or a class usable as NanoSLM(config) -class NanoSLM(torch.nn.Module): ... -``` - -`config` is the harness-owned `ModelConfig`: `vocab_size` (100278), `block_size` -(the TRAINING context, which the submission may set via a module-level -`BLOCK_SIZE`), `eval_block_size` (the SCORING context, always 8192 and never -negotiable), plus read-only budget hints. The returned module must implement: - -```python -def forward(self, idx): # idx: LongTensor [B, T] of token ids in [0, vocab_size) - return logits # FloatTensor [B, T, vocab_size] - # returning (logits, loss) is accepted, but the judge IGNORES any returned - # loss and computes cross-entropy itself for BOTH training and val_bpb. -``` - -The judge owns the loss so a model cannot report a fake low loss. The optimizer, -LR schedule, weight-decay grouping (no WD on 1-D params), data order and -wall-clock budget are all locked in `harness/train.py`; the training context is -the one training-side quantity the submission controls, and the evaluation -context is locked in `harness/data.py::val_windows`. - -## 5. Iso-wallclock protocol (the anti-gaming axis) - -- The judge trains the baseline and the submission for the **same fixed - wall-clock `T`** on the same H100, from the **same seed and data order** - (common random numbers), then evaluates both on the **same hidden held-out - validation set**. The scored (final/verifier) run measures baseline + submission - back-to-back in one job/GPU/process (no cache), mirroring - `nanowm`'s `run_pair(role="final")`. A cached baseline keyed by - `settings.config_fingerprint()` is used only for the cheap iterative - (agent-role) feedback path. Role is selected by the judge via - `FRONTIER_NANOSLM_ROLE` (`agent` = train submission only + cached baseline, ~T; - `final` = fresh baseline+submission CRN pair, ~2T; default `final`). GPU is - served on Modal from a CPU judge; a directly-attached H100 is - used only in local testing mode. -- The wall-clock cutoff is enforced **by the harness timer**, not by anything the - model can influence: training stops at the first optimizer step whose start - exceeds `T`. Model `forward` cost therefore trades directly against step count - — a more efficient architecture legitimately completes more useful steps. This - is the intended lever, and it is the mechanism by which a hybrid can win at - all: at ctx 8192 attention is estimated to dominate layer FLOPs (~78–84%, an - analytic estimate not yet confirmed on the judge GPU), so replacing most of it - with a cheaper mixer converts directly into extra optimizer steps. -- Determinism: fixed seeds; `torch.use_deterministic_algorithms(True, - warn_only=True)`, `cudnn.deterministic=True`, `benchmark=False`, TF32 - **disabled** during eval, `CUBLAS_WORKSPACE_CONFIG=:4096:8`. Iso-wallclock has - irreducible timing noise (step count varies run-to-run); the CRN pairing keeps - the *comparison* stable even so. The seed-to-seed noise floor itself is an open - calibration item, not yet measured on the judge GPU. - -## 6. Submission policy (validated before training; torch-free, unit-tested) - -`model.py` only (`.py`), ≤ 256 KB, no file deletion, safe path; the submission -must define `build_model(config)` or `class NanoSLM`. - -Two deny lists in `harness/policy.py`, scanned differently (representative, not -exhaustive): - -- `POLICY_DENY_TOKENS`, matched over the **full source** — a leak signal in a - comment is still a leak: - - **Escape / leakage:** `os.environ`, `os.getenv`, `putenv`, `subprocess`, - `socket`, `requests`, `urllib`, `httpx`, `FRONTIER_`/`JUDGE_`/`HARBOR_`/`MODAL_`, - `HF_TOKEN`, judge paths (`/judge`, `/opt/`, `/tests/`). - - **Pretrained-weight loading (must train from scratch):** `from_pretrained`, - `torch.load`, `load_state_dict`, `hf_hub`, `huggingface`, `safetensors`, `AutoModel`. - - **Filesystem reads:** `open(`, `Path(`, `np.load`, `np.fromfile`, `mmap`, `pickle.load`. - - **Timer / control-flow / concurrency:** `time.time`, `time.perf_counter`, - `time.sleep`, `while True`, `threading`, `multiprocessing`, `os.fork`, `ctypes`, - `exec(`, `__import__`. -- `POLICY_DENY_TOKENS_CODE`, matched over **code only** (comments and string - literals stripped) so a docstring *may* name the metric it targets: `val_ppl`, - `perplexity`, `val_bpb`, `bits_per_byte`, `holdout`, `val_data`, `val.bin`. - -`eval(` and `compile(` are deliberately **allowed** so `model.eval()` and -`torch.compile()` work — the sandbox and judge-owned loss cover the rest. - -The policy is a **static allow/deny gate**; §7 guards the residual dynamic -risks it cannot see. - -## 7. Dynamic guards (what the static scan can't catch) - -- **Judge owns the loss.** `val_bpb` is computed by the harness from the model's - `logits` via its own cross-entropy on **hidden** held-out bytes (never mounted - in `/app`), so a model cannot report a fake loss or memorize the val set. -- **Trained-from-scratch guard.** Params are fingerprinted at init and re-checked - after training; a model whose weights did not change (untrained / frozen - constant) or that produces a (near-)constant logit distribution over inputs is - scored 0. Blocks "return a cached distribution" degenerate submissions. -- **Resource caps.** Param count ≤ `PARAM_CAP` and peak activation memory within - the H100; OOM or over-cap → score 0 with a public message (no traceback). Stops - memory-bomb / timer-dodge attempts. -- **Sandboxing.** Submitted code is imported and run as an unprivileged user with - the evaluator source chmod-protected (same pattern as `erdos_demo`); evaluator - output returns public metrics and concise errors only — never raw submission - stdout/stderr or tracebacks (2.0 black-box safety rules). - -## 8. Scoring — ABSOLUTE bits-per-byte gain - -`harness/scoring.py` (shared with the public test). Lower `val_bpb` is better. -With `base_bpb` = locked-baseline held-out bits-per-byte and `sub_bpb` = the -submission's, both trained under the identical wall-clock budget with common -random numbers and both evaluated at the fixed 8192-token window: - -``` -gain = base_bpb - sub_bpb # THE MEASUREMENT -score = clip(100 * gain / bpb_score_scale, 0, 100) -``` - -Baseline-tying and worse-than-baseline both score 0. `score_unbounded` is the -un-clipped ratio and keeps rising past 100. - -### 8.1 `bpb_score_scale` is a display convention, NOT a calibration target - -The only measurement is `gain`, in bits per byte. `bpb_score_scale` exists solely -to map that gain onto the 0–100 range Harbor's `reward = score / 100` expects — a -raw gain of 0.05 bpb would otherwise surface as reward 0.0005. It is **not** a -calibrated definition of "a full win": no such number has been measured, and the -clip discards information (a submission at 2× the scale scores the same 100 as one -exactly at it), which is why `score_unbounded` stays un-clipped for an operator to -read. - -The one real requirement is **discrimination**: too large and everything pins at -0, too small and everything pins at 100 — checkable from a single real-corpus run, -no headroom study needed. Current value: `bpb_score_scale = 0.05` bpb. (The -predecessor `r_target` is dropped, not satisfied: it asked for a measured full-win -definition that was never well-posed.) - -## 9. Compute / infra - -- `tag: systems`; H100 served on Modal (one per environment), CPU judge — - identical shape to `nanowm_rollout_speedup` and `vllm_llm_serving_optimization`. -- `runtime.timeout_seconds`: long-horizon (agent iterates for hours); - per-submission training is capped at ≤ 1 h/train (single H100). -- Hidden judge assets (judge image only): the tokenized training shard, the - **held-out** validation byte stream, the locked baseline architecture, and the - cached baseline `val_bpb` (keyed by config fingerprint). None are mounted in - `/app`. \ No newline at end of file diff --git a/2.0/problems/nanoslm_hybrid_arch_design/config.yaml b/2.0/problems/nanoslm_hybrid_arch_design/config.yaml index a2d25e416..773880b46 100644 --- a/2.0/problems/nanoslm_hybrid_arch_design/config.yaml +++ b/2.0/problems/nanoslm_hybrid_arch_design/config.yaml @@ -3,7 +3,7 @@ runtime: # Submission is a single Python file, the model definition (/app/model.py). # `language: python` keeps the extension/CLI conventions standard. language: python - timeout_seconds: 21600 # long-horizon: the agent iterates for hours + timeout_seconds: 86400 environment: >- Full-freedom PyTorch model.py submitted against a locked training harness over a dolma2-BPE-tokenized FineWeb-Edu corpus; a single @@ -27,39 +27,34 @@ runtime: judge_pip_packages: - modal docker: - # Experimental images; build before a local Harbor trial. The agent image - # carries the harness + a starter model.py; the judge image additionally - # vendors the tokenized training shard, the HELD-OUT validation token stream - # and its manifest (which carries val_bytes, the bpb denominator), the - # locked baseline architecture, and the cached baseline metric. image: frontiercs/nanoslm-hybrid-arch-design-agent:experimental-v0 judge_image: frontiercs/nanoslm-hybrid-arch-design-judge:experimental-v0 environment: cpus: 8 memory_mb: 32768 storage_mb: 32768 - build_timeout_seconds: 5400 + build_timeout_seconds: 7200 evaluation: - # GPU served on Modal (one per environment); judge container is CPU-only. gpu: H100 - # dolma2 BPE (the OLMo-3 tokenizer). UNGATED -- no HF token required. tokenizer: allenai/dolma2-tokenizer - vocab_size: 100278 # > 65535, so token streams are uint32 + # 100352 = dolma2's 100278 real ids padded up (Olmo 3 convention; ids on + # disk stay < 100278). > 65535, so token streams are uint32. + vocab_size: 100352 dataset: HuggingFaceFW/fineweb-edu:sample-10BT - block_size: 8192 # default training context (agent may override) - eval_block_size: 8192 # fixed scoring context - train_seconds: 1800 # fixed wall-clock budget T per training run - max_train_seconds: 3300 # hard cap (< 1h) - val_tokens: 1048576 # held-out TARGET tokens scored - # The scored metric is bits per BYTE: nll_nats / (val_bytes * ln2). val_bytes + block_size: 8192 + eval_block_size: 8192 + train_seconds: 21600 # fixed wall-clock budget T per training run (6 h) + max_train_seconds: 43200 + val_tokens: 1048576 + # The scored metric is bits per byte: nll_nats / (val_bytes * ln2). val_bytes # is measured by docker/prep_assets.py and shipped in the judge image's - # manifest.json -- normalizing by the TOKEN count would make the metric + # manifest.json -- normalizing by the token count would make the metric # tokenizer-dependent. metric: val_bpb - # Scoring is on the ABSOLUTE val_bpb gain over the locked baseline: + # Scoring is on the absolute val_bpb gain over the locked baseline: # score = clip(100 * (base_bpb - sub_bpb) / bpb_score_scale, 0, 100) - bpb_score_scale: 0.05 # bpb gain that saturates the bounded score - param_cap: 400000000 # trainable-param cap; over-cap -> score 0 (matches harness/settings.py) + bpb_score_scale: 0.5 + param_cap: 400000000 seed: 1337 submission: kind: file diff --git a/2.0/problems/nanoslm_hybrid_arch_design/docker/prep_assets.py b/2.0/problems/nanoslm_hybrid_arch_design/docker/prep_assets.py index e7bfbc564..93865fd8e 100644 --- a/2.0/problems/nanoslm_hybrid_arch_design/docker/prep_assets.py +++ b/2.0/problems/nanoslm_hybrid_arch_design/docker/prep_assets.py @@ -29,7 +29,7 @@ While the tokenizer was byte-level, 1 token == 1 byte and normalizing by the token count was accidentally identical. Under BPE it is not: per-token CE/ln2 is a tokenizer-dependent quantity that is no longer comparable across setups, which -would destroy the reason bpb was chosen (DESIGN.md §2). Nothing at +would destroy the reason bpb was chosen. Nothing at eval time can recover the byte count from a token stream, so it is measured HERE -- by decoding the exact span of target tokens the harness scores and taking its UTF-8 length -- and recorded in the manifest as `val_bytes` alongside @@ -49,9 +49,13 @@ `val_tokens` TARGET tokens over non-overlapping windows, and the last target needs one more input token behind it. -`train.bin` defaults to 512M tokens (~2 GB on disk at uint32). A 30-minute H100 -run at batch 8 x ctx 8192 touches on the order of 10^7-10^8 tokens, so this -leaves ample headroom without the sampler ever exhausting the stream. +`train.bin` defaults to 2B tokens (~8 GB on disk at uint32), sized against the +6 h budget rather than the original 30-minute one: a 6 h H100 run at effective +batch 32 x ctx 8192 (262k tokens/step, ~10k steps) samples on the order of +2.5-5B tokens, so a 512M shard meant every window was drawn from ~5-10x-repeated +data -- and a fast short-context submission repeated it far more. 2B unique +tokens keeps the baseline at roughly 1-2 effective epochs. sample-10BT holds +~10B tokens, so there is headroom to raise this further if budgets grow. Usage: NANOSLM_ARCH_ASSETS=./assets python3 docker/prep_assets.py @@ -74,11 +78,15 @@ CORPUS_CONFIG = "sample-10BT" TOKENIZER = "allenai/dolma2-tokenizer" # matches TaskConfig.tokenizer_name -VOCAB_SIZE = 100278 # matches TaskConfig.vocab_size +# The TOKENIZER's true id space, used to range-check the written streams. +# Deliberately NOT TaskConfig.vocab_size, which is the model vocabulary padded +# up to 100352 (Olmo 3 convention): ids on disk must fit the tokenizer, and a +# stream that needed the padding rows would mean a corrupted corpus. +VOCAB_SIZE = 100278 TOKEN_DTYPE = np.uint32 # matches harness/data.TOKEN_DTYPE VAL_TOKENS = 1 << 20 # matches TaskConfig.val_tokens -TRAIN_TOKENS = 512_000_000 +TRAIN_TOKENS = 2_000_000_000 # see SIZING in the docstring (6 h budget) DOC_BATCH = 256 # docs per tokenizer call diff --git a/2.0/problems/nanoslm_hybrid_arch_design/evaluate.sh b/2.0/problems/nanoslm_hybrid_arch_design/evaluate.sh index 3e8e5a0ff..ab3fe4081 100755 --- a/2.0/problems/nanoslm_hybrid_arch_design/evaluate.sh +++ b/2.0/problems/nanoslm_hybrid_arch_design/evaluate.sh @@ -4,9 +4,6 @@ # Usage: # bash evaluate.sh /path/to/model.py # score a submission (needs torch+GPU) # bash evaluate.sh --selftest # torch-free policy/scoring/fingerprint tests -# -# Tip: FRONTIER_NANOSLM_SMOKE=1 shrinks the model/budget for a fast CPU wiring -# check (never used for real scoring). set -euo pipefail HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" diff --git a/2.0/problems/nanoslm_hybrid_arch_design/evaluator.py b/2.0/problems/nanoslm_hybrid_arch_design/evaluator.py index 0e9fa25f6..cdf9edf44 100644 --- a/2.0/problems/nanoslm_hybrid_arch_design/evaluator.py +++ b/2.0/problems/nanoslm_hybrid_arch_design/evaluator.py @@ -4,7 +4,7 @@ The submission is a single ``model.py`` (full architecture freedom). The judge trains it from scratch for a fixed wall-clock budget on a single H100 and scores held-out validation perplexity against a locked baseline architecture trained -under the identical budget. See DESIGN.md. +under the identical budget. Top-level imports are torch-free so this module loads (and self-tests) without a GPU; the training/eval path is lazy-imported inside :func:`evaluate`. @@ -26,14 +26,14 @@ from harness import policy, scoring, settings # noqa: E402 -# cuBLAS determinism MUST be configured before the first CUDA call, which means +# cuBLAS determinism must be configured before the first CUDA call, which means # before torch is imported anywhere. This module's top-level imports are # deliberately torch-free (see the docstring), so setting it here is early # enough on every path -- Modal, judge container, or a local GPU box. # -# TaskConfig has carried `cublas_workspace_config` since the problem was -# written, and DESIGN.md 5 lists CUBLAS_WORKSPACE_CONFIG=:4096:8 as part of the -# determinism story, but NOTHING READ IT. So every run so far had +# TaskConfig has carried `cublas_workspace_config` (CUBLAS_WORKSPACE_CONFIG= +# :4096:8) since the problem was written, as part of the determinism story, but +# nothing read it. So every run so far had # torch.use_deterministic_algorithms(True) set while cuBLAS remained free to be # non-deterministic -- which showed up as a UserWarning in the first GPU run and # quietly weakened the CRN pairing the score depends on. @@ -57,13 +57,13 @@ def _protect_evaluator_source() -> None: def _backend() -> str: - """MODAL (CPU judge -> Modal GPU) or LOCAL (directly-attached GPU). + """MODAL (CPU judge -> Modal GPU) or local (directly-attached GPU). The design specifies a CPU judge with the H100 served on Modal, the same shape as nanowm_rollout_* and vllm_llm_serving_optimization -- but until now nothing implemented it: `evaluate()` only called `_select_device()`, which in a CPU-only judge container returns "cpu" and would try to train a - ~190M model at ctx 8192 on CPU. Harbor invokes THIS module, not the manual + ~190M model at ctx 8192 on CPU. Harbor invokes this module, not the manual `modal_app.evaluate_remote` wrapper, so without this the problem is not runnable under Harbor at all. @@ -81,7 +81,7 @@ def _backend() -> str: def _run_pair_modal(solution_path: str, role: str): """Run both arms on a Modal GPU; score judge-side. - Only ARM METRICS cross the boundary -- scoring and the hidden constants + Only arm metrics cross the boundary -- scoring and the hidden constants (bpb_score_scale) stay in the judge, so a GPU worker never sees them and cannot hand back a score the judge did not compute. """ @@ -89,11 +89,7 @@ def _run_pair_modal(solution_path: str, role: str): source = Path(solution_path).read_text(encoding="utf-8", errors="replace") with app.run(): - res = run_pair_remote.remote( - solution_source=source, - smoke=settings.smoke_enabled(), - role=role, - ) + res = run_pair_remote.remote(solution_source=source, role=role) return res @@ -142,7 +138,7 @@ def _baseline_cache_put(cfg, fingerprint: str, ppl: float) -> None: def _load_submission_module(path: str): - """Import the submitted model.py AFTER the static policy gate has passed.""" + """Import the submitted model.py after the static policy gate has passed.""" spec = importlib.util.spec_from_file_location("submission_model", path) if spec is None or spec.loader is None: raise RuntimeError("could not load submission module") @@ -169,9 +165,23 @@ def evaluate(solution_path: str): except Exception as exc: # noqa: BLE001 - classify, never leak return (0.0, 0.0, f"environment_error: modal dispatch failed " f"({type(exc).__name__})", {"config_fingerprint": fp}) + # Guard rejections travel as a field, not an exception: GuardError is + # defined in harness.runner (which imports torch), so this CPU-only + # judge could not deserialize the raised form. The message is the + # public, classified reason -- same wording as the local path's. + if res.get("guard_error"): + return 0.0, 0.0, f"guard: {res['guard_error']}", { + "config_fingerprint": fp, "role": role, + "stack": res.get("stack"), + } b, c = res["baseline"], res["submission"] sr = scoring.score_from_bpb(b["val_bpb"], c["val_bpb"], cfg.bpb_score_scale) - note = ["smoke"] if settings.smoke_enabled() else [] + note = [] + if role == "agent": + # Mirrors the local path's note: says whether this feedback run + # paid ~T (cached baseline) or ~2T (fresh pair, cache populated). + note.append("iterative(cached-baseline)" + if res.get("baseline_cached") else "iterative") message = scoring.format_message( b["val_bpb"], c["val_bpb"], sr, steps=c["steps"], wall_seconds=c["wall_seconds"], @@ -180,7 +190,7 @@ def evaluate(solution_path: str): return sr.score, sr.score_unbounded, message, { "base_val_bpb": b["val_bpb"], "sub_val_bpb": c["val_bpb"], "abs_bpb_delta": sr.abs_bpb_delta, "sub_val_ppl": c["val_ppl"], - # abs_bpb_delta is the SCORED quantity; rel_improvement is kept + # abs_bpb_delta is the scored quantity; rel_improvement is kept # alongside it so an operator never has to recompute it. "rel_improvement": sr.rel_improvement, "bpb_score_scale": cfg.bpb_score_scale, @@ -193,7 +203,7 @@ def evaluate(solution_path: str): "config_fingerprint": fp, "stack": res.get("stack"), } - # LOCAL backend: torch + a directly-attached GPU. + # Local backend: torch + a directly-attached GPU. try: device = _select_device() except RuntimeError as exc: @@ -248,7 +258,7 @@ def evaluate(solution_path: str): ) sr = scoring.score_from_bpb(base_bpb, sub.val_bpb, cfg.bpb_score_scale) - note = ["smoke"] if settings.smoke_enabled() else [] + note = [] if role == "agent": note.append("iterative(cached-baseline)" if base_steps is None else "iterative") message = scoring.format_message( @@ -270,7 +280,7 @@ def evaluate(solution_path: str): "base_steps": base_steps, "sub_params": sub.n_params, "sub_wall_seconds": sub.wall_seconds, - # Agent-controlled TRAINING context vs the FIXED scoring window, both + # Agent-controlled training context vs the fixed scoring window, both # reported so a submission can see the steps-vs-extrapolation trade it # actually made rather than inferring it. "sub_train_block_size": sub.train_block_size, @@ -313,13 +323,29 @@ def prepare() -> dict: # --------------------------------------------------------------------------- # # Torch-free self-test: exercises policy + scoring + fingerprint without a GPU. # --------------------------------------------------------------------------- # +def _write_selftest_bin(dirpath: str, name: str, vocab_size: int, n: int, seed: int) -> str: + """Write a tiny real uint32 token stream for the self-test. + + The harness has no synthetic fallback, so the self-test manufactures its own + small real corpus (in-range ids, on disk) rather than relying on the loader + to fabricate one -- the same load path the judge uses, just tiny. + """ + import numpy as _np + + rng = _np.random.default_rng(seed) + arr = rng.integers(0, vocab_size, size=n).astype(_np.uint32) + path = os.path.join(dirpath, name) + arr.tofile(path) + return path + + def _selftest() -> int: ok = True def check(name, cond): nonlocal ok ok = ok and bool(cond) - print(f"[{'PASS' if cond else 'FAIL'}] {name}", file=sys.stderr) + print(f"[{'pass' if cond else 'fail'}] {name}", file=sys.stderr) # --- policy: accept reference + baseline, reject malicious variants --- ref_src = (_HERE / "reference.py").read_text() @@ -338,6 +364,24 @@ def check(name, cond): for name, src in malicious.items(): check(f"policy rejects {name}", not policy.check_source(src).ok) + # AST layer: substring evasions that motivated the strict allowlist. None + # of these contain a denied substring; all must die on the AST rules. + evasive = { + "import os (allowlist)": "import os\nclass NanoSLM: pass\n", + "bare getattr": "import torch\ng = getattr\nclass NanoSLM: pass\n", + "bare eval": "class NanoSLM: pass\ndef build_model(c):\n return eval('1')\n", + "dunder walk": "class NanoSLM: pass\ndef build_model(c):\n return ().__class__\n", + "aliased fromfile": "import numpy as q\nclass NanoSLM: pass\ndef build_model(c):\n return q.fromfile('x')\n", + "from-import load": "from torch import load\nclass NanoSLM: pass\n", + "wildcard import": "from torch import *\nclass NanoSLM: pass\n", + "relative import": "from . import secrets\nclass NanoSLM: pass\n", + "banned import segment (torch.hub)": "import torch.hub as h\nclass NanoSLM: pass\n", + "operator (attrgetter-by-string)": "import operator\nclass NanoSLM: pass\n", + "unparseable source": "def build_model(:\n", + } + for name, src in evasive.items(): + check(f"policy rejects {name}", not policy.check_source(src).ok) + # legitimate idioms must survive check( "policy allows model.eval()/torch.compile()", @@ -345,9 +389,19 @@ def check(name, cond): "class NanoSLM:\n def forward(self,x):\n self.eval(); return x\n" ).ok, ) + check( + "policy allows super().__init__ / torch.__version__ / triton", + policy.check_source( + "import torch\nimport triton\nimport triton.language as tl\n" + "class NanoSLM(torch.nn.Module):\n" + " def __init__(self):\n" + " super().__init__()\n" + " self.v = torch.__version__\n" + ).ok, + ) - # --- scoring: ABSOLUTE bpb gain. baseline->0, worse->0, scale->100 --- - # Cases are constructed by SUBTRACTING bits per byte, not by scaling + # --- scoring: Absolute bpb gain. baseline->0, worse->0, scale->100 --- + # Cases are constructed by subtracting bits per byte, not by scaling # base_bpb: the scored quantity is now `base_bpb - sub_bpb`, so a relative # construction would no longer test what it claims to. scale = 0.05 # bpb gain that saturates the bounded score @@ -365,22 +419,27 @@ def check(name, cond): check("unbounded keeps rewarding past 100", s_over.score_unbounded > 100.0) check("unbounded is exactly the un-clipped ratio", abs(s_over.score_unbounded - 200.0) < 1e-9) - # The score must depend on the ABSOLUTE gain only -- the same 0.05 bpb gain + # The score must depend on the absolute gain only -- the same 0.05 bpb gain # scores the same at any operating point. (This is the property the old - # relative form did NOT have, and the reason the operating point now - # matters: see DESIGN.md 8.) + # relative form did not have, and the reason the operating point now + # matters.) s_lowbase = scoring.score_from_bpb(1.5, 1.5 - scale, scale) check("score depends on absolute gain, not on base_bpb", abs(s_lowbase.score - s_sat.score) < 1e-9) - # ... while BOTH figures are still reported. + # ... while both figures are still reported. check("relative improvement still reported", abs(s_sat.rel_improvement - (scale / base_bpb)) < 1e-12 and abs(s_sat.abs_bpb_delta - scale) < 1e-12) - # A SCORING constant, not a training knob: changing it must not invalidate + # A scoring constant, not a training knob: changing it must not invalidate # a cached baseline, so it must stay out of the fingerprint. - check("bpb_score_scale is NOT fingerprinted", + check("bpb_score_scale is not fingerprinted", "bpb_score_scale" not in settings._FINGERPRINT_KEYS and "r_target" not in settings._FINGERPRINT_KEYS) + # Same for the plausibility floor (a guard constant), which must exist and + # sit far below any honest result at this scale. + check("min_plausible_bpb set, sane, and not fingerprinted", + 0.0 < settings.DEFAULT.min_plausible_bpb <= 0.5 + and "min_plausible_bpb" not in settings._FINGERPRINT_KEYS) # --- fingerprint: stable, and changes iff a locked knob changes --- from dataclasses import replace @@ -391,7 +450,7 @@ def check(name, cond): check("fingerprint stable", fp0 == fp_same) check("fingerprint changes on locked-knob change", fp0 != fp_diff) - # The EVAL window is fingerprinted (changing it changes what val_bpb means, + # The eval window is fingerprinted (changing it changes what val_bpb means, # so a cached baseline must not be reused) ... check( "fingerprint changes on eval_block_size change", @@ -399,11 +458,11 @@ def check(name, cond): replace(settings.DEFAULT, eval_block_size=4096) ), ) - # ... while the TRAINING context is NOT. It varies per submission now, so + # ... while the training context is not. It varies per submission now, so # fingerprinting it would give every submission a unique key and the cached # baseline would never hit -- doubling the cost of the agent-role path. check( - "fingerprint IGNORES the default training block_size", + "fingerprint ignores the default training block_size", fp0 == settings.config_fingerprint( replace(settings.DEFAULT, block_size=2048) ), @@ -450,10 +509,10 @@ def _mod(**attrs): rejected = True check(f"BLOCK_SIZE={bad!r} rejected", rejected) - # A model that sizes a positional table off the TRAINING context trains + # A model that sizes a positional table off the training context trains # fine and then fails at the 8192 eval -- the characteristic new failure # mode of the split contexts. It must come back as a GuardError naming - # BOTH contexts, not as a bare RuntimeError/IndexError. (torch raises + # Both contexts, not as a bare RuntimeError/IndexError. (torch raises # IndexError here, which an earlier `except RuntimeError` missed.) import torch import torch.nn as nn @@ -464,7 +523,7 @@ class _BadPos(nn.Module): def __init__(self, config): super().__init__() self.emb = nn.Embedding(config.vocab_size, 8) - # WRONG on purpose: should be config.eval_block_size. + # Wrong on purpose: should be config.eval_block_size. self.pos = nn.Embedding(config.block_size, 8) self.lin = nn.Linear(8, config.vocab_size) @@ -472,25 +531,72 @@ def forward(self, idx): p = self.pos(torch.arange(idx.shape[1], device=idx.device)) return self.lin(self.emb(idx) + p) - tiny = replace( - settings.DEFAULT, vocab_size=256, eval_block_size=512, batch_size=2, - train_seconds=0.5, max_train_seconds=10.0, val_tokens=1024, - train_tokens_path="", val_tokens_path="", manifest_path="", - ) - try: - _runner.run_arm( - _runner.load_factory(_mod(build_model=_BadPos)), - _TD(tiny), tiny, "cpu", 256, + import tempfile as _tempfile + + with _tempfile.TemporaryDirectory() as _cdir: + _tp = _write_selftest_bin(_cdir, "train.bin", 256, 16384, 1) + _vp = _write_selftest_bin(_cdir, "val.bin", 256, 16384, 2) + tiny = replace( + settings.DEFAULT, vocab_size=256, eval_block_size=512, batch_size=2, + train_seconds=0.5, max_train_seconds=10.0, val_tokens=1024, + train_tokens_path=_tp, val_tokens_path=_vp, manifest_path="", + val_bytes=1024 * 4, ) - got = "no error" - except _runner.GuardError as exc: - got = str(exc) + try: + _runner.run_arm( + _runner.load_factory(_mod(build_model=_BadPos)), + _TD(tiny), tiny, "cpu", 256, + ) + got = "no error" + except _runner.GuardError as exc: + got = str(exc) check( "eval-context failure -> GuardError naming both contexts", "evaluation failed at context 512" in got and "trained at 256" in got, ) - # --- THE CRUX: eval windows are 8192 whatever the training context is --- + # --- plausibility floor: a too-good bpb is rejected as leakage. The + # eval is stubbed to return an impossibly low bpb; run_arm must refuse + # it rather than score it. (Backstop for the static gate -- see + # settings.min_plausible_bpb.) --- + from harness.eval_ppl import EvalOutput as _EO + + class _TinyOK(nn.Module): + def __init__(self, config): + super().__init__() + self.emb = nn.Embedding(config.vocab_size, 8) + self.lin = nn.Linear(8, config.vocab_size) + + def forward(self, idx): + return self.lin(self.emb(idx)) + + with _tempfile.TemporaryDirectory() as _cdir: + _tp = _write_selftest_bin(_cdir, "train.bin", 256, 16384, 5) + _vp = _write_selftest_bin(_cdir, "val.bin", 256, 16384, 6) + tiny2 = replace( + settings.DEFAULT, vocab_size=256, eval_block_size=512, batch_size=2, + train_seconds=0.5, max_train_seconds=10.0, val_tokens=1024, + train_tokens_path=_tp, val_tokens_path=_vp, manifest_path="", + val_bytes=1024 * 4, + ) + _real_eval = _runner.evaluate_perplexity + _runner.evaluate_perplexity = ( + lambda *a, **k: _EO(0.01, 1.0, 0.01, 1024, 1.0, 4096.0) + ) + try: + _runner.run_arm( + _runner.load_factory(_mod(build_model=_TinyOK)), + _TD(tiny2), tiny2, "cpu", 256, + ) + floor_msg = "no error" + except _runner.GuardError as exc: + floor_msg = str(exc) + finally: + _runner.evaluate_perplexity = _real_eval + check("implausibly low bpb -> GuardError (leakage floor)", + "plausibility floor" in floor_msg) + + # --- The crux: eval windows are 8192 whatever the training context is --- # Numpy-only, no torch: exercise the window arithmetic val_windows uses. try: import numpy as np @@ -501,30 +607,38 @@ def forward(self, idx): TokenData = None if TokenData is not None: - smoke = replace( - settings.DEFAULT, val_tokens=32_768, train_tokens_path="", - val_tokens_path="", manifest_path="", val_bytes=0, - ) - widths, denoms = set(), set() - for train_ctx in (8192, 2048, 256): - d = TokenData(replace(smoke, block_size=train_ctx)) - # val_windows yields torch tensors, so replicate its slicing here - # rather than importing torch: same expression, no device. - eb = d.cfg.eval_block_size - n = min(d.val.size - 1, d.cfg.val_tokens) - n_windows = n // eb - widths.add((eb, n_windows)) - denoms.add(round(d.val_bytes_for(n_windows * eb), 6)) - # ... and the TRAINING batch really does follow the training ctx. - x = np.stack([d.train[i: i + train_ctx] for i in (0, 1)]) - check(f"train batch width == {train_ctx}", x.shape[1] == train_ctx) - check("eval window width/count independent of training ctx", len(widths) == 1) - check("bpb denominator independent of training ctx", len(denoms) == 1) - check("eval window width is 8192", widths.pop() == (8192, 4)) - - # --- baseline cache HITS across submissions with different train contexts --- + import tempfile as _tempfile + + with _tempfile.TemporaryDirectory() as _cdir: + _vocab = settings.DEFAULT.vocab_size + _tp = _write_selftest_bin(_cdir, "train.bin", _vocab, 16384, 3) + # val must hold > val_tokens ids so the window count is exercised. + _vp = _write_selftest_bin(_cdir, "val.bin", _vocab, 40_000, 4) + tiny_cfg = replace( + settings.DEFAULT, val_tokens=32_768, + train_tokens_path=_tp, val_tokens_path=_vp, manifest_path="", + val_bytes=32_768 * 4, + ) + widths, denoms = set(), set() + for train_ctx in (8192, 2048, 256): + d = TokenData(replace(tiny_cfg, block_size=train_ctx)) + # val_windows yields torch tensors, so replicate its slicing here + # rather than importing torch: same expression, no device. + eb = d.cfg.eval_block_size + n = min(d.val.size - 1, d.cfg.val_tokens) + n_windows = n // eb + widths.add((eb, n_windows)) + denoms.add(round(d.val_bytes_for(n_windows * eb), 6)) + # ... and the training batch really does follow the training ctx. + x = np.stack([d.train[i: i + train_ctx] for i in (0, 1)]) + check(f"train batch width == {train_ctx}", x.shape[1] == train_ctx) + check("eval window width/count independent of training ctx", len(widths) == 1) + check("bpb denominator independent of training ctx", len(denoms) == 1) + check("eval window width is 8192", widths.pop() == (8192, 4)) + + # --- baseline cache hits across submissions with different train contexts --- # This is the reason block_size left the fingerprint. The key is derived - # from the CONFIG only; a submission's BLOCK_SIZE is resolved per-arm inside + # from the config only; a submission's BLOCK_SIZE is resolved per-arm inside # run_arm and never reaches the key, so one cached baseline serves both. import tempfile @@ -537,7 +651,7 @@ def forward(self, idx): key = settings.config_fingerprint(cfg_c) _baseline_cache_put(cfg_c, key, 1.2345) # Two submissions, training at 8192 and at 2048, both look the - # baseline up under the SAME config -> the same key -> a hit. + # baseline up under the same config -> the same key -> a hit. hits = [ _baseline_cache_get(cfg_c, settings.config_fingerprint(cfg_c)) for _ in (8192, 2048) @@ -545,7 +659,7 @@ def forward(self, idx): check("baseline cache hits for both training contexts", hits == [1.2345, 1.2345]) check( - "cache MISSES when the eval window changes", + "cache misses when the eval window changes", _baseline_cache_get( cfg_c, settings.config_fingerprint( @@ -558,7 +672,7 @@ def forward(self, idx): else: os.environ["FRONTIER_NANOSLM_BASELINE_CACHE"] = prev - print(("SELFTEST OK" if ok else "SELFTEST FAILED"), file=sys.stderr) + print(("selftest OK" if ok else "selftest failed"), file=sys.stderr) return 0 if ok else 1 diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harbor/app/README.md b/2.0/problems/nanoslm_hybrid_arch_design/harbor/app/README.md index 3062eee82..8d18c885b 100644 --- a/2.0/problems/nanoslm_hybrid_arch_design/harbor/app/README.md +++ b/2.0/problems/nanoslm_hybrid_arch_design/harbor/app/README.md @@ -25,15 +25,26 @@ not a target, so simply maximize the gain. ## What you submit `/app/model.py`, defining `build_model(config)` or `class NanoSLM`, whose -`forward(idx)` returns logits `[B, T, 100278]` for `idx` of BPE token ids. +`forward(idx)` returns logits `[B, T, 100352]` for `idx` of BPE token ids. -`config` gives you `vocab_size` (always **100278** — the dolma2 BPE tokenizer, -`allenai/dolma2-tokenizer`), `block_size` (the context you are **trained** at), +`config` gives you `vocab_size` (always **100352** — the dolma2 BPE tokenizer +`allenai/dolma2-tokenizer`, whose 100278 real ids are padded up, the Olmo 3 +convention; actual ids stay `< 100278` but your logits must span the padded +width), `block_size` (the context you are **trained** at), `eval_block_size` (the context you are **scored** at — always 8192), and read-only hints. You control the architecture **and the training context length**; the optimizer, data, tokenizer, evaluation context and the wall-clock budget are fixed by the judge. +The static gate (`bash /app/public_test.sh` runs the judge's exact rules) is +strict: imports are limited to an allowlist — `torch`, `numpy`, `fla`, +`einops`, `triton` (custom kernels are fair game), `math` and a few +pure-computation stdlib modules — and dynamic-access primitives (bare +`eval`/`exec`/`getattr`/`__import__`/`open`, wildcard imports, dunder +attributes other than `__init__`/`__version__`) are rejected. `model.eval()` +and `torch.compile(...)` are fine. Run the gate locally before submitting; +its rejection reasons are exact. + ## Training context length — a real lever, with a real catch Declare it with a module-level integer in `model.py`: @@ -65,25 +76,28 @@ against `eval_block_size`, or build them lazily from the actual `T`. The baseline always trains at 8192, so you are trading against a fixed point. -Note the vocabulary is large: the embedding table alone is ~70M parameters at -d=704 (and the model ships it **tied**, one shared table), so how you handle it +Note the vocabulary is large: the embedding table alone is ~77M parameters at +d=768 (and the model ships it **tied**, one shared table), so how you handle it (tying, factorizing, resizing) is part of the design problem rather than a detail. ## The starting point -`/app/model.py` ships as the **3:1 Gated DeltaNet hybrid** from *Olmo Hybrid*: -three linear-recurrent layers per full-attention layer (25% attention), which -already beats the attention-only baseline. Your job is to push further. +`/app/model.py` ships as a **3:1 Gated DeltaNet hybrid** (following *Olmo Hybrid*, +arXiv:2604.03444): three linear-recurrent layers per full-attention layer (25% +attention), which already beats the attention-only baseline. Your job is to push +further. The 3:1 ratio is inherited, not tuned — it is the OLMo-3 sliding-window pattern `[4096, 4096, 4096, -1]`, and the recipe replaces exactly the sliding-window layers with GDN. It was chosen for a sliding window, not for a linear RNN, so there is no reason to believe it is optimal here. -The shipped model is also **parameter-matched** against the baseline the way -upstream does it (`REMOVE_HEADS`: d_model 768 -> 704, heads 12 -> 11, head_dim -64 preserved), which is why it is narrower than the shape its name suggests. +The shipped model is **not parameter-matched** to the baseline: it keeps the +baseline's shape (d=768, 12 heads, 12 layers) and lets the GDN mixer's wide +recurrent state run at its natural, larger size (~254M vs the baseline's +~190M). Capacity is bounded by the hard parameter cap, not by an artificial +per-arm match — how you spend the budget under that cap is yours to decide. Open questions it does not answer: the ratio (3:1 vs 5:1 vs 7:1), where the attention layers belong, recurrent state size, whether every layer should be diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harbor/app/model.py b/2.0/problems/nanoslm_hybrid_arch_design/harbor/app/model.py index 6170d9453..a9d2fa929 100644 --- a/2.0/problems/nanoslm_hybrid_arch_design/harbor/app/model.py +++ b/2.0/problems/nanoslm_hybrid_arch_design/harbor/app/model.py @@ -1,97 +1,87 @@ -"""Reference solution for nanoslm_hybrid_arch_design -- the Olmo-Hybrid recipe at 190M. +"""Reference solution for nanoslm_hybrid_arch_design -- a 3:1 GDN-to-attention hybrid built on the 190M olmo3_190M shape (254M params). -This is the STARTING POINT, not the ceiling. The locked baseline is a faithful -``olmo3_190M``; this file applies OLMo-core's own hybrid recipe -(``src/scripts/train/OLMo_hybrid/OLMo-hybrid-7B.py``) to it. +This is the starting point, not the ceiling. The locked baseline is a faithful +pure-attention ``olmo3_190M``; this file replaces most of its attention layers +with Gated DeltaNet (GDN) linear-recurrent layers. -WHERE THE 3:1 RATIO ACTUALLY COMES FROM ---------------------------------------- -It is NOT an arbitrary choice. ``olmo3_190M`` already carries:: +The 3:1 ratio +------------- +It is not an arbitrary choice. ``olmo3_190M`` already carries:: SlidingWindowAttentionConfig(pattern=[4096, 4096, 4096, -1], force_full_attention_on_first_layer=False, force_full_attention_on_last_layer=True) so 9 of its 12 layers are sliding-window and 3 (layers 3, 7, 11) are full -attention. The paper "replaces sliding window layers with Gated DeltaNet -layers", and upstream's recipe encodes exactly that:: +attention. This reference replaces each sliding-window layer with a GDN layer and +keeps the full-attention layers:: - config.block_pattern = ["gdn", "gdn", "gdn", "attn"] + block_pattern = ["gdn", "gdn", "gdn", "attn"] The three sliding-window layers of each 4-layer group become GDN; the -full-attention layer of each group survives untouched. The 3:1 ratio IS olmo3's -own sliding-window pattern. Note the consequence: the hybrid has NO sliding -window left anywhere -- every surviving attention layer sat at a ``-1`` position. +full-attention layer of each group survives untouched. The 3:1 ratio comes from +the base model's own sliding-window pattern. Note the consequence: the hybrid has +No sliding window left anywhere -- every surviving attention layer sat at a +``-1`` position. -PARAMETER MATCHING (upstream's REMOVE_HEADS practice) ------------------------------------------------------ -GDN's mixer is larger than an attention mixer, so upstream shrinks the model to -compensate rather than letting the hybrid quietly buy capacity:: +No parameter matching +--------------------- +A GDN mixer is larger than an attention mixer. This reference does not shrink the +hybrid to match the baseline's parameter count: it keeps the same d=768, h=12, +l=12 as the baseline and lets the hybrid run at its natural (larger) size. With +Tied embeddings:: - REMOVE_HEADS = 2 - config.d_model -= REMOVE_HEADS * 128 # 4096 -> 3840 - num_heads -= REMOVE_HEADS # 32 -> 30 - assert config.d_model / num_heads == 128 # head_dim preserved + baseline 190,354,176 (12 attention layers) + hybrid 254,430,936 (9 GDN + 3 attention layers) +33.7% total -At 190M with head_dim 64 the equivalent is ``REMOVE_HEADS = 1``:: +(Totals at the padded vocab 100352 -- dolma2's 100278 real ids plus 74 unused +rows on the tied table, the Olmo 3 padding convention.) - d_model 768 -> 704, n_heads 12 -> 11, 704 / 11 == 64 (head_dim preserved) +The hybrid is bigger -- the GDN mixer's wide recurrent state (head dim +ceil_128(0.75*d/h) = 128) costs params -- but well under the 400M ``param_cap``. +The hard cap -- not an artificial per-arm match -- is what bounds capacity, so +the scored question is whether the mixer choice pays off on quality within that +budget. -Matching is done on NON-EMBEDDING parameters, which is the quantity upstream -itself feeds to ``Duration.chinchilla_tokens(model_params=...)``. Both arms TIE -their single vocab table, which at this scale is ~40% of all parameters -- -including it would drown the mixer difference being matched. Measured -analytically: - - baseline non-embedding 113,283,840 - hybrid non-embedding 110,805,734 -2.19% - -which is the same tolerance upstream's own 7B pair achieves (+2.84%). One caveat -worth stating: because ``d_model`` shrinks 768 -> 704, the tied table shrinks too -(77.0M -> 70.6M), so the hybrid's TOTAL is 4.7% below the baseline's -(181,401,446 vs 190,297,344) even though its non-embedding count is within 2.2%. - -Note also that ``hidden_size`` stays 3072. Upstream mutates ``d_model`` in place -and never recomputes the feed-forward width, so the FFN keeps the width derived -from the ORIGINAL d_model. Reproduced here deliberately. +``hidden_size`` stays 3072 (the SwiGLU width for d_model 768), same as baseline. Everything else is identical to the baseline -- RMSNorm eps 1e-6, SwiGLU 3072, -reordered_norm blocks, qk_norm, RoPE theta 500_000, TIED embeddings (the baseline -ties, so this arm ties too for comparability) -- so the ONLY architectural +reordered_norm blocks, qk_norm, RoPE theta 500_000, tied embeddings (the baseline +ties, so this arm ties too for comparability) -- so the only architectural difference is the sequence mixer in those nine layers. -YOUR TASK: PUSH val_bpb FURTHER +Your task: Push val_bpb further ------------------------------- -Everything above is upstream's recipe, transposed to this scale. The open -questions it does NOT answer, each worth real bpb: +This 3:1 GDN hybrid is one point in a large design space. Open questions it does +Not settle, each worth real bpb: - * THE RATIO. 3:1 is inherited from olmo3's sliding-window pattern, which was - chosen for a sliding window, not for a linear RNN. 5:1 and 7:1 buy more - steps but give up more global context. - * PLACEMENT. Every 4th, or clustered (early layer for global context in, late + * The linear-to-attention ratio. 3:1 (GDN:attention) is inherited from the base model's sliding-window pattern, + which was chosen for a sliding window, not for a linear RNN. 5:1 and 7:1 buy + more steps but give up more global context. + * Layer placement. Every 4th, or clustered (early layer for global context in, late layer for read-out)? Same cost, different models. - * STATE SIZE. ``expand_v`` and ``num_v_heads`` set the recurrent state, the + * State size. ``expand_v`` and ``num_v_heads`` set the recurrent state, the main capacity knob of a linear RNN -- unlike a KV cache it does not grow with sequence length, so capacity is cheap at 8192. - * THE MIXER ITSELF. GDN is one choice; GLA, RetNet and Mamba2-style mixers are + * The mixer itself. GDN is one choice; GLA, RetNet and Mamba2-style mixers are all expressible in the same chunkwise matmul form. - * NON-UNIFORMITY. Nothing requires every GDN layer to be identical, or the + * Non-uniformity. Nothing requires every GDN layer to be identical, or the attention layers to be full-width. - * THE REST OF THE BLOCK. Norm placement, gating, MLP ratio, head count, + * The rest of the block. Norm placement, gating, MLP ratio, head count, embedding tying -- all still on the table, and all interact with the above. - * THE TRAINING CONTEXT. Declare a module-level ``BLOCK_SIZE`` int (a power of + * The training context. Declare a module-level ``BLOCK_SIZE`` int (a power of two in [256, 8192]) to train at a shorter context than the default 8192. Shorter steps are cheaper, so you complete more of them in the fixed - wall-clock budget -- but EVALUATION IS ALWAYS AT 8192, so the model must + wall-clock budget -- but **evaluation is always at 8192**, so the model must extrapolate to positions well beyond anything it trained on. How well it - does that is mostly a property of the POSITION ENCODING (plain RoPE degrades + does that is mostly a property of the position encoding (plain RoPE degrades sharply; NTK-aware/YaRN scaling and ALiBi hold up far better), which is a different question from mixer efficiency. This file declares no BLOCK_SIZE and therefore trains at 8192. Anything sized off ``config.block_size`` must still run at ``config.eval_block_size``. -Locked and not yours to change: optimizer, data, tokenizer, the EVALUATION -context (8192), and the wall-clock budget. Interface: ``build_model(config)`` / +Locked and not yours to change: optimizer, data, tokenizer, the evaluation +context, and the wall-clock budget. Interface: ``build_model(config)`` / ``NanoSLM(config)`` returning ``forward(idx) -> logits [B, T, vocab_size]``. """ @@ -104,36 +94,31 @@ import torch.nn.functional as F # --------------------------------------------------------------------------- # -# Shape: olmo3_190M with upstream's REMOVE_HEADS compensation applied. +# Architectural Details # --------------------------------------------------------------------------- # -REMOVE_HEADS = 1 - +# Attention hyperparameters +REMOVE_HEADS = 0 N_LAYER = 12 -N_HEAD = 12 - REMOVE_HEADS # 11 -HEAD_DIM = 64 # preserved, exactly as upstream asserts -N_EMBD = 768 - REMOVE_HEADS * HEAD_DIM # 704 -assert N_EMBD // N_HEAD == HEAD_DIM, "REMOVE_HEADS must preserve head_dim" +N_HEAD = 12 - REMOVE_HEADS # 12 +HEAD_DIM = 64 +N_EMBD = 768 - REMOVE_HEADS * HEAD_DIM # 768 +assert N_EMBD // N_HEAD == HEAD_DIM, "head_dim must be preserved" -# NOT recomputed from the reduced d_model -- upstream mutates d_model in place -# and leaves the feed-forward width at the value llama_like derived from the -# ORIGINAL 768. Reproduced deliberately. -HIDDEN_SIZE = 3072 +HIDDEN_SIZE = 3072 # SwiGLU width, same as baseline ROPE_THETA = 500_000.0 NORM_EPS = 1e-6 -# GatedDeltaNetConfig(n_heads=num_heads, head_dim=int(0.75*d_model/num_heads), -# allow_neg_eigval=True) with expand_v at its 2.0 default. -GDN_HEAD_DIM = int(0.75 * N_EMBD / N_HEAD) # 48 +# GDN hyperparameters +_gdn_raw = int(0.75 * N_EMBD / N_HEAD) +GDN_HEAD_DIM = ((_gdn_raw + 127) // 128) * 128 GDN_EXPAND_V = 2.0 -GDN_HEAD_V_DIM = int(GDN_HEAD_DIM * GDN_EXPAND_V) # 96 -GDN_KEY_DIM = N_HEAD * GDN_HEAD_DIM # 528 -GDN_VALUE_DIM = N_HEAD * GDN_HEAD_V_DIM # 1056 +GDN_ALLOW_NEG_EIGVAL = True +GDN_HEAD_V_DIM = int(GDN_HEAD_DIM * GDN_EXPAND_V) # 256 +GDN_KEY_DIM = N_HEAD * GDN_HEAD_DIM # 1536 +GDN_VALUE_DIM = N_HEAD * GDN_HEAD_V_DIM # 3072 GDN_CONV_SIZE = 4 -# config.block_pattern = ["gdn", "gdn", "gdn", "attn"] -- the three -# sliding-window layers of each group become GDN, the full-attention layer -# survives. Attention therefore lands at layers 3, 7, 11. PATTERN = ("gdn", "gdn", "gdn", "attn") assert N_LAYER % len(PATTERN) == 0 @@ -162,17 +147,7 @@ def _rope(x, theta: float = ROPE_THETA): class _Attention(nn.Module): - """Identical to the baseline's attention, always FULL causal. - - Full, not windowed, and that is not a simplification: the attention layers - the hybrid keeps are exactly the ones sitting at the ``-1`` positions of - olmo3's [4096, 4096, 4096, -1] pattern, so they were already full attention - before the GDN substitution. - - qk_norm spans the full n_heads*head_dim projection and is applied BEFORE the - head reshape -- upstream's behaviour when ``use_head_qk_norm`` is False, - which is the olmo3_190M default. - """ + """Identical to the baseline's attention, always full causal.""" def __init__(self, n_embd: int, n_head: int): super().__init__() @@ -198,123 +173,22 @@ def forward(self, x): return self.w_out(y.transpose(1, 2).contiguous().view(B, T, C)) -class _ChunkedDeltaNet(nn.Module): - """CPU-only chunkwise fallback, PARAMETER-IDENTICAL to fla's GatedDeltaNet. - - Exists so the CPU smoke path can build and report the same parameter count - the CUDA path will. Its module inventory is fla 0.5.1's exactly -- - q/k/v/a/b/g/o projections, three depthwise short convolutions, A_log, - dt_bias and the gated output norm -- so ``n_params`` is device-independent. - - The recurrence is a chunkwise GATED LINEAR ATTENTION, not the full delta - rule: it carries the decay and the gates but skips the delta-rule inverse. - That is a deliberate approximation. It is never scored -- CUDA raises rather - than falling back here (see ``_make_gdn``) -- and it exists to prove wiring, - not to reproduce GDN's numerics. - """ - - def __init__(self, chunk: int = 512): - super().__init__() - self.chunk = chunk - d, H = N_EMBD, N_HEAD - self.q_proj = nn.Linear(d, GDN_KEY_DIM, bias=False) - self.k_proj = nn.Linear(d, GDN_KEY_DIM, bias=False) - self.v_proj = nn.Linear(d, GDN_VALUE_DIM, bias=False) - self.a_proj = nn.Linear(d, H, bias=False) - self.b_proj = nn.Linear(d, H, bias=False) - self.g_proj = nn.Linear(d, GDN_VALUE_DIM, bias=False) - self.o_proj = nn.Linear(GDN_VALUE_DIM, d, bias=False) - self.A_log = nn.Parameter(torch.zeros(H)) - self.dt_bias = nn.Parameter(torch.zeros(H)) - # ShortConvolution weights: depthwise, (channels, 1, kernel), no bias. - self.q_conv1d = nn.Parameter(torch.zeros(GDN_KEY_DIM, 1, GDN_CONV_SIZE)) - self.k_conv1d = nn.Parameter(torch.zeros(GDN_KEY_DIM, 1, GDN_CONV_SIZE)) - self.v_conv1d = nn.Parameter(torch.zeros(GDN_VALUE_DIM, 1, GDN_CONV_SIZE)) - self.o_norm = nn.Parameter(torch.ones(GDN_HEAD_V_DIM)) - - @staticmethod - def _short_conv(x, w): - """Causal depthwise conv over time. x [B, T, C], w [C, 1, K].""" - C, K = w.shape[0], w.shape[2] - xt = F.pad(x.transpose(1, 2), (K - 1, 0)) - return F.conv1d(xt, w, groups=C).transpose(1, 2) - - def forward(self, x): - B, T, _ = x.shape - H, DK, DV, S = N_HEAD, GDN_HEAD_DIM, GDN_HEAD_V_DIM, self.chunk - - q = F.silu(self._short_conv(self.q_proj(x), self.q_conv1d)) - k = F.silu(self._short_conv(self.k_proj(x), self.k_conv1d)) - v = self._short_conv(self.v_proj(x), self.v_conv1d) - - q = q.view(B, T, H, DK).transpose(1, 2) - k = k.view(B, T, H, DK).transpose(1, 2) - v = v.view(B, T, H, DV).transpose(1, 2) - - # Per-head decay in (0, 1): the gate that makes this "gated". - dt = F.softplus(self.a_proj(x) + self.dt_bias) # [B, T, H] - g = torch.exp(-torch.exp(self.A_log) * dt) # [B, T, H] - g = g.permute(0, 2, 1).unsqueeze(-1) # [B, H, T, 1] - beta = torch.sigmoid(self.b_proj(x)).permute(0, 2, 1).unsqueeze(-1) - v = v * beta - - pad = (S - T % S) % S - if pad: - q, k, v = (F.pad(t, (0, 0, 0, pad)) for t in (q, k, v)) - g = F.pad(g, (0, 0, 0, pad), value=1.0) - nC = (T + pad) // S - rs = lambda t: t.view(B, H, nC, S, -1) # noqa: E731 - qc, kc, vc, gc = rs(q), rs(k), rs(v), rs(g) - - causal = torch.tril(torch.ones(S, S, device=x.device, dtype=torch.bool)) - intra = (qc @ kc.transpose(-1, -2)).masked_fill(~causal, 0.0) @ vc - - state = torch.zeros(B, H, DK, DV, device=x.device, dtype=x.dtype) - outs = [] - for c in range(nC): - outs.append(intra[:, :, c] + qc[:, :, c] @ state) - decay = gc[:, :, c].prod(dim=2, keepdim=True) - state = state * decay + kc[:, :, c].transpose(-1, -2) @ (vc[:, :, c] * gc[:, :, c]) - y = torch.stack(outs, dim=2).view(B, H, T + pad, DV)[:, :, :T] - - # Gated RMSNorm over head_v_dim, then merge heads. - yf = y.float() - y = (yf * torch.rsqrt(yf.pow(2).mean(-1, keepdim=True) + 1e-5)).to(x.dtype) - y = y * self.o_norm - y = y.transpose(1, 2).contiguous().view(B, T, GDN_VALUE_DIM) - y = y * F.silu(self.g_proj(x)) - return self.o_proj(y) - - def _make_gdn(): - """fla's fused GatedDeltaNet on CUDA; chunkwise only on CPU (smoke). - - WHY fla IS REQUIRED HERE, not merely preferred: PyTorch ships a fused kernel - for softmax attention (SDPA) but none for a linear/recurrent mixer, so a - hand-written GDN is unfused eager code. Scoring is fixed WALL-CLOCK, so an - unfused mixer loses on throughput regardless of architectural merit -- - measured at ctx 8192, the chunkwise fallback did 6% of attention's FLOPs and - still ran 3.5x slower. Both arms must be kernel-matched or the score - measures kernel quality, not architecture. - - HARD FAILURE ON CUDA, deliberately: a silent fallback here once produced a - scored-looking run where the reference "lost" 18 steps to 65 purely because - fla was missing. A missing kernel must raise, not quietly change what is - being measured. - """ + """fla's fused GatedDeltaNet. CUDA required.""" if not torch.cuda.is_available(): - return _ChunkedDeltaNet() - - from fla.layers import GatedDeltaNet # ImportError here is intentional + raise RuntimeError( + "reference hybrid requires a CUDA device: the GDN layers run only " + "through fla's fused Triton kernels, and there is deliberately no " + "CPU fallback" + ) - # head_dim MUST be passed: fla defaults it to 256, so omitting it builds - # num_heads*256-wide projections instead of the intended width -- a silent - # param blowup that trains a much larger model than the baseline. + from fla.layers import GatedDeltaNet return GatedDeltaNet( hidden_size=N_EMBD, num_heads=N_HEAD, head_dim=GDN_HEAD_DIM, expand_v=GDN_EXPAND_V, + allow_neg_eigval=GDN_ALLOW_NEG_EIGVAL, use_gate=True, use_short_conv=True, ) @@ -327,8 +201,8 @@ def __init__(self): def forward(self, x): out = self.impl(x) - # fla layers return (hidden_states, attentions, past_kv); the chunkwise - # fallback returns a tensor. Normalize so _Block need not care. + # fla layers return (hidden_states, attentions, past_kv) -- keep only + # the hidden states. return out[0] if isinstance(out, tuple) else out @@ -369,15 +243,12 @@ def __init__(self, config): kinds = [PATTERN[i % len(PATTERN)] for i in range(N_LAYER)] self.blocks = nn.ModuleList([_ReorderedNormBlock(N_EMBD, N_HEAD, k) for k in kinds]) self.norm_f = _RMSNorm(N_EMBD) - # TIED -- the baseline ties its embeddings (a deliberate deviation from - # upstream's untied default), so the reference ties too and the two arms - # stay comparable: the ONLY architectural difference between them is the - # sequence mixer in the GDN layers, not the output-head param budget. The - # tie is set AFTER init (below) so the shared table carries wte's init. + self.lm_head = nn.Linear(N_EMBD, config.vocab_size, bias=False) self.apply(self._init) + for name, p in self.named_parameters(): - if name.endswith("w_out.weight") or name.endswith("w2.weight"): + if name.endswith(("w_out.weight", "w2.weight", "o_proj.weight")): nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * N_LAYER)) self.lm_head.weight = self.wte.weight # weight tying @@ -420,14 +291,20 @@ def analytic_n_params(vocab_size: int) -> int: n_attn = sum(1 for i in range(N_LAYER) if PATTERN[i % len(PATTERN)] == "attn") n_gdn = N_LAYER - n_attn non_emb = n_attn * (attn_mix + block_norms + ffn) + n_gdn * (gdn_mix + block_norms + ffn) + d - # ONE vocab-sized table: lm_head is TIED to wte (see NanoSLM.__init__). + return non_emb + vocab_size * d -def _self_check(vocab_size: int = 100278) -> int: - from harness.model_config import ModelConfig +def _self_check(vocab_size: int = 100352) -> int: + # A minimal stand-in for the harness-provided config object + class _Cfg: + def __init__(self, vocab_size: int, block_size: int, + eval_block_size: int = 8192): + self.vocab_size = vocab_size + self.block_size = block_size + self.eval_block_size = eval_block_size - m = NanoSLM(ModelConfig(vocab_size=vocab_size, block_size=8192)) + m = NanoSLM(_Cfg(vocab_size, 8192)) got = sum(p.numel() for p in m.parameters() if p.requires_grad) want = analytic_n_params(vocab_size) assert got == want, f"param mismatch: built {got} != analytic {want}" diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harness/baseline_model.py b/2.0/problems/nanoslm_hybrid_arch_design/harness/baseline_model.py index eb2e1e2ad..c5ed765af 100644 --- a/2.0/problems/nanoslm_hybrid_arch_design/harness/baseline_model.py +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/baseline_model.py @@ -1,79 +1,10 @@ -"""Locked baseline architecture (hidden; the score-0 reference point). - -A plain-PyTorch reimplementation of OLMo-core's ``TransformerConfig.olmo3_190M``, -faithful in every respect EXCEPT that the word embeddings are TIED (see -"ON THE PARAMETER COUNT" below). Upstream unties them; this benchmark ties BOTH -arms (this baseline and ``reference.py``) instead, as a deliberate choice so the -only architectural difference between them is the sequence mixer, not the -output-head param budget. Tying is NOT forced by the ``param_cap`` (400M -- the -untied 267.3M model would fit); it is a comparability choice, and the single -deliberate deviation from upstream. - -Implements the exact interface the agent must also honor: ``build_model(config)`` -/ ``NanoSLM(config)`` returning a module whose ``forward(idx)`` yields logits -``[B, T, vocab_size]``. - -WHAT olmo3_190M ACTUALLY IS (verified against upstream, not reconstructed) -------------------------------------------------------------------------- -``olmo3_190M`` == ``olmo2_190M`` + sliding window + flash_2, and ``olmo2_190M`` -is ``llama_like(...)`` with:: - - d_model=768, n_layers=12, n_heads=12 -> head_dim 64 - hidden_size_multiplier=1.5 - block_name=TransformerBlockType.reordered_norm - qk_norm=True - rope_theta=500_000 - layer_norm_eps=1e-6 - -``llama_like`` then fixes the rest:: - - hidden_size = int(8 * d_model / 3) # 2048 - hidden_size = int(1.5 * hidden_size) # 3072 - hidden_size = ensure_multiple_of(., 256) # 3072 - bias = False everywhere - layer_norm = LayerNormType.rms -> RMSNorm (NOT LayerNorm) - -and ``olmo3_190M`` adds:: - - SlidingWindowAttentionConfig(pattern=[4096, 4096, 4096, -1], - force_full_attention_on_first_layer=False, - force_full_attention_on_last_layer=True) - -THE reordered_norm BLOCK IS NOT PRE-NORM. From upstream ``block.py``:: - - h = x + attention_norm(attention(x)) - out = h + feed_forward_norm(feed_forward(h)) - -The norm is applied to the residual branch's OUTPUT, not to its input. An -earlier revision of this file was a nanoGPT-style pre-norm block with only the -dimensions swapped -- LayerNorm, GELU 4x MLP, biases everywhere, no qk_norm, no -sliding window -- which measured 239,083,008 params and was not olmo3_190M in -any respect other than d/L/H. That is retracted. - -ON THE PARAMETER COUNT -- READ THIS BEFORE "FIXING" IT ------------------------------------------------------- -Upstream leaves ``tie_word_embeddings`` at its dataclass default of False, so a -faithful olmo3_190M carries TWO vocab-sized matrices and totals **267,310,848** -trainable params at vocab 100278. That fits comfortably under this task's -``param_cap`` of 400,000,000, so the tie is NOT a legality workaround. This -baseline (and ``reference.py``) TIE the embeddings so ``lm_head`` shares -``wte``'s weight, one 77,013,504-param table is removed, and the total drops to:: - - 12 blocks (113,283,072) + ONE shared embedding table (77,013,504) - + final norm (768) = 190,297,344 - -which is the "190M" the model's name refers to. Tying BOTH arms keeps them -differing only in the sequence mixer; the shared table is still ~40% of all -params -- at the 190M shape with a 100278-id vocabulary the embedding dominates; -at 7B the same table is ~11% and near-negligible. - -WHY THE BASELINE IS NOT A HYBRID --------------------------------- -``reference.py`` is the Olmo-Hybrid recipe applied to this exact model, and it -is the floor the agent starts from. For that reference to be a meaningful floor -it has to beat something, so the score-0 point is the pure-attention model it -improves on. The scored question is therefore "how much further past a -competent hybrid can you push val_bpb". +"""Baseline Transformer Architecture. CUDA-only, like everything scored here. + +A plain-PyTorch reimplementation of OLMo-core's ``TransformerConfig.olmo3_190M``. + +``reference.py`` is the Olmo-Hybrid recipe applied to the same scale, and it is +the floor the agent starts from. The agent will need to solve the problem of +how much further past a competent hybrid it can push val_bpb". """ from __future__ import annotations @@ -89,43 +20,19 @@ # --------------------------------------------------------------------------- # N_LAYER = 12 N_HEAD = 12 -N_EMBD = 768 # head_dim = 768 / 12 = 64 +N_EMBD = 768 HEAD_DIM = N_EMBD // N_HEAD - -# int(8*768/3) = 2048 -> int(1.5*2048) = 3072 -> ensure_multiple_of(3072, 256) HIDDEN_SIZE = 3072 - ROPE_THETA = 500_000.0 NORM_EPS = 1e-6 - -# SlidingWindowAttentionConfig(pattern=[4096,4096,4096,-1], -# force_full_attention_on_first_layer=False, -# force_full_attention_on_last_layer=True) SWA_PATTERN = (4096, 4096, 4096, -1) FORCE_FULL_FIRST = False FORCE_FULL_LAST = True - -# TRAINING CONTEXT OF THE LOCKED ARM, declared through the same module-level -# ``BLOCK_SIZE`` protocol a submission uses (runner.resolve_train_block_size). -# -# Stated EXPLICITLY rather than inheriting the task's ``block_size``: that field -# (8192 in the task config) is now only the submission-facing DEFAULT, and the -# score-0 reference point must not silently move if it is ever retuned. It also -# has to equal ``eval_block_size`` (also 8192) for the cached baseline to remain -# a valid comparison partner for submissions that train shorter -- the baseline -# never trades context for steps, so the trade a submission makes is measured -# against a fixed point. BLOCK_SIZE = 8192 def window_size_for_layer(layer_idx: int, n_layers: int) -> int: - """Port of ``SlidingWindowAttentionConfig._get_window_size``. -1 == full. - - Note the ``force_full_attention_on_first_layer`` branch also SHIFTS the - pattern index (upstream applies the pattern starting from the second layer - in that case). olmo3 sets it False, so there is no shift here -- but the - shift is reproduced anyway so the function is a faithful port. - """ + """Port of ``SlidingWindowAttentionConfig._get_window_size``. -1 == full.""" if FORCE_FULL_FIRST and layer_idx == 0: return -1 if FORCE_FULL_LAST and layer_idx == (n_layers - 1): @@ -134,9 +41,6 @@ def window_size_for_layer(layer_idx: int, n_layers: int) -> int: return SWA_PATTERN[eff % len(SWA_PATTERN)] -# For N_LAYER=12 this yields full attention at layers 3, 7, 11 and a 4096-wide -# window everywhere else -- the 9:3 split the reference's 3:1 GDN:attention -# ratio comes from. It is olmo3's own pattern, not an arbitrary choice. LAYER_WINDOWS = tuple(window_size_for_layer(i, N_LAYER) for i in range(N_LAYER)) @@ -156,38 +60,40 @@ def forward(self, x): # --------------------------------------------------------------------------- # -# Sliding-window mask cache. -# -# torch's SDPA has no native window flag, so a banded causal mask is built and -# handed to `attn_mask`. It is cached per (T, window, device): one bool mask at -# T=8192 is 64 MiB, so rebuilding it 9x per forward would be ruinous, and all -# nine windowed layers share the same window (4096) hence the same mask. -# -# KERNEL CAVEAT, stated plainly: passing an explicit `attn_mask` disqualifies -# SDPA's fused flash kernel and falls back to the memory-efficient backend, -# whereas upstream uses flash_2's native `window_size=(4095, 0)`. The windowed -# layers are therefore SLOWER here than upstream would be, despite doing fewer -# FLOPs. At ctx 8192 a 4096 window removes only ~25% of attention FLOPs anyway -# (full causal already averages T/2 = 4096 keys per query), so the window is a -# faithfulness feature here, not a speed feature. +# Windowed-attention kernel path: FlexAttention, CUDA-only. # --------------------------------------------------------------------------- # -_MASK_CACHE: dict = {} +_FLEX_BLOCK_MASKS: dict = {} +_FLEX_FN = None -def _sliding_causal_mask(T: int, window: int, device) -> torch.Tensor: - """Bool mask [T, T], True == attend. Keys in [i-window+1, i] inclusive. +def _windowed_attention(q, k, v, window: int) -> torch.Tensor: + global _FLEX_FN + if not q.is_cuda: + raise RuntimeError( + "baseline_model requires a CUDA device: sliding-window layers run " + "only through FlexAttention's compiled kernel (no CPU fallback)" + ) + from torch.nn.attention.flex_attention import ( + create_block_mask, + flex_attention, + ) - Matches upstream's flash window_size=(window-1, 0), documented there as - "window is [i - window_size[0], i + window_size[1]] inclusive". - """ - key = (T, window, str(device)) - m = _MASK_CACHE.get(key) - if m is None: - i = torch.arange(T, device=device) - q, k = i[:, None], i[None, :] - m = (k <= q) & (k > q - window) - _MASK_CACHE[key] = m - return m + if _FLEX_FN is None: + _FLEX_FN = torch.compile(flex_attention, dynamic=False) + + if torch.is_autocast_enabled("cuda"): + dt = torch.get_autocast_dtype("cuda") + q, k, v = q.to(dt), k.to(dt), v.to(dt) + T = q.shape[2] + key = (T, window, str(q.device)) + bm = _FLEX_BLOCK_MASKS.get(key) + if bm is None: + def mask_mod(b, h, qi, ki): + return (ki <= qi) & (ki > qi - window) + + bm = create_block_mask(mask_mod, None, None, T, T, device=str(q.device)) + _FLEX_BLOCK_MASKS[key] = bm + return _FLEX_FN(q, k, v, block_mask=bm) def _rope(x, theta: float = ROPE_THETA): @@ -206,14 +112,7 @@ def _rope(x, theta: float = ROPE_THETA): class _Attention(nn.Module): - """olmo_core AttentionConfig: bias=False, qk_norm=RMSNorm, RoPE, optional SWA. - - qk_norm is applied over the FULL n_heads*head_dim projection, BEFORE the - head reshape. That is upstream's behaviour when ``use_head_qk_norm`` is - False, which is the olmo3_190M default (``llama_like`` only forwards - ``use_head_qk_norm`` when explicitly asked, and olmo2_190M never asks). It - is NOT a per-head norm -- q_norm/k_norm are 768-wide, not 64-wide. - """ + """olmo_core AttentionConfig: bias=False, qk_norm=RMSNorm, RoPE, optional SWA.""" def __init__(self, n_embd: int, n_head: int, window: int): super().__init__() @@ -240,8 +139,7 @@ def forward(self, x): if self.window == -1: y = F.scaled_dot_product_attention(q, k, v, is_causal=True) else: - mask = _sliding_causal_mask(T, self.window, x.device) - y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask) + y = _windowed_attention(q, k, v, self.window) return self.w_out(y.transpose(1, 2).contiguous().view(B, T, C)) @@ -259,7 +157,7 @@ def forward(self, x): class _ReorderedNormBlock(nn.Module): - """TransformerBlockType.reordered_norm -- norm AFTER the residual branch.""" + """TransformerBlockType.reordered_norm -- norm after the residual branch.""" def __init__(self, n_embd: int, n_head: int, window: int): super().__init__() @@ -276,15 +174,18 @@ def forward(self, x): class NanoSLM(nn.Module): def __init__(self, config): super().__init__() + if not torch.cuda.is_available(): + raise RuntimeError( + "baseline_model requires a CUDA device: sliding-window layers " + "run only through FlexAttention's compiled kernel (no CPU " + "fallback)" + ) self.block_size = config.block_size self.wte = nn.Embedding(config.vocab_size, N_EMBD) self.blocks = nn.ModuleList( [_ReorderedNormBlock(N_EMBD, N_HEAD, w) for w in LAYER_WINDOWS] ) - # LMHeadConfig(layer_norm=rms, bias=False). Embeddings are TIED: the - # lm_head reuses wte's weight (a comparability choice so both arms differ - # only in the mixer -- see the module docstring; NOT forced by param_cap). - # The tie is set AFTER init so the shared tensor carries wte's init. + self.norm_f = _RMSNorm(N_EMBD) self.lm_head = nn.Linear(N_EMBD, config.vocab_size, bias=False) self.apply(self._init_weights) diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harness/data.py b/2.0/problems/nanoslm_hybrid_arch_design/harness/data.py index 761ccd1ce..019ef7a2c 100644 --- a/2.0/problems/nanoslm_hybrid_arch_design/harness/data.py +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/data.py @@ -1,22 +1,16 @@ """Token-level data loading (torch path). -Train/val data are flat ``uint32`` TOKEN-ID streams (``.bin``) baked into the +Train/val data are flat ``uint32`` token-ID streams (``.bin``) baked into the judge image by ``docker/prep_assets.py`` using the dolma2 BPE tokenizer -(``allenai/dolma2-tokenizer``, vocab 100278). The validation stream is HELD OUT -and never mounted in ``/app``. When the files are absent (CPU smoke / local dev) -a deterministic synthetic token stream is generated so the harness wiring runs -end-to-end without baked assets. - -WHY uint32 AND NOT uint16 -------------------------- -The dolma2 vocabulary is 100278 ids. ``uint16`` tops out at 65535, so the upper -~35k of the vocabulary would wrap silently into low ids -- a corrupted corpus -that still loads, still trains, and only shows up as an unexplained bpb floor. -The stream is therefore uint32 and ``_load_bin`` asserts every id is in range. - -BYTE ACCOUNTING +(``allenai/dolma2-tokenizer``, 100278 real ids; the model vocabulary is padded +to 100352, so streams never index the top 74 rows). The validation stream is held out +and never mounted in ``/app``. The corpus is required: when the ``.bin`` files +are absent or unusable the harness raises ``DataError`` rather than fabricating +tokens, so a run never reports a bits-per-byte computed from synthetic data. + +Byte accounting --------------- -``val_bytes_for(n)`` reports how many BYTES OF TEXT ``n`` scored target tokens +``val_bytes_for(n)`` reports how many bytes of text ``n`` scored target tokens cover. ``eval_ppl`` divides by that, not by the token count -- see ``settings.resolve_val_bytes_per_token`` for why that distinction is the whole point of the metric. @@ -30,7 +24,6 @@ import numpy as np from .settings import ( - SYNTHETIC_BYTES_PER_TOKEN, TaskConfig, resolve_val_bytes_per_token, ) @@ -39,33 +32,23 @@ class DataError(RuntimeError): - """Corpus is present but unusable (bad dtype, out-of-range ids, no byte count).""" - - -def _synthetic_tokens(n: int, seed: int, vocab_size: int) -> np.ndarray: - """Deterministic pseudo-text token stream for smoke runs (no baked data). - - Emits plausible BPE ids: a repeating "vocabulary" of a few thousand common - ids (real text is heavily Zipfian, so a smoke model can actually reduce - perplexity below uniform) plus occasional rare ids from the full range, so - the embedding table is exercised beyond its first rows. - """ - rng = np.random.default_rng(seed) - common = rng.integers(0, min(4096, vocab_size), size=max(1024, n // 8)) - stream = np.tile(common, (n // common.size) + 1)[:n].astype(np.int64) - rare_mask = rng.random(n) < 0.02 - stream[rare_mask] = rng.integers(0, vocab_size, size=int(rare_mask.sum())) - return stream.astype(TOKEN_DTYPE) + """Corpus is absent or unusable (missing/empty file, bad dtype, out-of-range ids, no byte count).""" -def _load_bin(path: str, *, fallback_n: int, seed: int, vocab_size: int) -> tuple[np.ndarray, bool]: - """Return ``(tokens, is_real)``; synthesize when the file is absent.""" +def _load_bin(path: str, *, vocab_size: int, label: str) -> np.ndarray: + """Load a ``uint32`` token stream, raising ``DataError`` if it is missing or unusable.""" try: arr = np.fromfile(path, dtype=TOKEN_DTYPE) - except OSError: - return _synthetic_tokens(fallback_n, seed, vocab_size), False + except OSError as exc: + raise DataError( + f"{label} corpus not found at {path!r} ({type(exc).__name__}). Stage the " + "train.bin/val.bin (docker/prep_assets.py, or mount the Modal corpus " + "volume)." + ) from exc if arr.size == 0: - return _synthetic_tokens(fallback_n, seed, vocab_size), False + raise DataError( + f"{label} corpus at {path!r} is empty -- re-run docker/prep_assets.py." + ) # A uint8 stream left over from the byte-level era reads as uint32 without # error -- it just produces garbage ids. Catch it here rather than as a # mystery bpb. (An out-of-range id would also index past the embedding @@ -76,19 +59,19 @@ def _load_bin(path: str, *, fallback_n: int, seed: int, vocab_size: int) -> tupl f"{path}: token id {hi} >= vocab_size {vocab_size} " f"(stale corpus, or written with the wrong dtype)" ) - return arr, True + return arr class TokenData: """Holds train/val token arrays and yields fixed-order training batches. - TWO CONTEXT LENGTHS -- READ BEFORE EDITING + Two context lengths -- Read before editing --------------------------------------------------------- ``val_windows`` is cut at ``cfg.eval_block_size`` (8192, judge-owned) and - NOTHING here may make it depend on the training context: the scored + Nothing here may make it depend on the training context: the scored bits-per-byte is only comparable across submissions because every arm is evaluated on the identical windows of the identical held-out stream. The - TRAINING path (``batch``) uses the per-arm training context instead, which a + Training path (``batch``) uses the per-arm training context instead, which a submission may set below 8192 to buy optimizer steps. """ @@ -99,36 +82,29 @@ def __init__(self, cfg: TaskConfig): # (runner.MAX_TRAIN_BLOCK), so sizing off the max keeps both paths safe # no matter what a submission asks for. self.max_block = max(cfg.block_size, cfg.eval_block_size) - self.train, train_real = _load_bin( - cfg.train_tokens_path, fallback_n=1 << 20, - seed=cfg.seed, vocab_size=cfg.vocab_size, + self.train = _load_bin( + cfg.train_tokens_path, vocab_size=cfg.vocab_size, label="train", ) - self.val, self.val_is_real = _load_bin( - cfg.val_tokens_path, fallback_n=cfg.val_tokens + cfg.eval_block_size + 1, - seed=cfg.seed + 999, vocab_size=cfg.vocab_size, + self.val = _load_bin( + cfg.val_tokens_path, vocab_size=cfg.vocab_size, label="val", ) if self.train.size < self.max_block + 1: - self.train = _synthetic_tokens(1 << 20, cfg.seed, cfg.vocab_size) - train_real = False - self.train_is_real = train_real + raise DataError( + f"train corpus at {cfg.train_tokens_path!r} has only " + f"{self.train.size} tokens; needs > {self.max_block} to fill one " + "window -- stage the full shard (docker/prep_assets.py)." + ) - # Bytes of text per scored target token. ASSERTED, never defaulted to 1: + # Bytes of text per scored target token. Asserted, never defaulted to 1: # a ratio of 1 would silently turn val_bpb back into per-token CE/ln2. ratio = resolve_val_bytes_per_token(cfg) if ratio is None: - if self.val_is_real: - raise DataError( - "held-out byte count unavailable: no val_bytes in " - f"{cfg.manifest_path!r}, no FRONTIER_NANOSLM_VAL_BYTES, and " - "TaskConfig.val_bytes is 0. bits-per-byte cannot be " - "computed from a token count -- re-run docker/prep_assets.py." - ) - # Synthetic stream: there is no underlying text, so use the measured - # dolma2 compression ratio to keep smoke bpb in a plausible range. - ratio = SYNTHETIC_BYTES_PER_TOKEN - self.val_bytes_estimated = True - else: - self.val_bytes_estimated = not self.val_is_real + raise DataError( + "held-out byte count unavailable: no val_bytes in " + f"{cfg.manifest_path!r}, no FRONTIER_NANOSLM_VAL_BYTES, and " + "TaskConfig.val_bytes is 0. bits-per-byte cannot be computed " + "from a token count -- re-run docker/prep_assets.py." + ) self.bytes_per_token = float(ratio) def val_bytes_for(self, n_target_tokens: int) -> float: @@ -138,11 +114,11 @@ def val_bytes_for(self, n_target_tokens: int) -> float: def batch(self, step: int, device: str, block_size: int | None = None): """Deterministic (step-keyed) training batch of token windows. - ``block_size`` is the TRAINING context of the arm being trained, which - may differ per submission; it defaults to the config's. The window START - INDICES are drawn against ``self.max_block`` rather than against the + ``block_size`` is the training context of the arm being trained, which + may differ per submission; it defaults to the config's. The window start + Indices are drawn against ``self.max_block`` rather than against the requested width, so two arms with different training contexts still - begin their windows at the SAME offsets -- a short-context arm reads a + begin their windows at the same offsets -- a short-context arm reads a prefix of what a long-context arm reads, instead of an unrelated slice. That keeps common random numbers as close to intact as differing window widths permit. @@ -170,12 +146,11 @@ def batch(self, step: int, device: str, block_size: int | None = None): def val_windows(self, device: str): """Yield non-overlapping (x, y) windows over the held-out token stream. - ALWAYS ``cfg.eval_block_size`` wide -- never the training context. A - submission that trained at 2048 is still scored on 8192-wide windows and - must extrapolate to those positions; that is the deliberate trade - described above, and the reason the bpb denominator (bytes per - scored target token x number of target tokens) is unaffected by the - training context. + Always ``cfg.eval_block_size`` wide. A submission that trained at 2048 + is still scored on 8192-wide windows and must extrapolate to those + positions; that is the deliberate trade described above, and the reason + the bpb denominator (bytes per scored target token x number of target + tokens) is unaffected by the training context. """ import torch diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harness/eval_ppl.py b/2.0/problems/nanoslm_hybrid_arch_design/harness/eval_ppl.py index 3a00dbbea..317324c80 100644 --- a/2.0/problems/nanoslm_hybrid_arch_design/harness/eval_ppl.py +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/eval_ppl.py @@ -1,38 +1,19 @@ """Held-out bits-per-byte evaluation (torch path). -The harness computes cross-entropy from the model's ``logits`` on the HELD-OUT +The harness computes cross-entropy from the model's ``logits`` on the held-out token stream and returns val_bpb = total_nll_nats / (total_val_BYTES * ln 2) Any loss returned by the model is ignored. -WHY bpb IS THE PRIMARY METRIC (not perplexity): +Why bpb is the primary metric (not perplexity): * It is tokenizer-independent, which is what makes it comparable across - architectures and ungameable -- see DESIGN.md §2. + architectures and ungameable. * It needs no ``exp``. ``val_ppl = exp(mean_ce)`` overflows to ``inf`` for a sufficiently bad submission (mean_ce > ~709), turning a merely-poor model into a non-finite number the guards then have to special-case. bpb is linear in the loss and cannot overflow. - -WHY THE DENOMINATOR IS BYTES AND NOT TOKENS -- READ BEFORE EDITING ------------------------------------------------------------------- -While the tokenizer was byte-level, 1 token == 1 byte and dividing the summed -NLL by the TOKEN count was accidentally identical to dividing by the byte count. -Under dolma2 BPE the two differ by the compression ratio (~4.4x), and per-token -CE/ln2 is a tokenizer-dependent number: a tokenizer that packs more text per -token raises it for the same model quality. Normalizing that way would destroy -exactly the property bpb was adopted for. So the denominator comes from -``data.val_bytes_for(...)``, which is backed by a byte count measured at -corpus-prep time and asserted present (``data.DataError`` otherwise) rather than -defaulted. - -``val_ppl`` is still derived and reported, for human readability only -- it is -no longer the scored quantity, and it remains a PER-TOKEN perplexity (so -``val_ppl != 2**val_bpb`` any more; that identity held only at byte level). It -is computed defensively so a diverged model reports ``inf`` rather than raising. -Determinism knobs (no TF32, deterministic algorithms) are set on the GPU path so -a cached/separate baseline is a valid common-random-numbers partner. """ from __future__ import annotations @@ -46,8 +27,8 @@ @dataclass class EvalOutput: - val_bpb: float # SCORED quantity: total_nll_nats / (bytes * ln 2) - val_ppl: float # derived, per-TOKEN, readability only + val_bpb: float # Scored quantity: total_nll_nats / (bytes * ln 2) + val_ppl: float # derived, per-token, readability only val_ce_nats: float # mean per-token CE n_tokens: int mean_abs_logit_std: float # degeneracy signal (near-constant logits -> ~0) @@ -97,7 +78,7 @@ def evaluate_perplexity(model, data: TokenData, cfg: TaskConfig, device: str) -> if total_tok == 0: return EvalOutput(float("inf"), float("inf"), float("inf"), 0, 0.0, 0.0) - # BYTES, not tokens -- see the module docstring. `total_ce` is the SUM of + # Bytes, not tokens. `total_ce` is the sum of # NLL in nats over every scored target token, so this is exactly # total_nll / (bytes * ln 2). total_bytes = data.val_bytes_for(total_tok) @@ -109,7 +90,7 @@ def evaluate_perplexity(model, data: TokenData, cfg: TaskConfig, device: str) -> val_bpb = total_ce / (total_bytes * math.log(2.0)) mean_ce = total_ce / total_tok - # Derived for readability only, and PER TOKEN. Guarded: exp overflows above + # Derived for readability only, and per token. Guarded: exp overflows above # mean_ce ~709, and a diverged submission can get there. try: val_ppl = math.exp(mean_ce) diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harness/modal_app.py b/2.0/problems/nanoslm_hybrid_arch_design/harness/modal_app.py index 61100f952..8401ec23a 100644 --- a/2.0/problems/nanoslm_hybrid_arch_design/harness/modal_app.py +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/modal_app.py @@ -5,26 +5,20 @@ Two things make this app much lighter than lm_arch_discovery's: - * NO EXTERNAL REPO AND NO CUDA SOURCE BUILD. The harness is self-contained + * No external repo and no CUDA source build. The harness is self-contained PyTorch; the only third-party imports across harness/*.py are torch and numpy. torch's own wheels ship the CUDA runtime, so a slim base works and the image builds in ~1 min rather than needing nvcc. - * NO DATASET DOWNLOAD AND NO GATED TOKENIZER. The corpus is PRE-TOKENIZED by - docker/prep_assets.py with dolma2 (allenai/dolma2-tokenizer, vocab 100278 -- - ungated, no HF token), so this image never tokenizes anything; and - harness/data.py falls back to a deterministic synthetic TOKEN stream when - the .bin paths are absent. So the wiring is provable before any corpus is - staged; a real shard is mounted later for calibration. - - CAVEAT while no shard is mounted here: with the synthetic stream there is no - underlying text, so the bits-per-byte denominator comes from - settings.SYNTHETIC_BYTES_PER_TOKEN (the measured dolma2 compression ratio) - rather than from a manifest. Smoke bpb is therefore in the right RANGE but - is not a measurement of anything. Staging the real train.bin/val.bin/ - manifest.json on this image is required before calibration. + * No dataset download and no gated tokenizer. The corpus is pre-tokenized by + docker/prep_assets.py with dolma2 (allenai/dolma2-tokenizer, 100278 ids, + model vocab padded to 100352 -- + ungated, no HF token), so this image never tokenizes anything. The real + train.bin/val.bin/manifest.json are mounted from the nanoslm-corpus Modal + Volume (below); the harness requires them and raises DataError rather than + synthesizing tokens, so every reported bits-per-byte is a real measurement. Rather than reimplement the training loop here, this calls evaluator.evaluate() --- the SAME entrypoint Harbor calls -- so there is no second code path to drift. +-- the same entrypoint Harbor calls -- so there is no second code path to drift. Parametrized via env vars: LMARCH_MODAL_GPU Modal GPU string (default "H100") @@ -41,13 +35,19 @@ APP_NAME = os.environ.get("LMARCH_MODAL_APP", "nanoslm-hybrid-arch-design") REMOTE_TASK = "/opt/nanoslm_arch/task" +# Function timeout must clear a full final-role pair: two arms trained back-to- +# back (~6 h each = 12 h normally, up to the 12 h hard cap each in the worst +# case), plus warmup and both evals. Set to Modal's 24 h maximum for the widest +# margin. (The old 6 h could not finish even one 6 h arm.) +FN_TIMEOUT_SECONDS = 24 * 60 * 60 + _PROBLEM_DIR = pathlib.Path(__file__).resolve().parent.parent def _ver(mod): - """Version string as a PLAIN str. + """Version string as a plain str. str() is required, not cosmetic: torch.__version__ is a - torch.torch_version.TorchVersion -- a str SUBCLASS defined inside torch -- + torch.torch_version.TorchVersion -- a str subclass defined inside torch -- so returning it raw makes the whole result unpicklable on a client without torch: DeserializationError: 'torch' module is not available. The judge is deliberately CPU-only and has no torch. @@ -67,24 +67,23 @@ def _ver(mod): # nvcc are needed -- the whole reason this image builds in ~1 min. .pip_install("torch==2.6.0", "numpy<2") - # flash-linear-attention: REQUIRED, not an optimization. PyTorch ships a + # flash-linear-attention: required, not an optimization. PyTorch ships a # fused kernel for softmax attention (SDPA/flash) but none for linear or - # recurrent mixers, so a hand-written GDN is unfused eager code. Measured - # at ctx 8192: chunkwise did 6% of attention's FLOPs and still ran 3.5x - # slower (11 optimizer steps vs 38). Under a fixed WALL-CLOCK budget that - # makes the comparison "fused kernel vs unfused kernel" rather than - # "architecture vs architecture", and a hybrid can never win. + # recurrent mixers, so a hand-written GDN is unfused eager code. Under a + # fixed wall-clock budget that would make the comparison "fused vs unfused + # kernel" rather than "architecture vs architecture", and a hybrid could + # never win. # - # torch 2.6.0 (not 2.5.1) because fla warns it needs Triton >= 3.2.0 and - # 2.5.1 ships 3.1.0 -- on 3.1.0 every GatedDeltaNet call hung 30+ min. + # torch 2.6.0 (not 2.5.1) for Triton >= 3.2.0, which fla requires; on the + # Triton 3.1.0 that torch 2.5.1 ships, GatedDeltaNet calls hung. # - # fla 0.5.1 (not 0.4.1): on 0.4.1 the FORWARD compiled fine but the - # BACKWARD did not -- + # fla 0.5.1 (not 0.4.1): on 0.4.1 the forward compiled fine but the + # Backward did not -- # fla/ops/gated_delta_rule/wy_fast.py:303 prepare_wy_repr_bwd # -> triton make_ttgir -> RuntimeError: PassManager::run failed # i.e. 0.4.1's backward kernels predate Triton 3.2's IR. This cost three - # wrong guesses (head count twice) because an early probe timed FORWARD - # ONLY and looked healthy; always exercise fwd+bwd when validating a + # wrong guesses (head count twice) because an early probe timed forward + # Only and looked healthy; always exercise fwd+bwd when validating a # fused kernel. .pip_install("flash-linear-attention==0.5.1", "transformers==4.46.3") # Ship the problem directory itself: harness/, evaluator.py, @@ -96,7 +95,7 @@ def _ver(mod): ) .env({ "PYTHONPATH": REMOTE_TASK, - # MUST be set before the first CUDA call, so it belongs in the image + # Must be set before the first CUDA call, so it belongs in the image # env -- setting it inside the function is too late. settings.py # declares this value but nothing was exporting it, so the Modal path # ran with torch.use_deterministic_algorithms(True) while cuBLAS was @@ -116,19 +115,19 @@ def _ver(mod): # Real corpus (FineWeb-Edu, dolma2-tokenized) staged on a Modal Volume and # mounted at the data path harness/settings.py expects, so the GPU arms train - # on REAL tokens instead of data.py's synthetic fallback. Populated once via - # `modal volume put nanoslm-corpus {train,val}.bin manifest.json /`. This is - # the "real shard mounted for calibration" the design leaves as a TODO. + # on real tokens. The harness has no synthetic fallback, so this mount is + # required. Populated once via + # `modal volume put nanoslm-corpus {train,val}.bin manifest.json /`. CORPUS_VOL = modal.Volume.from_name("nanoslm-corpus", create_if_missing=True) REMOTE_DATA = "/opt/nanoslm_arch/data" - @app.function(image=image, gpu=GPU, timeout=6 * 60 * 60, + @app.function(image=image, gpu=GPU, timeout=FN_TIMEOUT_SECONDS, volumes={REMOTE_DATA: CORPUS_VOL}) - def evaluate_remote(solution_source: str, smoke: bool = True, + def evaluate_remote(solution_source: str, role: str = "final", overrides: dict | None = None) -> dict: """Run the judge on `solution_source` and return its full result. - `solution_source` travels as TEXT, not a path: the submission is written + `solution_source` travels as text, not a path: the submission is written into a scratch file remotely, exactly as the judge receives it. """ import json @@ -138,7 +137,6 @@ def evaluate_remote(solution_source: str, smoke: bool = True, sys.path.insert(0, REMOTE_TASK) - os.environ["FRONTIER_NANOSLM_SMOKE"] = "1" if smoke else "0" os.environ["FRONTIER_NANOSLM_ROLE"] = role for k, v in (overrides or {}).items(): os.environ[k] = str(v) @@ -151,8 +149,8 @@ def evaluate_remote(solution_source: str, smoke: bool = True, # Dockerfile pins and what a warm container runs has already produced # two misleading results in this task's history. def _ver(mod): # noqa: F811 (module-level _ver is the shared one) - # str() is REQUIRED, not cosmetic. torch.__version__ is a - # torch.torch_version.TorchVersion (a str SUBCLASS defined inside + # str() is required, not cosmetic. torch.__version__ is a + # torch.torch_version.TorchVersion (a str subclass defined inside # torch), so returning it raw makes the whole result unpicklable on # any client without torch installed: # DeserializationError: 'torch' module is not available @@ -171,8 +169,8 @@ def _ver(mod): # noqa: F811 (module-level _ver is the shared one) sub = pathlib.Path(td) / "model.py" sub.write_text(solution_source) - # DEBUG PATH (operator-only, never reachable from a submission): - # rerun the arm OUTSIDE the evaluator's guard so the real traceback + # Debug path (operator-only, never reachable from a submission): + # rerun the arm outside the evaluator's guard so the real traceback # survives. evaluator.py deliberately collapses every training # exception to "training runtime error: RuntimeError" for black-box # safety -- correct for agents, useless for debugging the harness. @@ -186,7 +184,7 @@ def _ver(mod): # noqa: F811 (module-level _ver is the shared one) try: mod = evaluator._load_submission_module(str(sub)) runner.run_arm(runner.load_factory(mod), TokenData(cfg), cfg, "cuda") - # `stack` on the SUCCESS branch too. It was only on the + # `stack` on the success branch too. It was only on the # failure branch, so a passing run could not be attributed # to a version -- exactly the warm-container trap that # already produced one stale, misread result here. @@ -210,46 +208,85 @@ def _ver(mod): # noqa: F811 (module-level _ver is the shared one) "stack": stack, } - @app.function(image=image, gpu=GPU, timeout=6 * 60 * 60, + # Bump on every behavioral change to this file OR to the harness it ships + # (add_local_dir above). Two reasons, both learned the hard way: + # * A deploy does not evict warm containers, so the first call after + # `modal deploy` can still execute the old code -- this has produced + # misleading results twice in this task's history. Check this value + # before trusting a result. + # * It is part of the baseline-cache key. The config fingerprint hashes + # config VALUES only, so a harness behavior change (e.g. the LR + # schedule moving from step-based to wall-clock decay) changes the + # baseline's val_bpb without changing the fingerprint. Keying the cache + # by fingerprint alone would then pair fresh submissions against a + # stale baseline. + CODE_VERSION = "final-6h-agentcache-v8" + + # Baseline cache for the agent role, on the corpus Volume so it survives + # containers (the container FS is ephemeral; settings.baseline_cache_path + # is the local-backend location and is NOT mounted here). + BASELINE_CACHE = f"{REMOTE_DATA}/baseline/baseline_arm.json" + + @app.function(image=image, gpu=GPU, timeout=FN_TIMEOUT_SECONDS, volumes={REMOTE_DATA: CORPUS_VOL}) - def run_pair_remote(solution_source: str, smoke: bool = False, + def run_pair_remote(solution_source: str, role: str = "final") -> dict: - """Run BOTH arms on one GPU and return their metrics as plain dicts. + """Run the arms on one GPU and return their metrics as plain dicts. - This is the entrypoint the JUDGE uses. It deliberately returns ARM - METRICS, not a score: scoring stays judge-side so the scoring code and + This is the entrypoint the judge uses. It deliberately returns arm + Metrics, not a score: scoring stays judge-side so the scoring code and the hidden constants (bpb_score_scale) never ship to a GPU worker, and so the judge cannot be handed a score it did not compute. - Contrast with evaluate_remote(), which runs the whole evaluator - remotely -- useful for manual testing, but circular now that - evaluator.evaluate() dispatches here. + Role semantics (this is what makes iteration affordable): + + final fresh baseline + submission CRN pair on this GPU (~2T). + Authoritative. Opportunistically refreshes the baseline cache. + agent cheap iterative feedback: reuse the baseline cached on the + corpus Volume (keyed by config fingerprint + CODE_VERSION) and + train only the submission (~T). A cache miss falls back to the + full pair and populates the cache, so only the first agent + call after a deploy/config change pays ~2T. + + Guard rejections come back as {"guard_error": ...} rather than raised: + GuardError is defined in harness.runner, which imports torch, so the + CPU-only judge could not deserialize the exception -- and the public + guard message (param cap, BLOCK_SIZE range, eval-shape failure) is + exactly the feedback an iterating agent needs. Everything crossing this boundary must be a plain builtin: the judge container is CPU-only and has no torch, so e.g. torch.__version__ - (a str SUBCLASS defined in torch) would make the result unpicklable. + (a str subclass defined in torch) would make the result unpicklable. """ + import json import sys import tempfile sys.path.insert(0, REMOTE_TASK) - os.environ["FRONTIER_NANOSLM_SMOKE"] = "1" if smoke else "0" os.environ["FRONTIER_NANOSLM_ROLE"] = role import torch # noqa: F401 (ensures CUDA init before harness import) import evaluator # noqa: F401 (sets CUBLAS_WORKSPACE_CONFIG early) from harness import runner, settings as _s + from harness.data import TokenData cfg = _s.active_config() - with tempfile.TemporaryDirectory() as td: - sub = pathlib.Path(td) / "model.py" - sub.write_text(solution_source) - mod = evaluator._load_submission_module(str(sub)) - base, cand = runner.run_pair(mod, cfg, "cuda") + # The corpus identity is part of the key: the config fingerprint hashes + # dataset_name but not the staged shard's SIZE, and restaging a larger + # shard changes the sampled windows (data.batch draws offsets against + # train.size) -- so a baseline trained on the old shard must not be + # paired with submissions trained on the new one. The on-disk byte size + # is a cheap, sufficient stamp for that. + try: + corpus_stamp = str(os.path.getsize(cfg.train_tokens_path)) + except OSError: + corpus_stamp = "nocorpus" + cache_key = f"{_s.config_fingerprint(cfg)}:{corpus_stamp}:{CODE_VERSION}" + cache_path = pathlib.Path(BASELINE_CACHE) def _arm(a): - # Every value here must be a PLAIN builtin -- the judge is CPU-only + # Every value here must be a plain builtin -- the judge is CPU-only # and has no torch, so anything torch-defined is unpicklable there. return { "val_bpb": float(a.val_bpb), @@ -257,30 +294,88 @@ def _arm(a): "steps": int(a.steps), "wall_seconds": float(a.wall_seconds), "n_params": int(a.n_params), - # Context this arm TRAINED at. The baseline is always 8192; + # Context this arm trained at. The baseline is always 8192; # a submission may declare a shorter one via BLOCK_SIZE. Both - # arms are always EVALUATED at cfg.eval_block_size (8192). + # arms are always evaluated at cfg.eval_block_size (8192). "train_block_size": int(a.train_block_size), # Warmup is where Triton JIT is supposed to land. If # warmup_error is non-empty, warmup silently no-opped and its - # cost was deferred INTO the timed loop -- the exact failure - # that produced sub_steps=1 at 103.2s. + # cost was deferred into the timed loop -- a step-0 stall. These + # fields surface that regression. "warmup_seconds": float(a.warmup_seconds), "warmup_error": str(a.warmup_error), } + def _cache_read(): + """Cached baseline arm dict for this key, or None. Best-effort.""" + try: + CORPUS_VOL.reload() # see commits from other containers + except Exception: + pass + try: + entry = json.loads(cache_path.read_text()).get(cache_key) + # Sanity: a valid entry is an arm dict with a positive bpb. + if isinstance(entry, dict) and float(entry.get("val_bpb", 0)) > 0: + return entry + except Exception: + pass + return None + + def _cache_write(entry: dict) -> None: + """Best-effort persist; a failure only costs the next call ~T.""" + try: + cache_path.parent.mkdir(parents=True, exist_ok=True) + try: + data = json.loads(cache_path.read_text()) + if not isinstance(data, dict): + data = {} + except Exception: + data = {} + data[cache_key] = entry + cache_path.write_text(json.dumps(data)) + CORPUS_VOL.commit() + except Exception: + pass + + base_info = _cache_read() if role == "agent" else None + baseline_cached = base_info is not None + + with tempfile.TemporaryDirectory() as td: + sub = pathlib.Path(td) / "model.py" + sub.write_text(solution_source) + mod = evaluator._load_submission_module(str(sub)) + try: + if base_info is None: + # final role, or first agent call for this key: full CRN + # pair, then persist the baseline for future agent calls. + base, cand = runner.run_pair(mod, cfg, "cuda") + base_info = _arm(base) + _cache_write(base_info) + else: + # agent role, cache hit: submission arm only (~T). Resolve + # (and range-guard) BLOCK_SIZE before anything expensive, + # exactly as run_pair does. + sub_block = runner.resolve_train_block_size(mod, cfg) + data = TokenData(cfg) + cand = runner.run_arm( + runner.load_factory(mod), data, cfg, "cuda", sub_block + ) + except runner.GuardError as exc: + return { + "guard_error": str(exc), + "code_version": CODE_VERSION, + "stack": {"torch": str(torch.__version__), + "triton": _ver("triton"), "fla": _ver("fla")}, + } + return { - "baseline": _arm(base), + "baseline": base_info, "submission": _arm(cand), - # FIXED scoring window, reported so a result is self-describing: + "baseline_cached": baseline_cached, + # Fixed scoring window, reported so a result is self-describing: # both arms' val_bpb come from windows of exactly this width. "eval_block_size": int(cfg.eval_block_size), - # Bump on every behavioral change to this file. A deploy does NOT - # evict warm containers, so the first call after `modal deploy` can - # still execute the OLD code -- this has produced misleading results - # twice in this task's history. Check this value before trusting a - # result. - "code_version": "agent-train-ctx-v5", + "code_version": CODE_VERSION, "stack": {"torch": str(torch.__version__), "triton": _ver("triton"), "fla": _ver("fla")}, } diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harness/model_config.py b/2.0/problems/nanoslm_hybrid_arch_design/harness/model_config.py index 0cbd946ad..3d07fb692 100644 --- a/2.0/problems/nanoslm_hybrid_arch_design/harness/model_config.py +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/model_config.py @@ -6,12 +6,12 @@ informational (so a model can size itself to the budget) and carry no guarantees. -TWO CONTEXT LENGTHS +Two context lengths ------------------- -``block_size`` the context this model will be TRAINED at. It is the judge's +``block_size`` the context this model will be trained at. It is the judge's default (8192) unless model.py declares a module-level ``BLOCK_SIZE``, in which case it is that value. -``eval_block_size`` the context it will be SCORED at. ALWAYS 8192, whatever the +``eval_block_size`` the context it will be scored at. Always 8192, whatever the training context is. When those differ the model is asked for logits at positions it never saw during @@ -29,8 +29,8 @@ @dataclass(frozen=True) class ModelConfig: vocab_size: int - block_size: int # TRAINING context (per-submission) - eval_block_size: int = 8192 # SCORING context (fixed by the judge) + block_size: int # Training context (per-submission) + eval_block_size: int = 8192 # Scoring context (fixed by the judge) # Read-only budget hints (informational; not part of the scored contract). train_seconds_hint: float = 0.0 param_cap_hint: int = 0 diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harness/policy.py b/2.0/problems/nanoslm_hybrid_arch_design/harness/policy.py index 454c247ae..1f77f7e5a 100644 --- a/2.0/problems/nanoslm_hybrid_arch_design/harness/policy.py +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/policy.py @@ -2,12 +2,37 @@ Torch-free and unit-tested on CPU. Treat submission content as adversarial (2.0 black-box safety). This is the *static* layer; runtime guards that the -scan cannot see (trained-from-scratch check, judge-owned loss, resource caps) -live in the runner/evaluator. See DESIGN.md §6–§7. +scan cannot see (trained-from-scratch check, judge-owned loss, resource caps, +the bits-per-byte plausibility floor) live in the runner/evaluator. + +Three layers, applied in order, each fails closed: + + 1. Substring denylist over the full source (below). Catches dangerous words + in ANY spelling that contains them -- including comments, where a judge + path or an env var is a leak signal in itself. + 2. Metric/data-name denylist over code only (comments/strings stripped), so + a submission may *document* what it optimizes without tripping the gate. + 3. AST scan with a STRICT IMPORT ALLOWLIST. The substring layer alone is + evadable by construction -- ``getattr(os, 'envi' + 'ron')`` contains no + denied token -- so the AST layer removes the primitives such evasion + needs: only allowlisted module roots may be imported, and the + dynamic-access builtins (bare ``eval``/``exec``/``getattr``/ + ``__import__``/``open``/...), dunder attributes (``__class__`` and + friends are the classic sandbox-escape route), wildcard imports, and + raw-file-reader attribute names (``fromfile``, ``from_file``, + ``memmap``, ...) are rejected wherever they appear. ``model.eval()`` and + ``torch.compile(...)`` are attribute accesses, not bare names, and remain + allowed. + +Static scanning of Turing-complete code is defense in depth, not proof; the +runtime plausibility floor (runner.run_arm) backstops it by rejecting any +result too good to be honest. """ from __future__ import annotations +import ast + from dataclasses import dataclass MAX_BYTES = 256 * 1024 # 256 KB @@ -26,27 +51,27 @@ # --- filesystem reads (no data/weight/val access) --- "open(", "Path(", "np.load", "numpy.load", "np.fromfile", "mmap", "pickle.load", "joblib.load", - # NOTE: metric/data names are NOT here -- see POLICY_DENY_TOKENS_CODE. + # Note: metric/data names are not here -- see POLICY_DENY_TOKENS_CODE. # --- timer / control-flow short-circuits & concurrency tricks --- "time.time", "time.perf_counter", "perf_counter", "time.sleep", "while True", "threading", "multiprocessing", "os.fork", "ctypes", "exec(", "__import__", - # NOTE: "eval(" and "compile(" are intentionally NOT banned so that the + # Note: "eval(" and "compile(" are intentionally not banned so that the # legitimate `model.eval()` and `torch.compile()` idioms are allowed # (mirrors nanowm's allowlist). Sandboxing + judge-owned loss cover the rest. ) -# Scanned over CODE ONLY (comments and string literals stripped). +# Scanned over code only (comments and string literals stripped). # -# These names exist to stop a submission from READING the metric or the held-out +# These names exist to stop a submission from reading the metric or the held-out # data, which takes executable code. Scanning them over prose rejected any # submission that merely *documented its intent* -- "# should improve val_bpb" # -- with a confusing "forbidden token" error. Both the reference and the locked # baseline tripped it once their docstrings explained the task, which is how it # was found. A submission should be able to say what it optimizes. # -# Everything in POLICY_DENY_TOKENS above is still scanned over the FULL source: +# Everything in POLICY_DENY_TOKENS above is still scanned over the full source: # a judge path or an env var appearing in a comment is a leak signal in itself, # whereas the word "perplexity" in a docstring reads nothing. POLICY_DENY_TOKENS_CODE: tuple[str, ...] = ( @@ -55,10 +80,130 @@ ) +# --------------------------------------------------------------------------- # +# AST layer: strict import allowlist + dynamic-access primitives. +# --------------------------------------------------------------------------- # + +# Module ROOTS a submission may import. Everything a legitimate architecture +# needs, nothing that touches the filesystem, the environment, processes or +# the network: +# * torch / numpy / fla / einops / triton -- the modelling stack the Modal +# image ships (einops is an fla dependency; triton is allowed because a +# custom fused kernel is a legitimate architecture lever in this task -- +# fla itself exists for exactly that reason). +# * a handful of pure-computation stdlib modules. `operator` is deliberately +# absent: `operator.attrgetter("...")` is getattr-by-string, the exact +# primitive the banned-names rule below removes. `os`/`sys`/`io`/ +# `pathlib`/`importlib`/... are absent for the obvious reasons; +# `transformers` is absent because from_pretrained is a network/weights +# path (train-from-scratch rule). +POLICY_ALLOWED_IMPORT_ROOTS: frozenset = frozenset({ + "__future__", "abc", "collections", "contextlib", "copy", "dataclasses", + "enum", "functools", "itertools", "math", "typing", "warnings", + "torch", "numpy", "fla", "einops", "triton", +}) + +# Bare names whose *appearance* (any context -- call, alias, argument) is +# rejected. These are the primitives that turn computed strings into imports, +# attribute walks, or file reads; without them, string concatenation has no +# sink. `eval`/`compile` here are the BARE builtins: `model.eval()` and +# `torch.compile()` are Attribute nodes and unaffected. +POLICY_BANNED_NAMES: frozenset = frozenset({ + "eval", "exec", "compile", "__import__", "getattr", "setattr", "delattr", + "vars", "globals", "locals", "open", "input", "breakpoint", "exit", + "quit", "__builtins__", "__loader__", "__spec__", +}) + +# Attribute / submodule names rejected wherever they appear -- as an attribute +# access (aliasing a module defeats substring checks: `q = numpy; +# q.fromfile(...)`), as a dotted-path segment in an import, or as a name +# imported via `from X import Y` (`from torch import load` would create a +# BARE `load` that no substring catches). Raw-binary readers and +# process/env/compiler escape hatches reachable through the allowlisted +# roots: +# fromfile/from_file/memmap/open_memmap/... -- the readers that can open +# val.bin's raw uint32 stream. (np.load/torch.load stay usable on paper +# but require npy/pickle magic that a raw stream does not have, and the +# literal spellings are substring-denied above.) +# ctypeslib/f2py/distutils/cpp_extension/load_library -- compile-or-load- +# native-code paths inside numpy/torch. +# hub -- torch.hub downloads code and weights. +# os/sys/environ/popen/system -- e.g. `torch.os` re-exports the os module. +POLICY_BANNED_ATTRS: frozenset = frozenset({ + "os", "sys", "environ", "popen", "system", + "fromfile", "from_file", "tofile", "memmap", "open_memmap", + "loadtxt", "genfromtxt", "fromregex", "DataSource", + "ctypeslib", "f2py", "distutils", "cpp_extension", "load_library", + "hub", "attrgetter", +}) + +# `from X import Y` names additionally rejected: importing these unqualified +# strips the module prefix every substring rule keys on. +_BANNED_IMPORT_NAMES: frozenset = POLICY_BANNED_ATTRS | {"load", "loads"} + +# Dunder attributes allowed as Attribute access. Everything else dunder is +# rejected: `__class__`/`__globals__`/`__subclasses__`/`__dict__` are the +# standard object-graph escape routes. `__init__` is required by +# `super().__init__()`; `__version__` and `__name__` (the +# `type(exc).__name__` error-message idiom) are harmless strings that +# traverse nothing. +_ALLOWED_DUNDER_ATTRS: frozenset = frozenset( + {"__init__", "__version__", "__name__"} +) + + +def _check_ast(source: str) -> str: + """Return "" if the AST rules pass, else the (public) rejection reason.""" + try: + tree = ast.parse(source) + except SyntaxError as exc: + return f"submission does not parse (line {exc.lineno})" + + def _module_path_ok(path: str) -> str: + parts = path.split(".") + if parts[0] not in POLICY_ALLOWED_IMPORT_ROOTS: + return f"import of {parts[0]!r} is not in the import allowlist" + for seg in parts[1:]: + if seg in POLICY_BANNED_ATTRS: + return f"import path segment {seg!r} is not allowed" + return "" + + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + why = _module_path_ok(alias.name) + if why: + return why + elif isinstance(node, ast.ImportFrom): + if node.level and node.level > 0: + return "relative imports are not allowed" + why = _module_path_ok(node.module or "") + if why: + return why + for alias in node.names: + if alias.name == "*": + return "wildcard imports are not allowed" + if alias.name in _BANNED_IMPORT_NAMES: + return f"importing the name {alias.name!r} is not allowed" + if (alias.name.startswith("__") and alias.name.endswith("__") + and alias.name not in _ALLOWED_DUNDER_ATTRS): + return f"importing the name {alias.name!r} is not allowed" + elif isinstance(node, ast.Name): + if node.id in POLICY_BANNED_NAMES: + return f"use of the name {node.id!r} is not allowed" + elif isinstance(node, ast.Attribute): + if node.attr in POLICY_BANNED_ATTRS: + return f"attribute {node.attr!r} is not allowed" + if (node.attr.startswith("__") and node.attr.endswith("__") + and node.attr not in _ALLOWED_DUNDER_ATTRS): + return f"dunder attribute {node.attr!r} is not allowed" + return "" + + def _code_only(source: str) -> str: """Return `source` with comments and string literals removed. - Fails CLOSED: if the source will not tokenize, return it unchanged so the + Fails closed: if the source will not tokenize, return it unchanged so the scan still sees everything rather than waving a malformed file through. """ import io @@ -107,6 +252,11 @@ def check_source(source: str, *, size_bytes: int | None = None) -> PolicyResult: if tok in code: return PolicyResult(False, f"forbidden token in submission code: {tok!r}") + # AST layer: strict import allowlist + dynamic-access primitives. + why = _check_ast(source) + if why: + return PolicyResult(False, why) + # Must expose a model factory the harness can call. if ("def build_model" not in source) and ("class NanoSLM" not in source): return PolicyResult( diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harness/runner.py b/2.0/problems/nanoslm_hybrid_arch_design/harness/runner.py index b53a9d199..07230fe57 100644 --- a/2.0/problems/nanoslm_hybrid_arch_design/harness/runner.py +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/runner.py @@ -25,13 +25,13 @@ class GuardError(Exception): @dataclass class ArmMetrics: - val_bpb: float # SCORED quantity + val_bpb: float # Scored quantity val_ppl: float # derived, readability only steps: int wall_seconds: float n_params: int - train_block_size: int = 0 # context this arm actually TRAINED at - warmup_seconds: float = 0.0 # time absorbed OUTSIDE the timed window + train_block_size: int = 0 # context this arm actually trained at + warmup_seconds: float = 0.0 # time absorbed outside the timed window warmup_error: str = "" # "" when warmup completed cleanly eval: EvalOutput = field(repr=False, default=None) @@ -47,14 +47,14 @@ def load_factory(module) -> Callable: # --------------------------------------------------------------------------- # -# Agent-controlled TRAINING context. +# Agent-controlled training context. # # A submission declares it with a module-level ``BLOCK_SIZE`` int in model.py. # That is deliberately the whole protocol: no new required entrypoint, and the # existing build_model(config) / class NanoSLM contract is untouched, so a # submission that says nothing keeps training at the judge's 8192 default. # -# The EVALUATION context is never negotiable -- see settings.eval_block_size. +# The evaluation context is never negotiable -- see settings.eval_block_size. # --------------------------------------------------------------------------- # MIN_TRAIN_BLOCK = 256 @@ -72,7 +72,7 @@ def resolve_train_block_size(module, cfg: TaskConfig) -> int: and it keeps the space small enough to reason about; * >= 256 -- below that the context degenerates and a "language model" stops modelling anything the 8192-wide evaluation rewards; - * <= eval_block_size (8192) -- training LONGER than the scoring window + * <= eval_block_size (8192) -- training longer than the scoring window buys nothing measurable and only costs memory and steps. """ raw = getattr(module, "BLOCK_SIZE", None) @@ -95,8 +95,8 @@ def resolve_train_block_size(module, cfg: TaskConfig) -> int: def _model_config(cfg: TaskConfig, device: str) -> ModelConfig: return ModelConfig( vocab_size=cfg.vocab_size, - block_size=cfg.block_size, # per-arm TRAINING context - eval_block_size=cfg.eval_block_size, # fixed SCORING context + block_size=cfg.block_size, # per-arm training context + eval_block_size=cfg.eval_block_size, # fixed scoring context train_seconds_hint=cfg.train_seconds, param_cap_hint=cfg.param_cap, device_hint="cuda" if device.startswith("cuda") else device, @@ -115,59 +115,66 @@ def _param_snapshot(model): def _warmup(model, data: TokenData, cfg: TaskConfig, device: str): - """Compile/autotune kernels BEFORE the wall-clock timer starts. + """Compile/autotune kernels before the wall-clock timer starts. This is required for a fair fixed-wall-clock comparison, not an - optimization. Triton JIT-compiles per shape on first use: measured on an - H100, fla's GatedDeltaNet takes 15.5s for its first forward at T=8192 and - 74.5s at T=64, then 3ms warm. torch's SDPA path pays essentially none. + optimization. Triton JIT-compiles per shape on first use, so fla's + GatedDeltaNet pays a large first-forward cost while torch's SDPA path pays + essentially none. Left inside the timed window, an fla-based hybrid would spend a chunk of its budget T compiling while a pure-attention arm spent none -- so the score - would partly measure COMPILATION rather than throughput, and the smaller T + would partly measure compilation rather than throughput, and the smaller T is the worse the distortion. - THREE THINGS THIS MUST MATCH ABOUT ``train_model``, or it warms kernels the + Three things this must match about ``train_model``, or it warms kernels the training loop never calls and the JIT cost lands inside the timer anyway: - 1. DEVICE. The factories (``baseline_model.build_model``, the submission's) - construct on CPU; only ``train_model`` calls ``model.to(device)``. This - function used to run BEFORE that move, so ``model(x)`` with a CUDA batch - raised "Expected all tensors to be on the same device", which the old - bare ``except Exception: pass`` swallowed. Warmup was a silent no-op for - BOTH arms -- harmless for SDPA, fatal for fla, whose entire Triton - compile then landed in step 0 (measured: one 103.2s step against a 45s - budget). Hence ``model.to(device)`` here, and hence the error is now - REPORTED rather than swallowed. - 2. DTYPE. ``train_model`` runs under ``torch.autocast(bfloat16)``. Triton + 1. Device. The factories (``baseline_model.build_model``, the submission's) + construct on CPU; only ``train_model`` calls ``model.to(device)``. So + warmup moves the model to ``device`` here and reports errors rather than + swallowing them: an earlier version ran before the move and silently + no-op'd on a device mismatch, letting fla's entire Triton compile land + in step 0. + 2. Dtype. ``train_model`` runs under ``torch.autocast(bfloat16)``. Triton autotunes per dtype, so an fp32 warmup warms the wrong kernels. - 3. THE OPTIMIZER STEP. ``clip_grad_norm_`` and AdamW's foreach/fused + 3. The optimizer step. ``clip_grad_norm_`` and AdamW's foreach/fused kernels are first-call-compiled too, so warmup takes a real step and - then RESTORES the parameters, leaving training bit-identical to a run + then restores the parameters, leaving training bit-identical to a run that never warmed up. Two iterations, not one: Triton autotuning picks a config on the first call and can still recompile on the second. The second is ~ms once warm. - FOURTH THING, added with the agent-controlled training context (DESIGN.md - 15): the training and evaluation SHAPES can now differ. Warming only the + Fourth thing, added with the agent-controlled training context: the + training and evaluation shapes can now differ. Warming only the training shape leaves every eval-shape kernel cold for a submission that trained at, say, 2048 -- and Triton autotunes per shape, so that is a fresh - compile of tens of seconds. It does NOT currently land inside the scored + compile of tens of seconds. It does not currently land inside the scored window (``train_model`` stops its timer before ``evaluate_perplexity`` is called), so this is not a scoring bug today; it is warmed here anyway because (a) it is a real cost against the container/`max_train_seconds` budget, (b) it makes the property robust to any future reordering rather than accidental, and (c) it surfaces an eval-shape failure (a model whose - buffers were sized at the training context and cannot run at 8192) HERE, in - ``warmup_error``, instead of as a mystery exception after training. The - eval warm mirrors ``eval_ppl.evaluate_perplexity`` exactly: ``model.eval()``, - ``no_grad``, and NO autocast -- evaluation runs in fp32, so warming it under - bfloat16 would warm the wrong kernels again. - - Errors are captured and returned, never raised: a genuinely broken model - must fail in the real training loop where the guards can classify it. But it - is no longer INVISIBLE -- the message rides out in ArmMetrics.warmup_error. + buffers were sized at the training context and cannot run at 8192) here, + BEFORE the training budget is spent. The eval warm mirrors + ``eval_ppl.evaluate_perplexity`` exactly: ``model.eval()``, ``no_grad``, + and no autocast -- evaluation runs in fp32, so warming it under bfloat16 + would warm the wrong kernels again. + + Error handling is split by shape, deliberately: + + * Training-shape errors are captured and returned in + ``ArmMetrics.warmup_error``, never raised -- a genuinely broken model + must fail in the real training loop where the guards can classify it + (OOM vs shape error vs ...), and a warmup hiccup must not reject a + model the real loop could have trained. + * Eval-shape errors -- reached only when the training shape already ran + cleanly, so the model itself works -- are definitive: a model that + cannot run at ``eval_block_size`` can only ever score 0, and letting it + train first burns the full ``train_seconds`` of H100 to find that out. + These raise ``GuardError`` here, with the same classification the + post-training guard in ``run_arm`` would have produced. """ import time @@ -188,7 +195,7 @@ def _warmup(model, data: TokenData, cfg: TaskConfig, device: str): opt = torch.optim.AdamW(model.parameters(), lr=0.0) for i in range(2): - # cfg.block_size here is the ARM's training context (run_arm passes + # cfg.block_size here is the arm's training context (run_arm passes # a per-arm cfg), so this warms the shape train_model will use. x, y = data.batch(i, device, cfg.block_size) # same API train.py uses with autocast: @@ -207,12 +214,27 @@ def _warmup(model, data: TokenData, cfg: TaskConfig, device: str): p.copy_(s) model.zero_grad(set_to_none=True) del saved, opt + except Exception as exc: + # Training-shape failure: record and let the real training loop + # classify it (see the docstring). The eval-shape warm is skipped -- + # with the training shape already broken, its failure would say nothing + # about the eval context specifically. + err = f"{type(exc).__name__}: {exc}"[:300] + return round(time.monotonic() - t0, 2), err - # Eval-shape warm: [1, eval_block_size] in eval mode, no_grad, fp32 -- - # exactly how evaluate_perplexity will call the model. Uses the first - # real held-out window so the shape and dtype match to the element; it - # runs under no_grad and takes no optimizer step, so it cannot influence - # any parameter or leak the held-out stream into training. + # Eval-shape warm: [1, eval_block_size] in eval mode, no_grad, fp32 -- + # exactly how evaluate_perplexity will call the model. Uses the first + # real held-out window so the shape and dtype match to the element; it + # runs under no_grad and takes no optimizer step, so it cannot influence + # any parameter or leak the held-out stream into training. + # + # Failure here is a fast-fail (see the docstring): the training shape ran + # cleanly two iterations ago, so an exception now means the model cannot + # run at the fixed scoring context -- it can only ever score 0, and the + # only question is whether it finds that out now or after train_seconds of + # GPU time. Classified exactly like the post-training eval guard in + # run_arm, so the public message is identical either way. + try: model.eval() with torch.no_grad(): first = next(iter(data.val_windows(device)), None) @@ -227,8 +249,20 @@ def _warmup(model, data: TokenData, cfg: TaskConfig, device: str): if device.startswith("cuda"): torch.cuda.synchronize() torch.cuda.empty_cache() + except (GuardError, DataError): + raise except Exception as exc: - err = f"{type(exc).__name__}: {exc}"[:300] + msg = str(exc).lower() + if "out of memory" in msg: + raise GuardError( + f"evaluation ran out of GPU memory at context " + f"{cfg.eval_block_size}" + ) + raise GuardError( + f"evaluation failed at context {cfg.eval_block_size} " + f"(trained at {cfg.block_size}): {type(exc).__name__}. Buffers sized " + f"off config.block_size must still work at config.eval_block_size" + ) return round(time.monotonic() - t0, 2), err @@ -245,7 +279,7 @@ def run_arm(factory: Callable, data: TokenData, cfg: TaskConfig, device: str, train_block_size: int | None = None) -> ArmMetrics: """Instantiate -> guard params -> train -> eval -> guard training/degeneracy. - ``train_block_size`` is this arm's TRAINING context (default: cfg's). It is + ``train_block_size`` is this arm's training context (default: cfg's). It is threaded through as a per-arm ``cfg`` so that every consumer -- the model config handed to the factory, ``_warmup``, and ``train_model`` -- sees one consistent value. ``cfg.eval_block_size`` is untouched by it, so evaluation @@ -273,19 +307,19 @@ def run_arm(factory: Callable, data: TokenData, cfg: TaskConfig, device: str, if n_params > cfg.param_cap: raise GuardError(f"model exceeds param cap ({n_params} > {cfg.param_cap})") - # COMPILE BOTH ARMS, UNIFORMLY, IN THE HARNESS. + # Compile both arms, uniformly, in the harness. # # Two reasons, and the second is a correctness issue: # # 1. Throughput. The chunkwise linear mixer is unfused eager PyTorch and is # dominated by kernel-launch overhead, so it loses to SDPA's single fused - # flash-attention kernel despite doing ~6% of the FLOPs (measured: 11 - # optimizer steps vs 38 at ctx 8192). Compilation fuses the elementwise - # work and is what makes a linear mixer competitive at all here. + # flash-attention kernel despite doing a fraction of the FLOPs. + # Compilation fuses the elementwise work and is what makes a linear mixer + # competitive at all here. # # 2. Fairness. `torch.compile()` is explicitly allowed in submissions # (policy.py). If the harness did not compile, a submission could take - # the BASELINE ARCHITECTURE UNCHANGED, wrap it in torch.compile, gain + # the baseline architecture unchanged, wrap it in torch.compile, gain # wall-clock and score well with zero architectural insight -- the same # kernel-mismatch confound that fla created, by another route. Compiling # both arms here removes compilation as a lever and puts the comparison @@ -294,18 +328,18 @@ def run_arm(factory: Callable, data: TokenData, cfg: TaskConfig, device: str, # Failures fall back to eager: a submitted architecture may contain # something inductor cannot trace, and that should cost throughput, not be # a hard rejection. - # DEFAULT OFF. The fairness argument above is real and unresolved, but - # compilation measured WORSE than eager on the chunkwise mixer (1 step vs - # 11 at ctx 8192 -- inductor almost certainly graph-breaks on the chunk - # loop's data-dependent state), and its interaction with fla's Triton - # kernels is unvalidated. Re-enable only after measuring both arms with it. + # Default off. The fairness argument above is real and unresolved, but + # compilation measured worse than eager on the chunkwise mixer (inductor + # almost certainly graph-breaks on the chunk loop's data-dependent state), + # and its interaction with fla's Triton kernels is unvalidated. Re-enable + # only after measuring both arms with it. if getattr(cfg, "compile_model", False): try: model = torch.compile(model) except Exception: pass - # Warm kernels OUTSIDE the timed window -- see _warmup. This now also + # Warm kernels outside the timed window -- see _warmup. This now also # absorbs inductor's compilation, which is the whole reason it must precede # the timer. warmup_seconds, warmup_error = _warmup(model, data, cfg, device) @@ -313,8 +347,14 @@ def run_arm(factory: Callable, data: TokenData, cfg: TaskConfig, device: str, before = _param_snapshot(model) try: tout: TrainOutput = train_model(model, data, cfg, device) - except RuntimeError as exc: - # e.g. CUDA OOM or a shape error from a malformed architecture. + except (GuardError, DataError): + raise + except Exception as exc: + # `Exception`, not `RuntimeError`: a malformed architecture can raise + # anything mid-training -- torch raised ValueError for a kernel dtype + # mismatch, IndexError for an overrun table -- and an uncaught type + # here escapes classification entirely, surfacing as a misleading + # environment_error at the judge instead of a submission fault. msg = str(exc).lower() if "out of memory" in msg: raise GuardError("training ran out of GPU memory") @@ -329,8 +369,8 @@ def run_arm(factory: Callable, data: TokenData, cfg: TaskConfig, device: str, if mean_abs_delta < cfg.min_param_delta: raise GuardError("parameters did not change during training (not trained)") - # A NEW FAILURE MODE, since the training and evaluation contexts can differ: - # a model whose buffers were sized off the TRAINING context + # A new failure mode, since the training and evaluation contexts can differ: + # a model whose buffers were sized off the training context # trains happily and then blows up at eval, where T is always # eval_block_size. Left unclassified that surfaced as a bare # "submission_error: RuntimeError" with no hint about the cause, so it is @@ -339,7 +379,7 @@ def run_arm(factory: Callable, data: TokenData, cfg: TaskConfig, device: str, # positional embedding table indexed past its end, which torch raises as an # IndexError, not a RuntimeError. Catching only RuntimeError let exactly the # failure this guard exists for escape unclassified. DataError is re-raised - # untouched -- a missing held-out byte count is a JUDGE-side asset problem + # untouched -- a missing held-out byte count is a judge-side asset problem # and must not be reported as a submission fault. try: eout = evaluate_perplexity(model, data, cfg, device) @@ -364,6 +404,18 @@ def run_arm(factory: Callable, data: TokenData, cfg: TaskConfig, device: str, raise GuardError("evaluation produced a non-finite bits-per-byte") if eout.mean_abs_logit_std < 1e-4: raise GuardError("model produced near-constant logits (degenerate)") + # Plausibility floor (see settings.min_plausible_bpb): the static policy + # gate is defense in depth, not proof, and a submission that evaded it and + # memorized the held-out stream would land near zero. No honest model at + # this scale and budget can approach the floor, so below it the measurement + # is treated as leakage, not brilliance. Applied to both arms uniformly; + # the baseline sits several times above it. + floor = float(getattr(cfg, "min_plausible_bpb", 0.0)) + if floor > 0.0 and eout.val_bpb < floor: + raise GuardError( + f"bits-per-byte {eout.val_bpb:.4f} is below the plausibility " + f"floor ({floor}); flagged as presumed held-out leakage" + ) return ArmMetrics( val_bpb=eout.val_bpb, @@ -383,10 +435,10 @@ def baseline_factory() -> Callable: def baseline_block_size(cfg: TaskConfig) -> int: - """Training context of the LOCKED arm. + """Training context of the locked arm. Read through the same ``BLOCK_SIZE`` protocol a submission uses, from - ``baseline_model``, which declares 8192 EXPLICITLY rather than inheriting a + ``baseline_model``, which declares 8192 explicitly rather than inheriting a default. The locked arm's context must be unambiguous and must not move if ``TaskConfig.block_size`` (now only the submission-facing default) is ever retuned -- and it is what makes the fingerprint-keyed baseline cache valid @@ -397,7 +449,7 @@ def baseline_block_size(cfg: TaskConfig) -> int: def run_pair(submission_module, cfg: TaskConfig, device: str): """Scored path: train baseline + submission back-to-back on one GPU (CRN).""" - # Resolve (and range-guard) the submission's training context BEFORE + # Resolve (and range-guard) the submission's training context before # spending a baseline run on it: an out-of-range BLOCK_SIZE should cost the # judge nothing. sub_block = resolve_train_block_size(submission_module, cfg) diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harness/scoring.py b/2.0/problems/nanoslm_hybrid_arch_design/harness/scoring.py index ec4b25d1b..37b16fa5c 100644 --- a/2.0/problems/nanoslm_hybrid_arch_design/harness/scoring.py +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/scoring.py @@ -1,7 +1,7 @@ -"""Scoring: ABSOLUTE held-out bits-per-byte gain over the locked baseline. +"""Scoring: Absolute held-out bits-per-byte gain over the locked baseline. Torch-free and unit-tested on CPU. Shared by the evaluator and the public test -so the agent sees the same math the judge uses. See DESIGN.md §8. +so the agent sees the same math the judge uses. """ from __future__ import annotations @@ -14,7 +14,7 @@ class ScoreResult: score: float # bounded [0, 100] reported to Harbor score_unbounded: float # keeps rewarding past 100 rel_improvement: float # (base_bpb - sub_bpb) / base_bpb — reported only - abs_bpb_delta: float # base_bpb - sub_bpb — THE SCORED QUANTITY + abs_bpb_delta: float # base_bpb - sub_bpb — the scored quantity def _clip(x: float, lo: float, hi: float) -> float: @@ -28,7 +28,7 @@ def score_from_bpb( ) -> ScoreResult: """Map a submission's held-out bits-per-byte to a [0, 100] score. - Lower ``sub_bpb`` is better. The MEASUREMENT is the absolute gain:: + Lower ``sub_bpb`` is better. The measurement is the absolute gain:: gain = base_bpb - sub_bpb # bits per byte score = clip(100 * gain / bpb_score_scale, 0, 100) @@ -36,15 +36,15 @@ def score_from_bpb( The baseline architecture (``sub_bpb == base_bpb``) scores 0, and a submission worse than the baseline scores 0. - WHAT ``bpb_score_scale`` IS, AND WHAT IT IS NOT + What ``bpb_score_scale`` is, and what it is not ----------------------------------------------- - It is a DISPLAY CONVENTION and it is arbitrary by design. Its one necessary + It is a display convention and it is arbitrary by design. Its one necessary job is to map bits-per-byte into the 0-100 range Harbor expects: Harbor computes ``reward = score / 100``, so a raw ``gain`` of 0.05 bpb would surface as reward 0.0005 -- indistinguishable from zero for every submission. That is the whole reason the constant exists. - It is NOT a calibrated definition of "what counts as a full win". No such + It is not a calibrated definition of "what counts as a full win". No such number has been measured, and presenting an arbitrary constant as a target would (a) attach the score's meaning to a figure nobody determined and (b) throw away information at the clip, where a submission at 2x the scale @@ -52,19 +52,19 @@ def score_from_bpb( the un-clipped value and an operator should read it whenever the bounded score pins at 100. - The one weak requirement that does remain is DISCRIMINATION: set far too + The one weak requirement that does remain is discrimination: set far too large and every submission pins at 0, far too small and every submission pins at 100. That is a much weaker condition than calibration and can be set from a single real-corpus run. - WHY ABSOLUTE AND NOT RELATIVE + Why absolute and not relative ----------------------------- - Absolute bpb is the unit the literature quotes (e.g. Olmo Hybrid), so a + Absolute bpb is the unit the language-modelling literature quotes, so a result is directly comparable to published numbers; it is linear in cross-entropy, so equal absolute gains are equal information gains, and - gains compose roughly additively. The cost is that it is NOT scale-free: the + gains compose roughly additively. The cost is that it is not scale-free: the same absolute gain is roughly twice as hard at ``base_bpb`` 1.5 as at 2.9, - so the operating point matters -- see DESIGN.md §8. ``rel_improvement`` is + so the operating point matters. ``rel_improvement`` is still computed and reported so an operator can see the relative figure without recomputing it. """ @@ -97,13 +97,13 @@ def format_message( ) -> str: """Public feedback string — metrics only, no submission stdout/tracebacks. - Reports BOTH the scored absolute gain and the relative figure: only the - SCORED quantity changed to absolute, and an operator should not have to + Reports both the scored absolute gain and the relative figure: only the + Scored quantity changed to absolute, and an operator should not have to recompute the other one. """ msg = ( f"base_val_bpb={base_bpb:.5f}; sub_val_bpb={sub_bpb:.5f}; " - f"abs_bpb_delta={result.abs_bpb_delta:+.5f}; " # SCORED + f"abs_bpb_delta={result.abs_bpb_delta:+.5f}; " # Scored f"rel_improvement={result.rel_improvement:+.4%}; " f"steps={steps}; train_wall_s={wall_seconds:.1f}; " f"score={result.score:.4f}; score_unbounded={result.score_unbounded:.4f}" diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harness/settings.py b/2.0/problems/nanoslm_hybrid_arch_design/harness/settings.py index 5dae66924..333b88d2a 100644 --- a/2.0/problems/nanoslm_hybrid_arch_design/harness/settings.py +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/settings.py @@ -1,9 +1,7 @@ """Locked task settings + config fingerprint (torch-free). -All numbers here are read by the judge only. The agent never sees this file at -scoring time. Values tagged ``CALIBRATE`` are placeholders pending the -single-H100 calibration described in DESIGN.md §3 and must be frozen before the -task ships. +All numbers here are read by the judge only; the agent never sees this file at +scoring time. """ from __future__ import annotations @@ -20,41 +18,41 @@ @dataclass(frozen=True) class TaskConfig: # --- fixed model I/O contract (agent must honor these) --- - # dolma2 BPE (allenai/dolma2-tokenizer), the OLMo-3 tokenizer. UNGATED -- + # dolma2 BPE (allenai/dolma2-tokenizer), the OLMo-3 tokenizer. Ungated -- # no HF token is needed to download it, unlike lm_arch_discovery's Llama-2 - # tokenizer. 100278 ids > 65535, so token streams are uint32 on disk; uint16 - # would silently wrap the upper ~35k of the vocabulary. - vocab_size: int = 100278 - # HYBRID-DISCOVERY SETTING. 8192, not 1024: hybrids exist to escape - # attention's O(L^2), and at ctx 1024 attention is only ~31-40% of layer - # FLOPs, so a cheaper sequence mixer has little to win and the fixed - # wall-clock budget barely rewards it. At 8192 attention is ~78-84% of - # layer FLOPs, so the attention/recurrence trade is the dominant design - # question -- which is the point of this task. + # tokenizer. # - # TWO CONTEXT LENGTHS, AND THE DISTINCTION IS LOAD-BEARING. + # 100352 = the tokenizer's 100278 real ids PADDED up (Olmo 3's own + # convention, "padded for efficient embedding lookups"). The corpus is + # unaffected: ids on disk are always < 100278, the extra 74 embedding rows + # and logit columns are simply never indexed by data. Models must still + # produce logits over the full padded width. > 65535 either way, so token + # streams are uint32 on disk; uint16 would silently wrap the upper ~35k of + # the vocabulary. + vocab_size: int = 100352 + # 8192, not 1024: hybrids exist to escape attention's O(L^2), and attention + # only dominates layer FLOPs at long context, so the attention/recurrence + # trade is the dominant design question here -- the point of the task. # - # eval_block_size FIXED, judge-owned. Every submission is scored on + # Two context lengths, and the distinction is load-bearing. + # + # eval_block_size fixed, judge-owned. Every submission is scored on # non-overlapping windows of exactly this width, so - # val_bpb stays comparable across submissions. NEVER + # val_bpb stays comparable across submissions. Never # derive this from anything a submission controls. - # block_size DEFAULT training context. A submission may override it + # block_size default training context. A submission may override it # with a module-level ``BLOCK_SIZE`` int in model.py # (runner.resolve_train_block_size), because trading # context length for optimizer steps inside the fixed # wall-clock budget is now part of the design space -- # at the cost of having to extrapolate to 8192 at eval. - block_size: int = 8192 # DEFAULT training context (agent may override) - eval_block_size: int = 8192 # FIXED scoring window -- judge-owned + block_size: int = 8192 # Default training context (agent may override) + eval_block_size: int = 8192 # Fixed scoring window -- judge-owned # --- fixed optimization recipe (locked; agent designs architecture only) --- - # 2x16 (effective batch 32). MEASURED: micro-batch 8 at ctx 8192 OOMs on an - # H100 -- the [8, 8192, 100352] fp32 logits the harness materializes for CE - # are ~26GB, and the baseline transformer alone tried to allocate 24.48GiB - # with only 20.77GiB free. micro-batch 2 keeps the fp32-logits copy ~6.6GB - # and fits both arms (incl. the GDN hybrid) with margin; grad_accum 16 holds - # the effective batch at 32 so the optimization recipe stays comparable. - # (A fused/chunked cross-entropy would let micro-batch grow again -- TODO.) + # micro-batch 2 so the fp32 cross-entropy logits at ctx 8192 fit on one H100; + # grad_accum 16 holds the effective batch at 32 so the recipe stays + # comparable. (micro-batch 8 OOMs; a fused CE would let it grow again.) batch_size: int = 2 # sequences per micro-batch (per device) grad_accum: int = 16 # -> effective batch 32 sequences learning_rate: float = 3.0e-3 # AdamW peak LR (nanoGPT-speedrun class) @@ -65,65 +63,72 @@ class TaskConfig: grad_clip: float = 1.0 warmup_steps: int = 100 - # --- iso-wallclock budget (CALIBRATE: pick largest T leaving >=15min headroom) --- - train_seconds: float = 1800.0 # CALIBRATE wall-clock T per training run (30 min) - max_train_seconds: float = 3300.0 # hard cap (< 1h) — abort if exceeded + # --- iso-wallclock budget (matches config.yaml) --- + train_seconds: float = 21600.0 # wall-clock T per training run (6 h) + # Infrastructure backstop, enforced by the Modal function timeout (and the + # orchestrator), not by the training loop -- a between-step check could + # never fire before train_seconds and cannot interrupt a hung step. + max_train_seconds: float = 43200.0 # hard cap (12 h) # --- data / tokenizer (locked; hidden tokens live in the judge image) --- dataset_name: str = "HuggingFaceFW/fineweb-edu:sample-10BT" # provenance; matches config.yaml tokenizer_name: str = "allenai/dolma2-tokenizer" - token_dtype: str = "uint32" # vocab 100278 > 65535 -> uint16 is WRONG + token_dtype: str = "uint32" # ids up to 100277 > 65535 -> uint16 is wrong train_tokens_path: str = "/opt/nanoslm_arch/data/train.bin" - val_tokens_path: str = "/opt/nanoslm_arch/data/val.bin" # HELD OUT — never in /app + val_tokens_path: str = "/opt/nanoslm_arch/data/val.bin" # Held out — never in /app # Written by docker/prep_assets.py next to the streams; carries val_bytes. manifest_path: str = "/opt/nanoslm_arch/data/manifest.json" - val_tokens: int = 1_048_576 # held-out TOKENS scored (non-overlapping windows) - # BYTE length of the held-out text that `val_tokens` target tokens cover. + val_tokens: int = 1_048_576 # held-out tokens scored (non-overlapping windows) + # Byte length of the held-out text that `val_tokens` target tokens cover. # 0 means "resolve from the manifest / FRONTIER_NANOSLM_VAL_BYTES at load - # time" -- see resolve_val_bytes(). This must NEVER silently become a token - # count: bpb is normalized by BYTES so that it stays comparable across + # time" -- see resolve_val_bytes(). This must never silently become a token + # count: bpb is normalized by bytes so that it stays comparable across # architectures the way the byte-level tokenizer used to make it by # construction. val_bytes: int = 0 # --- iterative-role cached baseline (GPU via Modal; judge container CPU-only) --- # Mirrors nanowm's `baseline_cache`. The cheap agent-role feedback path trains - # only the submission (~T on a Modal H100) and reuses a baseline perplexity + # only the submission (~T on a Modal H100) and reuses a baseline result # cached by config fingerprint; the final/verifier path recomputes a fresh - # baseline+submission CRN pair (~2T) on one Modal GPU/process. In local testing - # mode the judge can use a directly-attached H100 instead of Modal. + # baseline+submission CRN pair (~2T) on one Modal GPU/process. + # + # On the Modal backend (production) the cache lives on the nanoslm-corpus + # Volume so it survives containers, and is keyed by fingerprint AND the + # staged train.bin's byte size AND the modal_app CODE_VERSION -- the + # fingerprint only hashes config values, so a harness behavior change + # (e.g. the LR schedule) must bump CODE_VERSION, and a corpus restage + # changes the size stamp automatically. See modal_app.run_pair_remote. + # The path below is only the local-backend (dev/testing) location. baseline_cache_path: str = "/opt/nanoslm_arch/baseline/baseline_ppl.json" # --- scoring --- - # ABSOLUTE bits-per-byte gain is the scored measurement: + # Absolute bits-per-byte gain is the scored measurement: # gain = base_bpb - sub_bpb # score = clip(100 * gain / bpb_score_scale, 0, 100) - # - # NOT A CALIBRATION TARGET, AND DELIBERATELY NOT LISTED AS ONE. - # `bpb_score_scale` is a DISPLAY CONVENTION whose only necessary job is to - # map bpb into the 0-100 range Harbor wants: Harbor computes - # `reward = score / 100`, so a raw gain of 0.05 bpb would arrive as reward - # 0.0005 -- zero, for every submission. It does NOT define "what counts as a - # full win"; no such number has been measured, and dressing an arbitrary - # constant as a measured target would put the score's meaning on a figure - # nobody determined. The only real requirement is DISCRIMINATION (too coarse - # and everyone pins at 0, too fine and everyone pins at 100), which is far - # weaker than calibration and settable from one real-corpus run. - # `score_unbounded` stays un-clipped so strong submissions remain visible. - # - # NOTE this replaced a RELATIVE `r_target` of the same numeric value, and - # the two are NOT equivalent: 0.05 relative was ~0.145 bpb absolute at the - # synthetic operating point (base_bpb ~2.9) and would be ~0.075-0.10 bpb - # once a real corpus puts the baseline nearer 1.5-2.0. So 0.05 ABSOLUTE is a - # STRICTER bar on real data than the relative form it replaced. - bpb_score_scale: float = 0.05 # bpb gain that saturates the bounded score + # `bpb_score_scale` is a display convention that maps the gain into Harbor's + # 0-100 reward, not a calibrated "full win" target; its only + # real requirement is discrimination, and `score_unbounded` stays un-clipped. + # 0.5 comes from a real-corpus run and is due a re-check at the 6 h budget. + bpb_score_scale: float = 0.5 # bpb gain that saturates the bounded score # --- guards / resource caps --- param_cap: int = 400_000_000 # max trainable params (over-cap -> score 0) min_param_delta: float = 1e-6 # trained-from-scratch guard: mean|Δparam| threshold + # Plausibility floor on the scored bits-per-byte. The static policy gate + # scans Turing-complete code and is defense in depth, not proof; a + # submission that evaded it and read the held-out stream at runtime would + # score a near-zero bpb. No honest result can approach this floor: the + # strongest published LLMs sit near ~0.55 bpb on web text, and a <=400M + # model trained 6 h from scratch lands far above that. A measurement below + # it is therefore treated as held-out leakage, not brilliance -> GuardError + # (score 0, public reason). A scoring/guard constant like bpb_score_scale, + # so it is deliberately NOT fingerprinted. + min_plausible_bpb: float = 0.4 seed: int = 1337 - # OFF pending validation -- see runner.run_arm. Closing the "compile the - # baseline architecture and win on wall-clock" loophole still needs doing. + # Off by default (see runner.run_arm): compiling the chunkwise mixer measured + # worse than eager, and the fairness argument for compiling both arms is real + # but unresolved. compile_model: bool = False # --- determinism knobs applied on the GPU path --- @@ -136,32 +141,28 @@ class TaskConfig: # --------------------------------------------------------------------------- # # Byte accounting for the held-out stream. # -# THIS IS THE SUBTLE PART OF THE BPE MIGRATION. At byte level 1 token == 1 byte, -# so `mean per-token CE / ln2` WAS bits-per-byte and dividing the total NLL by +# This is the subtle part of the BPE migration. At byte level 1 token == 1 byte, +# so `mean per-token CE / ln2` was bits-per-byte and dividing the total NLL by # the token count happened to be right. With a BPE tokenizer the two diverge by # the compression ratio (~4.4 bytes/token for dolma2 on English), and per-token -# CE/ln2 is a TOKENIZER-DEPENDENT quantity -- it is not comparable across setups +# CE/ln2 is a tokenizer-dependent quantity -- it is not comparable across setups # and it is not bits per byte. The whole reason bpb was chosen as the scored -# metric (DESIGN.md §2) is that it is tokenizer-independent, so the -# denominator must be BYTES. +# metric is that it is tokenizer-independent, so the +# denominator must be bytes. # # The byte count is produced at corpus-prep time (docker/prep_assets.py decodes # the scored token span and measures its UTF-8 length) and travels in -# manifest.json. eval_ppl.py ASSERTS it is available rather than falling back. +# manifest.json. eval_ppl.py asserts it is available rather than falling back; +# with no real corpus the harness raises DataError instead of estimating a ratio. # --------------------------------------------------------------------------- # -# Only used when no real corpus is staged (synthetic smoke stream). Roughly the -# measured dolma2 compression ratio on English web text, so a smoke bpb lands in -# a plausible range instead of being off by ~4.4x. -SYNTHETIC_BYTES_PER_TOKEN = 4.4 - def resolve_val_bytes_per_token(cfg: TaskConfig | None = None) -> float | None: - """Bytes of held-out TEXT per scored target token, or ``None`` if unknown. + """Bytes of held-out text per scored target token, or ``None`` if unknown. - A RATIO rather than a bare total, because the number of tokens actually + A ratio rather than a bare total, because the number of tokens actually scored is ``min(len(val)-1, cfg.val_tokens)`` rounded down to whole windows - -- the smoke config scores far fewer than the corpus contains. The ratio is + -- a test config may score fewer than the corpus contains. The ratio is exact for the production config (which scores the whole staged stream) and proportional otherwise, and it can never be confused for a token count. @@ -208,71 +209,45 @@ def _ratio(nbytes, ntok) -> float | None: return None -# --------------------------------------------------------------------------- # -# Smoke overrides: tiny, fast, CPU-friendly. Enabled by FRONTIER_NANOSLM_SMOKE=1. -# Used only to prove the harness wiring end-to-end without a GPU; never scored. -# --------------------------------------------------------------------------- # -def smoke_enabled() -> bool: - return os.environ.get("FRONTIER_NANOSLM_SMOKE", "") == "1" - - def active_config() -> TaskConfig: - if not smoke_enabled(): - return DEFAULT - # SMOKE SHRINKS THE BUDGET, NOT THE SHAPE. - # - # The obvious smoke config (block_size=64) is actively harmful here: fla's - # Triton kernels autotune per shape, and a 64-token sequence is far outside - # the regime they are tuned for. Measured on an H100, GatedDeltaNet's first - # forward costs 74.5s at T=64 but only 15.5s at T=8192 -- the tiny sequence - # is ~5x SLOWER to compile. A block_size=64 smoke therefore appears to hang - # for tens of minutes while compiling kernels no scored run will ever use. - # - # So the smoke keeps the REAL sequence length and cuts batch size, steps and - # eval tokens instead. It is slower than a toy config but it exercises the - # kernels the scored run actually uses, which is the point of a smoke test. - return TaskConfig( - block_size=8192, # REAL shape -- see above - eval_block_size=8192, # scoring window is NEVER shrunk by the smoke - batch_size=1, - grad_accum=1, - warmup_steps=1, - # 45s, not 3s. The score is a function of THROUGHPUT, so a smoke that - # yields one optimizer step per arm cannot validate the thing being - # measured -- at train_seconds=3.0 both arms did exactly 1 step and the - # reported +4.2% was noise. 45s gives tens of steps and a real - # steps-per-arm comparison, at the cost of a slower smoke. - train_seconds=45.0, - max_train_seconds=900.0, - val_tokens=32_768, - param_cap=400_000_000, - ) + """The one production configuration. + + (A FRONTIER_NANOSLM_SMOKE budget-shrinking mode used to live here for + cheap wiring checks; it was removed deliberately -- the harness runs one + configuration, the scored one. Ad-hoc testing configs belong in tests, + built with ``dataclasses.replace`` on ``DEFAULT``, as the evaluator + selftest does.) + """ + return DEFAULT # --------------------------------------------------------------------------- # # Config fingerprint: the cache key for the iterative-role cached baseline. -# A change to ANY locked knob invalidates a cached baseline rather than +# A change to any locked knob invalidates a cached baseline rather than # mispairing it (mirrors nanowm settings.config_fingerprint()). # --------------------------------------------------------------------------- # # -# NOTE `block_size` IS DELIBERATELY ABSENT, and `eval_block_size` replaces it. -# `block_size` is now only the DEFAULT training context: each submission may +# Note `block_size` is deliberately absent, and `eval_block_size` replaces it. +# `block_size` is now only the default training context: each submission may # pick its own (runner.resolve_train_block_size). Fingerprinting a per-arm value # would give every submission a distinct key, so the cached baseline would never # hit and the agent-role feedback path would silently cost ~2T instead of ~T -- -# the exact thing the cache exists to avoid. The BASELINE always trains at 8192 +# the exact thing the cache exists to avoid. The baseline always trains at 8192 # (baseline_model.BLOCK_SIZE) and is always scored at `eval_block_size`, so # neither of those depends on the submission and the cached number stays valid. -# `eval_block_size` DOES belong here: changing the scoring window changes what +# `eval_block_size` does belong here: changing the scoring window changes what # val_bpb means, so a baseline cached under the old one must not be reused. _FINGERPRINT_KEYS = ( "vocab_size", "eval_block_size", "batch_size", "grad_accum", "learning_rate", "min_lr", "weight_decay", "beta1", "beta2", "grad_clip", "warmup_steps", "train_seconds", "dataset_name", # tokenizer_name is fingerprinted alongside vocab_size: changing the - # tokenizer changes what val_bpb MEANS, so a baseline cached under the old + # tokenizer changes what val_bpb means, so a baseline cached under the old # one must not be reused. "tokenizer_name", "val_tokens", "seed", + # compile_model changes both arms' throughput, hence the baseline's step + # count and val_bpb -- flipping it must invalidate a cached baseline. + "compile_model", ) diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harness/train.py b/2.0/problems/nanoslm_hybrid_arch_design/harness/train.py index 260bcc14b..4ccbe755d 100644 --- a/2.0/problems/nanoslm_hybrid_arch_design/harness/train.py +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/train.py @@ -3,7 +3,7 @@ The optimizer, LR schedule, weight-decay grouping, gradient accumulation, data order, gradient clipping, and the wall-clock cutoff are all fixed here so the task measures *architecture*, not training tricks. The harness computes the -cross-entropy loss from the model's logits and IGNORES any loss the model +cross-entropy loss from the model's logits and ignores any loss the model returns. Each optimizer step accumulates ``cfg.grad_accum`` micro-batches of @@ -13,7 +13,10 @@ Timing note: wall-clock is measured by ``time.monotonic`` inside the harness (never by the model). Training stops at the first optimizer step whose start -exceeds ``cfg.train_seconds``. ``max_train_seconds`` is a hard safety abort. +exceeds ``cfg.train_seconds``. ``max_train_seconds`` is an infrastructure +backstop (the Modal function timeout / orchestrator), not a loop condition: a +between-step check could never fire before the ``train_seconds`` one, and it +could not interrupt a hung step anyway. """ from __future__ import annotations @@ -33,18 +36,32 @@ class TrainOutput: final_train_loss: float -def _lr_at(step: int, cfg: TaskConfig, total_steps_guess: int) -> float: +def _lr_at(step: int, elapsed: float, cfg: TaskConfig) -> float: + """Linear step-warmup, then cosine decay over the wall-clock budget. + + The decay is a function of ``elapsed / train_seconds``, not of a step + count. The stop condition is wall-clock, so a step horizon cannot be known + in advance, and any fixed guess hands different effective schedules to arms + with different throughput: a fast arm (short context, cheap mixer) overruns + the guess and trains most of its budget flat at ``min_lr``, while a slow + arm never finishes the decay -- which is known to cost real validation + loss. That would let the score partly measure schedule luck rather than + architecture, the exact confound the locked recipe exists to remove. + Time-based decay gives every architecture the identical schedule over the + identical budget. Warmup stays step-based: its job is stabilizing early + optimizer state, which is a per-step property, and at 100 steps it is a + negligible slice of the budget. + """ if step < cfg.warmup_steps: return cfg.learning_rate * (step + 1) / max(1, cfg.warmup_steps) - # cosine decay toward min_lr over an estimated horizon - denom = max(1, total_steps_guess - cfg.warmup_steps) - frac = min(1.0, (step - cfg.warmup_steps) / denom) + frac = min(1.0, max(0.0, elapsed / max(1e-9, cfg.train_seconds))) coeff = 0.5 * (1.0 + math.cos(math.pi * frac)) return cfg.min_lr + coeff * (cfg.learning_rate - cfg.min_lr) def _param_groups(model, cfg: TaskConfig): - """No weight decay on 1-D params (norms, biases, embeddings).""" + """No weight decay on 1-D params (norms, biases). 2-D params -- including + the (tied) embedding table -- are decayed, nanoGPT-style.""" decay, no_decay = [], [] for _, p in model.named_parameters(): if not p.requires_grad: @@ -75,9 +92,6 @@ def train_model(model, data: TokenData, cfg: TaskConfig, device: str) -> TrainOu lr=cfg.learning_rate, betas=(cfg.beta1, cfg.beta2), ) - # Rough horizon estimate for the cosine schedule; the true stop is wall-clock. - total_steps_guess = 10_000 - use_amp = device.startswith("cuda") autocast = ( torch.autocast(device_type="cuda", dtype=torch.bfloat16) @@ -92,17 +106,18 @@ def train_model(model, data: TokenData, cfg: TaskConfig, device: str) -> TrainOu elapsed = time.monotonic() - t0 if elapsed >= cfg.train_seconds: break - if elapsed >= cfg.max_train_seconds: # hard safety abort - break + # (max_train_seconds is enforced by the Modal function timeout, not + # here: train_seconds always breaks first between steps, and no + # between-step check can interrupt a hung step.) - lr = _lr_at(step, cfg, total_steps_guess) + lr = _lr_at(step, elapsed, cfg) for pg in opt.param_groups: pg["lr"] = lr # Gradient accumulation: grad_accum micro-batches of cfg.batch_size make # one optimizer step (effective batch = batch_size * grad_accum) while # peak activation memory stays at a single micro-batch. cfg.block_size is - # this arm's TRAINING context (run_arm passes a per-arm cfg); evaluation + # this arm's training context (run_arm passes a per-arm cfg); evaluation # is at cfg.eval_block_size regardless -- see data.val_windows. accum = max(1, int(cfg.grad_accum)) opt.zero_grad(set_to_none=True) diff --git a/2.0/problems/nanoslm_hybrid_arch_design/readme b/2.0/problems/nanoslm_hybrid_arch_design/readme index 49b6fac2f..e896bbd81 100644 --- a/2.0/problems/nanoslm_hybrid_arch_design/readme +++ b/2.0/problems/nanoslm_hybrid_arch_design/readme @@ -10,18 +10,30 @@ training + evaluation harness, trains it for a fixed wall-clock `T`, and scores its `val_bpb` against a **locked baseline architecture** trained under the identical budget. -The starting point is the 3:1 Gated DeltaNet hybrid from *Olmo Hybrid* -(arXiv:2604.03444); the question is how much further it can be pushed. The optimizer, learning-rate schedule, +The starting point is a 3:1 Gated DeltaNet hybrid (`reference.py`, following +*Olmo Hybrid*, arXiv:2604.03444); the question is how much further it can be +pushed. The optimizer, learning-rate schedule, weight-decay grouping, data, tokenizer, and the wall-clock budget are all **fixed by the judge**. You control the architecture — and, as of this revision, the **training context length** (see below). A more compute-efficient architecture legitimately completes more useful training steps within `T` — that is the intended lever. +Why a hybrid, specifically: the baseline is already a tuned modern transformer, +so the easy block-level wins (RMSNorm, rotary embeddings, QK-norm, SwiGLU) are +already in it — what is left on the table is the sequence-mixing direction. A +linear-recurrent mixer is cheaper per token at long context than attention, so +under a fixed wall-clock budget it can complete more optimizer steps, while a +few attention layers preserve the global context a pure linear RNN lacks. That +is the bet the reference hybrid makes; the design space it opens is the mixer +choice, the recurrence's internals, layer placement, the attention ratio, and +the training context length. + ## Metric The tokenizer is fixed by the judge: **dolma2 BPE** -(`allenai/dolma2-tokenizer`, vocabulary 100278) over a FineWeb-Edu corpus. The +(`allenai/dolma2-tokenizer`, 100278 real ids padded to a model vocabulary of +100352, the Olmo 3 convention) over a FineWeb-Edu corpus. The scored quantity is held-out **bits per byte**: ``` @@ -49,7 +61,9 @@ class NanoSLM(torch.nn.Module): ... `config` is provided by the harness and has: -- `config.vocab_size` — always `100278` (dolma2 BPE; do not change). +- `config.vocab_size` — always `100352` (dolma2 BPE padded, Olmo 3 + convention; do not change). Actual token ids are always `< 100278` — the + padding rows are never indexed, but your logits must span the full width. - `config.block_size` — the context length you will be **trained** at (8192 by default; see "Training context length" below). - `config.eval_block_size` — the context length you will be **scored** at. @@ -61,7 +75,8 @@ The returned module must implement: ```python def forward(self, idx): - # idx: LongTensor [B, T] of BPE token ids in [0, 100278) + # idx: LongTensor [B, T] of BPE token ids in [0, 100278); logits span + # the padded vocab_size (100352) return logits # FloatTensor [B, T, vocab_size] ``` @@ -75,8 +90,8 @@ The vocabulary is large enough that the embedding table is a first-class design concern — at `d=768` it is ~77M parameters on its own, so tying, factorizing or otherwise reshaping it is a real lever rather than a detail. -The 3:1 Gated DeltaNet hybrid `model.py` (the *Olmo Hybrid* recipe applied to the -`olmo3_190M` baseline) is provided as a starting point. +A 3:1 Gated DeltaNet hybrid `model.py` — GDN in most layers, full attention in +the rest, at the `olmo3_190M` baseline's shape — is provided as a starting point. ## Training context length — yours to choose, with a catch @@ -123,6 +138,17 @@ point. - Train from scratch: no loading pretrained weights, no reading files, no network, no environment access, no reading the clock, no timing short-circuits. Submissions containing such calls are rejected before running. +- Imports are restricted to an allowlist: `torch`, `numpy`, `fla`, `einops`, + `triton` (custom kernels are a legitimate lever), `math`, and a few + pure-computation stdlib modules (`typing`, `dataclasses`, `functools`, + `itertools`, `collections`, ...). Wildcard and relative imports, bare + `eval`/`exec`/`getattr`/`__import__`/`open`, dunder attributes other than + `__init__`/`__version__`, and raw-file-reader attributes (`fromfile`, + `memmap`, ...) are rejected statically. `model.eval()` and + `torch.compile(...)` remain allowed. +- A measured bits-per-byte below the plausibility floor (0.4 — far beyond any + honest result at this scale and budget) is rejected as presumed held-out + leakage, score 0. - Trainable parameters must be within the judge's param cap, and the model must fit and train within GPU memory. Over-cap or out-of-memory scores 0. - The model must actually train (its parameters must change) and must produce a @@ -167,3 +193,66 @@ bash /app/submit.sh The best successful iterative submission is kept if a later artifact is worse or you time out. + +## Problem structure + +``` +nanoslm_hybrid_arch_design/ +├── readme This file: problem statement, metric, scoring. +├── config.yaml Harbor problem config: runtime, images, GPU/eval +│ knobs (budget, param cap, bpb_score_scale). +├── evaluator.py Judge entrypoint. evaluate(model.py) -> (score, +│ score_unbounded, message, metrics): static policy +│ gate -> Modal (or local-GPU) dispatch -> scoring. +│ `--selftest` runs the torch-free test suite. +├── evaluate.sh Local CLI wrapper around evaluator.py. +├── reference.py Reference solution: the 3:1 GDN:attention hybrid +│ at the baseline's shape (~254M params). Also the +│ shipped starter (harbor/app/model.py is a copy). +├── harness/ The locked training+evaluation harness. +│ ├── settings.py TaskConfig: every locked knob (optimizer, budget, +│ │ data paths, caps), config fingerprint (baseline- +│ │ cache key), byte-accounting resolution. +│ ├── model_config.py ModelConfig: the read-only object handed to the +│ │ submission's build_model(config). +│ ├── policy.py Static gate: substring denylists + AST scan with +│ │ a strict import allowlist. Shipped verbatim to +│ │ the agent image so the rules cannot drift. +│ ├── data.py TokenData: loads the uint32 token streams, serves +│ │ CRN training batches (identical across arms) and +│ │ the fixed 8192-wide validation windows; converts +│ │ scored tokens to bytes for the bpb denominator. +│ ├── train.py Locked training loop: AdamW, step-warmup + +│ │ wall-clock cosine LR decay, grad accumulation, +│ │ harness-owned cross-entropy, wall-clock cutoff. +│ ├── eval_ppl.py Held-out evaluation: val_bpb (scored) + val_ppl +│ │ (readability) + degeneracy probe. +│ ├── runner.py Orchestration: BLOCK_SIZE resolution, kernel +│ │ warmup (incl. eval-shape fast-fail), dynamic +│ │ guards (param cap, OOM, trained-from-scratch, +│ │ plausibility floor), run_arm / run_pair (CRN). +│ ├── baseline_model.py The locked baseline: faithful plain-PyTorch +│ │ olmo3_190M (SWA via FlexAttention on CUDA). +│ └── modal_app.py Modal GPU app: run_pair_remote (role-aware -- +│ │ final = fresh CRN pair, agent = cached baseline +│ │ from the corpus Volume + submission-only train), +│ │ image pins (torch/triton/fla), corpus Volume. +├── docker/ +│ ├── build_images.sh Builds agent+judge images; stages the corpus on a +│ │ cold cache; leak-greps everything agent-visible. +│ ├── prep_assets.py Tokenizes FineWeb-Edu with dolma2 into +│ │ train.bin / val.bin / manifest.json (val_bytes, +│ │ the bpb denominator, measured here). +│ ├── agent/Dockerfile Agent workspace image: CPU-only, no corpus, no +│ │ torch; starter model.py + the policy gate. +│ └── judge/Dockerfile Judge image: CPU-only, holds the token streams + +│ manifest, dispatches GPU work to Modal. +└── harbor/app/ The agent's workspace (copied into /app). + ├── model.py Starter submission (copy of reference.py). + ├── README.md Agent-facing brief: interface, BLOCK_SIZE trade, + │ static-gate rules. + └── public_test.py/.sh Runs the judge's exact static gate locally. +``` + +The held-out `val.bin` and its manifest exist only in the judge image and the +Modal corpus volume — never in the agent workspace. diff --git a/2.0/problems/nanoslm_hybrid_arch_design/reference.py b/2.0/problems/nanoslm_hybrid_arch_design/reference.py index 6170d9453..a9d2fa929 100644 --- a/2.0/problems/nanoslm_hybrid_arch_design/reference.py +++ b/2.0/problems/nanoslm_hybrid_arch_design/reference.py @@ -1,97 +1,87 @@ -"""Reference solution for nanoslm_hybrid_arch_design -- the Olmo-Hybrid recipe at 190M. +"""Reference solution for nanoslm_hybrid_arch_design -- a 3:1 GDN-to-attention hybrid built on the 190M olmo3_190M shape (254M params). -This is the STARTING POINT, not the ceiling. The locked baseline is a faithful -``olmo3_190M``; this file applies OLMo-core's own hybrid recipe -(``src/scripts/train/OLMo_hybrid/OLMo-hybrid-7B.py``) to it. +This is the starting point, not the ceiling. The locked baseline is a faithful +pure-attention ``olmo3_190M``; this file replaces most of its attention layers +with Gated DeltaNet (GDN) linear-recurrent layers. -WHERE THE 3:1 RATIO ACTUALLY COMES FROM ---------------------------------------- -It is NOT an arbitrary choice. ``olmo3_190M`` already carries:: +The 3:1 ratio +------------- +It is not an arbitrary choice. ``olmo3_190M`` already carries:: SlidingWindowAttentionConfig(pattern=[4096, 4096, 4096, -1], force_full_attention_on_first_layer=False, force_full_attention_on_last_layer=True) so 9 of its 12 layers are sliding-window and 3 (layers 3, 7, 11) are full -attention. The paper "replaces sliding window layers with Gated DeltaNet -layers", and upstream's recipe encodes exactly that:: +attention. This reference replaces each sliding-window layer with a GDN layer and +keeps the full-attention layers:: - config.block_pattern = ["gdn", "gdn", "gdn", "attn"] + block_pattern = ["gdn", "gdn", "gdn", "attn"] The three sliding-window layers of each 4-layer group become GDN; the -full-attention layer of each group survives untouched. The 3:1 ratio IS olmo3's -own sliding-window pattern. Note the consequence: the hybrid has NO sliding -window left anywhere -- every surviving attention layer sat at a ``-1`` position. +full-attention layer of each group survives untouched. The 3:1 ratio comes from +the base model's own sliding-window pattern. Note the consequence: the hybrid has +No sliding window left anywhere -- every surviving attention layer sat at a +``-1`` position. -PARAMETER MATCHING (upstream's REMOVE_HEADS practice) ------------------------------------------------------ -GDN's mixer is larger than an attention mixer, so upstream shrinks the model to -compensate rather than letting the hybrid quietly buy capacity:: +No parameter matching +--------------------- +A GDN mixer is larger than an attention mixer. This reference does not shrink the +hybrid to match the baseline's parameter count: it keeps the same d=768, h=12, +l=12 as the baseline and lets the hybrid run at its natural (larger) size. With +Tied embeddings:: - REMOVE_HEADS = 2 - config.d_model -= REMOVE_HEADS * 128 # 4096 -> 3840 - num_heads -= REMOVE_HEADS # 32 -> 30 - assert config.d_model / num_heads == 128 # head_dim preserved + baseline 190,354,176 (12 attention layers) + hybrid 254,430,936 (9 GDN + 3 attention layers) +33.7% total -At 190M with head_dim 64 the equivalent is ``REMOVE_HEADS = 1``:: +(Totals at the padded vocab 100352 -- dolma2's 100278 real ids plus 74 unused +rows on the tied table, the Olmo 3 padding convention.) - d_model 768 -> 704, n_heads 12 -> 11, 704 / 11 == 64 (head_dim preserved) +The hybrid is bigger -- the GDN mixer's wide recurrent state (head dim +ceil_128(0.75*d/h) = 128) costs params -- but well under the 400M ``param_cap``. +The hard cap -- not an artificial per-arm match -- is what bounds capacity, so +the scored question is whether the mixer choice pays off on quality within that +budget. -Matching is done on NON-EMBEDDING parameters, which is the quantity upstream -itself feeds to ``Duration.chinchilla_tokens(model_params=...)``. Both arms TIE -their single vocab table, which at this scale is ~40% of all parameters -- -including it would drown the mixer difference being matched. Measured -analytically: - - baseline non-embedding 113,283,840 - hybrid non-embedding 110,805,734 -2.19% - -which is the same tolerance upstream's own 7B pair achieves (+2.84%). One caveat -worth stating: because ``d_model`` shrinks 768 -> 704, the tied table shrinks too -(77.0M -> 70.6M), so the hybrid's TOTAL is 4.7% below the baseline's -(181,401,446 vs 190,297,344) even though its non-embedding count is within 2.2%. - -Note also that ``hidden_size`` stays 3072. Upstream mutates ``d_model`` in place -and never recomputes the feed-forward width, so the FFN keeps the width derived -from the ORIGINAL d_model. Reproduced here deliberately. +``hidden_size`` stays 3072 (the SwiGLU width for d_model 768), same as baseline. Everything else is identical to the baseline -- RMSNorm eps 1e-6, SwiGLU 3072, -reordered_norm blocks, qk_norm, RoPE theta 500_000, TIED embeddings (the baseline -ties, so this arm ties too for comparability) -- so the ONLY architectural +reordered_norm blocks, qk_norm, RoPE theta 500_000, tied embeddings (the baseline +ties, so this arm ties too for comparability) -- so the only architectural difference is the sequence mixer in those nine layers. -YOUR TASK: PUSH val_bpb FURTHER +Your task: Push val_bpb further ------------------------------- -Everything above is upstream's recipe, transposed to this scale. The open -questions it does NOT answer, each worth real bpb: +This 3:1 GDN hybrid is one point in a large design space. Open questions it does +Not settle, each worth real bpb: - * THE RATIO. 3:1 is inherited from olmo3's sliding-window pattern, which was - chosen for a sliding window, not for a linear RNN. 5:1 and 7:1 buy more - steps but give up more global context. - * PLACEMENT. Every 4th, or clustered (early layer for global context in, late + * The linear-to-attention ratio. 3:1 (GDN:attention) is inherited from the base model's sliding-window pattern, + which was chosen for a sliding window, not for a linear RNN. 5:1 and 7:1 buy + more steps but give up more global context. + * Layer placement. Every 4th, or clustered (early layer for global context in, late layer for read-out)? Same cost, different models. - * STATE SIZE. ``expand_v`` and ``num_v_heads`` set the recurrent state, the + * State size. ``expand_v`` and ``num_v_heads`` set the recurrent state, the main capacity knob of a linear RNN -- unlike a KV cache it does not grow with sequence length, so capacity is cheap at 8192. - * THE MIXER ITSELF. GDN is one choice; GLA, RetNet and Mamba2-style mixers are + * The mixer itself. GDN is one choice; GLA, RetNet and Mamba2-style mixers are all expressible in the same chunkwise matmul form. - * NON-UNIFORMITY. Nothing requires every GDN layer to be identical, or the + * Non-uniformity. Nothing requires every GDN layer to be identical, or the attention layers to be full-width. - * THE REST OF THE BLOCK. Norm placement, gating, MLP ratio, head count, + * The rest of the block. Norm placement, gating, MLP ratio, head count, embedding tying -- all still on the table, and all interact with the above. - * THE TRAINING CONTEXT. Declare a module-level ``BLOCK_SIZE`` int (a power of + * The training context. Declare a module-level ``BLOCK_SIZE`` int (a power of two in [256, 8192]) to train at a shorter context than the default 8192. Shorter steps are cheaper, so you complete more of them in the fixed - wall-clock budget -- but EVALUATION IS ALWAYS AT 8192, so the model must + wall-clock budget -- but **evaluation is always at 8192**, so the model must extrapolate to positions well beyond anything it trained on. How well it - does that is mostly a property of the POSITION ENCODING (plain RoPE degrades + does that is mostly a property of the position encoding (plain RoPE degrades sharply; NTK-aware/YaRN scaling and ALiBi hold up far better), which is a different question from mixer efficiency. This file declares no BLOCK_SIZE and therefore trains at 8192. Anything sized off ``config.block_size`` must still run at ``config.eval_block_size``. -Locked and not yours to change: optimizer, data, tokenizer, the EVALUATION -context (8192), and the wall-clock budget. Interface: ``build_model(config)`` / +Locked and not yours to change: optimizer, data, tokenizer, the evaluation +context, and the wall-clock budget. Interface: ``build_model(config)`` / ``NanoSLM(config)`` returning ``forward(idx) -> logits [B, T, vocab_size]``. """ @@ -104,36 +94,31 @@ import torch.nn.functional as F # --------------------------------------------------------------------------- # -# Shape: olmo3_190M with upstream's REMOVE_HEADS compensation applied. +# Architectural Details # --------------------------------------------------------------------------- # -REMOVE_HEADS = 1 - +# Attention hyperparameters +REMOVE_HEADS = 0 N_LAYER = 12 -N_HEAD = 12 - REMOVE_HEADS # 11 -HEAD_DIM = 64 # preserved, exactly as upstream asserts -N_EMBD = 768 - REMOVE_HEADS * HEAD_DIM # 704 -assert N_EMBD // N_HEAD == HEAD_DIM, "REMOVE_HEADS must preserve head_dim" +N_HEAD = 12 - REMOVE_HEADS # 12 +HEAD_DIM = 64 +N_EMBD = 768 - REMOVE_HEADS * HEAD_DIM # 768 +assert N_EMBD // N_HEAD == HEAD_DIM, "head_dim must be preserved" -# NOT recomputed from the reduced d_model -- upstream mutates d_model in place -# and leaves the feed-forward width at the value llama_like derived from the -# ORIGINAL 768. Reproduced deliberately. -HIDDEN_SIZE = 3072 +HIDDEN_SIZE = 3072 # SwiGLU width, same as baseline ROPE_THETA = 500_000.0 NORM_EPS = 1e-6 -# GatedDeltaNetConfig(n_heads=num_heads, head_dim=int(0.75*d_model/num_heads), -# allow_neg_eigval=True) with expand_v at its 2.0 default. -GDN_HEAD_DIM = int(0.75 * N_EMBD / N_HEAD) # 48 +# GDN hyperparameters +_gdn_raw = int(0.75 * N_EMBD / N_HEAD) +GDN_HEAD_DIM = ((_gdn_raw + 127) // 128) * 128 GDN_EXPAND_V = 2.0 -GDN_HEAD_V_DIM = int(GDN_HEAD_DIM * GDN_EXPAND_V) # 96 -GDN_KEY_DIM = N_HEAD * GDN_HEAD_DIM # 528 -GDN_VALUE_DIM = N_HEAD * GDN_HEAD_V_DIM # 1056 +GDN_ALLOW_NEG_EIGVAL = True +GDN_HEAD_V_DIM = int(GDN_HEAD_DIM * GDN_EXPAND_V) # 256 +GDN_KEY_DIM = N_HEAD * GDN_HEAD_DIM # 1536 +GDN_VALUE_DIM = N_HEAD * GDN_HEAD_V_DIM # 3072 GDN_CONV_SIZE = 4 -# config.block_pattern = ["gdn", "gdn", "gdn", "attn"] -- the three -# sliding-window layers of each group become GDN, the full-attention layer -# survives. Attention therefore lands at layers 3, 7, 11. PATTERN = ("gdn", "gdn", "gdn", "attn") assert N_LAYER % len(PATTERN) == 0 @@ -162,17 +147,7 @@ def _rope(x, theta: float = ROPE_THETA): class _Attention(nn.Module): - """Identical to the baseline's attention, always FULL causal. - - Full, not windowed, and that is not a simplification: the attention layers - the hybrid keeps are exactly the ones sitting at the ``-1`` positions of - olmo3's [4096, 4096, 4096, -1] pattern, so they were already full attention - before the GDN substitution. - - qk_norm spans the full n_heads*head_dim projection and is applied BEFORE the - head reshape -- upstream's behaviour when ``use_head_qk_norm`` is False, - which is the olmo3_190M default. - """ + """Identical to the baseline's attention, always full causal.""" def __init__(self, n_embd: int, n_head: int): super().__init__() @@ -198,123 +173,22 @@ def forward(self, x): return self.w_out(y.transpose(1, 2).contiguous().view(B, T, C)) -class _ChunkedDeltaNet(nn.Module): - """CPU-only chunkwise fallback, PARAMETER-IDENTICAL to fla's GatedDeltaNet. - - Exists so the CPU smoke path can build and report the same parameter count - the CUDA path will. Its module inventory is fla 0.5.1's exactly -- - q/k/v/a/b/g/o projections, three depthwise short convolutions, A_log, - dt_bias and the gated output norm -- so ``n_params`` is device-independent. - - The recurrence is a chunkwise GATED LINEAR ATTENTION, not the full delta - rule: it carries the decay and the gates but skips the delta-rule inverse. - That is a deliberate approximation. It is never scored -- CUDA raises rather - than falling back here (see ``_make_gdn``) -- and it exists to prove wiring, - not to reproduce GDN's numerics. - """ - - def __init__(self, chunk: int = 512): - super().__init__() - self.chunk = chunk - d, H = N_EMBD, N_HEAD - self.q_proj = nn.Linear(d, GDN_KEY_DIM, bias=False) - self.k_proj = nn.Linear(d, GDN_KEY_DIM, bias=False) - self.v_proj = nn.Linear(d, GDN_VALUE_DIM, bias=False) - self.a_proj = nn.Linear(d, H, bias=False) - self.b_proj = nn.Linear(d, H, bias=False) - self.g_proj = nn.Linear(d, GDN_VALUE_DIM, bias=False) - self.o_proj = nn.Linear(GDN_VALUE_DIM, d, bias=False) - self.A_log = nn.Parameter(torch.zeros(H)) - self.dt_bias = nn.Parameter(torch.zeros(H)) - # ShortConvolution weights: depthwise, (channels, 1, kernel), no bias. - self.q_conv1d = nn.Parameter(torch.zeros(GDN_KEY_DIM, 1, GDN_CONV_SIZE)) - self.k_conv1d = nn.Parameter(torch.zeros(GDN_KEY_DIM, 1, GDN_CONV_SIZE)) - self.v_conv1d = nn.Parameter(torch.zeros(GDN_VALUE_DIM, 1, GDN_CONV_SIZE)) - self.o_norm = nn.Parameter(torch.ones(GDN_HEAD_V_DIM)) - - @staticmethod - def _short_conv(x, w): - """Causal depthwise conv over time. x [B, T, C], w [C, 1, K].""" - C, K = w.shape[0], w.shape[2] - xt = F.pad(x.transpose(1, 2), (K - 1, 0)) - return F.conv1d(xt, w, groups=C).transpose(1, 2) - - def forward(self, x): - B, T, _ = x.shape - H, DK, DV, S = N_HEAD, GDN_HEAD_DIM, GDN_HEAD_V_DIM, self.chunk - - q = F.silu(self._short_conv(self.q_proj(x), self.q_conv1d)) - k = F.silu(self._short_conv(self.k_proj(x), self.k_conv1d)) - v = self._short_conv(self.v_proj(x), self.v_conv1d) - - q = q.view(B, T, H, DK).transpose(1, 2) - k = k.view(B, T, H, DK).transpose(1, 2) - v = v.view(B, T, H, DV).transpose(1, 2) - - # Per-head decay in (0, 1): the gate that makes this "gated". - dt = F.softplus(self.a_proj(x) + self.dt_bias) # [B, T, H] - g = torch.exp(-torch.exp(self.A_log) * dt) # [B, T, H] - g = g.permute(0, 2, 1).unsqueeze(-1) # [B, H, T, 1] - beta = torch.sigmoid(self.b_proj(x)).permute(0, 2, 1).unsqueeze(-1) - v = v * beta - - pad = (S - T % S) % S - if pad: - q, k, v = (F.pad(t, (0, 0, 0, pad)) for t in (q, k, v)) - g = F.pad(g, (0, 0, 0, pad), value=1.0) - nC = (T + pad) // S - rs = lambda t: t.view(B, H, nC, S, -1) # noqa: E731 - qc, kc, vc, gc = rs(q), rs(k), rs(v), rs(g) - - causal = torch.tril(torch.ones(S, S, device=x.device, dtype=torch.bool)) - intra = (qc @ kc.transpose(-1, -2)).masked_fill(~causal, 0.0) @ vc - - state = torch.zeros(B, H, DK, DV, device=x.device, dtype=x.dtype) - outs = [] - for c in range(nC): - outs.append(intra[:, :, c] + qc[:, :, c] @ state) - decay = gc[:, :, c].prod(dim=2, keepdim=True) - state = state * decay + kc[:, :, c].transpose(-1, -2) @ (vc[:, :, c] * gc[:, :, c]) - y = torch.stack(outs, dim=2).view(B, H, T + pad, DV)[:, :, :T] - - # Gated RMSNorm over head_v_dim, then merge heads. - yf = y.float() - y = (yf * torch.rsqrt(yf.pow(2).mean(-1, keepdim=True) + 1e-5)).to(x.dtype) - y = y * self.o_norm - y = y.transpose(1, 2).contiguous().view(B, T, GDN_VALUE_DIM) - y = y * F.silu(self.g_proj(x)) - return self.o_proj(y) - - def _make_gdn(): - """fla's fused GatedDeltaNet on CUDA; chunkwise only on CPU (smoke). - - WHY fla IS REQUIRED HERE, not merely preferred: PyTorch ships a fused kernel - for softmax attention (SDPA) but none for a linear/recurrent mixer, so a - hand-written GDN is unfused eager code. Scoring is fixed WALL-CLOCK, so an - unfused mixer loses on throughput regardless of architectural merit -- - measured at ctx 8192, the chunkwise fallback did 6% of attention's FLOPs and - still ran 3.5x slower. Both arms must be kernel-matched or the score - measures kernel quality, not architecture. - - HARD FAILURE ON CUDA, deliberately: a silent fallback here once produced a - scored-looking run where the reference "lost" 18 steps to 65 purely because - fla was missing. A missing kernel must raise, not quietly change what is - being measured. - """ + """fla's fused GatedDeltaNet. CUDA required.""" if not torch.cuda.is_available(): - return _ChunkedDeltaNet() - - from fla.layers import GatedDeltaNet # ImportError here is intentional + raise RuntimeError( + "reference hybrid requires a CUDA device: the GDN layers run only " + "through fla's fused Triton kernels, and there is deliberately no " + "CPU fallback" + ) - # head_dim MUST be passed: fla defaults it to 256, so omitting it builds - # num_heads*256-wide projections instead of the intended width -- a silent - # param blowup that trains a much larger model than the baseline. + from fla.layers import GatedDeltaNet return GatedDeltaNet( hidden_size=N_EMBD, num_heads=N_HEAD, head_dim=GDN_HEAD_DIM, expand_v=GDN_EXPAND_V, + allow_neg_eigval=GDN_ALLOW_NEG_EIGVAL, use_gate=True, use_short_conv=True, ) @@ -327,8 +201,8 @@ def __init__(self): def forward(self, x): out = self.impl(x) - # fla layers return (hidden_states, attentions, past_kv); the chunkwise - # fallback returns a tensor. Normalize so _Block need not care. + # fla layers return (hidden_states, attentions, past_kv) -- keep only + # the hidden states. return out[0] if isinstance(out, tuple) else out @@ -369,15 +243,12 @@ def __init__(self, config): kinds = [PATTERN[i % len(PATTERN)] for i in range(N_LAYER)] self.blocks = nn.ModuleList([_ReorderedNormBlock(N_EMBD, N_HEAD, k) for k in kinds]) self.norm_f = _RMSNorm(N_EMBD) - # TIED -- the baseline ties its embeddings (a deliberate deviation from - # upstream's untied default), so the reference ties too and the two arms - # stay comparable: the ONLY architectural difference between them is the - # sequence mixer in the GDN layers, not the output-head param budget. The - # tie is set AFTER init (below) so the shared table carries wte's init. + self.lm_head = nn.Linear(N_EMBD, config.vocab_size, bias=False) self.apply(self._init) + for name, p in self.named_parameters(): - if name.endswith("w_out.weight") or name.endswith("w2.weight"): + if name.endswith(("w_out.weight", "w2.weight", "o_proj.weight")): nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * N_LAYER)) self.lm_head.weight = self.wte.weight # weight tying @@ -420,14 +291,20 @@ def analytic_n_params(vocab_size: int) -> int: n_attn = sum(1 for i in range(N_LAYER) if PATTERN[i % len(PATTERN)] == "attn") n_gdn = N_LAYER - n_attn non_emb = n_attn * (attn_mix + block_norms + ffn) + n_gdn * (gdn_mix + block_norms + ffn) + d - # ONE vocab-sized table: lm_head is TIED to wte (see NanoSLM.__init__). + return non_emb + vocab_size * d -def _self_check(vocab_size: int = 100278) -> int: - from harness.model_config import ModelConfig +def _self_check(vocab_size: int = 100352) -> int: + # A minimal stand-in for the harness-provided config object + class _Cfg: + def __init__(self, vocab_size: int, block_size: int, + eval_block_size: int = 8192): + self.vocab_size = vocab_size + self.block_size = block_size + self.eval_block_size = eval_block_size - m = NanoSLM(ModelConfig(vocab_size=vocab_size, block_size=8192)) + m = NanoSLM(_Cfg(vocab_size, 8192)) got = sum(p.numel() for p in m.parameters() if p.requires_grad) want = analytic_n_params(vocab_size) assert got == want, f"param mismatch: built {got} != analytic {want}" diff --git a/2.0/problems/nanoslm_hybrid_arch_design/repro/__init__.py b/2.0/problems/nanoslm_hybrid_arch_design/repro/__init__.py deleted file mode 100644 index ff7c20caa..000000000 --- a/2.0/problems/nanoslm_hybrid_arch_design/repro/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Repro rig for Olmo Hybrid Table 5 (190M Base-Easy BPB).""" diff --git a/2.0/problems/nanoslm_hybrid_arch_design/repro/modal_repro.py b/2.0/problems/nanoslm_hybrid_arch_design/repro/modal_repro.py deleted file mode 100644 index 9c8279e19..000000000 --- a/2.0/problems/nanoslm_hybrid_arch_design/repro/modal_repro.py +++ /dev/null @@ -1,483 +0,0 @@ -"""Modal app to reproduce Olmo Hybrid Table 5 (190M Base-Easy BPB). - -Two remote functions on one shared Volume: - - prep_data (CPU) stream the Olmo 3 mix as raw text, dolma2-tokenize it, and - write uint32 train.bin / val.bin (+ manifest with val_bytes) - to the Volume. Parameterized by target_tokens so a tiny - wiring test and a full 1x/8x-Chinchilla stage share code. - - train (H100) load the tokenized bins, build the paper-faithful 190M - transformer or GDN-3:1 hybrid (repro.paper_models), train - with a WSD-S schedule (repro.wsd), log train + held-out CE, - and checkpoint to the Volume. - -The DOWNSTREAM OlmoBaseEval-Easy bpb eval (the actual 0.950/0.891 metric) is a -separate stage (oe-eval) run on the produced checkpoints; this app produces the -checkpoints and a held-out-CE trajectory to sanity-check the training recipe. - -Run (from the driver venv, with ~/.modal.toml authed): - modal run repro/modal_repro.py::prep_data --target-tokens 20000000 # wiring - modal run repro/modal_repro.py::train --arm transformer --target-tokens 20000000 -""" - -from __future__ import annotations - -import os -import pathlib - -import modal - -APP_NAME = "nanoslm-repro" -DATA = "/data" -REPRO_REMOTE = "/root/repro" -_HERE = pathlib.Path(__file__).resolve().parent - -VOL = modal.Volume.from_name("nanoslm-repro", create_if_missing=True) - -# Same torch/fla/triton stack the benchmark's modal_app pins -- fla 0.5.1 needs -# Triton >= 3.2 (torch 2.6 ships it); 0.4.1's backward kernels predate it and -# hang/err. transformers carries the dolma2 tokenizer. -image = ( - modal.Image.debian_slim(python_version="3.11") - .pip_install( - "torch==2.6.0", "numpy<2", - "flash-linear-attention==0.5.1", "transformers==4.46.3", - "huggingface_hub>=0.25", "tokenizers>=0.20", "zstandard>=0.22", - "hf_transfer>=0.1.6", # HF_HUB_ENABLE_HF_TRANSFER=1 needs this present - ) - .add_local_dir(str(_HERE), REPRO_REMOTE, copy=True, - ignore=["__pycache__", "*.pyc"]) - .env({ - "CUBLAS_WORKSPACE_CONFIG": ":4096:8", - "TRITON_CACHE_DIR": "/tmp/triton-cache", - "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", - "HF_HUB_ENABLE_HF_TRANSFER": "1", - "PYTHONPATH": "/root", - }) -) - -app = modal.App(APP_NAME) - -TOKENIZER = "allenai/dolma2-tokenizer" # vocab 100278 (padded to 100352) -VOCAB_PADDED = 100352 -DEFAULT_DATASET = "allenai/dolma3_mix-150B-1025" # 150B sample of the Olmo 3 mix - - -# --------------------------------------------------------------------------- # -# Data prep: stream raw text -> dolma2 token ids -> uint32 bins on the Volume. -# --------------------------------------------------------------------------- # -@app.function(image=image, cpu=16.0, memory=32768, timeout=6 * 60 * 60, - volumes={DATA: VOL}) -def prep_data(target_tokens: int, val_tokens: int = 8_000_000, - dataset: str = DEFAULT_DATASET, text_field: str = "text", - tag: str = "smoke", seed: int = 0) -> dict: - """Tokenize `target_tokens` (+ val) tokens of `dataset` into the Volume. - - The dataset is 6k+ ``.jsonl.zst`` shards grouped by domain; the per-domain - shard COUNT already encodes the Olmo 3 mixture proportions. We list every - shard, DETERMINISTICALLY SHUFFLE it (so a small subset stays mixture- - representative rather than all-one-domain), then read shards via zstandard - until the token target is met. Train and val come from DISJOINT shards. - - Writes {DATA}/{tag}/train.bin, val.bin, manifest.json. manifest carries - val_bytes (UTF-8 length of the decoded val span) so bpb is normalized by - BYTES, not tokens. - """ - import io - import json - import time - - import numpy as np - import zstandard - from huggingface_hub import HfApi, hf_hub_download - from transformers import AutoTokenizer - - outdir = pathlib.Path(DATA) / tag - outdir.mkdir(parents=True, exist_ok=True) - tok = AutoTokenizer.from_pretrained(TOKENIZER) - eos = tok.eos_token_id if tok.eos_token_id is not None else 0 - - api = HfApi() - shards = sorted(f.rfilename for f in api.dataset_info(dataset).siblings - if f.rfilename.endswith(".jsonl.zst")) - order = np.random.default_rng(seed).permutation(len(shards)) - shards = [shards[i] for i in order] - print(f"{len(shards)} shards; shuffled seed={seed}") - - need = int(target_tokens) + int(val_tokens) - buf = np.empty(need + 4_000_000, dtype=np.uint32) - n = 0 - dctx = zstandard.ZstdDecompressor() - t0 = time.time() - - def append(texts): - nonlocal n - enc = tok(texts, add_special_tokens=False)["input_ids"] - for ids in enc: - m = len(ids) + 1 - if n + m > buf.shape[0]: - m = buf.shape[0] - n - if m <= 0: - return - buf[n:n + m - 1] = np.asarray(ids[:m - 1], dtype=np.uint32) - buf[n + m - 1] = eos - n += m - - shards_used = 0 - for path in shards: - if n >= need: - break - local = hf_hub_download(dataset, path, repo_type="dataset") - batch = [] - with open(local, "rb") as fh: - reader = io.TextIOWrapper(dctx.stream_reader(fh), encoding="utf-8") - for line in reader: - if not line.strip(): - continue - try: - t = json.loads(line).get(text_field) - except Exception: - continue - if not t: - continue - batch.append(t) - if len(batch) >= 1000: - append(batch); batch = [] - if n >= need: - break - if batch and n < need: - append(batch) - os.remove(local) # keep container disk bounded - shards_used += 1 - if shards_used % 10 == 0: - print(f" {shards_used} shards, {n/1e6:.1f}M/{need/1e6:.1f}M tok, " - f"{n/max(1e-9,time.time()-t0)/1e3:.0f}k tok/s") - - n = min(n, need) - ids = buf[:n] - hi = int(ids.max()) if n else 0 - assert hi < 100278, f"token id {hi} out of dolma2 vocab" - - val_ids = ids[-val_tokens:] - train_ids = ids[:-val_tokens] - train_ids.tofile(outdir / "train.bin") - val_ids.tofile(outdir / "val.bin") - - val_text = tok.decode(val_ids.tolist(), skip_special_tokens=True) - val_bytes = len(val_text.encode("utf-8")) - manifest = { - "tag": tag, "dataset": dataset, "tokenizer": TOKENIZER, - "shards_used": shards_used, "seed": seed, - "train_tokens": int(train_ids.size), "val_tokens": int(val_ids.size), - "val_bytes": int(val_bytes), - "val_bytes_per_token": val_bytes / max(1, val_ids.size), - "seconds": round(time.time() - t0, 1), - } - (outdir / "manifest.json").write_text(json.dumps(manifest, indent=2)) - VOL.commit() - print("prep_data done:", manifest) - return manifest - - -# --------------------------------------------------------------------------- # -# Training: paper-faithful 190M model, WSD-S schedule, checkpoint to Volume. -# --------------------------------------------------------------------------- # -@app.function(image=image, cpu=8.0, memory=16384, timeout=6 * 60 * 60, - volumes={DATA: VOL}) -def _tokenize_shards(args) -> dict: - """Worker: tokenize a list of shards -> one part_{idx}.bin on the Volume.""" - import io - import json - import time - - import numpy as np - import zstandard - from huggingface_hub import hf_hub_download - from transformers import AutoTokenizer - - idx, paths, dataset, text_field, tag = args - tok = AutoTokenizer.from_pretrained(TOKENIZER) - eos = tok.eos_token_id if tok.eos_token_id is not None else 0 - dctx = zstandard.ZstdDecompressor() - parts = pathlib.Path(DATA) / tag / "parts" - parts.mkdir(parents=True, exist_ok=True) - - chunks, ntok, t0 = [], 0, time.time() - for path in paths: - local = hf_hub_download(dataset, path, repo_type="dataset") - texts = [] - with open(local, "rb") as fh: - reader = io.TextIOWrapper(dctx.stream_reader(fh), encoding="utf-8") - for line in reader: - if not line.strip(): - continue - try: - t = json.loads(line).get(text_field) - except Exception: - continue - if t: - texts.append(t) - if len(texts) >= 1000: - enc = tok(texts, add_special_tokens=False)["input_ids"] - for ids in enc: - a = np.asarray(ids + [eos], dtype=np.uint32); chunks.append(a); ntok += a.size - texts = [] - if texts: - enc = tok(texts, add_special_tokens=False)["input_ids"] - for ids in enc: - a = np.asarray(ids + [eos], dtype=np.uint32); chunks.append(a); ntok += a.size - os.remove(local) - arr = np.concatenate(chunks) if chunks else np.empty(0, dtype=np.uint32) - hi = int(arr.max()) if arr.size else 0 - assert hi < 100278, f"token id {hi} out of dolma2 vocab" - outp = parts / f"part_{idx:04d}.bin" - arr.tofile(outp) - VOL.commit() - return {"idx": idx, "tokens": int(arr.size), "shards": len(paths), - "seconds": round(time.time() - t0, 1)} - - -@app.function(image=image, cpu=32.0, memory=65536, timeout=6 * 60 * 60, - volumes={DATA: VOL}) -def prep_data_parallel(target_tokens: int, val_tokens: int = 8_000_000, - dataset: str = DEFAULT_DATASET, tag: str = "run", - n_workers: int = 32, tokens_per_shard_est: int = 22_000_000, - headroom: float = 1.35, seed: int = 0) -> dict: - """Fan-out tokenization: split a shuffled shard subset across workers, then - concatenate the parts into train.bin / val.bin (+ manifest).""" - import json - import time - - import numpy as np - from huggingface_hub import HfApi - from transformers import AutoTokenizer - - outdir = pathlib.Path(DATA) / tag - outdir.mkdir(parents=True, exist_ok=True) - api = HfApi() - shards = sorted(f.rfilename for f in api.dataset_info(dataset).siblings - if f.rfilename.endswith(".jsonl.zst")) - order = np.random.default_rng(seed).permutation(len(shards)) - shards = [shards[i] for i in order] - - need = int(target_tokens) + int(val_tokens) - n_shards = min(len(shards), int(np.ceil(need / tokens_per_shard_est * headroom))) - picked = shards[:n_shards] - # Round-robin assign so each worker gets a mixture-representative slice. - buckets = [[] for _ in range(n_workers)] - for i, p in enumerate(picked): - buckets[i % n_workers].append(p) - jobs = [(i, b, dataset, "text", tag) for i, b in enumerate(buckets) if b] - print(f"{len(picked)} shards over {len(jobs)} workers (need {need/1e6:.0f}M tok)") - - t0 = time.time() - results = sorted(_tokenize_shards.map(jobs), key=lambda r: r["idx"]) - total = sum(r["tokens"] for r in results) - print(f"tokenized {total/1e6:.0f}M tok in {(time.time()-t0)/60:.1f}m across workers") - if total < need: - print(f"WARNING: only {total/1e6:.0f}M < need {need/1e6:.0f}M; increase headroom") - - # Concatenate parts in worker order into train.bin, then val.bin. - VOL.reload() - parts = sorted((outdir / "parts").glob("part_*.bin")) - train_target = int(target_tokens) - written_train = 0 - ftrain = open(outdir / "train.bin", "wb") - val_arrs, val_have = [], 0 - for pf in parts: - a = np.fromfile(pf, dtype=np.uint32) - if written_train < train_target: - take = min(a.size, train_target - written_train) - a[:take].tofile(ftrain); written_train += take - rest = a[take:] - else: - rest = a - if rest.size and val_have < val_tokens: - need_v = val_tokens - val_have - val_arrs.append(rest[:need_v]); val_have += min(rest.size, need_v) - ftrain.close() - val_ids = np.concatenate(val_arrs) if val_arrs else np.empty(0, dtype=np.uint32) - val_ids.tofile(outdir / "val.bin") - - tok = AutoTokenizer.from_pretrained(TOKENIZER) - val_text = tok.decode(val_ids.tolist(), skip_special_tokens=True) - val_bytes = len(val_text.encode("utf-8")) - for pf in parts: # free part files - pf.unlink() - manifest = { - "tag": tag, "dataset": dataset, "tokenizer": TOKENIZER, - "shards_used": len(picked), "n_workers": len(jobs), "seed": seed, - "train_tokens": int(written_train), "val_tokens": int(val_ids.size), - "val_bytes": int(val_bytes), - "val_bytes_per_token": val_bytes / max(1, val_ids.size), - "prep_minutes": round((time.time() - t0) / 60, 1), - } - (outdir / "manifest.json").write_text(json.dumps(manifest, indent=2)) - VOL.commit() - print("prep_data_parallel done:", manifest) - return manifest - - -@app.function(image=image, gpu="H100", memory=32768, timeout=12 * 60 * 60, - volumes={DATA: VOL}) -def train(arm: str = "transformer", target_tokens: int = 20_000_000, - tag: str = "smoke", seq_len: int = 4096, - # micro_batch defaults are ARM-AWARE (0 = auto): the transformer fits - # mb=8 (fp32 logits ~13GB), but the hybrid's fla GDN autotuner needs - # more headroom and OOMs at mb=8, so it uses mb=4. Effective batch is - # held at 128 sequences via grad_accum = 128 // micro_batch. - micro_batch: int = 0, grad_accum: int = 0, - peak_lr: float = 2.0e-3, warmup_frac: float = 0.02, - weight_decay: float = 0.1, beta1: float = 0.9, beta2: float = 0.95, - grad_clip: float = 1.0, log_every: int = 25, seed: int = 1337) -> dict: - """Train one arm for ~target_tokens tokens with WSD-S; log held-out CE.""" - import json - import sys - import time - - import numpy as np - import torch - import torch.nn.functional as F - - sys.path.insert(0, "/root") - from repro import paper_models, wsd - - torch.manual_seed(seed) - dev = "cuda" - ddir = pathlib.Path(DATA) / tag - # memmap the 15GB train stream (read-only) so it isn't loaded into RAM. - train_ids = np.memmap(ddir / "train.bin", dtype=np.uint32, mode="r") - val_ids = np.fromfile(ddir / "val.bin", dtype=np.uint32) # small (~32MB) - manifest = json.loads((ddir / "manifest.json").read_text()) - bytes_per_tok = float(manifest["val_bytes_per_token"]) - - class Cfg: - vocab_size = VOCAB_PADDED - block_size = seq_len - eval_block_size = seq_len - - build = paper_models.build_transformer if arm == "transformer" else paper_models.build_hybrid - model = build(Cfg()).to(dev) - n_params = sum(p.numel() for p in model.parameters() if p.requires_grad) - n_non_emb = n_params - model.wte.weight.numel() - - if not micro_batch: - micro_batch = 4 if arm == "hybrid" else 8 - if not grad_accum: - grad_accum = max(1, 128 // micro_batch) - tokens_per_step = micro_batch * grad_accum * seq_len - total_steps = max(1, int(target_tokens) // tokens_per_step) - warmup_steps = max(1, int(warmup_frac * total_steps)) - - decay, no_decay = [], [] - for p in model.parameters(): - (decay if p.dim() >= 2 else no_decay).append(p) - opt = torch.optim.AdamW( - [{"params": decay, "weight_decay": weight_decay}, - {"params": no_decay, "weight_decay": 0.0}], - lr=peak_lr, betas=(beta1, beta2), - ) - - rng = np.random.default_rng(seed) - hi = train_ids.size - seq_len - 1 - - def get_batch(nseq): - ix = rng.integers(0, hi, size=nseq) - x = np.stack([train_ids[i:i + seq_len] for i in ix]).astype(np.int64) - y = np.stack([train_ids[i + 1:i + 1 + seq_len] for i in ix]).astype(np.int64) - return (torch.from_numpy(x).pin_memory().to(dev, non_blocking=True), - torch.from_numpy(y).pin_memory().to(dev, non_blocking=True)) - - @torch.no_grad() - def eval_ce(max_windows=24): - model.eval() - n = min(val_ids.size - 1, max_windows * seq_len) - nw = n // seq_len - tot_ce, tot_tok = 0.0, 0 - for w in range(nw): - s = w * seq_len - x = torch.from_numpy(val_ids[s:s + seq_len][None].astype(np.int64)).to(dev) - y = torch.from_numpy(val_ids[s + 1:s + 1 + seq_len][None].astype(np.int64)).to(dev) - logits = model(x).float() - tot_ce += float(F.cross_entropy(logits.reshape(-1, logits.size(-1)), - y.reshape(-1), reduction="sum")) - tot_tok += y.numel() - model.train() - mean_ce = tot_ce / max(1, tot_tok) - val_bpb = tot_ce / (tot_tok * bytes_per_tok * np.log(2)) - return mean_ce, val_bpb - - logpath = ddir / f"train_{arm}.log" - log_lines = [] - - def log(msg): - print(msg, flush=True) - log_lines.append(msg) - - log(f"[{arm}] non_emb={n_non_emb/1e6:.1f}M total={n_params/1e6:.1f}M " - f"tokens/step={tokens_per_step} total_steps={total_steps} warmup={warmup_steps} " - f"target_tokens={target_tokens}") - - t0 = time.time() - for step in range(total_steps): - lr = wsd.lr_wsd(step, peak_lr=peak_lr, total_steps=total_steps, - warmup_steps=warmup_steps) - for pg in opt.param_groups: - pg["lr"] = lr - opt.zero_grad(set_to_none=True) - loss_acc = 0.0 - for _ in range(grad_accum): - x, y = get_batch(micro_batch) - with torch.autocast("cuda", dtype=torch.bfloat16): - logits = model(x) - loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)).float(), - y.reshape(-1)) / grad_accum - loss.backward() - loss_acc += float(loss) * grad_accum - torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip) - opt.step() - - if step % log_every == 0 or step == total_steps - 1: - toks = (step + 1) * tokens_per_step - dt = time.time() - t0 - tps = toks / max(1e-9, dt) - log(f"[{arm}] step {step:5d}/{total_steps} lr={lr:.2e} " - f"train_ce={loss_acc/grad_accum:.4f} tok={toks/1e6:.1f}M " - f"{tps/1e3:.0f}k tok/s wall={dt/60:.1f}m") - # Flush progress to the Volume so DETACHED runs are observable - # without a live client (see modal run --detach). - logpath.write_text("\n".join(log_lines)) - VOL.commit() - - mean_ce, val_bpb = eval_ce() - wall = time.time() - t0 - ckpt = ddir / f"ckpt_{arm}.pt" - torch.save({"model": model.state_dict(), "arm": arm, "cfg_vocab": VOCAB_PADDED, - "seq_len": seq_len, "target_tokens": target_tokens}, ckpt) - logpath.write_text("\n".join(log_lines)) - VOL.commit() - result = { - "arm": arm, "n_non_emb_M": round(n_non_emb / 1e6, 1), - "total_steps": total_steps, "tokens_trained_M": round(total_steps * tokens_per_step / 1e6, 1), - "final_train_ce": round(loss_acc / grad_accum, 4), - "val_ce_nats": round(mean_ce, 4), "val_bpb_lm": round(val_bpb, 4), - "wall_min": round(wall / 60, 1), - "tok_per_s_k": round(total_steps * tokens_per_step / wall / 1e3, 1), - "gpu": torch.cuda.get_device_name(0), - } - log(f"[{arm}] DONE {result}") - return result - - -@app.local_entrypoint() -def main(action: str = "prep", arm: str = "transformer", - target_tokens: int = 20_000_000, tag: str = "smoke"): - if action == "prep": - print(prep_data.remote(target_tokens=target_tokens, tag=tag)) - elif action == "prep_par": - print(prep_data_parallel.remote(target_tokens=target_tokens, tag=tag)) - elif action == "train": - print(train.remote(arm=arm, target_tokens=target_tokens, tag=tag)) - else: - raise SystemExit(f"unknown action {action!r}") diff --git a/2.0/problems/nanoslm_hybrid_arch_design/repro/paper_models.py b/2.0/problems/nanoslm_hybrid_arch_design/repro/paper_models.py deleted file mode 100644 index 5ac59d4b4..000000000 --- a/2.0/problems/nanoslm_hybrid_arch_design/repro/paper_models.py +++ /dev/null @@ -1,344 +0,0 @@ -"""Paper-faithful 190M ablation models for reproducing Olmo Hybrid Table 5. - -These are DELIBERATELY NOT the benchmark's ``baseline_model.py`` / ``reference.py``. -The Table 5 numbers (Transformer 0.950, GDN-3:1 0.891 Base-Easy BPB @ 190M) were -produced by the scaling-ladder ablation models described in arXiv:2604.03444 -Section D.3 + Table 22, which differ from the benchmark variants in three ways: - - * NORM PLACEMENT: pre-norm (RMSNorm BEFORE each sub-layer), per D.3 - ("RMSNorm applied before each sub-layer (pre-norm)"). The benchmark's - baseline_model.py uses reordered_norm (norm AFTER the residual branch). - * BASELINE ATTENTION: plain full multi-head causal attention. D.3 describes - no sliding window for the ablation ladder; the benchmark baseline uses a - 9:3 SWA (4096-window) pattern. - * HYBRID SIZING: the ablation hybrid keeps the SAME d=768, h=12, l=12 as the - transformer (non-embedding params grow 190M -> 254M). The benchmark's - reference.py instead param-matches via REMOVE_HEADS=1 (d704, h11). - -Shape (Table 22, 190M column): d=768, h=12, l=12, head_dim=64. -MLP: SwiGLU, hidden = round_up_256(1.5 * 8d/3) = 3072 (D.3). -Untied embeddings; RoPE theta 500_000; QK-norm; RMSNorm eps 1e-6. -Vocab: dolma2, 100352 padded (D.3). BPB is vocab-padding-invariant, but we use -100352 to match the paper's embedding shape exactly. - -GDN (D.3 + Appendix A.1): head sizing proportional to an attention head -- -d_k = 3/4 * head_dim = 48, d_v = 2 * d_k = 96 (expand_v = 2.0), allow_neg_eigval -= True, use_gate = True, use_short_conv = True. fla's GatedDeltaNet is REQUIRED -on CUDA (a hand-written mixer is unfused and would confound any timing); a -chunkwise CPU fallback exists only so param counts / wiring are checkable off-GPU. - -Placement (D.3): "every r-th layer is a full transformer block" with r=4 (3:1), -"and we additionally enforce the final layer be an attention layer". For l=12 -that is attention at layers 3, 7, 11 (0-indexed), i.e. pattern (gdn,gdn,gdn,attn). -""" - -from __future__ import annotations - -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F - -# --------------------------------------------------------------------------- # -# Shared shape (Table 22, 190M). -# --------------------------------------------------------------------------- # -N_LAYER = 12 -N_HEAD = 12 -N_EMBD = 768 -HEAD_DIM = N_EMBD // N_HEAD # 64 -HIDDEN_SIZE = 3072 # round_up_256(1.5 * 8*768/3) = 3072 -ROPE_THETA = 500_000.0 -NORM_EPS = 1e-6 - -# GDN head sizing (D.3): hGDN = ceil_128(0.75 * d/h), where ceil_128 rounds UP -# to the nearest multiple of 128. For 190M: 0.75 * 768/12 = 48 -> 128. Key dim -# = h * hGDN, value dim = h * 2*hGDN (expand_v = 2). This ceil-to-128 is what -# reference.py omitted (it used 48), and it is why the ablation hybrid is 254M -# non-embed, not ~201M -- the GDN mixer is ~9.5M/layer, not ~3.6M. -def _ceil_mult(x: int, m: int) -> int: - return ((x + m - 1) // m) * m - -GDN_HEAD_DIM = _ceil_mult(int(0.75 * HEAD_DIM), 128) # ceil_128(48) = 128 -GDN_EXPAND_V = 2.0 -GDN_HEAD_V_DIM = int(GDN_HEAD_DIM * GDN_EXPAND_V) # 256 -GDN_KEY_DIM = N_HEAD * GDN_HEAD_DIM # 1536 -GDN_VALUE_DIM = N_HEAD * GDN_HEAD_V_DIM # 3072 -GDN_CONV_SIZE = 4 - -# r=4 (3:1) interleave, final layer forced to attention -> attn at 3, 7, 11. -PATTERN = ("gdn", "gdn", "gdn", "attn") -assert N_LAYER % len(PATTERN) == 0 - - -class RMSNorm(nn.Module): - def __init__(self, d: int, eps: float = NORM_EPS): - super().__init__() - self.eps = eps - self.weight = nn.Parameter(torch.ones(d)) - - def forward(self, x): - dt = x.dtype - xf = x.float() - n = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps) - return n.to(dt) * self.weight - - -def _rope(x, theta: float = ROPE_THETA): - _, _, T, D = x.shape - half = D // 2 - freq = theta ** (-torch.arange(0, half, device=x.device, dtype=torch.float32) / half) - ang = torch.arange(T, device=x.device, dtype=torch.float32)[:, None] * freq[None, :] - cos, sin = ang.cos()[None, None], ang.sin()[None, None] - x1, x2 = x[..., :half].float(), x[..., half:].float() - return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1).to(x.dtype) - - -class Attention(nn.Module): - """Full causal MHA with QK-norm (full-width, pre-reshape) + RoPE. No SWA.""" - - def __init__(self, n_embd: int = N_EMBD, n_head: int = N_HEAD): - super().__init__() - assert n_embd % n_head == 0 - self.n_head, self.n_embd = n_head, n_embd - self.head_dim = n_embd // n_head - self.w_q = nn.Linear(n_embd, n_embd, bias=False) - self.w_k = nn.Linear(n_embd, n_embd, bias=False) - self.w_v = nn.Linear(n_embd, n_embd, bias=False) - self.w_out = nn.Linear(n_embd, n_embd, bias=False) - self.q_norm = RMSNorm(n_embd) - self.k_norm = RMSNorm(n_embd) - - def forward(self, x): - B, T, C = x.shape - q, k, v = self.w_q(x), self.w_k(x), self.w_v(x) - q, k = self.q_norm(q), self.k_norm(k) - q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2) - k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2) - v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2) - q, k = _rope(q), _rope(k) - y = F.scaled_dot_product_attention(q, k, v, is_causal=True) - return self.w_out(y.transpose(1, 2).contiguous().view(B, T, C)) - - -class _ChunkedDeltaNet(nn.Module): - """CPU-only chunkwise fallback, param-identical to fla's GatedDeltaNet. - - Numerically approximate (gated linear attention, no delta-rule inverse); - exists ONLY so param counts / wiring are checkable off-GPU. Never used on - CUDA -- _make_gdn raises there if fla is missing. - """ - - def __init__(self, chunk: int = 512): - super().__init__() - self.chunk = chunk - d, H = N_EMBD, N_HEAD - self.q_proj = nn.Linear(d, GDN_KEY_DIM, bias=False) - self.k_proj = nn.Linear(d, GDN_KEY_DIM, bias=False) - self.v_proj = nn.Linear(d, GDN_VALUE_DIM, bias=False) - self.a_proj = nn.Linear(d, H, bias=False) - self.b_proj = nn.Linear(d, H, bias=False) - self.g_proj = nn.Linear(d, GDN_VALUE_DIM, bias=False) - self.o_proj = nn.Linear(GDN_VALUE_DIM, d, bias=False) - self.A_log = nn.Parameter(torch.zeros(H)) - self.dt_bias = nn.Parameter(torch.zeros(H)) - self.q_conv1d = nn.Parameter(torch.zeros(GDN_KEY_DIM, 1, GDN_CONV_SIZE)) - self.k_conv1d = nn.Parameter(torch.zeros(GDN_KEY_DIM, 1, GDN_CONV_SIZE)) - self.v_conv1d = nn.Parameter(torch.zeros(GDN_VALUE_DIM, 1, GDN_CONV_SIZE)) - self.o_norm = nn.Parameter(torch.ones(GDN_HEAD_V_DIM)) - - @staticmethod - def _short_conv(x, w): - C, K = w.shape[0], w.shape[2] - xt = F.pad(x.transpose(1, 2), (K - 1, 0)) - return F.conv1d(xt, w, groups=C).transpose(1, 2) - - def forward(self, x): - B, T, _ = x.shape - H, DK, DV, S = N_HEAD, GDN_HEAD_DIM, GDN_HEAD_V_DIM, self.chunk - q = F.silu(self._short_conv(self.q_proj(x), self.q_conv1d)) - k = F.silu(self._short_conv(self.k_proj(x), self.k_conv1d)) - v = self._short_conv(self.v_proj(x), self.v_conv1d) - q = q.view(B, T, H, DK).transpose(1, 2) - k = k.view(B, T, H, DK).transpose(1, 2) - v = v.view(B, T, H, DV).transpose(1, 2) - dt = F.softplus(self.a_proj(x) + self.dt_bias) - g = torch.exp(-torch.exp(self.A_log) * dt).permute(0, 2, 1).unsqueeze(-1) - beta = torch.sigmoid(self.b_proj(x)).permute(0, 2, 1).unsqueeze(-1) - v = v * beta - pad = (S - T % S) % S - if pad: - q, k, v = (F.pad(t, (0, 0, 0, pad)) for t in (q, k, v)) - g = F.pad(g, (0, 0, 0, pad), value=1.0) - nC = (T + pad) // S - rs = lambda t: t.view(B, H, nC, S, -1) # noqa: E731 - qc, kc, vc, gc = rs(q), rs(k), rs(v), rs(g) - causal = torch.tril(torch.ones(S, S, device=x.device, dtype=torch.bool)) - intra = (qc @ kc.transpose(-1, -2)).masked_fill(~causal, 0.0) @ vc - state = torch.zeros(B, H, DK, DV, device=x.device, dtype=x.dtype) - outs = [] - for c in range(nC): - outs.append(intra[:, :, c] + qc[:, :, c] @ state) - decay = gc[:, :, c].prod(dim=2, keepdim=True) - state = state * decay + kc[:, :, c].transpose(-1, -2) @ (vc[:, :, c] * gc[:, :, c]) - y = torch.stack(outs, dim=2).view(B, H, T + pad, DV)[:, :, :T] - yf = y.float() - y = (yf * torch.rsqrt(yf.pow(2).mean(-1, keepdim=True) + 1e-5)).to(x.dtype) - y = y * self.o_norm - y = y.transpose(1, 2).contiguous().view(B, T, GDN_VALUE_DIM) - y = y * F.silu(self.g_proj(x)) - return self.o_proj(y) - - -def _make_gdn(): - if not torch.cuda.is_available(): - return _ChunkedDeltaNet() - from fla.layers import GatedDeltaNet # ImportError here is intentional - return GatedDeltaNet( - hidden_size=N_EMBD, - num_heads=N_HEAD, - head_dim=GDN_HEAD_DIM, - expand_v=GDN_EXPAND_V, - use_gate=True, - use_short_conv=True, - allow_neg_eigval=True, - ) - - -class GDNLayer(nn.Module): - def __init__(self): - super().__init__() - self.impl = _make_gdn() - - def forward(self, x): - out = self.impl(x) - return out[0] if isinstance(out, tuple) else out - - -class FeedForward(nn.Module): - """SwiGLU, hidden 3072, bias=False.""" - - def __init__(self, n_embd: int = N_EMBD, hidden: int = HIDDEN_SIZE): - super().__init__() - self.w1 = nn.Linear(n_embd, hidden, bias=False) - self.w3 = nn.Linear(n_embd, hidden, bias=False) - self.w2 = nn.Linear(hidden, n_embd, bias=False) - - def forward(self, x): - return self.w2(F.silu(self.w1(x)) * self.w3(x)) - - -class PreNormBlock(nn.Module): - """Pre-norm (D.3): norm BEFORE each sub-layer, residual add after.""" - - def __init__(self, kind: str): - super().__init__() - self.mixer = Attention() if kind == "attn" else GDNLayer() - self.mixer_norm = RMSNorm(N_EMBD) - self.ffn = FeedForward() - self.ffn_norm = RMSNorm(N_EMBD) - - def forward(self, x): - x = x + self.mixer(self.mixer_norm(x)) - return x + self.ffn(self.ffn_norm(x)) - - -class PaperModel(nn.Module): - """Paper-faithful ablation model. ``hybrid=False`` -> pure transformer.""" - - def __init__(self, config, hybrid: bool): - super().__init__() - self.block_size = config.block_size - self.wte = nn.Embedding(config.vocab_size, N_EMBD) - if hybrid: - kinds = [PATTERN[i % len(PATTERN)] for i in range(N_LAYER)] - else: - kinds = ["attn"] * N_LAYER - self.blocks = nn.ModuleList([PreNormBlock(k) for k in kinds]) - self.norm_f = RMSNorm(N_EMBD) - self.lm_head = nn.Linear(N_EMBD, config.vocab_size, bias=False) # untied - self.apply(self._init) - for name, p in self.named_parameters(): - if name.endswith("w_out.weight") or name.endswith("w2.weight"): - nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * N_LAYER)) - - @staticmethod - def _init(m): - if isinstance(m, nn.Linear): - nn.init.normal_(m.weight, mean=0.0, std=0.02) - if m.bias is not None: - nn.init.zeros_(m.bias) - elif isinstance(m, nn.Embedding): - nn.init.normal_(m.weight, mean=0.0, std=0.02) - - def forward(self, idx): - x = self.wte(idx) - for block in self.blocks: - x = block(x) - x = self.norm_f(x) - return self.lm_head(x) - - -def build_transformer(config): - return PaperModel(config, hybrid=False) - - -def build_hybrid(config): - return PaperModel(config, hybrid=True) - - -# --------------------------------------------------------------------------- # -# Analytic param counts (Table 22 convention: "non-embedding" = blocks + LM -# head, i.e. everything except the INPUT embedding table). At 190M this should -# reproduce Transformer 190 and GDN-3:1 254 (millions). -# --------------------------------------------------------------------------- # -def _attn_params(): - d = N_EMBD - return 4 * d * d + 2 * d # q/k/v/out + q/k norm - - -def _gdn_params(): - d, H = N_EMBD, N_HEAD - return ( - 2 * d * GDN_KEY_DIM # q_proj, k_proj - + 2 * d * GDN_VALUE_DIM # v_proj, g_proj - + 2 * d * H # a_proj, b_proj - + GDN_VALUE_DIM * d # o_proj - + 2 * H # A_log, dt_bias - + GDN_CONV_SIZE * (2 * GDN_KEY_DIM + GDN_VALUE_DIM) # short convs - + GDN_HEAD_V_DIM # o_norm - ) - - -def analytic_non_embed(hybrid: bool, vocab_size: int) -> int: - d = N_EMBD - block_norms = 2 * d - ffn = 3 * d * HIDDEN_SIZE - n_attn = sum(1 for i in range(N_LAYER) if PATTERN[i % len(PATTERN)] == "attn") if hybrid else N_LAYER - n_gdn = N_LAYER - n_attn - blocks = n_attn * (_attn_params() + block_norms + ffn) + n_gdn * (_gdn_params() + block_norms + ffn) - lm_head = vocab_size * d # counted as non-embedding (Table 22 convention) - return blocks + d + lm_head # + final norm - - -if __name__ == "__main__": - # Local self-check (CPU): param counts vs Table 22 (190 / 254 million). - class _Cfg: - vocab_size = 100352 - block_size = 256 - eval_block_size = 256 - - for hybrid, want in ((False, 190), (True, 254)): - m = PaperModel(_Cfg(), hybrid=hybrid) - total = sum(p.numel() for p in m.parameters() if p.requires_grad) - non_emb_built = total - m.wte.weight.numel() - non_emb_analytic = analytic_non_embed(hybrid, _Cfg.vocab_size) - assert non_emb_built == non_emb_analytic, (hybrid, non_emb_built, non_emb_analytic) - name = "hybrid GDN-3:1" if hybrid else "transformer" - print(f"{name:16s} non-embed(Table22 conv)={non_emb_built/1e6:6.1f}M " - f"(want ~{want}M) total={total/1e6:6.1f}M") - # tiny forward - idx = torch.randint(0, _Cfg.vocab_size, (2, 16)) - out = m(idx) - assert out.shape == (2, 16, _Cfg.vocab_size), out.shape - print("forward + param-count self-check OK") diff --git a/2.0/problems/nanoslm_hybrid_arch_design/repro/wsd.py b/2.0/problems/nanoslm_hybrid_arch_design/repro/wsd.py deleted file mode 100644 index f09539c02..000000000 --- a/2.0/problems/nanoslm_hybrid_arch_design/repro/wsd.py +++ /dev/null @@ -1,91 +0,0 @@ -"""WSD-S learning-rate schedule (warmup - stable - decay with periodic resets). - -Paper (arXiv:2604.03444, Section 4.1): a WSD-S schedule that is token-agnostic --- the LR at step t does not depend on the total token budget T. "We decay the -learning rate for 5% of the training tokens to 0 at each of the five Chinchilla -factors to obtain the trained checkpoint for that factor." Between factors the -LR resets back to the stable peak (the '-S' = periodic reset), so a single long -run yields checkpoints at 0.5x/1x/2x/4x/8x Chinchilla. - -Refs: Hu et al. 2024 (MiniCPM WSD); Wen et al. 2025b (river-valley view of WSD). - -Two schedule shapes here: - * ``lr_wsd``: single warmup -> stable -> terminal 5% decay-to-0. Used for a - fixed-budget run (e.g. the 1x-Chinchilla smoke). - * ``lr_wsd_s``: the periodic-reset variant. Given the ordered list of - decay-endpoint steps (the Chinchilla factors), it produces a checkpoint at - each: warmup once, hold at peak, and for the last ``decay_frac`` of the span - BEFORE each endpoint, decay to 0; immediately after an endpoint, jump back - to peak. Evaluate/snapshot AT each endpoint step. - -Decay shape: the WSD "river-valley" analysis (Wen et al.) favours a 1-sqrt decay -over linear; we use 1 - sqrt(progress), which spends more steps near the low LR. -Set ``decay_shape='linear'`` for a plain linear ramp to 0. -""" - -from __future__ import annotations - -import math - - -def _decay_mult(progress: float, shape: str) -> float: - """LR multiplier in [0,1] as decay progresses 0->1 (1 at start, 0 at end).""" - progress = min(1.0, max(0.0, progress)) - if shape == "linear": - return 1.0 - progress - # '1-sqrt' (river-valley): stays high then drops fast near the end. - return 1.0 - math.sqrt(progress) - - -def lr_wsd(step: int, *, peak_lr: float, total_steps: int, warmup_steps: int, - decay_frac: float = 0.05, decay_shape: str = "1-sqrt") -> float: - """Single warmup -> stable -> terminal decay-to-0 over the last ``decay_frac``.""" - if step < warmup_steps: - return peak_lr * (step + 1) / max(1, warmup_steps) - decay_steps = max(1, int(round(decay_frac * total_steps))) - decay_start = total_steps - decay_steps - if step < decay_start: - return peak_lr - progress = (step - decay_start) / decay_steps - return peak_lr * _decay_mult(progress, decay_shape) - - -def lr_wsd_s(step: int, *, peak_lr: float, endpoints: list[int], warmup_steps: int, - decay_frac: float = 0.05, decay_shape: str = "1-sqrt") -> float: - """Periodic-reset WSD-S. ``endpoints`` = sorted decay-endpoint steps. - - For each consecutive span (prev_endpoint, endpoint], the last ``decay_frac`` - of the span decays to 0; the rest holds at peak (after the one-time warmup). - Snapshot the model exactly AT each endpoint step to get that factor's ckpt. - """ - if step < warmup_steps: - return peak_lr * (step + 1) / max(1, warmup_steps) - prev = 0 - for end in endpoints: - if step <= end: - span = end - prev - decay_steps = max(1, int(round(decay_frac * span))) - decay_start = end - decay_steps - if step < decay_start: - return peak_lr - progress = (step - decay_start) / decay_steps - return peak_lr * _decay_mult(progress, decay_shape) - prev = end - return 0.0 # past the final endpoint - - -if __name__ == "__main__": - # Sanity checks. - P, T, W = 2e-3, 1000, 50 - assert abs(lr_wsd(0, peak_lr=P, total_steps=T, warmup_steps=W) - P / W) < 1e-12 - assert abs(lr_wsd(W, peak_lr=P, total_steps=T, warmup_steps=W) - P) < 1e-12 # stable - assert abs(lr_wsd(940, peak_lr=P, total_steps=T, warmup_steps=W) - P) < 1e-9 # still stable (decay=last 50) - assert lr_wsd(999, peak_lr=P, total_steps=T, warmup_steps=W) < P # decaying - assert abs(lr_wsd(T, peak_lr=P, total_steps=T, warmup_steps=W)) < 1e-9 # ~0 at end - # WSD-S: peak restored right after an endpoint. - eps = [500, 1000] - assert abs(lr_wsd_s(600, peak_lr=P, endpoints=eps, warmup_steps=W) - P) < 1e-9 # reset to peak after 500 - assert lr_wsd_s(500, peak_lr=P, endpoints=eps, warmup_steps=W) < P # decayed at endpoint 500 - print("WSD / WSD-S schedule self-check OK") - for s in (0, 25, 50, 500, 900, 950, 975, 1000): - print(f" step {s:4d} lr={lr_wsd(s, peak_lr=P, total_steps=T, warmup_steps=W):.3e}") diff --git a/adapters/frontier-cs-2.0/src/frontier_cs_2_0/adapter.py b/adapters/frontier-cs-2.0/src/frontier_cs_2_0/adapter.py index af2e05ae0..c5875385a 100644 --- a/adapters/frontier-cs-2.0/src/frontier_cs_2_0/adapter.py +++ b/adapters/frontier-cs-2.0/src/frontier_cs_2_0/adapter.py @@ -422,6 +422,12 @@ def _write_task_config(self, task_paths: "TaskPaths", problem: FrontierCS20Probl cpus=int(environment.get("cpus", 2)), memory_mb=int(environment.get("memory_mb", 4096)), storage_mb=int(environment.get("storage_mb", 4096)), + # Defaults to 0, so every existing problem still generates a + # byte-identical task.toml. Only a problem that explicitly sets + # `environment.gpus` in its config.yaml gets a GPU in the AGENT + # container -- needed when the task IS to run training, rather than + # to patch code the judge runs elsewhere. + gpus=int(environment.get("gpus", 0)), ) try: from harbor.models.task.config import TaskConfig diff --git a/adapters/frontier-cs-2.0/src/frontier_cs_2_0/task-template/task.toml b/adapters/frontier-cs-2.0/src/frontier_cs_2_0/task-template/task.toml index 6d721e89a..4cb2211bf 100644 --- a/adapters/frontier-cs-2.0/src/frontier_cs_2_0/task-template/task.toml +++ b/adapters/frontier-cs-2.0/src/frontier_cs_2_0/task-template/task.toml @@ -23,5 +23,5 @@ build_timeout_sec = {environment_build_timeout_sec} cpus = {cpus} memory_mb = {memory_mb} storage_mb = {storage_mb} -gpus = 0 +gpus = {gpus} allow_internet = true From 65d80cbe81f7f7e5b6e03860d6cb1169b37c9d25 Mon Sep 17 00:00:00 2001 From: Sijia Liu Date: Fri, 24 Jul 2026 15:38:49 -0700 Subject: [PATCH 5/5] [2.0] nanoslm: raw-bpb scoring, 30-min budget, full-horizon LR, wandb - score IS the submission's held-out val_bpb (lower better), no scale/clip; bpb_score_scale removed; failure paths return sentinel 9999 - train_seconds 21600 -> 1800 for fast iteration; fingerprinted lr_schedule_seconds=21600 keeps the full 6h cosine (a 30-min run is the prefix of a long run, not a compressed anneal) - opt-in wandb logging (operator-only env gate; scored paths byte-identical) incl. pre-clip grad_norm; stop shipping stale .assets into the GPU image - calibrated @30min (CRN pairs, full-horizon LR): baseline 1.33325, reference 1.31208; 6h numbers unchanged (0.98112 / 0.93795) Co-Authored-By: Claude Fable 5 --- .../nanoslm_hybrid_arch_design/config.yaml | 13 +- .../nanoslm_hybrid_arch_design/evaluator.py | 86 ++++----- .../harbor/app/README.md | 14 +- .../harness/modal_app.py | 36 +++- .../harness/runner.py | 12 +- .../harness/scoring.py | 93 +++------ .../harness/settings.py | 36 ++-- .../harness/train.py | 178 ++++++++++++------ .../nanoslm_hybrid_arch_design/readme | 40 ++-- 9 files changed, 292 insertions(+), 216 deletions(-) diff --git a/2.0/problems/nanoslm_hybrid_arch_design/config.yaml b/2.0/problems/nanoslm_hybrid_arch_design/config.yaml index 773880b46..8647f8564 100644 --- a/2.0/problems/nanoslm_hybrid_arch_design/config.yaml +++ b/2.0/problems/nanoslm_hybrid_arch_design/config.yaml @@ -43,17 +43,20 @@ evaluation: dataset: HuggingFaceFW/fineweb-edu:sample-10BT block_size: 8192 eval_block_size: 8192 - train_seconds: 21600 # fixed wall-clock budget T per training run (6 h) - max_train_seconds: 43200 + train_seconds: 1800 # fixed wall-clock budget T per training run (30 min) + # LR cosine horizon: the FULL 6h schedule. A 30-min run traverses only its + # first 1/12 (prefix-of-a-long-run behavior), not a compressed anneal. + lr_schedule_seconds: 21600 + max_train_seconds: 3600 val_tokens: 1048576 # The scored metric is bits per byte: nll_nats / (val_bytes * ln2). val_bytes # is measured by docker/prep_assets.py and shipped in the judge image's # manifest.json -- normalizing by the token count would make the metric # tokenizer-dependent. metric: val_bpb - # Scoring is on the absolute val_bpb gain over the locked baseline: - # score = clip(100 * (base_bpb - sub_bpb) / bpb_score_scale, 0, 100) - bpb_score_scale: 0.5 + # The score IS the submission's held-out val_bpb, raw and unscaled -- LOWER + # IS BETTER. No clipping, no baseline-relative mapping (bpb_score_scale was + # removed); the baseline's bpb and the gain are reported for context only. param_cap: 400000000 seed: 1337 submission: diff --git a/2.0/problems/nanoslm_hybrid_arch_design/evaluator.py b/2.0/problems/nanoslm_hybrid_arch_design/evaluator.py index cdf9edf44..edb1a3303 100644 --- a/2.0/problems/nanoslm_hybrid_arch_design/evaluator.py +++ b/2.0/problems/nanoslm_hybrid_arch_design/evaluator.py @@ -81,9 +81,8 @@ def _backend() -> str: def _run_pair_modal(solution_path: str, role: str): """Run both arms on a Modal GPU; score judge-side. - Only arm metrics cross the boundary -- scoring and the hidden constants - (bpb_score_scale) stay in the judge, so a GPU worker never sees them and cannot - hand back a score the judge did not compute. + Only arm metrics cross the boundary -- scoring stays in the judge, so a + GPU worker never hands back a score the judge did not compute. """ from harness.modal_app import app, run_pair_remote @@ -154,7 +153,8 @@ def evaluate(solution_path: str): # 1) Static policy gate (torch-free, adversarial-safe). pol = policy.check_file(solution_path) if not pol.ok: - return 0.0, 0.0, f"policy_rejected: {pol.reason}", {"config_fingerprint": fp} + return (scoring.FAILURE_SCORE, scoring.FAILURE_SCORE, + f"policy_rejected: {pol.reason}", {"config_fingerprint": fp}) # 2) Backend. On MODAL the judge never imports torch -- the arms run in a # Modal GPU container and only plain metrics come back. @@ -163,19 +163,20 @@ def evaluate(solution_path: str): try: res = _run_pair_modal(solution_path, role) except Exception as exc: # noqa: BLE001 - classify, never leak - return (0.0, 0.0, f"environment_error: modal dispatch failed " - f"({type(exc).__name__})", {"config_fingerprint": fp}) + return (scoring.FAILURE_SCORE, scoring.FAILURE_SCORE, + f"environment_error: modal dispatch failed " + f"({type(exc).__name__})", {"config_fingerprint": fp}) # Guard rejections travel as a field, not an exception: GuardError is # defined in harness.runner (which imports torch), so this CPU-only # judge could not deserialize the raised form. The message is the # public, classified reason -- same wording as the local path's. if res.get("guard_error"): - return 0.0, 0.0, f"guard: {res['guard_error']}", { + return scoring.FAILURE_SCORE, scoring.FAILURE_SCORE, f"guard: {res['guard_error']}", { "config_fingerprint": fp, "role": role, "stack": res.get("stack"), } b, c = res["baseline"], res["submission"] - sr = scoring.score_from_bpb(b["val_bpb"], c["val_bpb"], cfg.bpb_score_scale) + sr = scoring.score_from_bpb(b["val_bpb"], c["val_bpb"]) note = [] if role == "agent": # Mirrors the local path's note: says whether this feedback run @@ -193,7 +194,6 @@ def evaluate(solution_path: str): # abs_bpb_delta is the scored quantity; rel_improvement is kept # alongside it so an operator never has to recompute it. "rel_improvement": sr.rel_improvement, - "bpb_score_scale": cfg.bpb_score_scale, "sub_steps": c["steps"], "base_steps": b["steps"], "sub_params": c["n_params"], "sub_wall_seconds": c["wall_seconds"], "sub_train_block_size": c.get("train_block_size"), @@ -208,8 +208,8 @@ def evaluate(solution_path: str): device = _select_device() except RuntimeError as exc: return ( - 0.0, - 0.0, + scoring.FAILURE_SCORE, + scoring.FAILURE_SCORE, f"environment_error: {exc} (training path requires torch + GPU)", {"config_fingerprint": fp}, ) @@ -248,16 +248,17 @@ def evaluate(solution_path: str): from harness.runner import GuardError if isinstance(exc, GuardError): - return 0.0, 0.0, f"guard: {exc}", {"config_fingerprint": fp, "role": role} + return (scoring.FAILURE_SCORE, scoring.FAILURE_SCORE, + f"guard: {exc}", {"config_fingerprint": fp, "role": role}) # Never leak submission tracebacks/stdout (2.0 black-box safety). return ( - 0.0, - 0.0, + scoring.FAILURE_SCORE, + scoring.FAILURE_SCORE, f"submission_error: {type(exc).__name__}", {"config_fingerprint": fp, "role": role}, ) - sr = scoring.score_from_bpb(base_bpb, sub.val_bpb, cfg.bpb_score_scale) + sr = scoring.score_from_bpb(base_bpb, sub.val_bpb) note = [] if role == "agent": note.append("iterative(cached-baseline)" if base_steps is None else "iterative") @@ -275,7 +276,6 @@ def evaluate(solution_path: str): "abs_bpb_delta": sr.abs_bpb_delta, "sub_val_ppl": sub.val_ppl, # derived, readability only "rel_improvement": sr.rel_improvement, - "bpb_score_scale": cfg.bpb_score_scale, "sub_steps": sub.steps, "base_steps": base_steps, "sub_params": sub.n_params, @@ -400,39 +400,29 @@ def check(name, cond): ).ok, ) - # --- scoring: Absolute bpb gain. baseline->0, worse->0, scale->100 --- - # Cases are constructed by subtracting bits per byte, not by scaling - # base_bpb: the scored quantity is now `base_bpb - sub_bpb`, so a relative - # construction would no longer test what it claims to. - scale = 0.05 # bpb gain that saturates the bounded score + # --- scoring: the score IS the submission's raw val_bpb, lower is better --- base_bpb = 4.0 - s_tie = scoring.score_from_bpb(base_bpb, base_bpb, scale) - s_worse = scoring.score_from_bpb(base_bpb, base_bpb + 0.10, scale) - s_sat = scoring.score_from_bpb(base_bpb, base_bpb - scale, scale) - s_half = scoring.score_from_bpb(base_bpb, base_bpb - scale / 2, scale) - s_over = scoring.score_from_bpb(base_bpb, base_bpb - 2 * scale, scale) - check("baseline scores 0", abs(s_tie.score) < 1e-9) - check("worse-than-baseline scores 0", abs(s_worse.score) < 1e-9) - check("a gain of BPB_SCORE_SCALE saturates to 100", abs(s_sat.score - 100.0) < 1e-6) - check("half the scale ~= 50", abs(s_half.score - 50.0) < 1e-6) - check("bounded score caps at 100", abs(s_over.score - 100.0) < 1e-9) - check("unbounded keeps rewarding past 100", s_over.score_unbounded > 100.0) - check("unbounded is exactly the un-clipped ratio", - abs(s_over.score_unbounded - 200.0) < 1e-9) - # The score must depend on the absolute gain only -- the same 0.05 bpb gain - # scores the same at any operating point. (This is the property the old - # relative form did not have, and the reason the operating point now - # matters.) - s_lowbase = scoring.score_from_bpb(1.5, 1.5 - scale, scale) - check("score depends on absolute gain, not on base_bpb", - abs(s_lowbase.score - s_sat.score) < 1e-9) - # ... while both figures are still reported. - check("relative improvement still reported", - abs(s_sat.rel_improvement - (scale / base_bpb)) < 1e-12 - and abs(s_sat.abs_bpb_delta - scale) < 1e-12) - # A scoring constant, not a training knob: changing it must not invalidate - # a cached baseline, so it must stay out of the fingerprint. - check("bpb_score_scale is not fingerprinted", + s_tie = scoring.score_from_bpb(base_bpb, base_bpb) + s_worse = scoring.score_from_bpb(base_bpb, base_bpb + 0.10) + s_better = scoring.score_from_bpb(base_bpb, base_bpb - 0.05) + check("score is exactly the submission's val_bpb", + abs(s_better.score - (base_bpb - 0.05)) < 1e-12 + and abs(s_tie.score - base_bpb) < 1e-12 + and abs(s_worse.score - (base_bpb + 0.10)) < 1e-12) + check("score is unscaled and un-clipped (no [0,100] mapping)", + abs(s_worse.score - s_worse.score_unbounded) < 1e-12 + and 0.0 < s_better.score < 100.0 / 25.0) # a bpb, not a percentage + check("lower is better (worse model -> higher score value)", + s_worse.score > s_tie.score > s_better.score) + check("degenerate bpb maps to the failure sentinel, never a good score", + scoring.score_from_bpb(base_bpb, 0.0).score == scoring.FAILURE_SCORE + and scoring.score_from_bpb(base_bpb, float("inf")).score + == scoring.FAILURE_SCORE) + # The baseline comparison is still reported alongside, for context only. + check("gain figures still reported", + abs(s_better.abs_bpb_delta - 0.05) < 1e-12 + and abs(s_better.rel_improvement - (0.05 / base_bpb)) < 1e-12) + check("legacy scoring constants stay out of the fingerprint", "bpb_score_scale" not in settings._FINGERPRINT_KEYS and "r_target" not in settings._FINGERPRINT_KEYS) # Same for the plausibility floor (a guard constant), which must exist and diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harbor/app/README.md b/2.0/problems/nanoslm_hybrid_arch_design/harbor/app/README.md index 8d18c885b..cbc841b9b 100644 --- a/2.0/problems/nanoslm_hybrid_arch_design/harbor/app/README.md +++ b/2.0/problems/nanoslm_hybrid_arch_design/harbor/app/README.md @@ -16,11 +16,12 @@ bash /app/submit.sh # 3. enqueue for the judge ``` You do **not** train here — this container has no GPU and no torch. The judge -trains both your architecture and a locked baseline under the identical budget -and scores your **absolute bpb gain** over it, `base_bpb - sub_bpb`. Bits per -byte is the unit the literature quotes, so the gain is directly comparable to -published numbers; the constant that maps it onto 0–100 is a scaling convention, -not a target, so simply maximize the gain. +trains your architecture under a fixed wall-clock budget and **your score IS +your raw held-out bits per byte (`sub_bpb`) — LOWER IS BETTER**, with no +scaling and no clipping. A locked baseline is trained under the identical +budget and its bpb and your gain over it are reported for context only. Failed +or rejected runs score 9999 (worst). Bits per byte is the unit the literature +quotes, so your score is directly comparable to published numbers. ## What you submit @@ -53,7 +54,8 @@ Declare it with a module-level integer in `model.py`: BLOCK_SIZE = 2048 # power of two in [256, 8192]; omit it to train at 8192 ``` -Out-of-range or non-power-of-two values are rejected before training (score 0). +Out-of-range or non-power-of-two values are rejected before training +(failure score 9999 — the score is bits per byte, LOWER is better). **Evaluation is always at 8192**, whatever you train at. So: diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harness/modal_app.py b/2.0/problems/nanoslm_hybrid_arch_design/harness/modal_app.py index 8401ec23a..770f5db89 100644 --- a/2.0/problems/nanoslm_hybrid_arch_design/harness/modal_app.py +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/modal_app.py @@ -86,12 +86,19 @@ def _ver(mod): # Only and looked healthy; always exercise fwd+bwd when validating a # fused kernel. .pip_install("flash-linear-attention==0.5.1", "transformers==4.46.3") + # Operator-only observability (train._maybe_init_wandb). Inert unless a + # call passes FRONTIER_NANOSLM_WANDB=1 via `overrides`, so scored runs + # and the baseline cache are unaffected. + .pip_install("wandb") # Ship the problem directory itself: harness/, evaluator.py, # reference.py. The GPU function calls the real evaluator, so what runs - # remotely is byte-identical to what Harbor runs. + # remotely is byte-identical to what Harbor runs. `.assets` is the + # LOCAL dev staging (gigabytes of tokens, and a stale val split) -- the + # real corpus mounts from the Volume, so shipping it was pure deploy + # bloat plus an unguarded token-stream copy in the container. .add_local_dir( _PROBLEM_DIR, REMOTE_TASK, copy=True, - ignore=["__pycache__", "*.pyc", ".git", "docker"], + ignore=["__pycache__", "*.pyc", ".git", "docker", ".assets"], ) .env({ "PYTHONPATH": REMOTE_TASK, @@ -108,11 +115,20 @@ def _ver(mod): # Reduce allocator fragmentation for the large fp32-logits CE at # ctx 8192 (paired with the batch_size=2 fix in settings.py). "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", + # Keep wandb's scratch out of the shipped task tree. + "WANDB_DIR": "/tmp/wandb", }) ) app = modal.App(APP_NAME) + # WANDB_API_KEY for the opt-in observability path. The secret must exist + # for deploys to resolve (create once: + # modal secret create wandb WANDB_API_KEY=... + # ); a placeholder value just means wandb.init fails and logging silently + # stays off -- training is untouched either way. + WANDB_SECRET = modal.Secret.from_name("wandb") + # Real corpus (FineWeb-Edu, dolma2-tokenized) staged on a Modal Volume and # mounted at the data path harness/settings.py expects, so the GPU arms train # on real tokens. The harness has no synthetic fallback, so this mount is @@ -122,7 +138,7 @@ def _ver(mod): REMOTE_DATA = "/opt/nanoslm_arch/data" @app.function(image=image, gpu=GPU, timeout=FN_TIMEOUT_SECONDS, - volumes={REMOTE_DATA: CORPUS_VOL}) + volumes={REMOTE_DATA: CORPUS_VOL}, secrets=[WANDB_SECRET]) def evaluate_remote(solution_source: str, role: str = "final", overrides: dict | None = None) -> dict: """Run the judge on `solution_source` and return its full result. @@ -228,14 +244,15 @@ def _ver(mod): # noqa: F811 (module-level _ver is the shared one) BASELINE_CACHE = f"{REMOTE_DATA}/baseline/baseline_arm.json" @app.function(image=image, gpu=GPU, timeout=FN_TIMEOUT_SECONDS, - volumes={REMOTE_DATA: CORPUS_VOL}) + volumes={REMOTE_DATA: CORPUS_VOL}, secrets=[WANDB_SECRET]) def run_pair_remote(solution_source: str, - role: str = "final") -> dict: + role: str = "final", + overrides: dict | None = None) -> dict: """Run the arms on one GPU and return their metrics as plain dicts. This is the entrypoint the judge uses. It deliberately returns arm Metrics, not a score: scoring stays judge-side so the scoring code and - the hidden constants (bpb_score_scale) never ship to a GPU worker, and so the + no score constant ships to a GPU worker, and so the judge cannot be handed a score it did not compute. Role semantics (this is what makes iteration affordable): @@ -264,6 +281,10 @@ def run_pair_remote(solution_source: str, sys.path.insert(0, REMOTE_TASK) os.environ["FRONTIER_NANOSLM_ROLE"] = role + # Operator-only env passthrough (e.g. FRONTIER_NANOSLM_WANDB=1). The + # judge never passes overrides, so the scored path is unchanged. + for k, v in (overrides or {}).items(): + os.environ[k] = str(v) import torch # noqa: F401 (ensures CUDA init before harness import) @@ -358,7 +379,8 @@ def _cache_write(entry: dict) -> None: sub_block = runner.resolve_train_block_size(mod, cfg) data = TokenData(cfg) cand = runner.run_arm( - runner.load_factory(mod), data, cfg, "cuda", sub_block + runner.load_factory(mod), data, cfg, "cuda", sub_block, + run_label="submission", ) except runner.GuardError as exc: return { diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harness/runner.py b/2.0/problems/nanoslm_hybrid_arch_design/harness/runner.py index 07230fe57..e1c2d5483 100644 --- a/2.0/problems/nanoslm_hybrid_arch_design/harness/runner.py +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/runner.py @@ -276,7 +276,8 @@ def __exit__(self, *a): def run_arm(factory: Callable, data: TokenData, cfg: TaskConfig, device: str, - train_block_size: int | None = None) -> ArmMetrics: + train_block_size: int | None = None, + run_label: str = "") -> ArmMetrics: """Instantiate -> guard params -> train -> eval -> guard training/degeneracy. ``train_block_size`` is this arm's training context (default: cfg's). It is @@ -346,7 +347,8 @@ def run_arm(factory: Callable, data: TokenData, cfg: TaskConfig, device: str, before = _param_snapshot(model) try: - tout: TrainOutput = train_model(model, data, cfg, device) + tout: TrainOutput = train_model(model, data, cfg, device, + run_label=run_label) except (GuardError, DataError): raise except Exception as exc: @@ -454,6 +456,8 @@ def run_pair(submission_module, cfg: TaskConfig, device: str): # judge nothing. sub_block = resolve_train_block_size(submission_module, cfg) data = TokenData(cfg) - base = run_arm(baseline_factory(), data, cfg, device, baseline_block_size(cfg)) - sub = run_arm(load_factory(submission_module), data, cfg, device, sub_block) + base = run_arm(baseline_factory(), data, cfg, device, baseline_block_size(cfg), + run_label="baseline") + sub = run_arm(load_factory(submission_module), data, cfg, device, sub_block, + run_label="submission") return base, sub diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harness/scoring.py b/2.0/problems/nanoslm_hybrid_arch_design/harness/scoring.py index 37b16fa5c..c52b4265d 100644 --- a/2.0/problems/nanoslm_hybrid_arch_design/harness/scoring.py +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/scoring.py @@ -1,4 +1,4 @@ -"""Scoring: Absolute held-out bits-per-byte gain over the locked baseline. +"""Scoring: the submission's held-out bits-per-byte, reported raw. Torch-free and unit-tested on CPU. Shared by the evaluator and the public test so the agent sees the same math the judge uses. @@ -11,79 +11,48 @@ @dataclass class ScoreResult: - score: float # bounded [0, 100] reported to Harbor - score_unbounded: float # keeps rewarding past 100 + score: float # the submission's val_bpb — LOWER IS BETTER + score_unbounded: float # identical to score (kept for interface stability) rel_improvement: float # (base_bpb - sub_bpb) / base_bpb — reported only - abs_bpb_delta: float # base_bpb - sub_bpb — the scored quantity + abs_bpb_delta: float # base_bpb - sub_bpb — reported only -def _clip(x: float, lo: float, hi: float) -> float: - return max(lo, min(hi, x)) +# Failure sentinel: with a lower-is-better score, a failed/rejected run must +# NEVER return 0.0 (that would read as a perfect bpb). Finite rather than inf +# so every JSON serializer downstream stays standards-compliant; four orders of +# magnitude above any honest bpb, so it can never be confused for a +# measurement. +FAILURE_SCORE = 9999.0 def score_from_bpb( base_bpb: float, sub_bpb: float, - bpb_score_scale: float, ) -> ScoreResult: - """Map a submission's held-out bits-per-byte to a [0, 100] score. - - Lower ``sub_bpb`` is better. The measurement is the absolute gain:: - - gain = base_bpb - sub_bpb # bits per byte - score = clip(100 * gain / bpb_score_scale, 0, 100) - - The baseline architecture (``sub_bpb == base_bpb``) scores 0, and a - submission worse than the baseline scores 0. - - What ``bpb_score_scale`` is, and what it is not - ----------------------------------------------- - It is a display convention and it is arbitrary by design. Its one necessary - job is to map bits-per-byte into the 0-100 range Harbor expects: Harbor - computes ``reward = score / 100``, so a raw ``gain`` of 0.05 bpb would - surface as reward 0.0005 -- indistinguishable from zero for every - submission. That is the whole reason the constant exists. - - It is not a calibrated definition of "what counts as a full win". No such - number has been measured, and presenting an arbitrary constant as a target - would (a) attach the score's meaning to a figure nobody determined and - (b) throw away information at the clip, where a submission at 2x the scale - scores the same 100 as one exactly at it. ``score_unbounded`` is therefore - the un-clipped value and an operator should read it whenever the bounded - score pins at 100. - - The one weak requirement that does remain is discrimination: set far too - large and every submission pins at 0, far too small and every submission - pins at 100. That is a much weaker condition than calibration and can be set - from a single real-corpus run. - - Why absolute and not relative - ----------------------------- - Absolute bpb is the unit the language-modelling literature quotes, so a - result is directly comparable to published numbers; it is linear in - cross-entropy, so equal absolute gains are equal information gains, and - gains compose roughly additively. The cost is that it is not scale-free: the - same absolute gain is roughly twice as hard at ``base_bpb`` 1.5 as at 2.9, - so the operating point matters. ``rel_improvement`` is - still computed and reported so an operator can see the relative figure - without recomputing it. + """The score IS the submission's held-out bits-per-byte, unscaled. + + LOWER IS BETTER, and there is no clipping, no baseline-relative mapping, + and no display constant: ``score == sub_bpb``. The former + ``clip(100 * (base_bpb - sub_bpb) / bpb_score_scale, 0, 100)`` mapping was + removed deliberately -- the raw measurement is the quantity the + language-modelling literature quotes, and any rescaling constant attached + an arbitrary figure to the score's meaning. Consumers that rank runs must + sort ascending. + + The baseline comparison is still reported (``abs_bpb_delta``, + ``rel_improvement``) so an operator can read the gain without recomputing + it, but neither figure participates in the score. """ - # Invalid / degenerate values -> zero (defensive; runner also guards). + # Invalid / degenerate values -> the failure sentinel (worst; runner also + # guards). if not (base_bpb > 0.0) or not (sub_bpb > 0.0): - return ScoreResult(0.0, 0.0, 0.0, 0.0) + return ScoreResult(FAILURE_SCORE, FAILURE_SCORE, 0.0, 0.0) if not (base_bpb < float("inf")) or not (sub_bpb < float("inf")): - return ScoreResult(0.0, 0.0, 0.0, 0.0) - if bpb_score_scale <= 0.0: - raise ValueError("bpb_score_scale must be positive") + return ScoreResult(FAILURE_SCORE, FAILURE_SCORE, 0.0, 0.0) - abs_delta = base_bpb - sub_bpb # the scored measurement + abs_delta = base_bpb - sub_bpb # reported for readability only rel = abs_delta / base_bpb # reported for readability only - unbounded = 100.0 * abs_delta / bpb_score_scale - bounded = _clip(unbounded, 0.0, 100.0) - # No credit for tying or losing to the baseline. - if abs_delta <= 0.0: - bounded = 0.0 - return ScoreResult(bounded, unbounded, rel, abs_delta) + return ScoreResult(sub_bpb, sub_bpb, rel, abs_delta) def format_message( @@ -103,10 +72,10 @@ def format_message( """ msg = ( f"base_val_bpb={base_bpb:.5f}; sub_val_bpb={sub_bpb:.5f}; " - f"abs_bpb_delta={result.abs_bpb_delta:+.5f}; " # Scored + f"abs_bpb_delta={result.abs_bpb_delta:+.5f}; " f"rel_improvement={result.rel_improvement:+.4%}; " f"steps={steps}; train_wall_s={wall_seconds:.1f}; " - f"score={result.score:.4f}; score_unbounded={result.score_unbounded:.4f}" + f"score={result.score:.5f} (== sub_val_bpb; LOWER is better)" ) if extra: msg += f"; note={extra}" diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harness/settings.py b/2.0/problems/nanoslm_hybrid_arch_design/harness/settings.py index 333b88d2a..4cdfc74af 100644 --- a/2.0/problems/nanoslm_hybrid_arch_design/harness/settings.py +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/settings.py @@ -64,11 +64,21 @@ class TaskConfig: warmup_steps: int = 100 # --- iso-wallclock budget (matches config.yaml) --- - train_seconds: float = 21600.0 # wall-clock T per training run (6 h) + # 30 min (down from the original 6 h): chosen so an iterating agent sees + # signal quickly -- ~35-40 scored submissions fit in the 24 h session + # instead of ~4. The 6 h calibration numbers remain on record (baseline + # 0.98112 / reference ~0.938) but are NOT comparable to 30-min numbers. + train_seconds: float = 1800.0 # wall-clock T per training run (30 min) + # Cosine-decay horizon for the LR schedule -- deliberately the FULL 6 h + # schedule, decoupled from train_seconds: a 30-min run traverses only the + # first 1/12 of the cosine (LR still near peak at cutoff), behaving like + # the prefix of a long run instead of compressing the whole anneal into + # the short budget. See train._lr_at. + lr_schedule_seconds: float = 21600.0 # Infrastructure backstop, enforced by the Modal function timeout (and the # orchestrator), not by the training loop -- a between-step check could # never fire before train_seconds and cannot interrupt a hung step. - max_train_seconds: float = 43200.0 # hard cap (12 h) + max_train_seconds: float = 3600.0 # hard cap (1 h) # --- data / tokenizer (locked; hidden tokens live in the judge image) --- dataset_name: str = "HuggingFaceFW/fineweb-edu:sample-10BT" # provenance; matches config.yaml @@ -103,14 +113,13 @@ class TaskConfig: baseline_cache_path: str = "/opt/nanoslm_arch/baseline/baseline_ppl.json" # --- scoring --- - # Absolute bits-per-byte gain is the scored measurement: - # gain = base_bpb - sub_bpb - # score = clip(100 * gain / bpb_score_scale, 0, 100) - # `bpb_score_scale` is a display convention that maps the gain into Harbor's - # 0-100 reward, not a calibrated "full win" target; its only - # real requirement is discrimination, and `score_unbounded` stays un-clipped. - # 0.5 comes from a real-corpus run and is due a re-check at the 6 h budget. - bpb_score_scale: float = 0.5 # bpb gain that saturates the bounded score + # The score IS the submission's held-out bits-per-byte, raw and unscaled -- + # LOWER IS BETTER (see harness/scoring.py). The old + # clip(100*(base-sub)/bpb_score_scale) mapping and its constant were + # removed deliberately: the raw measurement is the quantity the literature + # quotes, and any rescaling constant attached an arbitrary figure to the + # score's meaning. The baseline is still trained and its gain figures + # reported for context; they no longer participate in the score. # --- guards / resource caps --- param_cap: int = 400_000_000 # max trainable params (over-cap -> score 0) @@ -122,8 +131,8 @@ class TaskConfig: # strongest published LLMs sit near ~0.55 bpb on web text, and a <=400M # model trained 6 h from scratch lands far above that. A measurement below # it is therefore treated as held-out leakage, not brilliance -> GuardError - # (score 0, public reason). A scoring/guard constant like bpb_score_scale, - # so it is deliberately NOT fingerprinted. + # (rejected, public reason). A guard constant, not a training knob, so it + # is deliberately NOT fingerprinted. min_plausible_bpb: float = 0.4 seed: int = 1337 # Off by default (see runner.run_arm): compiling the chunkwise mixer measured @@ -241,6 +250,9 @@ def active_config() -> TaskConfig: "vocab_size", "eval_block_size", "batch_size", "grad_accum", "learning_rate", "min_lr", "weight_decay", "beta1", "beta2", "grad_clip", "warmup_steps", "train_seconds", "dataset_name", + # The decay horizon changes every arm's effective LR trajectory, hence the + # baseline's val_bpb -- a cached baseline from another horizon is invalid. + "lr_schedule_seconds", # tokenizer_name is fingerprinted alongside vocab_size: changing the # tokenizer changes what val_bpb means, so a baseline cached under the old # one must not be reused. diff --git a/2.0/problems/nanoslm_hybrid_arch_design/harness/train.py b/2.0/problems/nanoslm_hybrid_arch_design/harness/train.py index 4ccbe755d..7b3c1423b 100644 --- a/2.0/problems/nanoslm_hybrid_arch_design/harness/train.py +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/train.py @@ -22,8 +22,9 @@ from __future__ import annotations import math +import os import time -from dataclasses import dataclass +from dataclasses import asdict, dataclass from .data import TokenData from .settings import TaskConfig @@ -37,24 +38,24 @@ class TrainOutput: def _lr_at(step: int, elapsed: float, cfg: TaskConfig) -> float: - """Linear step-warmup, then cosine decay over the wall-clock budget. - - The decay is a function of ``elapsed / train_seconds``, not of a step - count. The stop condition is wall-clock, so a step horizon cannot be known - in advance, and any fixed guess hands different effective schedules to arms - with different throughput: a fast arm (short context, cheap mixer) overruns - the guess and trains most of its budget flat at ``min_lr``, while a slow - arm never finishes the decay -- which is known to cost real validation - loss. That would let the score partly measure schedule luck rather than - architecture, the exact confound the locked recipe exists to remove. - Time-based decay gives every architecture the identical schedule over the - identical budget. Warmup stays step-based: its job is stabilizing early - optimizer state, which is a per-step property, and at 100 steps it is a - negligible slice of the budget. + """Linear step-warmup, then cosine decay over the FULL schedule horizon. + + The decay is a function of ``elapsed / lr_schedule_seconds``, not of a step + count (wall-clock, so every architecture sees the identical schedule over + identical time) and not of ``train_seconds`` (the stop time). The horizon + is deliberately decoupled from the budget: with a short iteration budget + (30 min) the run traverses only the first slice of the full-schedule + cosine -- i.e. it behaves like the prefix of a long training run, LR still + near peak at cutoff -- rather than compressing the entire anneal into the + short budget, which would over-weight the anneal and make short-budget + rankings unrepresentative of longer training. When + ``train_seconds == lr_schedule_seconds`` (the original 6 h config) the two + forms are identical. Warmup stays step-based: its job is stabilizing early + optimizer state, which is a per-step property. """ if step < cfg.warmup_steps: return cfg.learning_rate * (step + 1) / max(1, cfg.warmup_steps) - frac = min(1.0, max(0.0, elapsed / max(1e-9, cfg.train_seconds))) + frac = min(1.0, max(0.0, elapsed / max(1e-9, cfg.lr_schedule_seconds))) coeff = 0.5 * (1.0 + math.cos(math.pi * frac)) return cfg.min_lr + coeff * (cfg.learning_rate - cfg.min_lr) @@ -82,7 +83,51 @@ def _cross_entropy(logits, targets): ) -def train_model(model, data: TokenData, cfg: TaskConfig, device: str) -> TrainOutput: +def _maybe_init_wandb(model, cfg: TaskConfig, run_label: str): + """Opt-in observability sink; returns a wandb run or ``None``. + + Active only when the operator sets ``FRONTIER_NANOSLM_WANDB=1`` (plus a + valid ``WANDB_API_KEY`` in the environment); no scored path sets it, so + Harbor runs and the baseline cache are unaffected -- which is why this + needs no CODE_VERSION bump. Everything is best-effort: any failure + disables logging rather than touching the run. Init happens before the + budget timer starts, and the in-loop cost is an async enqueue every + ``_WANDB_LOG_EVERY`` steps (~1 min apart at production step times), so the + wall-clock training budget is unaffected. + """ + if os.environ.get("FRONTIER_NANOSLM_WANDB", "") != "1": + return None + try: + import wandb + + name = "-".join( + p for p in ( + os.environ.get("FRONTIER_NANOSLM_WANDB_NAME", ""), + os.environ.get("FRONTIER_NANOSLM_ROLE", ""), + run_label or type(model).__name__, + ) if p + ) + return wandb.init( + project=os.environ.get( + "FRONTIER_NANOSLM_WANDB_PROJECT", "nanoslm-hybrid-arch-design" + ), + group=os.environ.get("FRONTIER_NANOSLM_WANDB_GROUP") or None, + name=name or None, + config={ + **asdict(cfg), + "run_label": run_label, + "model_class": type(model).__name__, + }, + ) + except Exception: + return None + + +_WANDB_LOG_EVERY = 20 + + +def train_model(model, data: TokenData, cfg: TaskConfig, device: str, + run_label: str = "") -> TrainOutput: import torch model.to(device) @@ -99,45 +144,72 @@ def train_model(model, data: TokenData, cfg: TaskConfig, device: str) -> TrainOu else _nullcontext() ) + wb = _maybe_init_wandb(model, cfg, run_label) step = 0 last_loss = float("nan") t0 = time.monotonic() - while True: - elapsed = time.monotonic() - t0 - if elapsed >= cfg.train_seconds: - break - # (max_train_seconds is enforced by the Modal function timeout, not - # here: train_seconds always breaks first between steps, and no - # between-step check can interrupt a hung step.) - - lr = _lr_at(step, elapsed, cfg) - for pg in opt.param_groups: - pg["lr"] = lr - - # Gradient accumulation: grad_accum micro-batches of cfg.batch_size make - # one optimizer step (effective batch = batch_size * grad_accum) while - # peak activation memory stays at a single micro-batch. cfg.block_size is - # this arm's training context (run_arm passes a per-arm cfg); evaluation - # is at cfg.eval_block_size regardless -- see data.val_windows. - accum = max(1, int(cfg.grad_accum)) - opt.zero_grad(set_to_none=True) - step_ce = 0.0 - for micro in range(accum): - # Distinct deterministic micro-batch, identical across arms (CRN): - # keyed by (step * accum + micro) so an optimizer step's micro-batches - # never repeat and both arms see the same data. accum == 1 reproduces - # the pre-accumulation key exactly. - x, y = data.batch(step * accum + micro, device, cfg.block_size) - with autocast: - logits = _logits_only(model(x)) - loss = _cross_entropy(logits, y) / accum # harness-owned loss - loss.backward() - step_ce += float(loss.detach().float().item()) - torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip) - opt.step() - - last_loss = step_ce # sum of per-micro (already /accum) losses = mean CE - step += 1 + try: + while True: + elapsed = time.monotonic() - t0 + if elapsed >= cfg.train_seconds: + break + # (max_train_seconds is enforced by the Modal function timeout, not + # here: train_seconds always breaks first between steps, and no + # between-step check can interrupt a hung step.) + + lr = _lr_at(step, elapsed, cfg) + for pg in opt.param_groups: + pg["lr"] = lr + + # Gradient accumulation: grad_accum micro-batches of cfg.batch_size make + # one optimizer step (effective batch = batch_size * grad_accum) while + # peak activation memory stays at a single micro-batch. cfg.block_size is + # this arm's training context (run_arm passes a per-arm cfg); evaluation + # is at cfg.eval_block_size regardless -- see data.val_windows. + accum = max(1, int(cfg.grad_accum)) + opt.zero_grad(set_to_none=True) + step_ce = 0.0 + for micro in range(accum): + # Distinct deterministic micro-batch, identical across arms (CRN): + # keyed by (step * accum + micro) so an optimizer step's micro-batches + # never repeat and both arms see the same data. accum == 1 reproduces + # the pre-accumulation key exactly. + x, y = data.batch(step * accum + micro, device, cfg.block_size) + with autocast: + logits = _logits_only(model(x)) + loss = _cross_entropy(logits, y) / accum # harness-owned loss + loss.backward() + step_ce += float(loss.detach().float().item()) + # clip_grad_norm_ returns the pre-clip total norm; kept as a GPU + # tensor and only converted if the wandb branch logs it. + grad_norm = torch.nn.utils.clip_grad_norm_( + model.parameters(), cfg.grad_clip) + opt.step() + + last_loss = step_ce # sum of per-micro (already /accum) losses = mean CE + step += 1 + + if wb is not None and (step <= 5 or step % _WANDB_LOG_EVERY == 0): + try: + wb.log( + { + "train/loss": last_loss, + "train/lr": lr, + "train/grad_norm": float(grad_norm), + "train/elapsed_s": elapsed, + "train/tokens": step * accum * cfg.batch_size + * cfg.block_size, + }, + step=step, + ) + except Exception: + wb = None # a sick sink must never touch the run + finally: + if wb is not None: + try: + wb.finish() + except Exception: + pass if device.startswith("cuda"): import torch as _t diff --git a/2.0/problems/nanoslm_hybrid_arch_design/readme b/2.0/problems/nanoslm_hybrid_arch_design/readme index e896bbd81..f9e3023b3 100644 --- a/2.0/problems/nanoslm_hybrid_arch_design/readme +++ b/2.0/problems/nanoslm_hybrid_arch_design/readme @@ -103,7 +103,8 @@ BLOCK_SIZE = 2048 # power of two in [256, 8192]; omit it to train at 8192 ``` Omit it and you train at the default **8192**. A value outside `[256, 8192]`, or -one that is not a power of two, is rejected before training (score 0). +one that is not a power of two, is rejected before training (failure score +9999 — the score is bits per byte, so LOWER is better and 9999 is the worst). **You are always evaluated at 8192.** That is fixed for every submission — it is what makes `val_bpb` comparable across submissions — and it is the whole trade: @@ -126,7 +127,7 @@ Practical warning: anything you size or cache off `config.block_size` — RoPE tables, learned positional embeddings, causal or sliding-window mask buffers — must still work at `config.eval_block_size`. Build such buffers against `config.eval_block_size`, or lazily against the actual `T` you are handed. A -model that crashes or silently truncates at 8192 scores 0. +model that crashes or silently truncates at 8192 is rejected (failure score 9999). The baseline always trains at 8192, so this trade is measured against a fixed point. @@ -148,35 +149,36 @@ point. `torch.compile(...)` remain allowed. - A measured bits-per-byte below the plausibility floor (0.4 — far beyond any honest result at this scale and budget) is rejected as presumed held-out - leakage, score 0. + leakage, rejected (failure score 9999). - Trainable parameters must be within the judge's param cap, and the model must - fit and train within GPU memory. Over-cap or out-of-memory scores 0. + fit and train within GPU memory. Over-cap or out-of-memory is rejected + (failure score 9999). - The model must actually train (its parameters must change) and must produce a - non-degenerate output distribution. Untrained or constant-output models score 0. + non-degenerate output distribution. Untrained or constant-output models are + rejected (failure score 9999). ## Scoring Let `base_bpb` be the locked baseline architecture's held-out bits-per-byte and `sub_bpb` be your submission's, both trained for the same wall-clock `T` on the same H100 with the same data order (common random numbers), and both scored on -the same hidden held-out text at the same fixed 8192-token windows. The scored -measurement is your **absolute** bits-per-byte gain: +the same hidden held-out text at the same fixed 8192-token windows. **Your score IS +`sub_bpb` — your submission's raw held-out bits per byte. LOWER IS BETTER.** ``` -gain = base_bpb - sub_bpb # bits per byte -score = clip(100 * gain / BPB_SCORE_SCALE, 0, 100) +score = sub_bpb # raw bits per byte; no scaling, no clipping ``` -- The baseline architecture (tying its `val_bpb`) scores **0**. -- A submission worse than the baseline scores **0**. -- `BPB_SCORE_SCALE` is only a **scaling convention** — it exists to map bits per - byte onto the 0–100 range the platform expects, not to declare what counts as - a full result. Maximize `gain`; `score_unbounded` keeps rewarding improvements - past the cap. +- There is no baseline-relative mapping and no scaling constant: minimize your + own `val_bpb`, full stop. +- The baseline's `base_bpb` and your gain over it (`base_bpb - sub_bpb`, + absolute and relative) are still reported in the feedback for context — they + do not affect the score. +- Failed or rejected runs (policy, guards, crashes) receive the failure score + **9999** — unambiguously the worst possible value. -Bits per byte is the unit the language-modelling literature quotes, so a gain -here is directly comparable to published figures. A relative improvement is also -reported, for readability. +Bits per byte is the unit the language-modelling literature quotes, so your +score is directly comparable to published figures. Public feedback reports `base_val_bpb`, `sub_val_bpb`, the absolute gain and the relative improvement, the training context you used, the number of optimizer @@ -200,7 +202,7 @@ you time out. nanoslm_hybrid_arch_design/ ├── readme This file: problem statement, metric, scoring. ├── config.yaml Harbor problem config: runtime, images, GPU/eval -│ knobs (budget, param cap, bpb_score_scale). +│ knobs (budget, param cap). ├── evaluator.py Judge entrypoint. evaluate(model.py) -> (score, │ score_unbounded, message, metrics): static policy │ gate -> Modal (or local-GPU) dispatch -> scoring.