{"data":{"kind":"file","path":"README.md","version_id":"jolbu6urep3vmtarpflnmddf","entry":{"name":"README.md","path":"README.md","is_directory":false,"size":15240,"modified_at":"2026-09-05T10:39:16","content_hash":"03f6e6199d04efd04d3f8822319dda512a42f78595b651b33300e0503f5534f1"},"entries":[],"content":"# Scaling Book\n\n[Fork: `jahnclawdmonet/community-environments`](https://github.com/jahnclawdmonet/community-environments)\n\nAuthor: Jahn - Conatus AI\n\n`scaling-book` is a train-and-evaluation environment for first-principles accelerator and\ndistributed-LLM systems reasoning. It combines the book's worked problems, an authored graduate\nsystems exam bank, an unlimited deterministic cost-model generator, and executable implementation\nexercises.\n\nThe source is [jax-ml/scaling-book](https://github.com/jax-ml/scaling-book), licensed under the\n[MIT License](https://github.com/jax-ml/scaling-book/blob/ac72f020e4320b02e54b6f1b2324cdf1c579388c/LICENSE),\nat pinned commit [`ac72f020e4320b02e54b6f1b2324cdf1c579388c`](https://github.com/jax-ml/scaling-book/tree/ac72f020e4320b02e54b6f1b2324cdf1c579388c).\nSelected chapter Markdown is fetched from that commit inside `load_environment`; the package does\nnot ship a copy or a preprocessed dump of the book. The authored question banks contain only their\nown prompts, bounded source passages used to justify them, and reference answers.\n\n## Splits\n\n| Split | Default size | Grading | Purpose |\n|---|---:|---|---|\n| `book` | 79 | 40 numeric, 39 judge | Answer-bearing worked questions parsed from the pinned book |\n| `exam` | 121 | 71 numeric, 50 judge | Authored homework/exam questions with per-item validation |\n| `synthetic` | 256 | Numeric | Seeded cost-model training data; any positive size is accepted |\n| `impl` | 12 | Hidden subprocess tests | Python and CPU-only JAX implementation work |\n\nThe pinned source has 92 question markers, but only 79 have source-provided answers. Thirteen are\nexcluded instead of receiving invented references: sharding 6 and 8; applied-training 10 and 11;\ninference 4 and 6; applied-inference 13 through 16; profiling 2; and jax-stuff 3 and 4. In every\ncase, the reason is that the pinned chapter supplies no answer.\n\nThe two conclusion items from the earlier draft are intentionally omitted: the pinned conclusion\ncontains acknowledgements and a reading list, not source material that supports exam answers.\n\nThe generator covers matmul FLOPs and bytes, arithmetic intensity, crossover batch, roofline time,\nKV-cache size, transformer parameters, 6ND training and 2ND inference FLOPs, collectives, and FSDP\nversus tensor-parallel communication. Its hardware records are parsed at load time from the book's\nTPU v4p/v5e/v5p/v6e and H100/B200 tables.\n\n## Coverage\n\nThe map was derived from each chapter's headings and takeaways before the authored tasks were\nwritten. Counts below are links from tasks to key-content items; a task can cover more than one\nitem. Synthetic counts use the default 256 examples at seed 0. No mapped item is uncovered across\nall splits.\n\n| Chapter | Key content | Book | Exam | Synthetic | Impl |\n|---|---|---:|---:|---:|---:|\n| roofline | Compute, communication, and overlapped runtime bounds | 5 | 4 | 23 | 1 |\n| roofline | Arithmetic intensity and hardware crossover | 5 | 4 | 46 | 0 |\n| roofline | Exact matmul FLOPs, bytes, and token-batch crossover | 4 | 4 | 47 | 1 |\n| roofline | Inter-chip communication rooflines | 0 | 3 | 0 | 0 |\n| tpus | MXU, VPU, VMEM, and HBM roles | 3 | 3 | 0 | 0 |\n| tpus | Precision-dependent compute and MXU padding | 2 | 3 | 0 | 0 |\n| tpus | Host, tray, pod, and torus topology | 2 | 3 | 0 | 1 |\n| tpus | HBM, ICI, PCIe, and DCN bandwidth hierarchy | 3 | 3 | 0 | 0 |\n| tpus | v4p, v5e, v5p, and v6e hardware tables | 6 | 3 | 0 | 0 |\n| gpus | SM, CUDA core, Tensor Core, cache, and HBM organization | 5 | 2 | 0 | 0 |\n| gpus | H100 and B200 capacity, bandwidth, and compute tables | 7 | 2 | 0 | 0 |\n| gpus | NVLink and NVSwitch node bandwidth | 3 | 2 | 0 | 0 |\n| gpus | Cross-node fat-tree and bisection bandwidth | 2 | 2 | 0 | 0 |\n| gpus | Intra-node and cross-node collective costs | 6 | 4 | 0 | 0 |\n| gpus | Data, tensor, expert, and pipeline parallel rooflines | 3 | 2 | 0 | 0 |\n| gpus | Advertised versus attainable communication and SHARP | 8 | 2 | 0 | 0 |\n| sharding | Mesh and partition notation, including replication | 1 | 2 | 0 | 0 |\n| sharding | Local matmul without a sharded contracting dimension | 1 | 2 | 0 | 0 |\n| sharding | AllGather for one sharded contracting dimension | 2 | 3 | 0 | 0 |\n| sharding | Partial sums and AllReduce/ReduceScatter | 2 | 3 | 0 | 0 |\n| sharding | AllToAll for conflicting non-contracting shardings | 1 | 2 | 0 | 0 |\n| sharding | Collective throughput, latency, and overlap | 7 | 3 | 23 | 1 |\n| transformers | FLOP counting for tensor contractions | 2 | 2 | 0 | 0 |\n| transformers | Forward, backward, 2ND, and 6ND rules | 2 | 2 | 46 | 1 |\n| transformers | Gated MLP parameter and FLOP accounting | 3 | 2 | 24 | 0 |\n| transformers | MHA, GQA, and attention FLOP accounting | 4 | 2 | 0 | 0 |\n| transformers | Attention cost versus context length | 1 | 2 | 0 | 0 |\n| transformers | Mixture-of-experts total versus activated parameters | 1 | 2 | 0 | 0 |\n| transformers | KV cache dimensions and bytes | 1 | 3 | 0 | 1 |\n| training | Parameter, optimizer, and activation memory | 2 | 2 | 0 | 0 |\n| training | Data parallelism memory and gradient AllReduce | 1 | 2 | 0 | 0 |\n| training | FSDP parameter communication and critical batch | 1 | 2 | 0 | 0 |\n| training | Tensor-parallel activation communication | 1 | 2 | 0 | 0 |\n| training | Optimal hybrid FSDP and tensor parallelism | 1 | 3 | 23 | 2 |\n| training | Pipeline stages, bubbles, and microbatches | 0 | 2 | 0 | 0 |\n| training | Cross-pod data-parallel scaling | 0 | 2 | 0 | 0 |\n| applied-training | LLaMA 3 parameter, FLOP, and memory accounting | 3 | 1 | 0 | 0 |\n| applied-training | Training-duration estimates from topology and MFU | 2 | 1 | 0 | 0 |\n| applied-training | Practical hybrid sharding for LLaMA 3 | 5 | 1 | 0 | 0 |\n| inference | TTFT, per-token latency, throughput, and cost | 2 | 2 | 0 | 0 |\n| inference | Prefill/decode linear-operation rooflines | 1 | 2 | 0 | 0 |\n| inference | Prefill and decode attention rooflines | 1 | 2 | 0 | 0 |\n| inference | Model and KV-cache capacity | 4 | 2 | 24 | 1 |\n| inference | KV-cache architectural and runtime reductions | 2 | 2 | 0 | 0 |\n| inference | Prefill and generation parallelism | 4 | 2 | 0 | 0 |\n| inference | Disaggregation, continuous batching, and prefix caching | 0 | 3 | 0 | 0 |\n| applied-inference | LLaMA 3 KV and model memory sizing | 3 | 1 | 0 | 0 |\n| applied-inference | Decode latency and throughput by topology | 9 | 1 | 0 | 0 |\n| applied-inference | Communication-safe model parallelism | 1 | 1 | 0 | 0 |\n| applied-inference | Prefill/decode server balancing | 2 | 1 | 0 | 0 |\n| profiling | JAX to XLA to TPU software stack | 1 | 1 | 0 | 0 |\n| profiling | Trace-viewer operation and communication analysis | 1 | 1 | 0 | 0 |\n| profiling | Reading HLO shapes, layouts, and replica groups | 1 | 1 | 0 | 0 |\n| profiling | Peak memory and buffer-lifetime analysis | 1 | 1 | 0 | 0 |\n| jax-stuff | Meshes, PartitionSpec, and NamedSharding | 2 | 1 | 0 | 1 |\n| jax-stuff | jax.jit input/output shardings | 2 | 1 | 0 | 1 |\n| jax-stuff | Explicit per-shard computation with shard_map | 2 | 1 | 0 | 1 |\n| jax-stuff | JAX collective semantics | 2 | 1 | 0 | 0 |\n\n## Grading and validation\n\nNumeric answers are extracted from the last `Final answer:` line or last `\\boxed{}` expression,\nwhichever appears later, and compared by physical dimension. Decimal and binary byte units,\nnanoseconds through seconds, FLOPs/OPs and rates,\nbandwidth, percentages, scientific notation, and task-specific counts are normalized before\ncomparison. The default relative tolerance is 5% for `book` and `exam`, and 2% for `synthetic`.\nEvery required value must match; extra working values do not earn credit by themselves.\n\nDerivations and conceptual answers use a user-selected OpenAI-compatible judge against the stored\nreference. `judge_model` is mandatory for `book` or `exam` selections containing judge items. The\ngrader uses `vf.JudgeRubric` with a 256-token judge budget and\n`extra_body.chat_template_kwargs.enable_thinking=false`; the judge is never hard-coded. Each\nnumeric exam item has an item-specific executable oracle whose arguments are the raw quantities\nstated in the question. Every oracle performs at least two arithmetic operations; the 43 former\nsingle-operation items were rewritten as multi-step derivations, leaving zero. Tests recompute all\n71 numeric answers from those raw inputs. Conceptual references carry a chapter, an exact\npinned-source span, and a validation rationale. Fifteen oracle tests exercise package functions\nagainst genuine numbers printed in the book. Implementation reference solutions must pass every\nparent-held case.\n\nThe implementation executor is **best-effort process isolation and not a security boundary**.\nEvery case starts a new `python -I` child with CPU-only JAX, a scrubbed environment, a fresh empty\nworking directory and temporary `HOME`, a 15-second wall-clock timeout, a 10-second CPU limit, and\na hard address-space limit. The address-space allowance defaults to 2 GiB beyond the interpreter\nand required libraries. CPU and preliminary address-space limits are installed before exec; after\nany required JAX inputs are initialized, the child tightens the address-space limit before running\ncandidate code. The executor selects the strongest tier that its startup probe can actually run:\n\n- On macOS with `/usr/bin/sandbox-exec`, Seatbelt permits reads of the working directory, required\n  system runtime files, the Python standard library, and installed JAX/NumPy modules. It denies\n  reads of the environment package and sibling temporary directories, all network access, and\n  writes outside the working directory.\n- On Linux where `bwrap` can create its namespaces, bubblewrap mounts the host root read-only,\n  replaces `/tmp` with a private tmpfs containing the empty working directory and fresh `HOME`,\n  creates a network namespace with no external network access, and uses `--die-with-parent`.\n  Non-temporary host files, including installed libraries and package files, remain readable\n  through the read-only root; bubblewrap prevents changing them, but does not hide them.\n- On Linux without `bwrap`, or where its startup probe fails (for example because user namespaces\n  are unavailable), the executor uses the plain child. The empty working directory, temporary\n  `HOME`, scrubbed environment, isolated-mode Python, wall timeout, and pre-exec CPU and\n  address-space limits still apply. This tier does not enforce filesystem isolation: the child can\n  read host files and can write anywhere allowed by the grader account. It also has no kernel-level\n  network or process namespace; the Python socket guard and AST deny-list are only defense in depth.\n\nIf the selected child Python process itself cannot be started, grading fails closed with an error.\n\nThe child receives only candidate source and one case's inputs over stdin and returns the computed\nvalue over stdout. Expected outputs and exception expectations exist only in the parent process,\nwhich performs every comparison. There is no hidden-test runner file, test source, expected value,\npass token, or sentinel in the child's filesystem, environment, command line, or input. The AST\ndeny-list remains defense in depth for dangerous process, network, and filesystem calls while\nallowing ordinary standard-library imports such as `typing` and `functools`.\n\n## Installation and use\n\nPython 3.11 through 3.13 is supported. From the repository root:\n\n```bash\nuv run vf-install scaling-book\nuv run vf-eval scaling-book -m <model> -n 5 -r 3 -s\n```\n\nThe default command selects `synthetic`, so it needs only the endpoint configuration already used\nby `vf-eval`. Select another split with environment arguments:\n\n```bash\nuv run vf-eval scaling-book -m <model> -a '{\"split\":\"impl\"}'\nuv run vf-eval scaling-book -m <model> -a '{\"split\":\"synthetic\",\"num_synthetic_examples\":1024,\"seed\":7}'\nuv run vf-eval scaling-book -m <model> -a '{\"split\":\"exam\",\"chapters\":\"roofline,tpus\",\"judge_model\":\"<judge-model>\"}'\n```\n\nFor a local OpenAI-compatible judge, add `\"judge_base_url\":\"http://127.0.0.1:8089/v1\"` to the\nenvironment arguments. Otherwise `OPENAI_BASE_URL` is used. The judge API key is read from\n`OPENAI_API_KEY` by default; `judge_api_key_var` changes only the environment-variable name.\nBefore trusting a smoke run, grade a stored reference against itself through that endpoint and\nrequire reward 1. For thinking-capable local models, also pass generation sampling arguments under\n`extra_body.chat_template_kwargs`, not as a top-level `enable_thinking` field, and use at least\n1024 completion tokens.\n\n`load_environment` arguments are `split`, `chapters` (comma-separated string or sequence),\n`num_synthetic_examples`, `seed`, `tolerance`, `judge_model`, `judge_base_url`, and\n`judge_api_key_var`. Source download failure is fatal and explicit; there is no mutable-source or\ncached-book fallback.\n\n## Design choices\n\n- Pinning and fetching the source at load time preserves provenance without distributing the book.\n- Answerless source questions are excluded because fabricated references would undermine fidelity.\n- Worked prompts include the pinned specification rows for every chip named by the question or\n  answer, even when the problem appears outside the TPU/GPU chapters. Context is bounded to 6000\n  characters and has an 80-character minimum carrying at least its source location.\n- Numeric targets use deterministic code rather than a judge; judges are reserved for answers that\n  genuinely require semantic comparison.\n- The FSDP/TP generator and implementations use the book's forward-pass convention: communicated\n  bytes are `2*element_bytes*D*F/Y` for FSDP and `2*element_bytes*B*D/X` for tensor parallelism.\n- The generated split interleaves family groups; its first eight seed-0 tasks span roofline, KV\n  cache, transformer, collective, and FSDP/TP work before the 11-type cycle repeats. Shapes and\n  hardware are sampled deterministically from the seed.\n- Implementation case inputs are generated in the parent, and reference solutions ship separately\n  from prompts. CPU-only JAX keeps the exercises runnable without accelerator access. Use an\n  independently hardened container or VM for hostile submissions or multi-tenant RL training.\n\n## Known limitations\n\n- Loading requires network access to `raw.githubusercontent.com` for the pinned chapters.\n- The book parser follows the question and Liquid-details markup at the pinned commit. A future\n  source version requires a new pin and parser review rather than silently changing the dataset.\n- The judge splits require a configured model and inherit that model's semantic-grading variance.\n- On Linux without usable bubblewrap, implementation grading uses the weaker plain-child tier\n  described above. Use an independently hardened container or VM for hostile submissions.\n- Book hardware tables mix advertised and measured figures; each task states which figure it uses.\n- The H100 chapter alternates between 3.35 and 3.4 TB/s HBM figures, and the B200 worked answer\n  reports 281 FLOPs/byte while the printed 2.3 PFLOP/s and 8.0 TB/s table entries yield 287.5. Book\n  tasks retain the worked answers; table-derived synthetic tasks use the table entries.\n","encoding":"utf-8","truncated":false,"total_bytes":15240},"status":null}