nanoscope-lab 0.3.0__tar.gz

nanoscope_lab-0.3.0/.gitignore

@@ -0,0 +1,16 @@
+.venv/
+__pycache__/
+*.py[cod]
+.pytest_cache/
+.ruff_cache/
+.coverage
+htmlcov/
+.env
+wandb/
+runs/
+checkpoints/
+experiments/**/artifacts/
+*.tmp
+.kaggle/
+.vscode/
+graphify-out/

nanoscope_lab-0.3.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 almajd3713
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

nanoscope_lab-0.3.0/Makefile

@@ -0,0 +1,30 @@
+.PHONY: install test test-all test-blocks test-curricula test-curricula-full openapi lint typecheck check
+
+install:  ## everything, including dev tools
+	uv sync --all-extras
+
+test:  ## fast offline tests
+	uv run pytest -m "not network and not gpu"
+
+test-all:  ## also the first-notebook timing test (downloads TinyStories once) and GPU tests
+	uv run pytest
+
+test-blocks:  ## the block library: blocks, references, graph, describe, block stats
+	uv run pytest tests/test_blocks.py tests/test_reference.py tests/test_graph.py tests/test_describe.py tests/test_blockstats.py
+
+test-curricula:  ## every lesson's starter fails and its solution passes (offline, CPU)
+	uv run pytest tests/test_curricula.py tests/test_checks.py tests/test_learn.py tests/test_gating.py -m "not network and not gpu"
+
+test-curricula-full:  ## also train every lesson's solution on real TinyStories (minutes; needs network once)
+	uv run pytest tests/test_curricula.py
+
+openapi:  ## regenerate docs/openapi.json from the API (commit the result)
+	uv run python -m nanoscope.server.openapi docs/openapi.json
+
+lint:
+	uv run ruff check nanoscope tests
+
+typecheck:
+	uv run pyright nanoscope
+
+check: lint typecheck test  ## what CI runs

nanoscope_lab-0.3.0/PKG-INFO

@@ -0,0 +1,144 @@
+Metadata-Version: 2.5
+Name: nanoscope-lab
+Version: 0.3.0
+Summary: See what your language model learns
+Project-URL: Homepage, https://github.com/almajd3713/nanoscope
+Project-URL: Source, https://github.com/almajd3713/nanoscope
+Project-URL: Issues, https://github.com/almajd3713/nanoscope/issues
+Author: almajd3713
+License-Expression: MIT
+License-File: LICENSE
+Keywords: education,interpretability,language-models,pytorch,research,transformers
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Education
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: <3.13,>=3.10
+Requires-Dist: datasets<5,>=3.0
+Requires-Dist: huggingface-hub<2,>=0.27
+Requires-Dist: jsonschema<5,>=4.20
+Requires-Dist: matplotlib<4,>=3.8
+Requires-Dist: numpy<3,>=1.26
+Requires-Dist: scipy<2,>=1.11
+Requires-Dist: tiktoken<1,>=0.8
+Requires-Dist: tokenizers<1,>=0.20
+Requires-Dist: tomli-w>=1.0
+Requires-Dist: tomli>=2; python_full_version < '3.11'
+Requires-Dist: torch<3,>=2.2
+Provides-Extra: dev
+Requires-Dist: httpx<1,>=0.27; extra == 'dev'
+Requires-Dist: hypothesis<7,>=6.100; extra == 'dev'
+Requires-Dist: jsonschema<5,>=4.20; extra == 'dev'
+Requires-Dist: libcst<2,>=1.1; extra == 'dev'
+Requires-Dist: nbformat<6,>=5.10; extra == 'dev'
+Requires-Dist: pyright>=1.1.390; extra == 'dev'
+Requires-Dist: pytest-cov<8,>=6; extra == 'dev'
+Requires-Dist: pytest<9,>=8.3; extra == 'dev'
+Requires-Dist: ruff<1,>=0.9; extra == 'dev'
+Provides-Extra: graph
+Requires-Dist: libcst<2,>=1.1; extra == 'graph'
+Provides-Extra: server
+Requires-Dist: fastapi<1,>=0.115; extra == 'server'
+Requires-Dist: libcst<2,>=1.1; extra == 'server'
+Requires-Dist: ruff<1,>=0.9; extra == 'server'
+Requires-Dist: uvicorn[standard]<1,>=0.30; extra == 'server'
+Requires-Dist: watchfiles<2,>=0.24; extra == 'server'
+Provides-Extra: wandb
+Requires-Dist: wandb<1,>=0.19; extra == 'wandb'
+Description-Content-Type: text/markdown
+
+# nanoscope
+
+See what your language model learns. You write an `nn.Module`; nanoscope handles the
+data, the training loop, evaluation, checkpoints and comparison against baselines.
+
+It has two levels that share one core. Learners call `run()`. Researchers write a `Study`
+with seeds, budgets, parameter matching and preregistration. The level changes what you
+see, never which code runs.
+
+## Learn
+
+```bash
+pip install git+https://github.com/almajd3713/nanoscope
+```
+
+```python
+import torch.nn as nn
+from nanoscope import compare, run
+
+class Bigram(nn.Module):
+    def __init__(self, vocab_size: int, d_model: int = 32):
+        super().__init__()
+        self.token_embedding = nn.Embedding(vocab_size, d_model)
+        self.head = nn.Linear(d_model, vocab_size, bias=False)
+
+    def forward(self, idx):
+        return self.head(self.token_embedding(idx))
+
+result = run(Bigram, preset="tinystories-5min")   # about a minute on a laptop CPU
+result.plot()
+compare(result, "gpt2")                          # against a shipped 3-seed baseline
+```
+
+Work through the notebooks in order:
+
+1. [`01-first-model`](notebooks/01-first-model.ipynb): write a bigram model and train it.
+2. [`02-gpt2`](notebooks/02-gpt2.ipynb): a real transformer.
+3. [`03-modern-block`](notebooks/03-modern-block.ipynb): RoPE, RMSNorm, SwiGLU, GQA, QK-norm, z-loss.
+4. [`04-ablations`](notebooks/04-ablations.ipynb): which part matters, with seeds and confidence intervals.
+
+Token data downloads from the Hub (`RedhouaneLazib/nanoscope-tokens`) when a preset has
+it, and is tokenized locally otherwise. Everything is cached under `~/.nanoscope/data`.
+
+## Research
+
+See the [research guide](docs/research.md). The research program itself is in
+[`docs/project-nanoscope.md`](docs/project-nanoscope.md).
+
+## Build models from blocks
+
+`GPT2` and `Modern` are short compositions of the blocks in `nanoscope.blocks`, and so can your
+own models. See [docs/blocks.md](docs/blocks.md).
+
+## Learn
+
+Guided paths build the models step by step, with checks that say why. See
+[docs/learn.md](docs/learn.md): `nanoscope learn list`, `learn start`, `learn check`.
+
+## Command line
+
+```bash
+nanoscope presets                                   # available presets
+nanoscope run nanoscope/models/gpt2.py:GPT2 --seeds 3
+nanoscope compare modern gpt2 --preset tinystories-5min
+nanoscope study studies/m1_ablation.py --devices cuda:0
+nanoscope report studies/m1_ablation.py             # writes experiments/<name>/
+nanoscope bench modern --compile reduce-overhead    # speed, and whether CPU or GPU is the limit
+nanoscope status runs                               # what is running and how far along
+nanoscope describe nanoscope/models/modern.py:Modern  # shapes, params, FLOPs, memory per module
+nanoscope graph my_model.py                         # a model file's architecture, without running it
+nanoscope blocks                                    # the blocks models are composed from
+nanoscope prepare-data tinystories-5min             # download or tokenize now
+nanoscope publish-data tinystories-5min user/repo   # upload tokens to a Hub dataset
+```
+
+## Develop
+
+```bash
+uv sync --all-extras
+make test       # offline tests
+make lint
+make typecheck
+make check      # lint, typecheck, then test: what CI runs
+```
+
+## Data credits
+
+Tokens are derived from [TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories)
+(CDLA-Sharing-1.0) and [FineWeb-Edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu)
+(ODC-By 1.0).

nanoscope_lab-0.3.0/README.md

@@ -0,0 +1,90 @@
+# nanoscope
+
+See what your language model learns. You write an `nn.Module`; nanoscope handles the
+data, the training loop, evaluation, checkpoints and comparison against baselines.
+
+It has two levels that share one core. Learners call `run()`. Researchers write a `Study`
+with seeds, budgets, parameter matching and preregistration. The level changes what you
+see, never which code runs.
+
+## Learn
+
+```bash
+pip install git+https://github.com/almajd3713/nanoscope
+```
+
+```python
+import torch.nn as nn
+from nanoscope import compare, run
+
+class Bigram(nn.Module):
+    def __init__(self, vocab_size: int, d_model: int = 32):
+        super().__init__()
+        self.token_embedding = nn.Embedding(vocab_size, d_model)
+        self.head = nn.Linear(d_model, vocab_size, bias=False)
+
+    def forward(self, idx):
+        return self.head(self.token_embedding(idx))
+
+result = run(Bigram, preset="tinystories-5min")   # about a minute on a laptop CPU
+result.plot()
+compare(result, "gpt2")                          # against a shipped 3-seed baseline
+```
+
+Work through the notebooks in order:
+
+1. [`01-first-model`](notebooks/01-first-model.ipynb): write a bigram model and train it.
+2. [`02-gpt2`](notebooks/02-gpt2.ipynb): a real transformer.
+3. [`03-modern-block`](notebooks/03-modern-block.ipynb): RoPE, RMSNorm, SwiGLU, GQA, QK-norm, z-loss.
+4. [`04-ablations`](notebooks/04-ablations.ipynb): which part matters, with seeds and confidence intervals.
+
+Token data downloads from the Hub (`RedhouaneLazib/nanoscope-tokens`) when a preset has
+it, and is tokenized locally otherwise. Everything is cached under `~/.nanoscope/data`.
+
+## Research
+
+See the [research guide](docs/research.md). The research program itself is in
+[`docs/project-nanoscope.md`](docs/project-nanoscope.md).
+
+## Build models from blocks
+
+`GPT2` and `Modern` are short compositions of the blocks in `nanoscope.blocks`, and so can your
+own models. See [docs/blocks.md](docs/blocks.md).
+
+## Learn
+
+Guided paths build the models step by step, with checks that say why. See
+[docs/learn.md](docs/learn.md): `nanoscope learn list`, `learn start`, `learn check`.
+
+## Command line
+
+```bash
+nanoscope presets                                   # available presets
+nanoscope run nanoscope/models/gpt2.py:GPT2 --seeds 3
+nanoscope compare modern gpt2 --preset tinystories-5min
+nanoscope study studies/m1_ablation.py --devices cuda:0
+nanoscope report studies/m1_ablation.py             # writes experiments/<name>/
+nanoscope bench modern --compile reduce-overhead    # speed, and whether CPU or GPU is the limit
+nanoscope status runs                               # what is running and how far along
+nanoscope describe nanoscope/models/modern.py:Modern  # shapes, params, FLOPs, memory per module
+nanoscope graph my_model.py                         # a model file's architecture, without running it
+nanoscope blocks                                    # the blocks models are composed from
+nanoscope prepare-data tinystories-5min             # download or tokenize now
+nanoscope publish-data tinystories-5min user/repo   # upload tokens to a Hub dataset
+```
+
+## Develop
+
+```bash
+uv sync --all-extras
+make test       # offline tests
+make lint
+make typecheck
+make check      # lint, typecheck, then test: what CI runs
+```
+
+## Data credits
+
+Tokens are derived from [TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories)
+(CDLA-Sharing-1.0) and [FineWeb-Edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu)
+(ODC-By 1.0).

nanoscope_lab-0.3.0/docs/blocks.md

@@ -0,0 +1,144 @@
+# Blocks
+
+`nanoscope.blocks` is a library of short modules, one idea each, that models are composed
+from. `GPT2` and `Modern` are themselves compositions of these blocks, so everything you read
+here is what the shipped models run.
+
+## Compose a model
+
+```python
+from nanoscope.blocks import Attention, Block, Decoder, RMSNorm, RoPE, SwiGLU
+
+
+class MyLM(Decoder):
+    def __init__(self, vocab_size: int, context_length: int = 256):
+        super().__init__(
+            vocab_size, context_length, d_model=128, n_layers=4,
+            block=Block(norm=RMSNorm(), mlp=SwiGLU(hidden=344),
+                        attn=Attention(n_heads=4, n_kv_heads=2, pos=RoPE(), qk_norm=True)),
+            final_norm=RMSNorm(),
+            tie_weights=True, z_loss=1e-4,
+        )
+```
+
+`run(MyLM)` trains it like any model. Calling a block with only its options (`SwiGLU(hidden=344)`)
+gives a *spec*, a description with no weights. `Decoder` builds the spec once per layer, so
+layers never share parameters. A spec knows its arguments, so `spec.to_dict()` is what the graph
+and `describe` show.
+
+Useful `Decoder` options:
+
+| Option | Meaning |
+|---|---|
+| `block` | One block spec used for every layer. |
+| `pattern=[a, b]` | Instead of `block`: repeat a list of specs through the layers (`a, b, a, b, ...`), e.g. sliding-window and global attention. |
+| `final_norm` | A norm after the last layer. |
+| `pos_emb` | A position embedding added after the token embedding (`LearnedPosition()`); leave it out when attention carries positions (`RoPE`). |
+| `tie_weights` | Share the output matrix with the token embedding. |
+| `z_loss` | Add a penalty that keeps the softmax normaliser near 1. The model then returns `(logits, aux_loss)`. |
+
+Parameters start as N(0, 0.02), biases at zero, and each layer's output projections at
+`0.02 / sqrt(2 * n_layers)`.
+
+## The representable subset
+
+The architecture graph and `nanoscope graph` read your file with `ast`, without importing or
+running it. A class is shown as a graph when it is a `Decoder` (or `Composite`) subclass whose
+`__init__` is one `super().__init__(...)` call. Each argument is one of:
+
+- a literal (`128`, `True`, `1e-4`),
+- one of the class's own `__init__` parameters (`vocab_size`),
+- a call to a registered block, with its options by keyword,
+- a list of those (`pattern=[...]`).
+
+A call to anything else becomes an *opaque* node: it is drawn, and `describe` reports its
+shapes, but its inside is not editable. Any other expression is kept as text. Any other
+*statement* in `__init__` (an assignment, a loop, an `if`) makes the class **code-only**: it
+still appears, with the reason and the line, and you edit it as source.
+
+```
+$ nanoscope graph my_model.py:MyLM
+MyLM(Decoder)  my_model.py:5
+  vocab_size = <vocab_size>
+  context_length = <context_length>
+  d_model = 128
+  block = Block
+    norm = RMSNorm
+    ...
+```
+
+Edits go back into the file with `blocks.graph.emit(graph, source)`, which changes only the
+arguments that changed (keeping comments and formatting) and adds a missing import when a new
+block is used. An unedited round trip is byte-identical.
+
+## Blocks and their references
+
+Every block has a typed constructor, a docstring, `flops_per_token(context_length)` (6 FLOPs per
+weight used, plus the two attention matmuls) and a naive reference in `nanoscope/reference/`
+that its test compares it with. The references import only `torch` and `math`.
+
+| Family | Blocks | Reference |
+|---|---|---|
+| embedding | `TokenEmbedding`, `LearnedPosition`, `Head` (tied or untied) | `embed_one_hot`, `add_learned_position`, `tied_head` |
+| positional | `RoPE`, `NoPE` | `naive_rope` (complex rotation; scores depend only on distance) |
+| norm | `LayerNorm`, `RMSNorm` | `layer_norm`, `rms_norm` |
+| attention | `Attention(n_heads, n_kv_heads, pos, qk_norm, window, bias)`: MHA, GQA, MQA and sliding window | `naive_causal_attention` (loops over heads and positions) |
+| mlp | `GELUMLP(hidden, bias)`, `SwiGLU(hidden)` | `gelu`, `swiglu` |
+| structure | `Block(norm, attn, mlp, order)` (pre or post norm), `Decoder` | causality and an untrained loss near `ln(vocab)` |
+| primitive | `Linear`, `Activation`, `CausalMask`, `ScaledDotScores`, `Softmax`, `WeightedSum`, `SplitHeads`, `MergeHeads`, `Residual` | one formula each |
+
+`nanoscope blocks` prints the same table with every option, and `nanoscope blocks --json` prints
+it as `blocks.v1`. *Primitives* are the smallest pieces: build your own attention from them to
+see how the bigger blocks work. `Attention.attention_weights(x)` shows the softmax the fused
+kernel never exposes.
+
+## Templates
+
+A `Composite` is a block made of other blocks in named slots. A lesson template is a
+`Composite` with empty slots for you to fill:
+
+```python
+from nanoscope.blocks import Composite
+
+
+class PreNormAttention(Composite):
+    SLOTS = ("norm", "attn")
+
+    def forward(self, x):
+        return x + self.attn(self.norm(x))
+
+
+block = PreNormAttention(norm=RMSNorm(), attn=Attention(n_heads=4))
+```
+
+Each slot is a spec, built once per instance. A `Composite` has slots and a `forward`, not an
+`__init__`.
+
+## Register your own block
+
+```python
+from nanoscope.blocks import register_block
+
+def naive_gate(x, weight):
+    return x * torch.sigmoid(weight)
+
+@register_block(reference=naive_gate, family="mlp")
+class Gate(nn.Module):
+    """x times sigmoid of a learned vector."""
+    def __init__(self, d_model, context_length, scale=1.0): ...
+```
+
+The constructor takes `d_model` and `context_length` first, like every block; the rest are the
+options a model file writes as `Gate(scale=2.0)`. `reference` is a plain function that computes
+the same thing slowly and obviously. A block without one is listed as uncertified.
+`nanoscope blocks --workspace my_folder` finds registered blocks by reading the files.
+
+## Inspect
+
+| Command | Shows |
+|---|---|
+| `nanoscope graph file.py[:Class] [--json]` | the architecture graph, without running the file |
+| `nanoscope describe file.py:Class [--preset p] [--json] [--set k=v ...]` | shapes, parameters, FLOPs per token and memory per module, traced on the `meta` device (it imports your file). A shape error names the module and the line. |
+| `nanoscope blocks [--json] [--workspace dir]` | the palette |
+| `run(Model, block_stats=True)` | at each eval step, per block: activation size, gradient norm, update-to-weight ratio and attention entropy, appended to `blockstats.jsonl` |
+| `nanoscope status --blocks <run ref>` | the latest block statistics of a run |