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/README.md b/2.0/README.md index 381cf6cb7..08c2c3db7 100644 --- a/2.0/README.md +++ b/2.0/README.md @@ -127,6 +127,18 @@ 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 + +This task formulates the hybrid language model architecture design as a scored task +at 200M scale (capped at 400M parameters), 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 and scores the +**absolute** reduction in held-out **bits-per-byte** (`val_bpb`, normalized by +bytes so it is tokenizer-independent) against a baseline pure-attention +`olmo3_190M` baseline. + ## Structured-LWE Public Witness Recovery This cryptanalysis task publishes 200 structured-LWE instances spanning ten 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..8647f8564 --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/config.yaml @@ -0,0 +1,65 @@ +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: 86400 + 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: + 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: 7200 +evaluation: + gpu: H100 + tokenizer: allenai/dolma2-tokenizer + # 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 + eval_block_size: 8192 + 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 + # 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: + 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..93865fd8e --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/docker/prep_assets.py @@ -0,0 +1,302 @@ +"""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. 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 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 + 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 +# 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 = 2_000_000_000 # see SIZING in the docstring (6 h budget) +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..ab3fe4081 --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/evaluate.sh @@ -0,0 +1,16 @@ +#!/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 +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..edb1a3303 --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/evaluator.py @@ -0,0 +1,688 @@ +"""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. + +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` (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. +# `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 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 + + source = Path(solution_path).read_text(encoding="utf-8", errors="replace") + with app.run(): + res = run_pair_remote.remote(solution_source=source, 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 (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. + role = _role() + if _backend() == "modal": + try: + res = _run_pair_modal(solution_path, role) + except Exception as exc: # noqa: BLE001 - classify, never leak + 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 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"]) + 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"], + 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, + "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 ( + scoring.FAILURE_SCORE, + scoring.FAILURE_SCORE, + 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 (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 ( + 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) + note = [] + 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, + "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 _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) + + # --- 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) + + # 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()", + policy.check_source( + "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: 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) + 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 + # 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 + + 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) + + 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, + ) + 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, + ) + + # --- 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 + + 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: + 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 + # 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..cbc841b9b --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/harbor/app/README.md @@ -0,0 +1,113 @@ +# 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 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 + +`/app/model.py`, defining `build_model(config)` or `class NanoSLM`, whose +`forward(idx)` returns logits `[B, T, 100352]` for `idx` of BPE token ids. + +`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`: + +```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 +(failure score 9999 — the score is bits per byte, LOWER is better). + +**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 ~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 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 **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 +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..a9d2fa929 --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/harbor/app/model.py @@ -0,0 +1,311 @@ +"""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 +pure-attention ``olmo3_190M``; this file replaces most of its attention layers +with Gated DeltaNet (GDN) linear-recurrent layers. + +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. This reference replaces each sliding-window layer with a GDN layer and +keeps the full-attention layers:: + + 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 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. + +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:: + + baseline 190,354,176 (12 attention layers) + hybrid 254,430,936 (9 GDN + 3 attention layers) +33.7% total + +(Totals at the padded vocab 100352 -- dolma2's 100278 real ids plus 74 unused +rows on the tied table, the Olmo 3 padding convention.) + +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. + +``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 +difference is the sequence mixer in those nine layers. + +Your task: Push val_bpb further +------------------------------- +This 3:1 GDN hybrid is one point in a large design space. Open questions it does +Not settle, each worth real bpb: + + * 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 + 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, 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 + +# --------------------------------------------------------------------------- # +# Architectural Details +# --------------------------------------------------------------------------- # +# Attention hyperparameters +REMOVE_HEADS = 0 +N_LAYER = 12 +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" + +HIDDEN_SIZE = 3072 # SwiGLU width, same as baseline + +ROPE_THETA = 500_000.0 +NORM_EPS = 1e-6 + +# 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_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 + +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.""" + + 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)) + + +def _make_gdn(): + """fla's fused GatedDeltaNet. CUDA required.""" + if not torch.cuda.is_available(): + 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" + ) + + 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, + ) + + +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) -- keep only + # the hidden states. + 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) + + 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", "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 + + @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 + + return non_emb + vocab_size * d + + +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(_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}" + 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..c5ed765af --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/baseline_model.py @@ -0,0 +1,215 @@ +"""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 + +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 = N_EMBD // N_HEAD +HIDDEN_SIZE = 3072 +ROPE_THETA = 500_000.0 +NORM_EPS = 1e-6 +SWA_PATTERN = (4096, 4096, 4096, -1) +FORCE_FULL_FIRST = False +FORCE_FULL_LAST = True +BLOCK_SIZE = 8192 + + +def window_size_for_layer(layer_idx: int, n_layers: int) -> int: + """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): + return -1 + eff = layer_idx - 1 if FORCE_FULL_FIRST else layer_idx + return SWA_PATTERN[eff % len(SWA_PATTERN)] + + +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 + + +# --------------------------------------------------------------------------- # +# Windowed-attention kernel path: FlexAttention, CUDA-only. +# --------------------------------------------------------------------------- # +_FLEX_BLOCK_MASKS: dict = {} +_FLEX_FN = None + + +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, + ) + + 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): + """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.""" + + 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: + y = _windowed_attention(q, k, v, self.window) + 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__() + 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] + ) + + 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..019ef7a2c --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/data.py @@ -0,0 +1,172 @@ +"""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``, 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 +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 ( + TaskConfig, + resolve_val_bytes_per_token, +) + +TOKEN_DTYPE = np.uint32 + + +class DataError(RuntimeError): + """Corpus is absent or unusable (missing/empty file, bad dtype, out-of-range ids, no byte count).""" + + +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 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: + 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 + # 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 + + +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 = _load_bin( + cfg.train_tokens_path, vocab_size=cfg.vocab_size, label="train", + ) + self.val = _load_bin( + cfg.val_tokens_path, vocab_size=cfg.vocab_size, label="val", + ) + if self.train.size < self.max_block + 1: + 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: + # 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: + 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: + """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. 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..317324c80 --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/eval_ppl.py @@ -0,0 +1,101 @@ +"""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. + * 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. +""" + +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. `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..770f5db89 --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/modal_app.py @@ -0,0 +1,407 @@ +"""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, 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. + +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" + +# 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. + + 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. 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) 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/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") + # 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. `.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", ".assets"], + ) + .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", + # 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 + # 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=FN_TIMEOUT_SECONDS, + 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. + + `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_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, + } + + # 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}, secrets=[WANDB_SECRET]) + def run_pair_remote(solution_source: str, + 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 + 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): + + 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. + """ + import json + import sys + import tempfile + + 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) + + 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() + # 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 + # 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 -- 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, + run_label="submission", + ) + 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": base_info, + "submission": _arm(cand), + "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), + "code_version": CODE_VERSION, + "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..3d07fb692 --- /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..1f77f7e5a --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/policy.py @@ -0,0 +1,274 @@ +"""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, +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 + +# 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", +) + + +# --------------------------------------------------------------------------- # +# 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 + 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}") + + # 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( + 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..e1c2d5483 --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/runner.py @@ -0,0 +1,463 @@ +"""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, 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 + 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)``. 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 + 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: 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, + 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 + + 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 + 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. + # + # 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) + 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 (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" + ) + + 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, + 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 + 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 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 + # 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 (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, + run_label=run_label) + 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") + 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)") + # 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, + 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), + 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 new file mode 100644 index 000000000..c52b4265d --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/scoring.py @@ -0,0 +1,82 @@ +"""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. +""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass +class ScoreResult: + 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 — reported only + + +# 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, +) -> ScoreResult: + """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 -> the failure sentinel (worst; runner also + # guards). + if not (base_bpb > 0.0) or not (sub_bpb > 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(FAILURE_SCORE, FAILURE_SCORE, 0.0, 0.0) + + abs_delta = base_bpb - sub_bpb # reported for readability only + rel = abs_delta / base_bpb # reported for readability only + return ScoreResult(sub_bpb, sub_bpb, 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}; " + f"rel_improvement={result.rel_improvement:+.4%}; " + f"steps={steps}; train_wall_s={wall_seconds:.1f}; " + f"score={result.score:.5f} (== sub_val_bpb; LOWER is better)" + ) + 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..4cdfc74af --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/settings.py @@ -0,0 +1,271 @@ +"""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. +""" + +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. + # + # 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. + # + # 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) --- + # 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) + 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 (matches config.yaml) --- + # 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 = 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 + tokenizer_name: str = "allenai/dolma2-tokenizer" + 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 + # 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 result + # cached by config fingerprint; the final/verifier path recomputes a fresh + # 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 --- + # 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) + 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 + # (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 + # 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 --- + 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 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; +# with no real corpus the harness raises DataError instead of estimating a ratio. +# --------------------------------------------------------------------------- # + + +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 + -- 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. + + 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 + + +def active_config() -> TaskConfig: + """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 +# 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", + # 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. + "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", +) + + +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..7b3c1423b --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/harness/train.py @@ -0,0 +1,235 @@ +"""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 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 + +import math +import os +import time +from dataclasses import asdict, 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, elapsed: float, cfg: TaskConfig) -> float: + """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.lr_schedule_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). 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: + 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 _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) + model.train() + opt = torch.optim.AdamW( + _param_groups(model, cfg), + lr=cfg.learning_rate, + betas=(cfg.beta1, cfg.beta2), + ) + use_amp = device.startswith("cuda") + autocast = ( + torch.autocast(device_type="cuda", dtype=torch.bfloat16) + if use_amp + else _nullcontext() + ) + + wb = _maybe_init_wandb(model, cfg, run_label) + step = 0 + last_loss = float("nan") + t0 = time.monotonic() + 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 + + _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..f9e3023b3 --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/readme @@ -0,0 +1,260 @@ +# 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 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`, 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**: + +``` +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 `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. + 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); logits span + # the padded vocab_size (100352) + 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. + +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 + +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 (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: + +- 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 is rejected (failure score 9999). + +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. +- 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, 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 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 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. **Your score IS +`sub_bpb` — your submission's raw held-out bits per byte. LOWER IS BETTER.** + +``` +score = sub_bpb # raw bits per byte; no scaling, no clipping +``` + +- 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 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 +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. + +## 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). +├── 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 new file mode 100644 index 000000000..a9d2fa929 --- /dev/null +++ b/2.0/problems/nanoslm_hybrid_arch_design/reference.py @@ -0,0 +1,311 @@ +"""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 +pure-attention ``olmo3_190M``; this file replaces most of its attention layers +with Gated DeltaNet (GDN) linear-recurrent layers. + +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. This reference replaces each sliding-window layer with a GDN layer and +keeps the full-attention layers:: + + 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 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. + +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:: + + baseline 190,354,176 (12 attention layers) + hybrid 254,430,936 (9 GDN + 3 attention layers) +33.7% total + +(Totals at the padded vocab 100352 -- dolma2's 100278 real ids plus 74 unused +rows on the tied table, the Olmo 3 padding convention.) + +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. + +``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 +difference is the sequence mixer in those nine layers. + +Your task: Push val_bpb further +------------------------------- +This 3:1 GDN hybrid is one point in a large design space. Open questions it does +Not settle, each worth real bpb: + + * 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 + 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, 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 + +# --------------------------------------------------------------------------- # +# Architectural Details +# --------------------------------------------------------------------------- # +# Attention hyperparameters +REMOVE_HEADS = 0 +N_LAYER = 12 +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" + +HIDDEN_SIZE = 3072 # SwiGLU width, same as baseline + +ROPE_THETA = 500_000.0 +NORM_EPS = 1e-6 + +# 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_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 + +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.""" + + 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)) + + +def _make_gdn(): + """fla's fused GatedDeltaNet. CUDA required.""" + if not torch.cuda.is_available(): + 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" + ) + + 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, + ) + + +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) -- keep only + # the hidden states. + 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) + + 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", "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 + + @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 + + return non_emb + vocab_size * d + + +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(_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}" + return got 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