Delete docs/v0.3.0-design.md and strip all references to it
CI / Format (ruff format) (push) Failing after 28s
CI / Lint (ruff check) (push) Successful in 29s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 33s
CI / Type check (ty) (push) Successful in 37s
CI / Format (ruff format) (pull_request) Failing after 37s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 37s
CI / Tests (pull_request) Successful in 2m49s
CI / Tests (push) Successful in 2m55s

The design doc and its followups doc are no longer needed as a live
reference now that the v0.3.0 redesign is implemented — comments and
docstrings across the codebase cited it extensively (file path, "design
doc §X.Y", "decision N", or bare "§X.Y" section numbers) as design
rationale. Removed docs/ and edited every citing comment/docstring to
drop the now-dangling reference while keeping the substantive
explanation next to it. CLAUDE.md's v0.3.0 roadmap bullet loses its
trailing pointer to the deleted file.

Verified: no remaining "docs/v0.3.0", "design doc", "decision N", or
"§N.N" references (repo-wide grep); ruff and ty clean; full test suite
on the heaviest-touched modules (network, sample, rollout, migration,
config, train) passes.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
2026-08-10 11:19:02 +02:00
parent f46628141d
commit 878e9ddca3
41 changed files with 263 additions and 1346 deletions
+1 -1
View File
@@ -91,6 +91,6 @@ GIANT is a conditional generative surrogate for the Geant4 step function. It rep
A sampling-calorimeter (multi-material) dataset is still a planned future direction, not yet built. See the knowledge base (`/home/lars/knowledge-base/meta/roadmap.md`).
**v0.3.0 — Stage-2 autoregressive redesign (designed, not implemented; branch `v0.3.0-stage2-autoregressive`):** the 2026-08-03 WGAN rollout benchmark failed specifically at the secondary-species level (zero photon secondaries, ~4M hallucinated `-14` muon antineutrinos). The agreed response pivots Stage 2 to **autoregressive generation** in descending-energy order with teacher forcing, and switches the particle-type representation back to **categorical** (top N1 by training-set count + an "other" bucket), reversing the 2026-07-17 continuous `(log-mass, charge)` target. This requires a config break: `[conditioning]` / `[stage1_model]` / `[stage2_model]` / `[train]` blocks replace the single global `train.mode` + `[model]`, so per-stage generators (`stage1 = flow` + `stage2 = wgan`), stage-2-only training, and one-shot-vs-autoregressive comparison are all expressible. `network.py` is refactored from ten permutation classes into composable parts (encoder × trunk × objective), which also makes routed WGAN work for the first time. **Full design contract, with every config option documented: `docs/v0.3.0-design.md` — read it before touching `giant/config.py` or `giant/model/network.py`.**
**v0.3.0 — Stage-2 autoregressive redesign (designed, not implemented; branch `v0.3.0-stage2-autoregressive`):** the 2026-08-03 WGAN rollout benchmark failed specifically at the secondary-species level (zero photon secondaries, ~4M hallucinated `-14` muon antineutrinos). The agreed response pivots Stage 2 to **autoregressive generation** in descending-energy order with teacher forcing, and switches the particle-type representation back to **categorical** (top N1 by training-set count + an "other" bucket), reversing the 2026-07-17 continuous `(log-mass, charge)` target. This requires a config break: `[conditioning]` / `[stage1_model]` / `[stage2_model]` / `[train]` blocks replace the single global `train.mode` + `[model]`, so per-stage generators (`stage1 = flow` + `stage2 = wgan`), stage-2-only training, and one-shot-vs-autoregressive comparison are all expressible. `network.py` is refactored from ten permutation classes into composable parts (encoder × trunk × objective), which also makes routed WGAN work for the first time.
**Condor-submitted GPU training/rollout (in progress, `condor-gpu-train-rollout` branch, not yet merged):** moves `giant train`/`giant rollout` off the shared portal GPU dev machines (see Compute environment) onto remote-GPU HTCondor submission on TOpAS/NEMO2 (`giant/condor.py`). Partway between "needs major features" and feature-complete — not ready to merge yet.
File diff suppressed because it is too large Load Diff
-53
View File
@@ -1,53 +0,0 @@
# v0.3.0 — post-implementation audit: open discrepancies
**Status:** steps 17 of `docs/v0.3.0-design.md` §12 are implemented (branch
`v0.3.0-stage2-autoregressive`, commits `eb6dd27`..`200c6d2`). This document
tracked discrepancies found between that implementation and the design contract
during a 2026-08-07 audit, as concrete work items. **All items (1-9) are now
resolved** — either implemented (1-5, 7-9) or explicitly deferred into
`docs/v0.3.0-design.md` §11.2 (6: `stage2_model.generator = "ddpm"`). Step 8
(`estimate_batch_size` recalibration) was intentionally still outstanding per
§12 and was never tracked here.
---
## Confirmed correct during the audit (no action needed)
For reference — these were explicitly checked against the design doc and
match it, including two spots the doc itself flagged as likely stale that
turned out fine:
- Config schema, `DEFAULT_CONFIG`, `migrate_config` table (§4), `save_config`/
`merge_cli_overrides` recursion, `default_out_dir_name`, `Conditioning` enum,
`n_sec.mode = "stop_token"` error (§9, §11.2).
- `network.py`'s full class decomposition (§5.3), dict-returning
`build_models`/`build_critics` (§5.4), `ExpertTrunk` separate in/out dims,
ST-Gumbel wiring (§2.1), AR token layout (§6.1), Markov/Attention history
encoders (§6.2).
- Shared PDG top-N type map (one map, not two, per §8), `other_policy`
sample/modal/drop (§11.1), embedding L1-nearest decode, and the L1-distance
diagnostic surfaced in `giant analyze` (§11.3).
- `analysis/render.py`/`analysis/router_gating.py` correctly branch
old-flat vs new-nested `model_config["router"]` location — doc flagged this
as a likely stale spot (§10) but it's actually fine.
- `giant/training/`'s per-stage trainers, mixed flow+wgan runs, WGAN critic cadence,
per-stage router auxiliary losses, stage-prefixed metrics, stage-2-only
training via ground-truth `x1_s1` (§7).
## Behavior change: WGAN best-checkpoint selection
The `giant/training/` split fixed a dead guard in WGAN validation scoring. A
WGAN stage was *meant* to contribute its marginal KL to the `val_loss` that
drives `best.pt`, but the guard `if n not in val_loss_per_stage` could never
fire (every stage was pre-seeded to `0.0`), so the stage contributed a flat
`0.0` and the KL was written to `metrics.csv` without ever being used.
`WGANStageTrainer.val_objective` now returns the KL as intended.
**Consequence:** any checkpoint selected before this commit under a config
with a WGAN stage — including the v0.3.0 default (`stage2_model.generator =
"wgan"`) — picked its best epoch on the non-adversarial stages alone. Measured
on the test harness's default flow+wgan config, `val_loss` went from `2.182`
(stage 1 only) to `15.137` (stage 1 + KL `12.954`), and which epoch won
changed. Do not compare `val/loss` or `best.pt` choice across this commit.
- `pipeline.py`'s deleted wgan+router rejection, per-stage `centers_init`
seeding, removed stale expert-size warning (§9, §10).
+1 -1
View File
@@ -72,7 +72,7 @@ class Bundle:
r_phys: pl.LazyFrame # rollout, physical steps only
t_phys: pl.LazyFrame # reference, physical steps only
checkpoint: str | None = None # from the rollout YAML; router_gating only
# §11.3 diagnostic pre-aggregated at rollout time (giant.rollout.
# Diagnostic pre-aggregated at rollout time (giant.rollout.
# L1DistCollector.summary()) — from the rollout YAML, type_embedding_l1_distance
# only. Unlike checkpoint/router_gating, this needs no live model: it's
# already a finished histogram, just passed through.
+1 -1
View File
@@ -83,7 +83,7 @@ _PLOT_META_KEYS = (
"best_val_loss",
"training_config",
"training_meta",
# §11.3 diagnostic — only present when giant rollout ran under
# Diagnostic — only present when giant rollout ran under
# stage2_model.particle_type.target="embedding" (see giant/cli.py's
# rollout command and giant.rollout.L1DistCollector); absent otherwise,
# which the type_embedding_l1_distance PlotSpec (catalog.py) reads as
+2 -2
View File
@@ -51,8 +51,8 @@ def _figure_params_v2(mc: dict, run_meta: dict) -> dict:
"""`_figure_params` for a new-shape (nested) `model_config` — has a
`stage1_model` key. Reports stage 1's architecture (the headline
generator); stage 2's generator is only added (`mode_s2`) when it
differs from stage 1's, since a mixed run (docs/v0.3.0-design.md's
`stage1=flow` + `stage2=wgan` case) is the interesting exception, not
differs from stage 1's, since a mixed run (the `stage1=flow` +
`stage2=wgan` case) is the interesting exception, not
the common case."""
s1 = mc["stage1_model"]
s2 = mc.get("stage2_model") or {}
+2 -2
View File
@@ -71,8 +71,8 @@ def _conditioning_axes(model_cfg: dict, default: str = "embedding") -> tuple[str
`giant.data.transforms.build_cond_features` from either a v0.2
checkpoint's flat `model_config["conditioning"]` (one shared string, same
for both axes) or a new-format one (independent
`model_config["conditioning"]["particle"/"material"]["type"]`
docs/v0.3.0-design.md §3.1: the two axes may differ). Mirrors
`model_config["conditioning"]["particle"/"material"]["type"]` the two
axes may differ). Mirrors
`giant.cli._conditioning_axes`."""
raw = model_cfg.get("conditioning", default)
if isinstance(raw, dict):
+1 -1
View File
@@ -1,4 +1,4 @@
"""Secondary-type embedding-distance diagnostic (docs/v0.3.0-design.md §11.3).
"""Secondary-type embedding-distance diagnostic.
Unlike every other diagnostic in this package, the data isn't derivable from
a rollout/reference parquet at all it's the L1 distance between each
+7 -9
View File
@@ -78,9 +78,8 @@ def _conditioning_axes(model_cfg: dict, default: str = "embedding") -> tuple[str
`giant.data.transforms.build_cond_features`/`build_features` from
either a v0.2 checkpoint's flat `model_config["conditioning"]` (one
shared string, same for both axes) or a new-format one (independent
`model_config["conditioning"]["particle"/"material"]["type"]`
docs/v0.3.0-design.md §3.1: the two axes are configured independently and
may differ)."""
`model_config["conditioning"]["particle"/"material"]["type"]` the two
axes are configured independently and may differ)."""
raw = model_cfg.get("conditioning", default)
if isinstance(raw, dict):
return (
@@ -390,7 +389,7 @@ def train(
typer.Option(
"--stage1-generator",
help="Stage 1's generative objective — overrides --mode for stage 1 "
"only (see docs/v0.3.0-design.md decision 7)",
"only",
),
] = None,
stage1_hidden_dim: Annotated[
@@ -725,8 +724,7 @@ def train(
# --hidden-dim/--n-blocks/--dropout are stage-1-only shorthands kept for
# backward compatibility (they predate stage2_model having its own
# flags); --stage1-*/--stage2-* below are the explicit, discoverable
# per-stage flags docs/v0.3.0-design.md decision 7 calls for, and take
# precedence when both are given.
# per-stage flags, and take precedence when both are given.
cli_stage1_model: dict[str, object] = {
k: v
for k, v in {
@@ -1360,7 +1358,7 @@ def predict(
)
if coord == Coord.global_:
# A fresh v0.3.0 Stage1Model owns no n_sec_head (decision 1)
# A fresh v0.3.0 Stage1Model owns no n_sec_head —
# sample_stage1 returns n_sec_pred=None then, so ask stage 2.
n_sec_pred = resolve_n_sec(
model, sec_decoder, cc, ck, stage1_norm, n_sec_pred
@@ -1764,7 +1762,7 @@ def rollout(
writer = pq.ParquetWriter(out, table.schema)
writer.write_table(table)
# Only meaningful under particle_type.target="embedding" (§11.3) — a
# Only meaningful under particle_type.target="embedding" — a
# no-op collector otherwise, cheaper than branching the call itself.
l1_dist_collector = L1DistCollector()
@@ -1820,7 +1818,7 @@ def rollout(
"rollout_seed": seed,
"n_rows": summary["n_rows"],
"termination_reason_counts": summary["termination_reason_counts"],
# §11.3 diagnostic — only present under
# Diagnostic — only present under
# stage2_model.particle_type.target="embedding"; omitted (not
# written as null) otherwise, so giant.analysis can tell "not
# applicable to this checkpoint" apart from "collector empty".
+6 -8
View File
@@ -325,8 +325,7 @@ def auto_device() -> torch.device:
# batch_size * hidden_dim * n_blocks (the ResBlock stack dominates), so this
# is a rough estimate rather than a guaranteed bound.
# NOTE: not yet recalibrated for the v0.3.0 autoregressive stage-2 trunk —
# see docs/v0.3.0-design.md §9/§11.3, deliberately last in the implementation
# order.
# deliberately last in the implementation order.
_REF_BYTES = 7683 * 1024**2
_REF_BATCH_SIZE = 29696
_REF_HIDDEN_DIM = 1024
@@ -528,10 +527,9 @@ def migrate_config(cfg: dict) -> dict:
Called on every config.toml load (see merge_cli_overrides) so old
training configs on disk keep working under new code without hand-
editing (decision 3, docs/v0.3.0-design.md §4). `[meta].config_version ==
CONFIG_VERSION` marks a dict as already-v0.3; its absence is read as
"this is v0.2" (the design doc's stated rule), so an already-migrated
dict is returned unchanged (deep-copied).
editing. `[meta].config_version == CONFIG_VERSION` marks a dict as
already-v0.3; its absence is read as "this is v0.2", so an
already-migrated dict is returned unchanged (deep-copied).
Only keys actually present in `cfg` are translated `cfg` may be a
partial file (e.g. `[train]\\nepochs = 5\\n` with no [model] section at
@@ -543,7 +541,7 @@ def migrate_config(cfg: dict) -> dict:
to set.
Operates on the config.toml shape. A checkpoint's `model_config` dict
(which additionally carries n_sec_head ownership, §4.1, and needs
(which additionally carries n_sec_head ownership and needs
`network.build_models`'s cooperation) is a separate migration surface,
deferred to the network.py refactor.
"""
@@ -729,7 +727,7 @@ def validate_config(cfg: dict) -> None:
raise ValueError(
"stage2_model.n_sec.mode = 'truth' is invalid for a "
"rollout-capable checkpoint (both stage1_model.active and "
"stage2_model.active = true — see docs/v0.3.0-design.md §9): "
"stage2_model.active = true): "
"'truth' takes n_sec from ground truth, which giant rollout "
"doesn't have. 'truth' is for standalone stage-2 evaluation "
"only — set stage1_model.active = false for that, or use "
+1 -2
View File
@@ -44,7 +44,6 @@ class StreamingStepsDataset(IterableDataset):
where:
cond_cont: (B, COND_DIM) float32
cond_cat: (B, 2/3/4) int64 width 2 unless conditioning="onehot"
(see docs/v0.3.0-design.md decision 4)
target_s1: (B, 9) float32 normalised Stage-1 primary target
n_sec: (B,) int64 true secondary count per step
sec_cont: (B, k_max, SEC_SLOT_DIM) float32 [stick_logit,
@@ -60,7 +59,7 @@ class StreamingStepsDataset(IterableDataset):
("onehot", "embedding"); zeros (unused) otherwise
`k_max` (constructor arg, default the module constant) should match
`stage2_model.k_max` (docs/v0.3.0-design.md §8) it sets the padded
`stage2_model.k_max` it sets the padded
width of `sec_cont`/`sec_type_idx` above.
"""
+5 -6
View File
@@ -176,7 +176,7 @@ def iter_file_chunks(
"""Yield one parquet row-group at a time so a large file never fully loads.
`k_max` sets the padded width of the sec_*_list columns (should match
`stage2_model.k_max` see docs/v0.3.0-design.md §8); defaults to the
`stage2_model.k_max`); defaults to the
module constant for callers that don't care (e.g. Stage-1-only reads)."""
pf = pq.ParquetFile(path)
for i in range(pf.num_row_groups):
@@ -277,8 +277,7 @@ def _topn_plus_other_map(counts: dict, n_classes: int) -> tuple[dict, dict]:
Returns `(class_map, other_members)` `other_members` is `{key: count}`
for every key bucketed into "other" (the empirical within-bucket
distribution, for `other_policy = "sample"` at rollout see
docs/v0.3.0-design.md §8).
distribution, for `other_policy = "sample"` at rollout).
"""
ranked = sorted(counts, key=lambda k: counts[k], reverse=True)
keep = ranked[: max(n_classes - 1, 0)]
@@ -324,7 +323,7 @@ def build_topn_map_from_files(
`_topn_plus_other_map`), generalized over the source column and key type.
Used for the material axis (`column="material"`, `cast=str`, matching
`mat_map`'s key type) — see docs/v0.3.0-design.md §8. The PDG axis uses
`mat_map`'s key type). The PDG axis uses
`build_pdg_topn_map_from_files` instead (it needs to pool two columns,
which this single-column form can't express). Also records
`other_members` (the empirical within-"other" distribution), needed
@@ -341,8 +340,8 @@ def build_pdg_topn_map_from_files(files: list[Path], n_classes: int) -> TopNMap:
plays in this dataset: a step's own primary particle (`pdg` column) and
an emitted secondary's species (`sec_pdg_list`, exploded) — shared by
`conditioning.particle.type = "onehot"` and
`stage2_model.particle_type.target = "onehot"` (docs/v0.3.0-design.md
§8). Pooling both is what keeps a species that's common as a secondary
`stage2_model.particle_type.target = "onehot"`. Pooling both is what
keeps a species that's common as a secondary
but rare as a primary (or vice versa) from being pushed into "other"
just because one role's count alone looks small — the meeting's failure
mode (zero photon secondaries, hallucinated antineutrinos) was
+3 -3
View File
@@ -104,7 +104,7 @@ def normalizer_key(
) -> str:
# .6g avoids float-repr drift (e.g. 0.1 vs 0.10000000000000002) causing
# spurious cache misses between runs with the "same" val_fraction. The two
# conditioning axes are independent (docs/v0.3.0-design.md §3.1) and both
# conditioning axes are independent and both
# affect which cond_cont columns are computed for real vs. zero-filled
# (giant.data.transforms._physical_cond_columns), so both must be part of
# the key or two mixed-axis runs could collide on the same cache entry.
@@ -114,7 +114,7 @@ def normalizer_key(
)
# Top-N-map axes (docs/v0.3.0-design.md §8): "pdg" keys match pdg_map's int
# Top-N-map axes: "pdg" keys match pdg_map's int
# keys (shared by conditioning.particle.type="onehot" and
# stage2_model.particle_type.target="onehot" — one map for both), "material"
# keys match mat_map's str keys.
@@ -190,7 +190,7 @@ class SetupCache:
proc_maps: dict[int, dict[str, int]] = field(default_factory=dict)
normalizers: dict[str, NormalizerEntry] = field(default_factory=dict)
topn_maps: dict[str, TopNMap] = field(default_factory=dict)
"""Keyed by `topn_key(axis, n_classes)` — see docs/v0.3.0-design.md §8."""
"""Keyed by `topn_key(axis, n_classes)`."""
@classmethod
def empty(cls, files: list[Path]) -> "SetupCache":
+16 -19
View File
@@ -561,9 +561,9 @@ def encode_secondary_type_idx(
`class_map` is either a top-N-plus-other map's `class_map`
(`stage2_model.particle_type.target = "onehot"`, see
`giant.data.loader.build_pdg_topn_map_from_files`) or the dense `pdg_map`
(`target = "embedding"`). Not used at all for `target = "physical"`
(see docs/v0.3.0-design.md decision 1) that target keeps using
`encode_secondaries`'s (log_mass, charge) columns unchanged.
(`target = "embedding"`). Not used at all for `target = "physical"`
that target keeps using `encode_secondaries`'s (log_mass, charge)
columns unchanged.
Padding slots get index 0 (their looked-up value is discarded downstream
by the `sec_valid`/`n_sec` mask regardless, so any in-vocabulary dummy
@@ -599,7 +599,7 @@ def decode_secondary_cont(
stick-breaking energy split and local->world direction generator/
`particle_type.target`-independent, since every target (`"physical"`,
`"onehot"`, `"embedding"`) shares the same `CONT_SLOT_DIM`-wide
(stick_logit, dir) prefix (docs/v0.3.0-design.md §6.1) and differs only
(stick_logit, dir) prefix and differs only
in what follows it. `decode_secondaries` (target="physical") is the
original all-in-one form built on top of this; `target` in `("onehot",
"embedding")` decodes their type slice separately via
@@ -727,9 +727,8 @@ def _physical_cond_columns(
) -> np.ndarray:
"""(N, PARTICLE_PHYS_DIM + MATERIAL_PHYS_DIM) physical conditioning columns.
The particle and material blocks are gated independently
(docs/v0.3.0-design.md §3.1: "configured independently and may mix
freely e.g. material `physical` with particle `embedding`"), so e.g.
The particle and material blocks are gated independently and may mix
freely e.g. material `physical` with particle `embedding` so e.g.
`particle_conditioning="embedding"` + `material_conditioning="physical"`
zero-fills only the particle columns and computes the material ones for
real.
@@ -802,16 +801,15 @@ def build_cond_features(
) -> tuple[np.ndarray, np.ndarray]:
"""Build conditioning arrays only — no target, no post-step variables.
`particle_conditioning`/`material_conditioning` are independent
(docs/v0.3.0-design.md §3.1) e.g. `particle_conditioning="embedding"` +
`particle_conditioning`/`material_conditioning` are independent
e.g. `particle_conditioning="embedding"` +
`material_conditioning="physical"` is a valid mix.
`pdg_topn_map`/`mat_topn_map` (a top-N-plus-other `class_map`, see
`giant.data.loader.build_topn_map_from_files`) append extra `cond_cat`
columns read by `ConditionEncoder`'s `"onehot"` mode
(docs/v0.3.0-design.md decision 4): pdg topN index at column 2 (iff
`pdg_topn_map` given), material topN index at column 3 (iff
`mat_topn_map` given, after column 2 if both are). Only ever given when
columns read by `ConditionEncoder`'s `"onehot"` mode: pdg topN index at
column 2 (iff `pdg_topn_map` given), material topN index at column 3
(iff `mat_topn_map` given, after column 2 if both are). Only ever given when
the corresponding axis is `"onehot"`; `cond_cat` stays `(N, 2)` otherwise.
"""
cond_cont = np.column_stack(
@@ -936,9 +934,8 @@ def build_features(
the secondary's real physical identity (from its ground-truth
PDG code), a fixed regression target, not a learned/snapped one.
Always computed the same way regardless of
`stage2_model.particle_type.target` (docs/v0.3.0-design.md
decision 1/3) only actually used downstream under `target =
"physical"`.
`stage2_model.particle_type.target` only actually used
downstream under `target = "physical"`.
proc_idx: (N,) integer process-class label (ProcessRouter supervision only
never conditioning). Zeros when `proc_map` is None or the loaded
data has no "process" column (e.g. pre-conversion parquet files).
@@ -966,7 +963,7 @@ def build_features(
dense `pdg_map` for `target = "embedding"` (pass `pdg_map` itself).
`None` for `target = "physical"`.
k_max: should match `stage2_model.k_max` (docs/v0.3.0-design.md §8)
k_max: should match `stage2_model.k_max`
overridden internally by `data["sec_E_list"]`'s own padded width when
present (the loader already padded it to some k_max; that width is
authoritative), so this only actually matters when secondary list
@@ -1015,7 +1012,7 @@ def build_features(
cat_cols.append(_vectorized_map_lookup(data["pdg"], pdg_topn_map))
if mat_topn_map is not None:
cat_cols.append(_vectorized_map_lookup(data["material"], mat_topn_map))
cond_cat = np.column_stack(cat_cols) # (N, 2/3/4) — see decision 4
cond_cat = np.column_stack(cat_cols) # (N, 2/3/4)
n_sec_raw = data["n_sec"].astype(
np.int64
@@ -1029,7 +1026,7 @@ def build_features(
# The loader already padded sec_*_list to some k_max (see
# giant.data.loader.iter_file_chunks); that padded width is
# authoritative over whatever this call happened to pass in, so the
# two can never drift apart (docs/v0.3.0-design.md §8).
# two can never drift apart.
k_max = sec_E_list.shape[1]
# Clamp the classification label to k_max: the head only has k_max+1
+47 -50
View File
@@ -21,7 +21,7 @@ from giant.constants import (
)
# ---------------------------------------------------------------------------
# Building blocks (docs/v0.3.0-design.md §5.2/§5.3)
# Building blocks
# ---------------------------------------------------------------------------
@@ -47,7 +47,7 @@ def cat_col_layout(
particle_type: str, material_type: str
) -> tuple[int | None, int | None]:
"""`cond_cat` column indices for each axis's top-N-onehot index, or
`None` if that axis isn't `"onehot"` (docs/v0.3.0-design.md decision 4).
`None` if that axis isn't `"onehot"`.
Columns 0/1 are always the dense pdg/material vocab index. The particle
top-N column (if any) comes next, then the material top-N column (if
@@ -92,9 +92,9 @@ class ConditionEncoder(nn.Module):
"""Fuses continuous conditioning with particle/material identity.
The particle and material axes are configured independently
(`particle_cfg`/`material_cfg`, each `{"type", "emb_dim", "n_layers"}`
see docs/v0.3.0-design.md §3.1) and may mix freely, e.g. material
"physical" with particle "embedding". Three modes per axis:
(`particle_cfg`/`material_cfg`, each `{"type", "emb_dim", "n_layers"}`)
and may mix freely, e.g. material "physical" with particle "embedding".
Three modes per axis:
- "embedding": a learned `nn.Embedding` lookup, indexed by `cond_cat`'s
dense training-vocab index. Memorizes the training menu.
- "physical": an `n_layers`-deep MLP over the axis's raw physical
@@ -104,8 +104,7 @@ class ConditionEncoder(nn.Module):
- "onehot": a fixed, unlearned one-hot vector over a top-N-plus-other
class map (`giant.data.loader.build_topn_map_from_files`/
`build_pdg_topn_map_from_files`), read from `cond_cat`'s extra
top-N-index column(s) see `_cat_col_layout` and
docs/v0.3.0-design.md decision 4.
top-N-index column(s) see `_cat_col_layout`.
"""
def __init__(
@@ -193,7 +192,7 @@ class ContextAdapter(nn.Module):
fixed-width context vector for a downstream stage's conditioning —
`stage2_model.context_dim`. Was `SecondaryConditionEncoder.stage1_proj`
(+ its `tanh`) in v0.2; pulled out as its own module in v0.3.0 since
`SecondaryConditionEncoder` as a wrapper class disappears (design doc §5.2)."""
`SecondaryConditionEncoder` as a wrapper class disappears."""
def __init__(self, in_dim: int, context_dim: int) -> None:
super().__init__()
@@ -223,7 +222,7 @@ class ResBlock(nn.Module):
# ---------------------------------------------------------------------------
# Routers — carried over unchanged from v0.2 (docs/v0.3.0-design.md §5.3)
# Routers — carried over unchanged from v0.2
# ---------------------------------------------------------------------------
@@ -615,7 +614,7 @@ def _build_router_from_cfg(
# ---------------------------------------------------------------------------
# Trunks (docs/v0.3.0-design.md §5.2 (b))
# Trunks
# ---------------------------------------------------------------------------
@@ -783,7 +782,7 @@ def build_trunk(
# ---------------------------------------------------------------------------
# History encoders — stage-2 autoregressive only (docs/v0.3.0-design.md §6.2)
# History encoders — stage-2 autoregressive only
# ---------------------------------------------------------------------------
@@ -791,8 +790,8 @@ class HistoryEncoder(nn.Module):
"""Interface for stage-2 autoregressive per-token history summaries:
`forward(feat, has_prev) -> (B, K, out_dim)`, a single parallel pass over
a full (teacher-forced) token sequence used by training. `MarkovHistory`
and `AttentionHistory` (docs/v0.3.0-design.md §6.2) are the two
implementations. Inference (`giant/sample.py`) generates one token at a
and `AttentionHistory` are the two implementations. Inference
(`giant/sample.py`) generates one token at a
time and cannot afford `forward`'s per-step cost to be O(K) (attention
would then be O(K^2) over a rollout's k_max loop); encoders that need
incremental state for that path additionally implement `init_cache`/
@@ -806,13 +805,13 @@ class HistoryEncoder(nn.Module):
class MarkovHistory(HistoryEncoder):
"""Summarizes the previous secondary's own `(energy_fraction, direction,
type_representation)` through one small MLP the "markov" history
(docs/v0.3.0-design.md §6.2): token i+1 only ever sees token i plus the
running scalars (`remaining_frac`/`slot_idx`, fused in separately by
type_representation)` through one small MLP the "markov" history:
token i+1 only ever sees token i plus the running scalars
(`remaining_frac`/`slot_idx`, fused in separately by
`Stage2Autoregressive._token_cond`), not the full prefix.
At slot 0 (`has_prev` False) substitutes a learned start vector rather
than zeros a reasonable default, not specified by the design doc.
than zeros a reasonable default.
"""
def __init__(self, in_dim: int, out_dim: int) -> None:
@@ -879,9 +878,9 @@ class _CausalAttnBlock(nn.Module):
class AttentionHistory(HistoryEncoder):
"""Causal self-attention over the emitted-token prefix
(docs/v0.3.0-design.md §6.2) the more expressive alternative to
`MarkovHistory`'s fixed previous-token-only summary. `feat`/`has_prev`
"""Causal self-attention over the emitted-token prefix — the more
expressive alternative to `MarkovHistory`'s fixed previous-token-only
summary. `feat`/`has_prev`
follow the same shifted-by-one convention `MarkovHistory` and
`Stage2Autoregressive._token_cond` use: `feat[:, i]` is token `i - 1`'s
own `(energy_fraction, direction, type_representation)`, with a learned
@@ -893,9 +892,9 @@ class AttentionHistory(HistoryEncoder):
`forward` is the parallel training path (one pass over the whole
teacher-forced sequence); `init_cache`/`step` are the incremental
inference path `giant/sample.py` uses, one new token per call, to avoid
re-encoding the whole prefix from scratch every slot (docs/v0.3.0-design.md
§10's "KV cache" note) — `step` must be called exactly once per slot (its
cache-extension is not idempotent), so a slot's output must be reused for
re-encoding the whole prefix from scratch every slot `step` must be
called exactly once per slot (its cache-extension is not idempotent),
so a slot's output must be reused for
every model call within that slot (`forward`'s ODE substeps, or a separate
`predict_type` call) rather than re-derived see
`Stage2Autoregressive.history_step`.
@@ -947,7 +946,7 @@ class AttentionHistory(HistoryEncoder):
# ---------------------------------------------------------------------------
# Stage models (docs/v0.3.0-design.md §5.3)
# Stage models
# ---------------------------------------------------------------------------
@@ -955,7 +954,7 @@ def stage2_type_dim(particle_type_cfg: dict, emb_dim: int) -> int:
"""Width of a single secondary slot's type slice —
`PARTICLE_PHYS_DIM` (log_mass, charge) for `target = "physical"`, else
`emb_dim` (both `"onehot"` class logits and `"embedding"` vectors are
`conditioning.particle.emb_dim` wide docs/v0.3.0-design.md §3.3)."""
`conditioning.particle.emb_dim` wide)."""
target = particle_type_cfg.get("target", "physical")
return PARTICLE_PHYS_DIM if target == "physical" else emb_dim
@@ -963,15 +962,15 @@ def stage2_type_dim(particle_type_cfg: dict, emb_dim: int) -> int:
def stage2_trunk_sec_dim(
particle_type_cfg: dict, generator: str, k_max: int, emb_dim: int
) -> int:
"""`Stage2OneShot`'s trunk output width (docs/v0.3.0-design.md decision 2).
"""`Stage2OneShot`'s trunk output width.
`target = "physical"` is untouched from v0.2/today (decision 1):
`target = "physical"` is untouched from v0.2/today:
`k_max * SEC_SLOT_DIM`, the type slice folded into the same
flow-matched/WGAN vector as the continuous stick/dir slots.
`target` in `("onehot", "embedding")`: under `generator == "wgan"` the
type slice is still folded in (adversarial for onehot via ST-Gumbel,
already-continuous for embedding §2.1), just `emb_dim` wide instead of
already-continuous for embedding), just `emb_dim` wide instead of
`PARTICLE_PHYS_DIM` wide: `k_max * (CONT_SLOT_DIM + emb_dim)`. Under
`generator in ("flow", "ddpm")` the type slice isn't part of this vector
at all it's `Stage2OneShot.type_head`'s job instead so the trunk
@@ -986,14 +985,14 @@ def stage2_trunk_sec_dim(
class Stage1Model(nn.Module):
"""Predicts the 9D primary post-step vector. No `n_sec_head` — decision 1
(docs/v0.3.0-design.md §2) moves it to stage 2, except for a migrated
v0.2 checkpoint (`n_sec_head_k_max` given), where it stays attached here
"""Predicts the 9D primary post-step vector. No `n_sec_head` — fresh runs
move it to stage 2, except for a migrated v0.2 checkpoint
(`n_sec_head_k_max` given), where it stays attached here
since that's where its weights live and what conditioning it was trained
against (see `_migrate_legacy_model_config`).
`cond_enc`, if given, is used in place of building a fresh
`ConditionEncoder` `conditioning.share_stages = true` (§3.1): `build_models`
`ConditionEncoder` `conditioning.share_stages = true`: `build_models`
constructs one shared instance and passes it to both stages, halving the
conditioning parameter count and forcing a common representation."""
@@ -1060,11 +1059,11 @@ class Stage1Model(nn.Module):
) -> torch.Tensor:
"""Return n_sec logits (B, K_MAX+1) from conditioning alone. Only
valid on a migrated v0.2 checkpoint's Stage1Model — fresh v0.3.0
configs predict n_sec from Stage2OneShot instead (decision 1)."""
configs predict n_sec from Stage2OneShot instead."""
if self.n_sec_head is None:
raise RuntimeError(
"this Stage1Model has no n_sec_head — n_sec now lives on "
"stage 2 by default (decision 1); this method only exists "
"stage 2 by default; this method only exists "
"for a migrated v0.2 checkpoint (legacy_owner='stage1')"
)
c_emb = self.cond_enc(cond_cont, cond_cat)
@@ -1073,17 +1072,16 @@ class Stage1Model(nn.Module):
class Stage2OneShot(nn.Module):
"""Predicts all `k_max` secondary slots simultaneously — v0.2 behaviour,
reproduced exactly (see docs/v0.3.0-design.md §12 step 2 acceptance
criterion; `decoder = "autoregressive"` is `Stage2Autoregressive`,
reproduced exactly (`decoder = "autoregressive"` is `Stage2Autoregressive`,
step 4/5, not implemented yet).
Owns `n_sec_head` by default (decision 1) unless `build_n_sec_head=False`
Owns `n_sec_head` by default unless `build_n_sec_head=False`
(a migrated v0.2 checkpoint, whose n_sec_head instead attaches to
Stage1Model see `_migrate_legacy_model_config`).
`particle_type_cfg["target"]` (default `"physical"`) selects the
secondary-type mechanism (docs/v0.3.0-design.md decision 2):
`"physical"` keeps the type slice folded into the trunk's own
secondary-type mechanism: `"physical"` keeps the type slice folded into
the trunk's own
flow-matched/WGAN output, unchanged from v0.2 (`sec_dim` computed by
the caller via `stage2_trunk_sec_dim` already reflects this). Under
`"onehot"`/`"embedding"` with `generator in ("flow", "ddpm")`, the type
@@ -1225,8 +1223,8 @@ class Stage2OneShot(nn.Module):
class Stage2Autoregressive(nn.Module):
"""Emits secondaries one at a time in descending-energy order
(docs/v0.3.0-design.md §6), instead of `Stage2OneShot`'s simultaneous
"""Emits secondaries one at a time in descending-energy order, instead
of `Stage2OneShot`'s simultaneous
k_max-slot prediction. `history` selects `MarkovHistory` or
`AttentionHistory` (`attn_n_heads`/`attn_n_layers`, attention only).
`teacher_forcing` handling lives entirely in the trainer
@@ -1555,7 +1553,7 @@ class CriticModel(nn.Module):
# ---------------------------------------------------------------------------
# v0.2 -> v0.3 checkpoint migration (docs/v0.3.0-design.md §4.1, §4.3)
# v0.2 -> v0.3 checkpoint migration
# ---------------------------------------------------------------------------
@@ -1568,14 +1566,14 @@ def _migrate_legacy_model_config(model_config: dict) -> dict:
Sets `stage2_model.n_sec.legacy_owner = "stage1"` so the n_sec_head
weights a v0.2 checkpoint carries on its Stage-1 module keep loading
there (design doc §4.1) instead of the new default location
there instead of the new default location
(`Stage2OneShot`) the n_sec head was trained against Stage 1's own
`ConditionEncoder` output, so it has to stay attached to Stage 1's
module, not just be labeled as such.
Only the monolithic (non-routed) trunk shape is exercised by the step-2
migration test (docs/v0.3.0-design.md §4.3); a routed v0.2 checkpoint
still builds correctly here (the router config passes through), but its
migration test; a routed v0.2 checkpoint still builds correctly here
(the router config passes through), but its
state dict isn't covered by `migrate_legacy_state_dict` below.
"""
m = model_config
@@ -1599,7 +1597,7 @@ def _migrate_legacy_model_config(model_config: dict) -> dict:
"hidden_dim/n_res_blocks), so this checkpoint's routed experts "
"have a different width/depth than the monolith — silently "
"dropping these keys would resize the experts instead of "
"refusing (docs/v0.3.0-design.md §4.2). This checkpoint can "
"refusing. This checkpoint can "
"only be loaded by v0.2 code."
)
router_cfg.setdefault("enabled", False)
@@ -1651,8 +1649,7 @@ def migrate_legacy_state_dict(
`(Stage1Model, Stage2OneShot)` module structure produced by
`build_models(_migrate_legacy_model_config(model_config))`.
Only the monolithic (non-routed) trunk shape is handled see
docs/v0.3.0-design.md §4.3's migration test scope.
Only the monolithic (non-routed) trunk shape is handled.
"""
def _trunk_prefix(k: str) -> str:
@@ -1682,7 +1679,7 @@ def migrate_legacy_state_dict(
# ---------------------------------------------------------------------------
# Factories (docs/v0.3.0-design.md §5.4)
# Factories
# ---------------------------------------------------------------------------
@@ -1698,7 +1695,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
(v0.2's actual — probably accidental — behaviour: two routers built from
one config with no semantic relationship between them).
`conditioning.share_stages = true` (§3.1) builds one `ConditionEncoder`
`conditioning.share_stages = true` builds one `ConditionEncoder`
instance here and passes it to both stages (`Stage1Model`/`Stage2OneShot`/
`Stage2Autoregressive`'s `cond_enc` param), instead of each stage
building its own halving the conditioning parameter count and forcing a
+3 -3
View File
@@ -90,8 +90,8 @@ def flow_matching_loss_secondary(
"physical"`'s width and the only case this function handled before
v0.3.0 step 4. `0` means no type slice is in `x1` at all (`target`
in `("onehot", "embedding")` under `generator in ("flow", "ddpm")`
see docs/v0.3.0-design.md decision 2, `Stage2OneShot.type_head`
handles the type loss separately in that case).
`Stage2OneShot.type_head` handles the type loss separately in that
case).
Only valid-slot dimensions contribute to the loss; padded slots are zeroed
before averaging, so the loss is not diluted by empty slots.
@@ -156,7 +156,7 @@ def flow_matching_loss_secondary_ar(
model call signature differs enough (four extra per-token conditioning
tensors) that merging would need an awkward shape-flag + closure.
Under teacher forcing (docs/v0.3.0-design.md §6.2 point 3) this is still a
Under teacher forcing this is still a
single parallel pass over all K_MAX tokens `x1`/`history_feat`/etc. are
already built from ground truth for every slot by the caller
(`giant.training.stage2_inputs._assemble_stage2_ar_inputs`/`_assemble_stage2_ar_target`).
+10 -10
View File
@@ -8,8 +8,8 @@ nuclide table doesn't cover — confirmed necessary for ~32% of the nuclear code
actually present in the multi-material dataset
(`0932fb02-f2ce-43ca-a4ef-60a2b1221bbc.parquet`).
Also holds the v0.3.0 stage-2 categorical-type rollout decode (§3.3/§8/§11.3
of docs/v0.3.0-design.md): `decode_topn_class`/`decode_embedding_nearest` turn
Also holds the v0.3.0 stage-2 categorical-type rollout decode:
`decode_topn_class`/`decode_embedding_nearest` turn
`Stage2Autoregressive`/`Stage2OneShot`'s `"onehot"`/`"embedding"` type
predictions back into concrete PDG codes, the one place a secondary's
categorical/continuous type representation is ever discretized (its
@@ -145,8 +145,8 @@ def decode_topn_class(
rng: np.random.Generator | None = None,
) -> np.ndarray:
"""`conditioning.particle.type` / `stage2_model.particle_type.target =
"onehot"` inference decode (docs/v0.3.0-design.md §3.3): per-row top-N
class index -> concrete PDG code.
"onehot"` inference decode: per-row top-N class index -> concrete PDG
code.
class_idx: int array, any shape, values in `[0, n_classes)`.
topn_map: the `TopNMap` (`giant.data.loader.build_pdg_topn_map_from_files`)
@@ -205,9 +205,9 @@ def decode_embedding_nearest(
emb_weight: np.ndarray,
idx_to_pdg: dict[int, int],
) -> tuple[np.ndarray, np.ndarray]:
"""`stage2_model.particle_type.target = "embedding"` inference decode
(docs/v0.3.0-design.md §3.3): L1-nearest row of the conditioning's own
particle embedding table, since a generative model's continuous output
"""`stage2_model.particle_type.target = "embedding"` inference decode:
L1-nearest row of the conditioning's own particle embedding table, since
a generative model's continuous output
essentially never lands within float tolerance of a table row (the exact-
match form is only valid as a round-trip test assertion, never here).
@@ -220,9 +220,9 @@ def decode_embedding_nearest(
idx_to_pdg: `invert_dense_map(pdg_map)` embedding row index -> PDG.
Returns `(pdg, l1_dist)`, both shaped like `vectors.shape[:-1]`. `l1_dist`
is the §11.3 diagnostic: a heavy tail means the decoder is emitting
vectors off the embedding manifold, the direct analogue of the species-
collapse symptom this redesign exists to fix.
is a diagnostic: a heavy tail means the decoder is emitting vectors off
the embedding manifold, the direct analogue of the species-collapse
symptom this redesign exists to fix.
"""
emb_dim = vectors.shape[-1]
flat = np.asarray(vectors, dtype=np.float64).reshape(-1, emb_dim)
+10 -10
View File
@@ -68,7 +68,7 @@ def _seed_energy_router(
"""Mutate `router_cfg["centers_init"]` in place from real data quantiles,
when this stage's router is an enabled EnergyRouter. Shared by both
stages' router configs — each seeded independently, since v0.3.0 stages
may have entirely different router configs (see docs/v0.3.0-design.md)."""
may have entirely different router configs."""
active = router_cfg.get("enabled") and router_cfg.get("type") == "energy"
if not active:
return
@@ -107,7 +107,7 @@ def run_setup_stage(
`cfg` is the full merged v0.3 config (`conditioning`/`stage1_model`/
`stage2_model`), already passed through `giant.config.validate_config`.
`conditioning.particle.type` and `conditioning.material.type` are
independent (docs/v0.3.0-design.md §3.1) and may differ.
independent and may differ.
Reads from and writes to the `giant.data.setup_cache` sidecar when
`cache_setup` is set (`rebuild_setup_cache` ignores but still
@@ -197,8 +197,8 @@ def run_setup_stage(
if cache is not None:
cache.proc_maps[n_experts] = proc_map
# Top-N-plus-other maps for onehot conditioning/type axes
# (docs/v0.3.0-design.md §8). The PDG axis is shared by
# Top-N-plus-other maps for onehot conditioning/type axes.
# The PDG axis is shared by
# conditioning.particle.type="onehot" and
# stage2_model.particle_type.target="onehot" (both key off
# conditioning.particle.emb_dim), so at most one PDG scan is needed even
@@ -403,11 +403,11 @@ def run_train_job(
setup.n_train_steps,
)
# cond_cat's onehot columns (docs/v0.3.0-design.md decision 4) are
# present per-axis, independently, under that axis's own
# conditioning.{particle,material}.type == "onehot" (§3.1: the two axes
# may mix freely). run_setup_stage builds each map whenever its own axis
# is "onehot" (see its own particle_cfg["type"]/material_cfg["type"]
# cond_cat's onehot columns are present per-axis, independently, under
# that axis's own conditioning.{particle,material}.type == "onehot"
# (the two axes may mix freely). run_setup_stage builds each map
# whenever its own axis is "onehot" (see its own
# particle_cfg["type"]/material_cfg["type"]
# checks), so they're guaranteed non-None here — asserted, not just
# assumed, so a future wiring bug fails loudly instead of silently
# dropping the onehot columns.
@@ -422,7 +422,7 @@ def run_train_job(
# The secondary type-index map depends on stage2_model.particle_type.target,
# independently of conditioning's own onehot/embedding choice above
# (docs/v0.3.0-design.md §3.3 — physical stays untouched/None).
# (physical stays untouched/None).
particle_type_target = (
cfg["stage2_model"].get("particle_type", {}).get("target", "physical")
)
+16 -19
View File
@@ -53,14 +53,14 @@ if TYPE_CHECKING:
class L1DistCollector:
"""Accumulates the §11.3 L1-distance diagnostic across a whole rollout
"""Accumulates the L1-distance diagnostic across a whole rollout
run: the L1 distance between each emitted secondary's raw predicted
embedding vector and the nearest table row it snapped to (only
meaningful under `particle_type.target = "embedding"`
`giant.particles.decode_embedding_nearest`). A heavy tail means the
decoder is emitting vectors off the embedding manifold the direct
analogue of the species-collapse symptom the v0.3.0 redesign exists to
fix (docs/v0.3.0-design.md §11.3).
fix.
Not folded into `rollout()`'s own return value (which is shape-typed as
step records, see `_RECORD_KEYS`/`RolloutSummary`) passed in and read
@@ -124,8 +124,7 @@ def decode_secondary_identity(
np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray | None
]:
"""Decode Stage 2's raw (sec_cont, sec_type) output into physical
secondary attributes, branching on `sec_decoder.particle_type_cfg`
(docs/v0.3.0-design.md §3.3):
secondary attributes, branching on `sec_decoder.particle_type_cfg`:
- `"physical"`: unchanged v0.2 path `sec_type` already *is* (log_mass,
charge), used as the secondary's identity as-is (no snapping).
@@ -138,8 +137,8 @@ def decode_secondary_identity(
- `"embedding"`: `sec_type` is a raw vector in the conditioning's own
embedding space `giant.particles.decode_embedding_nearest` L1-snaps
it to the nearest table row for the PDG (+ physics via
`particle_phys_array`), and also returns the L1 distance (§11.3
diagnostic see `giant/rollout.py`'s L1-distance accumulator).
`particle_phys_array`), and also returns the L1 distance (see this
module's `L1DistCollector`).
Returns (sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg,
sec_type_l1_dist) the last is `None` except under `"embedding"`.
@@ -480,24 +479,22 @@ def rollout(
`n_events * max_steps * avg_tracks_per_event`.
There is no `mode` parameter each stage's generative objective is read
directly off the model instance's own `generator_kind`
(docs/v0.3.0-design.md decision 2: stage 1 and stage 2 objectives are
independent, e.g. `stage1_model.generator="flow"` +
`stage2_model.generator="wgan"`), and the decoder (one-shot vs
directly off the model instance's own `generator_kind` (stage 1 and
stage 2 objectives are independent, e.g. `stage1_model.generator="flow"`
+ `stage2_model.generator="wgan"`), and the decoder (one-shot vs
autoregressive) is inferred from `sec_decoder`'s own class — see
`sample_stage1`/`sample_stage2` (giant.sample).
`pdg_topn_map`/`mat_topn_map` serve two independent purposes that happen
to share `pdg_topn_map` (docs/v0.3.0-design.md §8 one PDG map, not
two): they're required whenever `particle_conditioning`/
`material_conditioning` is `"onehot"` (feeds `build_cond_features`'s
extra `cond_cat` top-N columns §3.1), and `pdg_topn_map`/`other_policy`
are additionally read under `stage2_model.particle_type.target =
"onehot"` (§3.3, secondary-species decode). `seed` seeds the
`other_policy = "sample"` draw only (torch/numpy sampling itself is
seeded by the caller, same as today).
to share `pdg_topn_map` (one PDG map, not two): they're required
whenever `particle_conditioning`/`material_conditioning` is `"onehot"`
(feeds `build_cond_features`'s extra `cond_cat` top-N columns), and
`pdg_topn_map`/`other_policy` are additionally read under
`stage2_model.particle_type.target = "onehot"` (secondary-species
decode). `seed` seeds the `other_policy = "sample"` draw only
(torch/numpy sampling itself is seeded by the caller, same as today).
`l1_dist_collector`, if given, accumulates the §11.3 embedding-distance
`l1_dist_collector`, if given, accumulates the embedding-distance
diagnostic across the whole run see `L1DistCollector`. Only populated
under `particle_type.target = "embedding"`; a no-op otherwise.
"""
+23 -27
View File
@@ -10,10 +10,9 @@ def _predict_n_sec_if_owned(
model: torch.nn.Module, cond_cont: torch.Tensor, cond_cat: torch.Tensor
) -> torch.Tensor | None:
"""Stage-1 `n_sec_head` is only present on a migrated v0.2 checkpoint
(docs/v0.3.0-design.md decision 1 moves it to stage 2 for fresh runs
see `Stage1Model`'s docstring). `None` here means "ask stage 2 instead",
which every caller (`giant/rollout.py`, `giant/cli.py`) must do for a
fresh checkpoint."""
(fresh runs move it to stage 2 see `Stage1Model`'s docstring). `None`
here means "ask stage 2 instead", which every caller (`giant/rollout.py`,
`giant/cli.py`) must do for a fresh checkpoint."""
if getattr(model, "n_sec_head", None) is None:
return None
logits = model.predict_n_sec(cond_cont, cond_cat)
@@ -124,7 +123,7 @@ def _stage2_flat_width(sec_decoder: torch.nn.Module) -> int:
(continuous + type) under `particle_type.target = "physical"` or
`generator = "wgan"`, continuous-only otherwise (the type slice then
comes from `predict_type` instead see `stage2_trunk_sec_dim`'s
docstring, docs/v0.3.0-design.md decision 2)."""
docstring)."""
return stage2_trunk_sec_dim(
sec_decoder.particle_type_cfg,
sec_decoder.generator_kind,
@@ -147,9 +146,9 @@ def _decode_stage2_flat(
n_sec_pred: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Reshape a flat `(B, flat_width)` `Stage2OneShot` output into per-slot
tensors, generator/`particle_type.target`-agnostic (docs/v0.3.0-design.md
decision 2/3): shared by `sample_secondaries`/`sample_secondaries_wgan`,
which differ only in how `x` was produced.
tensors, generator/`particle_type.target`-agnostic: shared by
`sample_secondaries`/`sample_secondaries_wgan`, which differ only in how
`x` was produced.
Returns (sec_cont, sec_type, sec_valid):
sec_cont: (B, k_max, CONT_SLOT_DIM) [stick_logit, local_dir]
@@ -241,23 +240,22 @@ def sample_secondaries_ar(
n_sec_pred: torch.Tensor,
steps: int = 10,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""`Stage2Autoregressive` inference loop (docs/v0.3.0-design.md §6.4):
one token at a time, in descending-energy slot order, `k_max` sequential
calls. Unlike training (teacher forcing, §6.2 point 3 a single
parallel pass over ground-truth tokens, see
`giant.training.stage2_inputs._assemble_stage2_ar_inputs`), there is no ground truth at
inference: each token's conditioning is built free-running, from the
PREVIOUS TOKEN'S OWN just-generated output — the train/inference gap
§6.2 point 4 explicitly flags as the cost of markov history's
"""`Stage2Autoregressive` inference loop: one token at a time, in
descending-energy slot order, `k_max` sequential calls. Unlike training
(teacher forcing a single parallel pass over ground-truth tokens, see
`giant.training.stage2_inputs._assemble_stage2_ar_inputs`), there is no
ground truth at inference: each token's conditioning is built
free-running, from the PREVIOUS TOKEN'S OWN just-generated output — the
train/inference gap that is the cost of markov history's
expressiveness.
A `{flow,ddpm}` token costs `steps` ODE substeps; `wgan` costs one pass
§6.4's "K sequential forwards" cost note applies per-token here, not
the "K sequential forwards" cost applies per-token here, not
once, so a flow/ddpm AR run costs ~`k_max * steps` model calls per
physics step.
Under `history="attention"` the history encoding is computed once per
slot via `Stage2Autoregressive.history_step` (a KV-cache append, §10)
slot via `Stage2Autoregressive.history_step` (a KV-cache append)
rather than re-derived by every model call inside that slot so an ODE
loop's `steps` substeps, and the separate `predict_type` call when the
type slice isn't folded into the trunk output, all reuse the SAME `hist`
@@ -389,9 +387,9 @@ def sample_secondaries_ar(
# ---------------------------------------------------------------------------
# Per-stage dispatch — shared by giant/rollout.py and giant/cli.py's
# `predict` command, since both need "given a stage model, produce a
# sample" without hand-picking the sampler themselves (docs/v0.3.0-design.md
# decision 2: each stage's generative objective is independent, read off the
# model's own `generator_kind`, not a caller-supplied `mode` string).
# sample" without hand-picking the sampler themselves (each stage's
# generative objective is independent, read off the model's own
# `generator_kind`, not a caller-supplied `mode` string).
# ---------------------------------------------------------------------------
@@ -450,11 +448,10 @@ def resolve_n_sec(
) -> torch.Tensor:
"""`n_sec_pred` is already populated when `stage1_model` owns a legacy
`n_sec_head` (a migrated v0.2 checkpoint see `Stage1Model`'s
docstring); otherwise ask stage 2, which owns it by default under
decision 1 (docs/v0.3.0-design.md §2). Raises if neither stage owns a
head at all the only way that happens is `stage2_model.n_sec.mode`
other than `"head"` (`"truth"`/`"stop_token"`), neither of which is a
valid rollout-/predict-capable checkpoint (§3.3, §9)."""
docstring); otherwise ask stage 2, which owns it by default. Raises if
neither stage owns a head at all the only way that happens is
`stage2_model.n_sec.mode` other than `"head"` (`"truth"`/`"stop_token"`),
neither of which is a valid rollout-/predict-capable checkpoint."""
if n_sec_pred is not None:
return n_sec_pred
if getattr(sec_decoder, "n_sec_head", None) is None:
@@ -462,7 +459,6 @@ def resolve_n_sec(
"checkpoint has no n_sec_head on either stage — needs "
"stage2_model.n_sec.mode = 'head' (the default); 'truth' is "
"standalone-evaluation-only and 'stop_token' isn't implemented"
"(docs/v0.3.0-design.md §3.3/§9)"
)
logits = sec_decoder.predict_n_sec(cond_cont, cond_cat, stage1_out)
return logits.argmax(dim=-1)
+3 -4
View File
@@ -72,9 +72,8 @@ class _GracefulShutdown:
def _try_validate_marginals(trainer: StageTrainer, val_loader, device, **kwargs):
"""Runs `validate_marginals` on `trainer`'s sampling model (EMA model if
present, else the raw model). `validate_marginals` itself dispatches
through `giant.sample.sample_stage1`/`sample_stage2`/`resolve_n_sec`
(docs/v0.3.0-design.md §10), so this is generator- and
one-shot-vs-autoregressive-agnostic."""
through `giant.sample.sample_stage1`/`sample_stage2`/`resolve_n_sec`, so
this is generator- and one-shot-vs-autoregressive-agnostic."""
model = trainer.sampling_model()
return validate_marginals(model, val_loader, device=device, **kwargs)
@@ -124,7 +123,7 @@ def train(
wandb_log_every: int = 50,
) -> None:
"""Train whichever of stage1/stage2 are active, each through its own
`StageTrainer` (design doc §7). `models`/`critics` are the dicts
`StageTrainer`. `models`/`critics` are the dicts
`giant.model.network.build_models`/`build_critics` return a `None`
entry means that stage is `active = false`.
"""
+14 -14
View File
@@ -4,7 +4,7 @@ Pure functions, no optimizer/model state: they turn a batch's ground-truth
secondary tensors into the per-token targets and autoregressive conditioning
inputs `giant.training.trainers` feeds to `Stage2OneShot` /
`Stage2Autoregressive`. Split out of the trainers so the (target, generator,
decoder) width rules the fiddliest part of docs/v0.3.0-design.md §2.1/§6
decoder) width rules the fiddliest part of this codebase
live in one place and stay unit-testable on their own.
"""
@@ -65,8 +65,7 @@ def _assemble_stage2_ar_target(
analogue of `_assemble_stage2_real` (defined below in terms of this),
matching whatever width `Stage2Autoregressive`'s (or `Stage2OneShot`'s)
own trunk produces for this (target, generator) combination
(`giant.model.network.stage2_trunk_sec_dim`; docs/v0.3.0-design.md
decision 2/3):
(`giant.model.network.stage2_trunk_sec_dim`):
- `target = "physical"`: unchanged from v0.2 `sec_cont` (stick_logit,
dir, log_mass, charge) as-is.
@@ -116,7 +115,7 @@ def _stick_fraction(sec_cont: torch.Tensor) -> torch.Tensor:
def _remaining_energy_fraction(fraction: torch.Tensor) -> torch.Tensor:
"""(B, K_MAX) — fraction of the original e_sec budget unclaimed entering
slot i: `1.0` at `i=0`, `prod_{j<i}(1-fraction_j)` for `i>=1`
(docs/v0.3.0-design.md §6.3 "no re-derivation needed": the existing
("no re-derivation needed": the existing
stick-breaking encoding is already scale-free, so this is derivable from
the batch's ground-truth stick logits alone, no `e_sec` required)."""
cumprod = torch.cumprod(1.0 - fraction, dim=1)
@@ -166,7 +165,7 @@ def _assemble_stage2_ar_inputs(
) -> dict[str, torch.Tensor]:
"""Ground-truth per-token AR conditioning tensors — all `(B, K_MAX, ...)`
or `(B, K_MAX)`, built in one vectorized pass (teacher forcing means
every token's input is ground truth, docs/v0.3.0-design.md §6.2 point 3).
every token's input is ground truth).
Keys match `Stage2Autoregressive.forward`'s trailing kwargs."""
device = sec_cont.device
B, K = sec_cont.shape[0], sec_cont.shape[1]
@@ -187,9 +186,10 @@ def _stage2_tf_prob(
mode: str, p_start: float, p_end: float, epoch: int, total_epochs: int
) -> float:
"""P(condition slot k+1 on the TRUE token k rather than the model's own
prediction), for the current epoch (docs/v0.3.0-design.md §3.3
`stage2_model.autoregressive.teacher_forcing`). `"always"`/`"never"` are
the two degenerate constants; `"scheduled"` linearly interpolates
prediction), for the current epoch
(`stage2_model.autoregressive.teacher_forcing`).
`"always"`/`"never"` are the two degenerate constants; `"scheduled"`
linearly interpolates
`p_start` (epoch 0) to `p_end` (the final epoch) standard scheduled
sampling (Bengio et al. 2015)."""
if mode == "always":
@@ -240,7 +240,7 @@ def _assemble_stage2_ar_inputs_scheduled(
sample_steps: int,
) -> dict[str, torch.Tensor]:
"""Scheduled-sampling counterpart of `_assemble_stage2_ar_inputs`
(docs/v0.3.0-design.md §3.3 `teacher_forcing` = "scheduled"/"never"):
(`teacher_forcing` = "scheduled"/"never"):
each slot's history is the TRUE previous token with probability `p_tf`
(an independent per-example, per-slot Bernoulli draw) and the model's own
free-running prediction otherwise closing the train/inference gap that
@@ -252,8 +252,8 @@ def _assemble_stage2_ar_inputs_scheduled(
`giant.sample.sample_secondaries_ar` under `torch.no_grad()` not a
cheap one-step proxy, so building it costs the same `k_max` (`* steps`
for flow) sequential forwards `sample.py` pays at inference, EVERY batch
this is called on (§6.4's cost note, paid at train time too whenever
teacher_forcing != "always"). Fully detached: gradient only ever flows
this is called on (paid at train time too whenever teacher_forcing !=
"always"). Fully detached: gradient only ever flows
through the "real" target path each stage trainer already uses
(`_assemble_stage2_ar_target`), never through this self-sample.
"""
@@ -302,7 +302,7 @@ def _relax_onehot_type_slice(
) -> torch.Tensor:
"""Straight-through Gumbel-softmax relaxation of the per-slot type slice
inside a flattened `(B, k_max * (cont_dim + type_dim))` WGAN generator
output decision 5 (docs/v0.3.0-design.md §2.1): the forward pass is a
output: the forward pass is a
hard one-hot (matching what the critic sees from real data), the
backward pass flows smooth gradient. Continuous slots (stick/dir, and
the type slice itself under `target = "embedding"`, which never calls
@@ -311,8 +311,8 @@ def _relax_onehot_type_slice(
`grad_probe`, if given, gets `["cont"]`/`["type"]` populated with the L2
norm of the gradient reaching this split point during the next
`.backward()` call that touches it a backward hook, not a second
backward pass. This is the §11.4 differentiability validation-obligation
instrumentation (docs/v0.3.0-design.md): the trunk-gradient contribution
backward pass. This is the differentiability validation-obligation
instrumentation: the trunk-gradient contribution
from the type slice vs. the continuous slices, for
`particle_type.target="onehot"` + `generator="wgan"`. Only ever populated
on a `did_g_step` batch the critic step backprops through
+17 -19
View File
@@ -166,7 +166,7 @@ class StageSpec:
tf_p_end=ar_cfg.get("tf_p_end", 1.0),
# AR self-sampling under scheduled/never teacher forcing reuses
# train.validate_steps as its flow-matching ODE step count — no
# dedicated config key for this (docs/v0.3.0-design.md §3.3 lists
# dedicated config key for this (the autoregressive config lists
# tf_p_start/tf_p_end/attn_n_heads/attn_n_layers only).
ar_sample_steps=t["validate_steps"],
ddpm_n_steps=stage_cfg.get("ddpm", {}).get("n_steps", 1000),
@@ -184,12 +184,11 @@ class StageTrainer:
Reads only the shared batch tuple `(cond_cont, cond_cat, x1_s1, n_sec,
sec_cont, proc_idx, sec_type_idx)` stage 2 always conditions on the
ground-truth `x1_s1` (`stage2_model.stage1_context = "truth"`,
stage-level teacher forcing; `"sampled"` is not implemented see
docs/v0.3.0-design.md §3.3), so stage trainers never need each other's
output at train time. This means "stage-2-only training is a cheap
ablation, not new plumbing" (design doc §7) falls out for free: a trainer
only exists for active stages, and inactive stages are simply never
constructed.
stage-level teacher forcing; `"sampled"` is not implemented), so stage
trainers never need each other's output at train time. This means
"stage-2-only training is a cheap ablation, not new plumbing" falls out
for free: a trainer only exists for active stages, and inactive stages
are simply never constructed.
Grad-norm clipping is per-stage here v0.2's single shared optimizer
clipped both stages' gradients jointly; splitting per stage is a small,
@@ -381,7 +380,7 @@ class StageTrainer:
non-adversarial and WGAN trainers.
Gated on `n_sec_head is None`, not on `n_sec.mode`: a future
`mode="stop_token"` model (design doc §11.2, currently rejected in
`mode="stop_token"` model (currently rejected in
`validate_config`) carries no head and would train its EOS signal in
the generator/AR loss path instead, so this correctly stays zero.
"""
@@ -458,8 +457,7 @@ class FlowDDPMStageTrainer(StageTrainer):
raise NotImplementedError(
f"stage2_model.generator={spec.generator!r} is accepted by the "
"schema but not implemented in v0.3.0 for stage 2 (only "
"'flow' and 'wgan' have a stage-2 secondary-decoder loss"
"see docs/v0.3.0-design.md §11.2)"
"'flow' and 'wgan' have a stage-2 secondary-decoder loss)"
)
super().__init__(spec, model, device)
self.particle_type_lambda = self.particle_type_cfg.get("lambda", 1.0)
@@ -562,7 +560,7 @@ class FlowDDPMStageTrainer(StageTrainer):
):
"""CE (`target="onehot"`) or MSE (`target="embedding"`) loss for the
stage-2 model's `type_head` — the non-adversarial counterpart to
WGANStageTrainer's ST-Gumbel-into-the-critic path (decision 2/5).
WGANStageTrainer's ST-Gumbel-into-the-critic path.
Zero when this stage has no `type_head` (stage 1, or
`particle_type.target = "physical"`)."""
l_type = torch.zeros((), device=device)
@@ -774,8 +772,8 @@ class WGANStageTrainer(StageTrainer):
"grad_norm_g",
]
if self.is_stage2 and self.particle_type_cfg.get("target") == "onehot":
# §11.4 differentiability instrumentation — only meaningful when
# the type slice is a straight-through Gumbel relaxation.
# Differentiability instrumentation — only meaningful when the
# type slice is a straight-through Gumbel relaxation.
train_keys += ["grad_norm_type_slice", "grad_norm_cont_slice"]
self.train_metrics = [train_metric(key) for key in train_keys]
self.val_metrics = []
@@ -860,11 +858,11 @@ class WGANStageTrainer(StageTrainer):
)
if self.particle_type_cfg.get("target", "physical") == "onehot":
# Straight-through Gumbel-softmax relaxation of the type
# slice only (decision 5) — the critic must see a hard
# one-hot forward (matching what "real" data looks like)
# while gradient still flows smoothly to the generator.
# grad_probe captures the §11.4 gradient-magnitude
# instrumentation — see _relax_onehot_type_slice's docstring.
# slice only — the critic must see a hard one-hot forward
# (matching what "real" data looks like) while gradient
# still flows smoothly to the generator. grad_probe captures
# the gradient-magnitude instrumentation — see
# _relax_onehot_type_slice's docstring.
tau = _gumbel_tau(
global_step,
self.total_steps,
@@ -898,7 +896,7 @@ class WGANStageTrainer(StageTrainer):
)
# On a non-generator-step batch with no n_sec_head on this stage
# (n_sec now defaults to stage 2, decision 1), there's nothing for
# (n_sec now defaults to stage 2), there's nothing for
# the generator optimizer to do this batch — g_loss would otherwise
# be a graph-less zero tensor, which .backward() rejects outright.
skip_g_step = not did_g_step and self.model.n_sec_head is None
+1 -1
View File
@@ -85,7 +85,7 @@ def validate_marginals(
When `sec_decoder` is given, also validates Stage 2 via
`giant.sample.sample_stage2`/`resolve_n_sec` (generator- and
one-shot-vs-autoregressive-agnostic, docs/v0.3.0-design.md §10): n_sec
one-shot-vs-autoregressive-agnostic): n_sec
distribution (+ classification accuracy), per-slot energy-fraction
marginals, and a particle-type marginal whose shape depends on
`sec_decoder.particle_type_cfg["target"]` restricted to each side's own
+4 -6
View File
@@ -1,6 +1,5 @@
"""Portal-machine follow-up for v0.3.0 step 2 (docs/v0.3.0-design.md §4.3):
diff a real v0.2 checkpoint's outputs against the new `build_models` on the
same input batch.
"""Portal-machine follow-up for v0.3.0 step 2: diff a real v0.2 checkpoint's
outputs against the new `build_models` on the same input batch.
`tests/test_migration_v02_v03.py` already proves this bit-identical with
synthetic random weights, but that test can't run where it matters (no
@@ -16,7 +15,7 @@ Usage (from the repo root, on a portal machine):
uv run python scripts/check_migration_v02_v03.py /ceph/lbogner/.../best.pt --batch 32 --seed 1
Run it once against a flow (or ddpm) checkpoint and once against a wgan
checkpoint (design doc §4.3's "one flow checkpoint and one WGAN checkpoint").
checkpoint ("one flow checkpoint and one WGAN checkpoint").
A routed checkpoint (`model_config["router"]["enabled"]`) is only checked for
successful construction `giant.model.network.migrate_legacy_state_dict`
doesn't yet remap routed (Expert-per-router) state dicts, so the
@@ -115,8 +114,7 @@ def main() -> int:
print(
" routed checkpoint: migrate_legacy_state_dict only handles the "
"monolithic trunk shape — verifying construction only, skipping "
"the bit-identical weight/output comparison. See "
"docs/v0.3.0-design.md §2.4's scope note."
"the bit-identical weight/output comparison."
)
print("PASS (construction only, routed checkpoint)")
return 0
+1 -1
View File
@@ -536,7 +536,7 @@ def warm_cache(
"--material-conditioning",
help="Must match the `giant train` run(s)' conditioning.material.type "
"to warm for — independent of --particle-conditioning "
"(docs/v0.3.0-design.md §3.1: the two axes may differ)",
"(the two axes may differ)",
),
] = Conditioning.physical,
router: Annotated[
+2 -2
View File
@@ -31,8 +31,8 @@ def run_warm_setup_cache(
select the normalizer cache entry
(`giant.data.setup_cache.normalizer_key`) pass the same values a later
`giant train` invocation will use so it hits this warmed entry. The two
conditioning axes are independent (docs/v0.3.0-design.md §3.1) and may
differ. `router_enabled`/`router_type`/`n_experts` only matter for
conditioning axes are independent and may differ.
`router_enabled`/`router_type`/`n_experts` only matter for
`router_type == "process"` (warms that `n_experts`'s process map); the
energy-router quantile summary is always collected regardless, so a
later `--router-type energy` run never needs to rescan just to seed
+4 -4
View File
@@ -1,13 +1,13 @@
"""Frozen snapshot of `giant/model/network.py` as it stood at the v0.3.0
"step 1" commit (eb6dd27), i.e. the last commit before the step-2 §5
decomposition (see `docs/v0.3.0-design.md`).
"step 1" commit (eb6dd27), i.e. the last commit before the step-2
composable-parts decomposition.
This is a deliberate verbatim copy, not an import of the live module the
whole point is that this file's classes keep behaving exactly as v0.2 did
even after `giant/model/network.py` itself is rewritten, so
`tests/test_migration_v02_v03.py` has a stable "old" side to diff the new
`build_models`/`Stage1Model`/`Stage2OneShot` against (design doc §4.3's
bit-identical acceptance test). Do not edit this file to track future
`build_models`/`Stage1Model`/`Stage2OneShot` against (the bit-identical
acceptance test). Do not edit this file to track future
`network.py` changes it exists specifically to stop tracking them.
"""
+2 -3
View File
@@ -1,6 +1,5 @@
"""Tests for `giant train`'s stage-prefixed CLI flags (docs/v0.3.0-design.md
decision 7 / docs/v0.3.0-followups.md item 2): --stage1-*/--stage2-* must
independently override each stage's config block, and must take precedence
"""Tests for `giant train`'s stage-prefixed CLI flags: --stage1-*/--stage2-*
must independently override each stage's config block, and must take precedence
over the older shared flags (--mode/--hidden-dim/--n-critic/... ) that still
apply the same value to both stages for backward compatibility."""
+4 -4
View File
@@ -585,8 +585,8 @@ def test_validate_config_embedding_target_passes_with_embedding_conditioning():
def test_validate_config_mixed_particle_material_conditioning_is_valid():
"""docs/v0.3.0-design.md §3.1: the particle and material conditioning
axes are configured independently and may mix freely e.g. material
"""The particle and material conditioning axes are configured
independently and may mix freely e.g. material
"physical" with particle "embedding" and the data pipeline
(giant/data/transforms.py) now implements that end-to-end, so
validate_config must not reject it."""
@@ -638,8 +638,8 @@ def test_validate_config_stop_token_not_implemented():
def test_validate_config_n_sec_truth_rejected_for_rollout_capable_checkpoint():
"""docs/v0.3.0-design.md §9: 'n_sec.mode = "truth" is invalid for a
rollout-capable checkpoint' — both stages active means giant rollout
"""'n_sec.mode = "truth" is invalid for a rollout-capable checkpoint'
both stages active means giant rollout
could load this checkpoint, but 'truth' has no ground truth to draw
n_sec from at rollout time."""
cfg = _cfg_with(
+1 -2
View File
@@ -210,8 +210,7 @@ def test_build_topn_map_from_files_fewer_values_than_n_classes(tmp_path):
def test_build_pdg_topn_map_from_files_pools_primary_and_secondary_pdg(tmp_path):
"""A species that's rare as a primary but common as a secondary must
still rank by its pooled (primary + secondary) count, not just its
primary-role count alone the whole point of pooling both roles
(docs/v0.3.0-design.md §8)."""
primary-role count alone the whole point of pooling both roles."""
path = tmp_path / "a.parquet"
# primary pdg: mostly 11 (electron), one lone 22 (photon)
pdg = [11] * 5 + [22] * 1
+13 -15
View File
@@ -1,13 +1,12 @@
"""Migration acceptance test for v0.3.0 step 2 (docs/v0.3.0-design.md §4.3,
§12 step 2): "load a v0.2 checkpoint through migrate_config + the new
build_models, and diff its outputs against v0.2 code on the same input
batch bit-identical, or the refactor has changed something it should not
have."
"""Migration acceptance test for v0.3.0 step 2: "load a v0.2 checkpoint
through migrate_config + the new build_models, and diff its outputs against
v0.2 code on the same input batch bit-identical, or the refactor has
changed something it should not have."
No `/ceph` access on this machine (see CLAUDE.md's Compute environment
section), so a real trained checkpoint can't be used here — see
docs/v0.3.0-design.md's plan for the separate portal-machine follow-up with a
real checkpoint. This test is the synthetic stand-in: build a v0.2-shaped
section), so a real trained checkpoint can't be used here — a separate
portal-machine follow-up with a real checkpoint is planned instead. This
test is the synthetic stand-in: build a v0.2-shaped
model from the frozen `tests/legacy/network_v02_snapshot.py` classes with
fixed-seed random weights (playing the role of "a v0.2 checkpoint"), migrate
its config and remap its state dict onto the new `build_models` output, and
@@ -134,7 +133,7 @@ def _run_migration_check(mode: str, conditioning: str) -> None:
assert isinstance(new_stage1, net.Stage1Model)
assert isinstance(new_stage2, net.Stage2OneShot)
# legacy_owner="stage1": n_sec lives on stage1, not stage2, for a
# migrated v0.2 checkpoint (design doc §4.1).
# migrated v0.2 checkpoint.
assert new_stage1.n_sec_head is not None
assert new_stage2.n_sec_head is None
@@ -199,12 +198,11 @@ def test_migrate_legacy_model_config_shape():
def test_migrate_legacy_model_config_nonzero_expert_dims_raises():
"""docs/v0.3.0-followups.md item 8 regression: a v0.2 checkpoint's
model_config carrying a non-default expert_hidden_dim/expert_n_blocks
must fail loudly through this path too (§4.2) not just
giant.config.migrate_config's parallel TOML-load path. Silently dropping
these keys (build_router's kwarg filtering) would resize the experts
instead of refusing."""
"""Regression: a v0.2 checkpoint's model_config carrying a non-default
expert_hidden_dim/expert_n_blocks must fail loudly through this path too
not just giant.config.migrate_config's parallel TOML-load path.
Silently dropping these keys (build_router's kwarg filtering) would
resize the experts instead of refusing."""
legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical")
legacy_cfg["router"] = {
"enabled": True,
+9 -9
View File
@@ -84,7 +84,7 @@ def test_stage1_model_gradients_flow():
def test_stage1_model_no_n_sec_head_by_default():
"""Fresh v0.3.0 construction (no n_sec_head_k_max) has no n_sec head —
decision 1 (docs/v0.3.0-design.md §2) moves it to stage 2."""
it moves to stage 2."""
model = Stage1Model(
pdg_vocab=3, mat_vocab=2, particle_cfg=PARTICLE_CFG, material_cfg=MATERIAL_CFG
)
@@ -206,7 +206,7 @@ def test_condition_encoder_onehot_is_a_true_one_hot_vector():
assert torch.all(pdg_e.sum(dim=-1) == 1.0)
# --- Stage2OneShot particle_type architecture (docs/v0.3.0-design.md decision 2) --
# --- Stage2OneShot particle_type architecture --------------------------------
def _build_stage2(target: str, generator: str, emb_dim: int = 6) -> Stage2OneShot:
@@ -300,7 +300,7 @@ def test_stage2_oneshot_forward_shape_onehot_flow_excludes_type():
assert out.shape == (B, k_max * CONT_SLOT_DIM)
# --- MarkovHistory (docs/v0.3.0-design.md §6.2) -----------------------------
# --- MarkovHistory -----------------------------------------------------------
def test_markov_history_shape():
@@ -314,8 +314,8 @@ def test_markov_history_shape():
def test_markov_history_uses_start_vector_when_no_prev():
"""Slot 0's own raw feature must be ignored — a learned start vector is
substituted there instead (a reasonable default not specified by the
design doc, see Stage2Autoregressive's docstring)."""
substituted there instead (a reasonable default, see
Stage2Autoregressive's docstring)."""
hist = MarkovHistory(in_dim=4, out_dim=6)
B, K = 2, 3
has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
@@ -328,7 +328,7 @@ def test_markov_history_uses_start_vector_when_no_prev():
assert torch.allclose(out_a[:, 1:], out_b[:, 1:])
# --- AttentionHistory (docs/v0.3.0-design.md §6.2, v0.3.0 step 7) ----------
# --- AttentionHistory (v0.3.0 step 7) ---------------------------------------
def test_attention_history_shape():
@@ -392,7 +392,7 @@ def test_attention_history_step_matches_forward():
assert torch.allclose(stepped, expected, atol=1e-5)
# --- Stage2Autoregressive (docs/v0.3.0-design.md §6, v0.3.0 step 5) ---------
# --- Stage2Autoregressive (v0.3.0 step 5) -----------------------------------
def _build_stage2_ar(
@@ -671,8 +671,8 @@ def test_build_models_share_stages_true_shared_params_are_in_both_stage_paramete
"""The shared encoder's parameters must actually appear in both stages'
own `.parameters()` that's what makes each stage's independent
optimizer include (and update) them, which is the actual mechanism behind
"shared weights, forced common representation" (docs/v0.3.0-design.md
§3.1), not just object identity on `.cond_enc`."""
"shared weights, forced common representation", not just object identity
on `.cond_enc`."""
built = build_models(_minimal_model_config(share_stages=True))
stage1, stage2 = built["stage1"], built["stage2"]
assert stage1 is not None and stage2 is not None
+5 -5
View File
@@ -155,7 +155,7 @@ def test_run_train_job_second_run_hits_cache(tmp_path, data, monkeypatch):
def test_run_train_job_builds_caches_and_persists_pdg_topn_map(tmp_path, data):
"""DEFAULT_CONFIG's stage2_model.particle_type.target defaults to
"onehot" (docs/v0.3.0-design.md §3.3/§8) a plain _tiny_cfg() run must
"onehot" a plain _tiny_cfg() run must
build the shared pdg top-N map, cache it in the setup-cache sidecar, and
persist it into the checkpoint, with no extra config needed."""
echo1 = _run(data, tmp_path / "out1")
@@ -242,7 +242,7 @@ def test_run_train_job_new_val_fraction_is_partial_hit(tmp_path, data, monkeypat
def test_run_train_job_custom_k_max_end_to_end(tmp_path, data):
"""docs/v0.3.0-followups.md item 3 regression: stage2_model.k_max other
"""Regression: stage2_model.k_max other
than the K_MAX module constant's default (15) must not produce a shape
mismatch between the data pipeline (loader.py/transforms.py padding) and
the model (network.py's trunks, sized from this same config value)."""
@@ -254,9 +254,9 @@ def test_run_train_job_custom_k_max_end_to_end(tmp_path, data):
def test_run_train_job_mixed_particle_material_conditioning_end_to_end(tmp_path, data):
"""docs/v0.3.0-followups.md item 4 regression: conditioning.particle.type
"""Regression: conditioning.particle.type
and conditioning.material.type are configured independently and may mix
freely (docs/v0.3.0-design.md §3.1) e.g. particle "embedding" with
freely e.g. particle "embedding" with
material "physical" end-to-end through the real data pipeline, not
just accepted by validate_config."""
cfg = _tiny_cfg()
@@ -286,7 +286,7 @@ def test_run_train_job_mixed_particle_material_conditioning_end_to_end(tmp_path,
def test_run_train_job_share_stages_end_to_end(tmp_path, data):
"""docs/v0.3.0-followups.md item 5 regression: conditioning.share_stages
"""Regression: conditioning.share_stages
= true must actually train (not raise NotImplementedError), and the
resulting checkpoint's two stages must reload into a single shared
ConditionEncoder instance rather than two independent ones."""
+10 -10
View File
@@ -373,7 +373,7 @@ def _models_v3(
particle_cfg = {"type": conditioning, "emb_dim": emb_dim, "n_layers": 1}
material_cfg = {"type": conditioning, "emb_dim": emb_dim, "n_layers": 1}
# A fresh v0.3.0 Stage1Model — no n_sec_head_k_max, unlike _models() above
# (decision 1 moves n_sec ownership to stage 2 by default).
# (n_sec ownership moves to stage 2 by default).
s1 = Stage1Model(
pdg_vocab=3,
mat_vocab=2,
@@ -470,9 +470,9 @@ def _run_v3(
def test_rollout_stage2_owns_n_sec_when_stage1_has_no_head(fake_material_props):
"""A fresh v0.3.0 Stage1Model has no n_sec_head (decision 1) — n_sec
must come from Stage2's own head instead, and the run must still
complete and conserve energy."""
"""A fresh v0.3.0 Stage1Model has no n_sec_head — n_sec must come from
Stage2's own head instead, and the run must still complete and
conserve energy."""
s1, s2 = _models_v3()
rec = _run_v3(s1, s2)
assert len(rec["event_id"]) > 0
@@ -501,7 +501,7 @@ def test_rollout_physical_target_decoder_generator_matrix(
):
"""Every (decoder, stage2 generator) combination under
particle_type.target="physical" must run to completion and conserve
energy the matrix docs/v0.3.0-design.md §7 calls out for comparison."""
energy."""
s1, s2 = _models_v3(decoder=decoder, generator2=generator2)
rec = _run_v3(s1, s2)
assert len(rec["event_id"]) > 0
@@ -555,10 +555,10 @@ def test_rollout_onehot_target_missing_topn_map_raises(fake_material_props):
_run_v3(s1, s2, pdg_topn_map=None)
# --- conditioning.{particle,material}.type = "onehot" (docs/v0.3.0-followups.md
# item 7) — a separate axis from stage2_model.particle_type.target above: this
# is what feeds cond_cat's extra top-N columns for ConditionEncoder's own
# "onehot" mode, not the secondary-species decode. ---------------------------
# --- conditioning.{particle,material}.type = "onehot" — a separate axis from
# stage2_model.particle_type.target above: this is what feeds cond_cat's
# extra top-N columns for ConditionEncoder's own "onehot" mode, not the
# secondary-species decode. ---------------------------------------------
COND_PDG_TOPN_MAP = TopNMap(class_map=dict(PDG_MAP), other_members={})
COND_MAT_TOPN_MAP = TopNMap(class_map={"G4_AIR": 0, "G4_PbWO4": 1}, other_members={})
@@ -653,7 +653,7 @@ def test_rollout_embedding_target_end_to_end(decoder):
def test_l1_dist_collector_populated_only_for_embedding_target():
"""§11.3: the L1-distance diagnostic only makes sense under
"""The L1-distance diagnostic only makes sense under
particle_type.target="embedding" a physical-target run must leave the
collector empty rather than silently accumulating garbage."""
s1, s2 = _models_v3(target="physical")
+4 -4
View File
@@ -896,8 +896,8 @@ def test_build_models_routed_pair_composed_router_is_drop_in_for_sample_flow():
cond_cont, cond_cat = _cond(B, pdg=3, mat=2)
stage1_norm, n_sec_pred = sample_flow(stage1, cond_cont, cond_cat, steps=2)
assert stage1_norm.shape == (B, X_DIM)
# A fresh v0.3.0 Stage1Model has no n_sec_head (decision 1 moves it to
# stage 2) — sample_flow returns n_sec_pred=None here, and n_sec must be
# A fresh v0.3.0 Stage1Model has no n_sec_head (it moves to stage 2) —
# sample_flow returns n_sec_pred=None here, and n_sec must be
# asked of stage2 instead, using the just-sampled stage1_norm as context.
assert n_sec_pred is None
n_sec_pred = stage2.predict_n_sec(cond_cont, cond_cat, stage1_norm).argmax(dim=-1)
@@ -1092,8 +1092,8 @@ def test_build_models_routed_pair_is_drop_in_for_sample_flow():
cond_cont, cond_cat = _cond(B, pdg=3, mat=2)
stage1_norm, n_sec_pred = sample_flow(stage1, cond_cont, cond_cat, steps=2)
assert stage1_norm.shape == (B, X_DIM)
# A fresh v0.3.0 Stage1Model has no n_sec_head (decision 1 moves it to
# stage 2) — sample_flow returns n_sec_pred=None here, and n_sec must be
# A fresh v0.3.0 Stage1Model has no n_sec_head (it moves to stage 2) —
# sample_flow returns n_sec_pred=None here, and n_sec must be
# asked of stage2 instead, using the just-sampled stage1_norm as context.
assert n_sec_pred is None
n_sec_pred = stage2.predict_n_sec(cond_cont, cond_cat, stage1_norm).argmax(dim=-1)
+3 -3
View File
@@ -1,6 +1,6 @@
"""Tests for giant/sample.py's v0.3.0 stage-model sampling — the AR loop
(`sample_secondaries_ar`) and non-"physical" `particle_type.target` coverage
for the one-shot samplers (docs/v0.3.0-design.md step 6)."""
for the one-shot samplers."""
import pytest
import torch
@@ -38,7 +38,7 @@ def _cond(B: int, pdg: int = 3, mat: int = 2) -> tuple[torch.Tensor, torch.Tenso
def _conditioning_for(target: str) -> str:
# target="embedding" regresses against the conditioning's own embedding
# table (docs/v0.3.0-design.md §3.3) — only meaningful when the
# table — only meaningful when the
# conditioning axis is itself "embedding".
return "embedding" if target == "embedding" else "physical"
@@ -103,7 +103,7 @@ def _expected_type_dim(target: str, emb_dim: int) -> int:
return PARTICLE_PHYS_DIM if target == "physical" else emb_dim
# ── Stage-1 n_sec ownership (decision 1) ────────────────────────────────────
# ── Stage-1 n_sec ownership ──────────────────────────────────────────────────
def test_sample_flow_returns_none_n_sec_when_stage1_owns_no_head():
+6 -6
View File
@@ -90,7 +90,7 @@ def test_wandb_run_config_handles_missing_model_config():
assert wcfg["model_config"] == {}
# --- AR helper functions (v0.3.0 step 5, docs/v0.3.0-design.md §6) ---------
# --- AR helper functions (v0.3.0 step 5) ---------
def test_stick_fraction_matches_sigmoid_of_logit():
@@ -119,7 +119,7 @@ def test_ar_has_prev_false_only_at_slot_zero():
assert has_prev.tolist() == [[False, True, True, True, True]]
# --- _stage2_tf_prob (docs/v0.3.0-design.md §3.3, v0.3.0 step 7) -----------
# --- _stage2_tf_prob (v0.3.0 step 7) -----------
def test_stage2_tf_prob_always_is_constant_one():
@@ -569,7 +569,7 @@ def test_metrics_csv_columns_are_stage_prefixed():
def test_wgan_stage_trainer_skips_generator_step_when_no_grad_this_batch():
"""Regression test: on a non-generator-step batch, if this stage's model
has no n_sec_head (n_sec defaults to stage 2, decision 1), g_loss is a
has no n_sec_head (n_sec defaults to stage 2), g_loss is a
graph-less zero .backward() must not be called on it."""
cfg = _base_cfg()
cfg["stage1_model"]["generator"] = "wgan"
@@ -672,11 +672,11 @@ def test_train_end_to_end_ar_attention_history_scheduled_teacher_forcing(
def test_ar_wgan_onehot_grad_norm_instrumentation_populates_metrics():
"""§11.4 differentiability validation-obligation instrumentation: the
"""Differentiability validation-obligation instrumentation: the
trunk-gradient-norm-by-slice columns must appear and actually fire for
generator='wgan' + particle_type.target='onehot' under decoder=
'autoregressive' (added at v0.3.0 step 5 per the design doc's
instruction to accrue evidence during the architecture comparison)."""
'autoregressive' (added at v0.3.0 step 5 to accrue evidence during the
architecture comparison)."""
cfg = _base_cfg()
cfg["stage2_model"]["decoder"] = "autoregressive"
cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0}
+1 -1
View File
@@ -1,5 +1,5 @@
"""Tests for the secondary-type embedding-distance diagnostic
(giant.analysis.type_embedding_distance) the §11.3 diagnostic."""
(giant.analysis.type_embedding_distance)."""
from __future__ import annotations
+2 -2
View File
@@ -11,8 +11,8 @@ _K_MAX = 5
def _tiny_models(particle_type_cfg: dict | None = None):
"""A fresh v0.3.0 pair: Stage1Model owns no n_sec_head (decision 1), so
n_sec always comes from Stage2OneShot."""
"""A fresh v0.3.0 pair: Stage1Model owns no n_sec_head, so n_sec always
comes from Stage2OneShot."""
s1 = Stage1Model(
pdg_vocab=3,
mat_vocab=2,
Generated
+1 -1
View File
@@ -534,7 +534,7 @@ wheels = [
[[package]]
name = "giant"
version = "0.2.0"
version = "0.3.0"
source = { editable = "." }
dependencies = [
{ name = "numpy" },