{"data":{"kind":"file","path":"README.md","version_id":"n5fp4ygpif459b76cv1hrfxg","entry":{"name":"README.md","path":"README.md","is_directory":false,"size":6709,"modified_at":"2026-08-07T06:30:03.648000","content_hash":"bf580db3881f3b98963b5989f69c7480ce658f5c42b9ee4ceb637adae45f1f03"},"entries":[],"content":"# scaling-env\r\n\r\nRL environment for the **quantitative side of scaling ML models on accelerators**, on the\r\ntopics of Google DeepMind's free book *[How To Scale Your\r\nModel](https://jax-ml.github.io/scaling-book/)*.\r\n\r\n- **Tasks:** 35 across 5 categories (roofline 7, flops 8, memory 6, sharding 7, inference 7)\r\n- **Dataset:** [`scaling-tasks-v1`](https://huggingface.co/datasets/eltociear/scaling-tasks-v1)\r\n- **Reward:** binary, from a row comparison at a 1e-9 relative tolerance\r\n- **Dependencies in the sandbox:** none. Every task is answered with `math`.\r\n\r\n```bash\r\nprime env install eltociear/scaling-env\r\n```\r\n\r\n## Original problems, not the book's\r\n\r\nNothing here is copied. Every task is written fresh on the topics the book teaches — a\r\ncopyright decision and the stronger evaluation choice, since a model can recall a published\r\nworked example but cannot recall an original problem with an executed answer key.\r\n\r\n## Every hardware number is stated in the prompt\r\n\r\nNo task requires knowing a real chip's peak FLOPs or HBM bandwidth. That is deliberate, and it\r\nis the main design decision in this environment.\r\n\r\nAn environment that required recalling a spec sheet would be measuring **memorisation of\r\nnumbers that change with every accelerator generation** — and worse, its answer keys would\r\nsilently become *wrong* as hardware moves, while still looking authoritative. Stating the\r\nparameters in the task makes every problem self-contained, permanently valid, and a test of\r\nwhether the model can actually reason about rooflines rather than whether it remembers a\r\ndatasheet.\r\n\r\nThe conventions are stated too, for the same reason — the book's own results depend on them:\r\n`6ND` training FLOPs, bf16 at 2 bytes, a ring collective moving `V·(N-1)/N` bytes per device,\r\nAdam state as 2 + 4 + 4 + 4 bytes per parameter.\r\n\r\n## What it measures\r\n\r\nThe arithmetic that decides real scaling choices, and where the answers are counterintuitive:\r\n\r\n- the same matmul is **compute-bound at batch 1024** (intensity 683) and **memory-bound at\r\n  batch 1** (intensity 0.9995) — the crossover is the whole reason decode and prefill behave\r\n  differently;\r\n- attention is only **16%** of a layer's FLOPs at `S = 8192` with `d_ff = 3.5·d_model`, so\r\n  \"attention is quadratic\" does not mean attention dominates;\r\n- Adam optimizer state is **6× the weights** — a 70B model needs 912 GiB and therefore\r\n  **11 chips** with 96 GiB each just to be held;\r\n- a 64-way data-parallel gradient all-reduce on a 70B model takes **6.1× longer than the\r\n  compute step it overlaps**, which makes that configuration communication-bound;\r\n- at batch 1 a 70B model decodes at **11.4 tokens/s** on 1.6 TB/s of bandwidth — a number\r\n  fixed by memory bandwidth alone, with peak FLOPs playing no part;\r\n- once the KV cache is counted, it is **38%** of everything read per decode step.\r\n\r\n## Grading, and why there is a tolerance here\r\n\r\nFloats are compared at a **relative tolerance of 1e-9**; strings, booleans and `None` are\r\ncompared exactly.\r\n\r\nThis is the opposite call from the [`csapp-env`](../csapp_env/) sibling, which is graded with\r\nno tolerance at all. The difference is the domain: CS:APP answers are exact bit patterns, while\r\nthese are *computed* floats where the order of operations legitimately varies —\r\n`6·N·D/(chips·peak·mfu)` and `((6N)/chips)·(D/peak)/mfu` are the same answer with different\r\nlast bits.\r\n\r\n1e-9 absorbs exactly that and nothing else. The comparator's own unit checks\r\n(`test_compare.py`, 17/17) assert both directions: a reassociated computation **passes**, while\r\na ring all-reduce missing its `(N-1)/N` factor, a factor of 2, and an error of one part in a\r\nmillion all **fail**. A wrong `'compute'`/`'memory'` classification can never be laundered by\r\nthe tolerance, and a string `\"250.0\"` is never accepted for the number `250.0`.\r\n\r\n## Verify it yourself\r\n\r\n```bash\r\npython environments/scaling_env/build_tasks.py --verify   # 35/35\r\npython environments/scaling_env/test_compare.py           # 17/17\r\n```\r\n\r\n`--verify` re-checks every task: that it runs, that it is deterministic across two fresh\r\nexecutions, that it returns JSON-safe primitives, that no value is NaN or infinite, that no\r\nvalue is exactly `0.0` (almost always a formula that cancelled), and that it survives the\r\nserialisation round-trip.\r\n\r\n## Environment arguments\r\n\r\n`load_environment()` deliberately exposes very little: the program's guidance is that an\r\nenvironment should have one correct way to be run, so nothing about the prompts, the parsing or\r\nthe grading is configurable.\r\n\r\n| Argument | Default | Meaning |\r\n|---|---|---|\r\n| `split` | `'train'` | Dataset split to load. |\r\n| `dataset_name` | `'eltociear/scaling-tasks-v1'` | Hugging Face dataset of tasks. Change only to point at a fork. |\r\n| `max_turns` | `5` | Tool-use turns the model gets before the rollout ends. |\r\n| `**kwargs` | — | Passed through to the underlying `SandboxEnv`. |\r\n\r\n## Reward rubric\r\n\r\n| Reward function | Weight | What it returns |\r\n|---|---|---|\r\n| `correctness` | 1.0 | Binary: 1.0 when the answer matches the reference, else 0.0. |\r\n\r\nThere is no LLM judge and no partial credit. The score is computed on the host in\r\n`post_rollout` and read back by the rubric, so the reward is a deterministic function of the\r\nvalues the model left in `result`. A **harness** failure (dead sandbox, unreadable read-back)\r\nalso scores 0.0 but additionally sets `state[\"scoring_error\"]`, so an eval run can tell a\r\nbroken harness from a wrong answer instead of blaming the model.\r\n\r\n## Dependencies\r\n\r\n`datasets>=4.1.0`, `verifiers>=0.1.8`\r\n\r\n## Sample `vf-eval` usage\r\n\r\n```bash\r\nuv run vf-install scaling-env\r\nuv run vf-eval -s scaling-env -m gpt-4.1 -n 5 -r 3\r\nuv run vf-tui                      # inspect the outputs/ folder it writes\r\n```\r\n\r\n## Known limitation\r\n\r\nThe Docker **sandbox transport** has not been executed — `docker run`, `pip install`, and the\r\nmodel's code running inside the container — because that needs a Docker runtime and an\r\ninference provider key, neither available where this was authored.\r\n\r\nEverything else is exercised. `environments/test_scoring_path.py` runs this environment's\r\n`post_rollout` and `Rubric` with the sandbox mocked out, asserting that a correct answer scores\r\n1.0, a well-formed wrong answer scores 0.0, and a dead sandbox scores 0.0 *and* sets\r\n`scoring_error` so a harness failure is never mistaken for a bad model. It also checks that\r\nwhat `build_tasks.py` emits is exactly what the scorer expects. The environment mirrors the\r\nstructure of `polars_env` (already accepted into the Environments Program) and imports cleanly\r\nagainst `verifiers` 0.2.1.\r\n","encoding":"utf-8","truncated":false,"total_bytes":6709},"status":null}