feat: add WGAN + AR stop-token config variant

Branches off configs/baseline.toml with both stages on WGAN-GP and
stage-2 n_sec.mode = stop_token, 30 epochs — combines two unbenchmarked
roadmap axes (post-v0.3.0 WGAN, and the AR stop-token multiplicity mode)
into one variant that stays a single edit away from baseline for
attribution.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01KzPghrmFcAJYrWUvHApY9N

.bumpversion.toml

@@ -1,5 +1,5 @@
 [tool.bumpversion]
-current_version = "0.3.11"
+current_version = "0.3.16"
 parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
 serialize = ["{major}.{minor}.{patch}"]
 search = "{current_version}"
@@ -8,7 +8,7 @@ regex = false
 allow_dirty = false
 commit = true
 tag = false
-message = "chore: bump version {current_version} -> {new_version} [skip ci]"
+message = "chore: bump version {current_version} -> {new_version}"
 pre_commit_hooks = ["uv lock", "git add uv.lock"]
 
 [[tool.bumpversion.files]]

.gitea/workflows/ci.yml

@@ -2,10 +2,9 @@ name: CI
 
 "on":
   push:
-    branches: ["**"]
+    branches: ["master"]
     tags: ["**"]
-  pull_request:
+  pull_request: {}
-    branches: [master]
 
 env:
   UV_CACHE_DIR: /uv-cache
@@ -156,7 +155,7 @@ jobs:
             uv run git-cliff --tag "$TAG" --unreleased --prepend CHANGELOG.md
             git add CHANGELOG.md
             if ! git diff --cached --quiet -- CHANGELOG.md; then
-              git commit -m "chore: update changelog for $TAG [skip ci]"
+              git commit -m "chore: update changelog for $TAG"
             else
               git restore --staged CHANGELOG.md
             fi
@@ -193,7 +192,7 @@ jobs:
             git config user.name "gitea-actions"
             git config user.email "actions@git.larsbogner.de"
             git add pyproject.toml uv.lock
-            git commit -m "chore: sync project version to tag ${GITHUB_REF_NAME} [skip ci]"
+            git commit -m "chore: sync project version to tag ${GITHUB_REF_NAME}"
             git push origin HEAD:master
             git push origin ":refs/tags/${GITHUB_REF_NAME}"
             git tag -f "${GITHUB_REF_NAME}" HEAD
@@ -201,3 +200,33 @@ jobs:
           else
             echo "Tag version matches project version ($CURRENT_VERSION)"
           fi
+
+  publish-package:
+    name: Publish package to Gitea package registry
+    needs: [ruff-check, ruff-format, type-check, test, sync-version-on-tag]
+    if: startsWith(github.ref, 'refs/tags/')
+    runs-on: ubuntu-latest
+    container:
+      image: docker.gitea.com/runner-images:ubuntu-latest
+      volumes:
+        - /srv/act-runner-cache/uv:/uv-cache
+    steps:
+      # Check out by tag name (not the triggering SHA) since sync-version-on-tag
+      # may have force-moved the tag to a version-corrected commit.
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ github.ref_name }}
+      - uses: astral-sh/setup-uv@v5
+        with:
+          enable-cache: false
+      - run: |
+          echo "UV_CACHE_DIR=/uv-cache" >> "$GITHUB_ENV"
+          echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
+      - run: uv build
+      # CI_TOKEN needs write:package scope (in addition to write:repository,
+      # used elsewhere) for this upload to authenticate.
+      - run: |
+          uv publish \
+            --publish-url "https://git.larsbogner.de/api/packages/lars/pypi" \
+            --username gitea-actions \
+            --password "${{ secrets.CI_TOKEN }}"

CHANGELOG.md

@@ -1,5 +1,39 @@
 # Changelog
 
+## [0.3.16] - 2026-08-31
+
+### Changed
+
+- Docs: record analysis_341dfb14 baseline rollout benchmark results
+
+- Feat: add eval-cost benchmark — Geant4 reference vs surrogate rollout timing
+
+## [0.3.15] - 2026-08-28
+
+### Changed
+
+- Perf: defer heavy imports in giant/dwarf CLIs until commands run
+
+## [0.3.14] - 2026-08-28
+
+### Changed
+
+- Ci: give automated commits visible checks, scope CI triggers, publish releases
+
+- Ci: fix pull_request trigger not registering
+
+## [0.3.13] - 2026-08-28
+
+### Added
+
+- Add inference-time model_config overrides with a sampling-key allowlist [gitea #87](https://git.larsbogner.de/lars/giant/issues/87)
+
+## [0.3.12] - 2026-08-28
+
+### Added
+
+- Add sampled n_sec under n_sec.mode = 'head' [gitea #86](https://git.larsbogner.de/lars/giant/issues/86)
+
 ## [0.3.11] - 2026-08-26
 
 ### Changed

CLAUDE.md

@@ -113,6 +113,8 @@ Secondary energies are a **stick-breaking partition of the `e_sec` budget** from
 
 **v0.3.0 — Stage-2 autoregressive redesign (implemented, released; on `master` since 2026-08-13):** motivated by the 2026-08-03 WGAN rollout benchmark, which failed specifically at the secondary-species level (zero photon secondaries, ~4M hallucinated `-14` muon antineutrinos). Stage 2 became autoregressive in descending-energy order with teacher forcing, and the particle-type representation went back to **categorical** (`particle_type.target = "onehot"`), reversing the 2026-07-17 continuous `(log-mass, charge)` target. The config break (`[conditioning]`/`[stage1_model]`/`[stage2_model]`/`[train]` replacing the flat `train.mode` + `[model]`) makes per-stage generators, stage-2-only training, and one-shot-vs-autoregressive comparison all expressible, and the `network.py` refactor into composable parts (encoder × trunk × objective) also makes routed WGAN work for the first time.
 
+**Baseline benchmark (done, 2026-08-26):** `configs/baseline.toml`'s first full rollout-vs-Geant4 validation (`analysis_341dfb14`, checkpoint `20260814_1743_s2-flow_h512_s2h512_bs36864_ep50/best.pt`, epoch 50/50). Confirms the v0.3.0 pivot fixed the species collapse — zero photon secondaries / hallucinated `-14` muon antineutrinos are both gone (γ at 95% of truth, no `-14` in the top species) — and rules out `conditioning.*.type = "physical"` as the cause, since this checkpoint pairs it with `flow`/no-router and still doesn't collapse. Bulk shower observables are close to Geant4 (total deposited energy +1.9%, containment depth-90%/95% both 0.986×), but steps/event now *over*-shoots by 1.32× (the opposite sign from every pre-v0.3.0 checkpoint), no hadronic/nuclear secondaries are produced at all, and event-to-event energy variance is ~16× too narrow. Writeup: `/home/lars/knowledge-base/experiments/giant-baseline-flow-ar-rollout-validation.md`.
+
 v0.2 configs and checkpoints are auto-migrated (`config.migrate_config`, `model._legacy._migrate_legacy_model_config`, both drawing on shared facts in `giant/_migration.py`). **v0.2 checkpoint-loading support has no expiry decided yet**: `/ceph` still holds pre-v0.3.0 checkpoints and analysis runs referencing them, so don't delete or substantially alter either migration function or `tests/legacy/network_v02_snapshot.py` (the frozen v0.2 snapshot they're tested against) without an explicit decision to do so first.
 
 **Faster-eval architectures — both implemented, neither validated.** Target is a ~10× native-Geant4 eval budget; no eval-latency number exists for any configuration yet, so that budget is unverified across the board.

cliff.toml

@@ -37,7 +37,7 @@ commit_preprocessors = [
 protect_breaking_commits = false
 commit_parsers = [
     { message = "^Merge ", skip = true },
-    { message = "\\[skip ci\\]", skip = true },
+    { message = "^chore: (bump version|update changelog|sync project version)", skip = true },
     { message = "^Add", group = "<!-- 0 -->Added" },
     { message = "^(Fix|Clamp|Clip)", group = "<!-- 1 -->Fixed" },
     { message = "^(Remove|Drop|Deprecate)", group = "<!-- 2 -->Removed" },

configs/baseline.toml

@@ -29,11 +29,21 @@
 #     capacity overfitting is not the binding constraint, and every recent
 #     run used 0.0.
 #
-# Known weak spots this baseline is expected to *exhibit* (they are the
+# Known weak spots, now measured against this exact config rather than
-# reason for the comparisons, not a reason to retune this file): every model
+# extrapolated from the pre-v0.3 field (analysis_341dfb14, best.pt @ epoch
-# on record under-produces steps per event by ~2x (rollout ~7e4 vs Geant4
+# 50/50, full writeup: knowledge-base/experiments/
-# ~1.4e5) and secondaries per event by 2-3.5x (~2-3e4 vs 7.2e4), and n_sec
+# giant-baseline-flow-ar-rollout-validation.md). Unlike every pre-v0.3
-# head accuracy sits at 0.863-0.867 regardless of size or objective.
+# checkpoint (which under-produced steps/event by 1.6-5x), this baseline
+# OVER-produces steps/event by 1.32x (1.86e5 vs Geant4 1.41e5) and
+# under-produces secondaries/event by 0.84x (5.97e4 vs 7.14e4) — the sign on
+# steps flipped with the v0.3 autoregressive pivot, so don't assume it still
+# undershoots. Secondary-species hallucination (zero photons, hallucinated
+# `-14` muon antineutrinos) that broke every prior checkpoint is gone; the
+# remaining species gap is a total absence of hadronic/nuclear secondaries
+# (protons, neutrons, ion recoils), not miscalibration of the ones produced.
+# Total deposited energy/event is +1.9% high but its event-to-event spread is
+# ~16x too narrow (31 MeV vs Geant4's 491 MeV). Per-step deposited energy is
+# the worst per-step marginal (KS 0.179 vs 0.004-0.071 for the others).
 
 [meta]
 # REQUIRED. Without it config.migrate_config reads this file as v0.2 and

configs/wgan_ar_stoptoken.toml

@@ -0,0 +1,122 @@
+# GIANT WGAN-GP + AR stop-token variant of configs/baseline.toml.
+#
+# Two roadmap axes, combined into one run: WGAN-GP generators for both
+# stages (unbenchmarked since the 2026-08-03 pre-v0.3.0 failure, which was
+# secondary-species mode collapse — the failure v0.3.0's AR/categorical
+# pivot exists to fix) and the AR stop-token multiplicity mode
+# (stage2_model.n_sec.mode = "stop_token", never benchmarked at all).
+# Everything else is byte-identical to baseline.toml so a rollout compared
+# against baseline's analysis_341dfb14 is attributable to these two axes
+# alone: conditioning (physical/physical), hidden_dim 512 / n_res_blocks 6 /
+# dropout 0.0 per stage, k_max 15, history "markov", teacher_forcing
+# "always", particle_type.target "onehot" (n_classes 32, other_policy
+# "sample"), lr 3e-4, warmup_epochs 3, weight_decay 0.01, ema_decay 0.9999,
+# val_fraction 0.1, num_workers 4, seed 0, validate_steps 10, W&B on.
+#
+# No [stage1_model.wgan] / [stage2_model.wgan] block: the dataclass defaults
+# (noise_dim 64, n_critic 5, gp_weight 10.0, critic_lr 0.0 = inherit
+# train.lr, critic_hidden_dim/critic_n_res_blocks 0 = inherit the stage's
+# 512/6, stage 2's gumbel_tau_start/_end 1.0/0.1) are what the earlier WGAN
+# runs used — writing them out would add keys that don't vary.
+#
+# particle_type.class_weighting stays "none" (the default): config.py's
+# validate_config rejects any other value under stage2_model.generator =
+# "wgan", since that path feeds the type slice to the critic via a
+# straight-through Gumbel relaxation instead of a weighted cross-entropy.
+#
+# Prior WGAN writeup (pre-v0.3.0, describes the failure this run re-tests):
+# /home/lars/knowledge-base/experiments/giant-wgan-physical-rollout-validation.md
+
+[meta]
+# REQUIRED. Without it config.migrate_config reads this file as v0.2 and
+# rewrites it from V02_FIXED_FACTS — silently forcing decoder = "one_shot",
+# particle_type.target = "physical" and the v0.2 default sizes, while still
+# passing validate_config.
+config_version = 3
+
+[conditioning]
+# Physical-property MLPs rather than learned vocab embeddings: computable for
+# any PDG code / material, which is what the held-out-species and
+# held-out-material generalization comparisons need.
+out_dim      = 128
+share_stages = false
+
+# n_layers = 2 rather than the v0.3 default of 1: v0.2's conditioning MLP was
+# always 2 deep (see _migration.V02_FIXED_FACTS), so this keeps the encoder
+# identical to baseline.toml.
+[conditioning.particle]
+type     = "physical"
+emb_dim  = 16
+n_layers = 2
+
+[conditioning.material]
+type     = "physical"
+emb_dim  = 16
+n_layers = 2
+
+[stage1_model]
+generator    = "wgan"
+hidden_dim   = 512
+n_res_blocks = 6
+dropout      = 0.0
+
+[stage2_model]
+# Autoregressive in descending-energy order, as baseline.toml — this variant
+# only swaps the generator (flow -> wgan) and the multiplicity mode
+# (head -> stop_token), not the decoder shape.
+decoder      = "autoregressive"
+generator    = "wgan"
+hidden_dim   = 512
+n_res_blocks = 6
+dropout      = 0.0
+k_max        = 15
+
+[stage2_model.autoregressive]
+history         = "markov"
+teacher_forcing = "always"
+
+[stage2_model.n_sec]
+# EOS-style per-slot stop head on the AR secondary decoder, replacing the
+# n_sec classifier entirely (mutually exclusive — see NSecConfig's
+# docstring in giant/config.py). Requires decoder = "autoregressive" and
+# owner = "stage2" (both already true above/by default); validate_config
+# enforces this.
+mode = "stop_token"
+
+[stage2_model.particle_type]
+target = "onehot"
+# Decoupled from conditioning.particle.emb_dim (gitea #29). 32 classes + the
+# "other" bucket keeps essentially all real secondary species out of "other"
+# without making the head expensive.
+n_classes    = 32
+other_policy = "sample"
+
+[train]
+epochs = 30
+# Halved from baseline's 36864. That figure came from a measured linear fit
+# of the *flow-AR* training step (peak reserved MiB = 0.9736 * batch_size +
+# 115); WGAN invalidates it twice over — each stage gains a critic that by
+# default inherits the stage's own 512/6 body, and gradient_penalty
+# (giant/model/wgan.py, forced fp32 internally) runs a double-backward every
+# batch. 18432 is a conservative choice pending a real memory measurement on
+# this exact config, not a re-derived fit. Throughput is already flat above
+# bs~4096 on the 4070, so this costs occupancy on the L40S, not step
+# efficiency.
+batch_size    = 18432
+lr            = 3e-4
+warmup_epochs = 3
+weight_decay  = 0.01
+ema_decay     = 0.9999
+val_fraction  = 0.1
+num_workers   = 4
+seed          = 0
+# Tightened from baseline's 10: WGANStageTrainer.supports_val_loss = False,
+# and with both stages adversarial there is no per-epoch val loss at all, so
+# validate_every's marginal-KL pass (giant/training/trainers.py's
+# val_objective) is the only comparable-across-epochs best-checkpoint
+# selection signal available. 5 gives 6 evaluations over 30 epochs instead
+# of baseline's 3, at ~6x5000s of extra walltime.
+validate_every = 5
+validate_steps = 10
+wandb          = true
+wandb_project  = "giant"

giant/analysis/catalog.py

@@ -45,6 +45,7 @@ import numpy as np
 import polars as pl
 
 from giant.analysis.context import Context
+from giant.analysis.geant4_reference import GEANT4_REFERENCE, geant4_per_step_us
 from giant.analysis.grouping import (
     energy_bin_labels,
     event_energy_bins,
@@ -125,6 +126,7 @@ class Bundle:
                 phys=physical_steps(r_all, Side.rollout),
                 checkpoint=rs.checkpoint,
                 type_embedding_l1_dist=rs.type_embedding_l1_dist,
+                timing=rs.timing,
             )
         return cls(ctx=ctx, rollouts=sides, t_all=t_all, t_phys=physical_steps(t_all, Side.reference))
 
@@ -973,6 +975,80 @@ def _sec_cos_angle_finalize(parts: list[dict], ctx: Context) -> Reduced:
     )
 
 
+# ---------------------------------------------------------------------------
+# eval cost (not chunked — metadata-only, no row scan)
+# ---------------------------------------------------------------------------
+
+_EVAL_COST_LABELS = ["sampling / simulation", "parquet write / convert", "total"]
+_EVAL_COST_NOTE = (
+    "no rollout in this run carries a `timing` block — re-run `giant rollout` "
+    "(timing instrumentation added after this checkpoint's rollout run) to "
+    "populate this plot"
+)
+
+
+def _eval_cost_per_step(b: Bundle) -> Reduced:
+    """Per-rollout µs/physical-step vs the measured Geant4 reference.
+
+    ``timing`` (``giant.cli``'s ``rollout`` command) is metadata carried on
+    the rollout YAML, not derived from the row data, so this needs no chunked
+    scan — same shape as the router diagnostics above.
+    """
+    series: dict[str, list[float]] = {}
+    speedup: dict[str, float] = {}
+    for name, rs in b.rollouts.items():
+        t = rs.timing
+        if not t or t.get("us_per_step") is None:
+            continue
+        sample_us = t["us_per_step"]
+        write_us = t.get("write_us_per_step") or 0.0
+        series[name] = [sample_us, write_us, sample_us + write_us]
+
+    if not series:
+        return Reduced(
+            id="eval_cost_per_step",
+            family="cost",
+            kind="unavailable",
+            title="Eval cost per step: surrogate vs Geant4",
+            xlabel="n/a",
+            payload={"note": _EVAL_COST_NOTE},
+        )
+
+    g4 = geant4_per_step_us()
+    reference = [g4["sim_us_per_step"], g4["convert_us_per_step"], g4["total_us_per_step"]]
+    for name, vals in series.items():
+        speedup[name] = reference[-1] / vals[-1] if vals[-1] else float("inf")
+
+    return Reduced(
+        id="eval_cost_per_step",
+        family="cost",
+        kind="bar",
+        title="Eval cost per step: surrogate vs Geant4",
+        xlabel="phase",
+        payload={
+            "labels": _EVAL_COST_LABELS,
+            "series": series,
+            "reference": reference,
+            "ylabel": "µs per physical step",
+            "log_y": True,
+        },
+        meta={
+            "speedup_vs_geant4_total": speedup,
+            "geant4_provenance": GEANT4_REFERENCE["provenance"],
+            "caveat": (
+                "The Geant4 reference is measured single-threaded on one CPU core "
+                "(see giant.analysis.geant4_reference); a rollout's timing is "
+                "whatever device it actually ran on (see each series' device in "
+                "run_meta.json's plot_meta). This is a deployment-speedup ratio, "
+                "not a same-hardware or per-FLOP comparison."
+            ),
+        },
+    )
+
+
+_eval_cost_per_step_partial, _eval_cost_per_step_finalize = _unchunkable(_eval_cost_per_step)
+
+
 # ---------------------------------------------------------------------------
 # router diagnostics (not chunked — already bounded/subsampled)
 # ---------------------------------------------------------------------------
@@ -1160,6 +1236,13 @@ def build_catalog() -> list[PlotSpec]:
             compute_partial=_sec_cos_angle_partial,
             finalize=_sec_cos_angle_finalize,
         ),
+        PlotSpec(
+            "eval_cost_per_step",
+            "cost",
+            compute_partial=_eval_cost_per_step_partial,
+            finalize=_eval_cost_per_step_finalize,
+            chunkable=False,
+        ),
         PlotSpec(
             "router_gating",
             "model",

giant/analysis/condor.py

@@ -82,6 +82,7 @@ _PLOT_META_KEYS = (
     "rollout_seed",
     "n_rows",
     "termination_reason_counts",
+    "timing",
     "model_config",
     "training_epoch",
     "best_val_loss",
@@ -329,9 +330,9 @@ def compute_reduced(
 ) -> Path:
     """Core: run one (plot, chunk)'s partial reduction against explicit paths.
 
-    ``rollouts``: ``[{"name", "path", "checkpoint"?, "type_embedding_l1_dist"?},
+    ``rollouts``: ``[{"name", "path", "checkpoint"?, "type_embedding_l1_dist"?,
-    ...]``, one per rollout series (insertion order preserved through to every
+    "timing"?}, ...]``, one per rollout series (insertion order preserved
-    plot's ``Reduced.payload["series"]``).
+    through to every plot's ``Reduced.payload["series"]``).
 
     Writes a ``Partial`` JSON — the raw, not-yet-merged output of
     ``PlotSpec.compute_partial`` — never a finished ``Reduced``; ``merge_one``
@@ -352,6 +353,7 @@ def compute_reduced(
             source=r["path"],
             checkpoint=r.get("checkpoint"),
             type_embedding_l1_dist=r.get("type_embedding_l1_dist"),
+            timing=r.get("timing"),
         )
         for r in rollouts
     ]
@@ -377,6 +379,7 @@ def compute_one(spec_id: str, run_dir: str | Path, chunk_index: int = 0) -> Path
             "path": ro["path"],
             "checkpoint": ro["plot_meta"].get("checkpoint"),
             "type_embedding_l1_dist": ro["plot_meta"].get("type_embedding_l1_dist"),
+            "timing": ro["plot_meta"].get("timing"),
         }
         for ro in meta.rollouts
     ]

giant/analysis/geant4_reference.py

@@ -0,0 +1,76 @@
+"""Measured Geant4 (miniCaloSim) per-step eval cost — the reference line for
+``eval_cost_per_step`` in ``catalog.py``.
+
+Mirrors the precedent set by ``runtime_estimate.py``'s ``_COST_MODEL``: a
+constant table measured once on a specific machine and pasted in, with the
+methodology and provenance recorded in this docstring rather than derived at
+runtime (there is no live Geant4 install on the machines that run
+``giant analyze``, and re-measuring per invocation would be both slow and
+noisy — see the module docstring precedent).
+
+**Methodology** (``scratchpad/bench_geant4.py``, a one-off, not a `dwarf`
+subcommand): ``run_pbwo4`` (the default homogeneous-PbWO4 miniCaloSim
+executable, see ``~/Programming/minicalosim``) was timed at 3 beam energies
+(1/10/50 GeV) and **4 event counts each**, converting each run's ROOT output
+to Parquet with ``giant.tools.steps_to_parquet.convert_steps_to_parquet``
+immediately after. Event counts were scaled down as energy rose (100/400/
+1000/2000 at 1 GeV, 30/100/200/300 at 10 GeV, 10/25/45/60 at 50 GeV) to keep
+every run's row count under ~8.1M — a naive 50/200 pair at 50 GeV produces
+~27M steps and OOM'd the conversion step on a 14GB laptop. Per-energy linear
+fits (``t = intercept + slope * n``) separate Geant4's one-time init (physics
+tables, geometry construction) from its true marginal per-event cost — the
+slope, not a naive ``t / n_events`` from a single run, is what feeds
+``sim_us_per_step`` below. The per-step denominator is the produced
+``Steps``-tree/Parquet row count, matching the "physical step" unit
+``giant rollout``'s ``timing.n_physical_rows`` uses on the surrogate side.
+Both stages ran single-threaded (default Geant4 threading), pinned to one
+CPU core.
+
+``sim_us_per_step``/``convert_us_per_step``/``sim_ms_per_event`` below are
+the mean across the 3 energies. With 4 event-count points per energy (up
+from an initial 2-point pass, which had ~80% spread and nonsensical negative
+fitted intercepts at 10/50 GeV — an artifact of extrapolating a 2-point
+line), both quantities are now energy-flat as physically expected:
+``sim_us_per_step`` spread ~5%, ``convert_us_per_step`` spread ~13.5%. Treat
+these as reliable to about that precision.
+
+**Caveat — hardware asymmetry**: this reference is single-core CPU. A
+surrogate rollout's ``timing`` block will typically be measured on a batched
+GPU. The resulting ratio in ``eval_cost_per_step`` is a *deployment* speedup
+(what you'd actually see swapping Geant4 for the surrogate in a production
+pipeline), not a same-hardware or per-FLOP comparison — state this whenever
+quoting the number.
+
+**Staleness**: re-run ``scratchpad/bench_geant4.py`` (and update this file)
+if measured on different hardware, after a miniCaloSim/Geant4 version bump,
+or if this reference is more than a year or two stale.
+"""
+
+from __future__ import annotations
+
+GEANT4_REFERENCE: dict = {
+    "sim_us_per_step": 11.2903,
+    "convert_us_per_step": 11.1014,
+    "sim_ms_per_event": 609.6848,
+    "provenance": {
+        "cpu": "AMD Ryzen 7 PRO 4750U with Radeon Graphics",
+        "geant4_version": "11.4.1",
+        "minicalosim_sha": "ea917da",
+        "measured": "2026-08-31",
+        "energies_gev": [1.0, 10.0, 50.0],
+        "spread_pct_sim": 4.96,
+        "spread_pct_convert": 13.52,
+        "threads": 1,
+    },
+}
+
+
+def geant4_per_step_us() -> dict[str, float]:
+    """Sim / convert / total microseconds per physical step, from ``GEANT4_REFERENCE``."""
+    sim = GEANT4_REFERENCE["sim_us_per_step"]
+    convert = GEANT4_REFERENCE["convert_us_per_step"]
+    return {
+        "sim_us_per_step": sim,
+        "convert_us_per_step": convert,
+        "total_us_per_step": sim + convert,
+    }

giant/analysis/render.py

@@ -274,6 +274,8 @@ def _render_bar(r: Reduced, params: dict):
     ax.set_xticks(x)
     ax.set_xticklabels(labels, rotation=45, ha="right")
     ax.set_ylabel(r.payload.get("ylabel", "value"))
+    if r.payload.get("log_y"):
+        ax.set_yscale("log")
     ps.style_legend(ax, title="source")
     return fig
 

giant/analysis/runtime_estimate.py

@@ -96,6 +96,10 @@ _COST_MODEL: dict[str, tuple[float, float]] = {
     "sec_count_per_species": (0.0, 4.963e-07),
     "sec_energy": (0.0, 4.727e-07),
     "sec_cos_angle": (0.0, 2.749e-06),
+    # Metadata-only (YAML-carried `timing`, no row scan) — same shape as the
+    # router diagnostics' fixed cost, just cheaper since there's no live
+    # torch checkpoint to load.
+    "eval_cost_per_step": (0.0, 0.0),
 }
 
 

giant/analysis/sources.py

@@ -110,6 +110,7 @@ class RolloutSpec:
     source: str | Path | pl.LazyFrame
     checkpoint: str | None = None
     type_embedding_l1_dist: dict | None = None
+    timing: dict | None = None
 
 
 @dataclass
@@ -124,6 +125,10 @@ class RolloutSide:
     # only. Unlike checkpoint, this needs no live model: it's already a
     # finished histogram, just passed through.
     type_embedding_l1_dist: dict | None = None
+    # Wall-clock cost of this rollout run (giant.cli's rollout command),
+    # from the rollout YAML — eval_cost_per_step only. None on rollout runs
+    # that predate timing instrumentation.
+    timing: dict | None = None
 
 
 def _check_rollout_metadata(path: Path) -> None:

giant/checkpoint_io.py

@@ -14,11 +14,18 @@ directly and imported from non-CLI code (`giant.analysis.router_gating`,
 lazily — see that module's docstring for why). Failures raise
 `CheckpointCompatibilityError` with the same wording the CLI has always
 shown; the CLI layer catches it and does the `typer.echo`/`Exit(1)`.
+
+`load_for_inference`'s `config_overrides` (gitea #87) lets a caller change a
+checkpoint's `model_config` at load time, restricted to
+`giant.config.INFERENCE_OVERRIDES` — the allowlist of keys that only affect
+sampling, never module construction/shapes or the preprocessing normalizers/
+vocab maps were fit under.
 """
 
 from __future__ import annotations
 
-from dataclasses import dataclass
+import copy
+from dataclasses import dataclass, field
 from pathlib import Path
 
 import torch
@@ -29,13 +36,45 @@ from giant.constants import K_MAX
 from giant.data.loader import TopNMap
 from giant.data.setup_cache import topnmap_from_json
 from giant.data.transforms import Normalizer
-from giant.model.network import build_models
+from giant.model.network import _migrate_legacy_model_config, build_models
 
 
 class CheckpointCompatibilityError(Exception):
     """Checkpoint is missing something `load_for_inference` needs."""
 
 
+def apply_config_overrides(model_cfg: dict, overrides: dict[str, object] | None) -> dict:
+    """Deep-merge dotted-path *overrides* into a checkpoint's `model_config`,
+    validated against `giant.config.INFERENCE_OVERRIDES` — the allowlist of
+    keys that only affect sampling, not module construction/shapes or the
+    preprocessing normalizers/vocab maps were fit under (gitea #87).
+
+    Migrates a v0.2 flat `model_config` to the nested v0.3 shape first: a
+    dotted path like "stage1_model.ddpm.n_steps" would otherwise silently
+    write into a dict that `build_models` still reads as flat (it decides
+    v0.2-vs-v0.3 by `"stage1_model" in model_config`), suppressing migration.
+
+    Raises `CheckpointCompatibilityError` — never a bare `ValueError` or a
+    downstream `load_state_dict` size mismatch — for an unknown/disallowed
+    path or a value that fails its allowlisted check.
+    """
+    if not overrides:
+        return model_cfg
+    cfg = model_cfg if "stage1_model" in model_cfg else _migrate_legacy_model_config(model_cfg)
+    cfg = copy.deepcopy(cfg)
+    for path, value in overrides.items():
+        spec = gconfig.INFERENCE_OVERRIDES.get(path)
+        if spec is None:
+            allowed = ", ".join(sorted(gconfig.INFERENCE_OVERRIDES))
+            raise CheckpointCompatibilityError(f"{path!r} is not an inference-safe override — allowed paths: {allowed}")
+        try:
+            spec.check(path, value)
+        except ValueError as exc:
+            raise CheckpointCompatibilityError(str(exc)) from exc
+        gconfig._set_path(cfg, path, value)
+    return cfg
+
+
 def conditioning_axes(model_cfg: dict, default: str = "embedding") -> tuple[str, str]:
     """(particle_conditioning, material_conditioning) for
     `giant.data.transforms.build_cond_features`/`build_features` — from
@@ -128,6 +167,7 @@ class InferenceContext:
     model_config: dict
     epoch: int | None
     best_val_loss: float | None
+    config_overrides: dict[str, object] = field(default_factory=dict)
 
 
 def load_for_inference(
@@ -136,6 +176,7 @@ def load_for_inference(
     command_name: str,
     weights: str = "raw",
     require_stage2: bool = True,
+    config_overrides: dict[str, object] | None = None,
 ) -> InferenceContext:
     """Load *checkpoint* and reconstruct everything `predict`/`rollout` need
     to run it forward, on *device*, in `eval()` mode.
@@ -148,6 +189,13 @@ def load_for_inference(
     both stages) or an acceptable `stage2 = None` result — kept as a real
     parameter since `stage{1,2}_model.active` is a real, if currently
     stage1+stage2-only-in-practice, config option.
+
+    *config_overrides* deep-merges dotted `model_config` paths (e.g.
+    `{"stage2_model.n_sec.sampling": "sample"}`) before anything is
+    derived from `model_config` or built — see `apply_config_overrides` for
+    the allowlist and validation. Every derived `InferenceContext` field
+    (`other_policy`, `stage{1,2}_ddpm_steps`, the built modules, ...)
+    reflects the overridden config.
     """
     ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)
     for key in ("model_config", "sec_decoder"):
@@ -159,7 +207,7 @@ def load_for_inference(
 
     gconfig.warn_if_checkpoint_config_mismatch(checkpoint)
 
-    model_cfg = ckpt["model_config"]
+    model_cfg = apply_config_overrides(ckpt["model_config"], config_overrides)
     particle_conditioning, material_conditioning = conditioning_axes(model_cfg)
     pdg_topn_map = load_pdg_topn_map(ckpt)
     mat_topn_map = load_mat_topn_map(ckpt)
@@ -231,4 +279,5 @@ def load_for_inference(
         model_config=model_cfg,
         epoch=ckpt.get("epoch"),
         best_val_loss=ckpt.get("best_val_loss"),
+        config_overrides=dict(config_overrides) if config_overrides else {},
     )

giant/cli.py

@@ -1,21 +1,19 @@
+from __future__ import annotations
+
 from collections import Counter
 from datetime import datetime, timezone
 from enum import Enum
 import math
 from pathlib import Path
 import re
-from typing import Optional
+from typing import TYPE_CHECKING, Optional, cast
 import uuid as uuid_mod
 
-import numpy as np
-import yaml
-import torch
 import typer
 from typing_extensions import Annotated
 
-import pyarrow as pa
+if TYPE_CHECKING:
-import pyarrow.parquet as pq
+    import numpy as np
-from tqdm import tqdm
 
 from giant import config as gconfig
 from giant.constants import (
@@ -25,30 +23,11 @@ from giant.constants import (
     PREDICT_SCHEMA_VERSION_KEY,
     ROLLOUT_COORD_VALUE,
 )
-from giant.data.loader import (
+
-    event_id_offset,
+# giant.materials only pulls in numpy (no torch/pandas), and MATERIAL_PROPERTIES
-    find_parquet_files,
+# is needed at decoration time below (a Typer option default), so it can't be
-    iter_file_chunks,
+# deferred into a command body like the rest of this module's heavy imports.
-    iter_cond_chunks,
-)
-from giant.data.transforms import (
-    build_features,
-    build_cond_features,
-    energy_simplex_decode,
-    inv_local_frame_rotation,
-    inv_log_transform,
-    reconstruct_post_pos,
-)
-from giant.checkpoint_io import CheckpointCompatibilityError, load_for_inference
-from giant.geometry import GeometryOracle
 from giant.materials import MATERIAL_PROPERTIES
-from giant.pipeline import run_train_job
-from giant.rollout import (
-    L1DistCollector,
-    decode_secondary_identity,
-    rollout as run_rollout,
-)
-from giant.sample import resolve_n_sec, sample_stage1, sample_stage2
 
 app = typer.Typer(no_args_is_help=True)
 
@@ -141,6 +120,23 @@ def _parse_router_axis_flags(specs: list[str]) -> dict[str, object]:
     return out
 
 
+def _parse_set_flags(specs: Optional[list[str]]) -> dict[str, object]:
+    """Parse repeated `--set dotted.path=value` flags into a dict, typing
+    each value with `_coerce_scalar` the same way a TOML file's native types
+    would arrive. Validation against the inference-safe allowlist happens
+    downstream in `giant.checkpoint_io.apply_config_overrides` — this only
+    parses syntax.
+    """
+    out: dict[str, object] = {}
+    for spec in specs or []:
+        path, sep, val = spec.partition("=")
+        if not sep:
+            typer.echo(f"error: --set {spec!r} must be 'dotted.path=value'", err=True)
+            raise typer.Exit(1)
+        out[path] = _coerce_scalar(val)
+    return out
+
+
 def _router_cli_overrides(
     router: bool | None,
     router_type: str | None,
@@ -193,6 +189,8 @@ def _write_prediction_ref(
     comment: str | None = None,
 ) -> Path:
     """Write a YAML sidecar in the checkpoint directory and return its path."""
+    import yaml
+
     ref = {
         "prediction_id": pred_uuid,
         "output": str(out),
@@ -207,6 +205,45 @@ def _write_prediction_ref(
     return ref_path
 
 
+def _build_rollout_timing(
+    *,
+    setup_s: float,
+    rollout_s: float,
+    write_s: float,
+    n_rows: int,
+    termination_reason_counts: dict[str, int],
+    n_seed_events: int,
+    device: str,
+    torch_threads: int,
+) -> dict:
+    """Assemble ``giant rollout``'s ``timing`` sidecar block.
+
+    ``n_physical_rows`` excludes the synthetic termination rows (escape/
+    unknown-pdg/energy-cutoff/max-steps markers `giant.rollout` emits but
+    Geant4 never does) so ``us_per_step`` is comparable to
+    ``giant.analysis.geant4_reference``'s per-step Geant4 measurement — see
+    ``giant/analysis/catalog.py``'s ``eval_cost_per_step`` spec.
+    """
+    from giant.analysis.sources import SYNTHETIC_TERMINATION_REASONS
+
+    sample_s = rollout_s - write_s
+    n_synthetic_rows = sum(termination_reason_counts.get(reason, 0) for reason in SYNTHETIC_TERMINATION_REASONS)
+    n_physical_rows = n_rows - n_synthetic_rows
+    return {
+        "setup_s": setup_s,
+        "rollout_s": rollout_s,
+        "write_s": write_s,
+        "sample_s": sample_s,
+        "n_rows": n_rows,
+        "n_physical_rows": n_physical_rows,
+        "us_per_step": (sample_s / n_physical_rows * 1e6) if n_physical_rows else None,
+        "write_us_per_step": (write_s / n_physical_rows * 1e6) if n_physical_rows else None,
+        "ms_per_event": (rollout_s / n_seed_events * 1e3) if n_seed_events else None,
+        "device": device,
+        "torch_threads": torch_threads,
+    }
+
+
 @app.callback()
 def _main() -> None:
     """GIANT — Geant4 step-function surrogate."""
@@ -622,6 +659,10 @@ def train(
     ] = None,
 ) -> None:
     """Train the GIANT surrogate model."""
+    import torch
+
+    from giant.pipeline import run_train_job
+
     batch_size_auto = False
     batch_size_value: Optional[int] = None
     if batch_size is not None:
@@ -1020,8 +1061,36 @@ def predict(
             help="Free-text note recorded in the prediction's YAML sidecar",
         ),
     ] = None,
+    set_: Annotated[
+        Optional[list[str]],
+        typer.Option(
+            "--set",
+            help="Override a sampling-only model_config key on this checkpoint, "
+            "'dotted.path=value' (repeatable) — see giant.config.INFERENCE_OVERRIDES "
+            "for the allowlist, e.g. --set stage2_model.n_sec.sampling=sample",
+        ),
+    ] = None,
 ) -> None:
     """Run trained model on a parquet file and save predictions."""
+    import numpy as np
+    import pyarrow as pa
+    import pyarrow.parquet as pq
+    import torch
+    from tqdm import tqdm
+
+    from giant.checkpoint_io import CheckpointCompatibilityError, load_for_inference
+    from giant.data.loader import event_id_offset, find_parquet_files, iter_cond_chunks, iter_file_chunks
+    from giant.data.transforms import (
+        build_cond_features,
+        build_features,
+        energy_simplex_decode,
+        inv_local_frame_rotation,
+        inv_log_transform,
+        reconstruct_post_pos,
+    )
+    from giant.rollout import decode_secondary_identity
+    from giant.sample import resolve_n_sec, sample_stage1, sample_stage2
+
     batch_size_auto = False
     batch_size_value: Optional[int] = None
     if batch_size.strip().lower() == "auto":
@@ -1040,8 +1109,11 @@ def predict(
     typer.echo(f"device: {_device}")
 
     # --- Load checkpoint ---
+    config_overrides = _parse_set_flags(set_)
     try:
-        ctx = load_for_inference(checkpoint, _device, "predict", weights=weights.value)
+        ctx = load_for_inference(
+            checkpoint, _device, "predict", weights=weights.value, config_overrides=config_overrides
+        )
     except CheckpointCompatibilityError as exc:
         typer.echo(f"error: {exc}", err=True)
         raise typer.Exit(1)
@@ -1314,6 +1386,10 @@ def _seed_from_data(files: list[Path], n_events: int | None) -> dict[str, np.nda
     the codebase's convention for the primary (a secondary always carries less
     energy than its parent). See giant/analysis/reduce.py:entry_axis.
     """
+    import numpy as np
+
+    from giant.data.loader import event_id_offset, iter_cond_chunks
+
     best_E: dict[int, float] = {}
     best: dict[int, tuple] = {}
     for file_idx, path in enumerate(files):
@@ -1407,8 +1483,32 @@ def rollout(
         Optional[int],
         typer.Option("--seed", help="Torch/numpy seed for reproducibility"),
     ] = None,
+    set_: Annotated[
+        Optional[list[str]],
+        typer.Option(
+            "--set",
+            help="Override a sampling-only model_config key on this checkpoint, "
+            "'dotted.path=value' (repeatable) — see giant.config.INFERENCE_OVERRIDES "
+            "for the allowlist, e.g. --set stage2_model.n_sec.sampling=sample",
+        ),
+    ] = None,
 ) -> None:
     """Roll the surrogate forward into full showers (autoregressive)."""
+    import time
+
+    import numpy as np
+    import pyarrow as pa
+    import pyarrow.parquet as pq
+    import torch
+    import yaml
+
+    from giant.checkpoint_io import CheckpointCompatibilityError, load_for_inference
+    from giant.data.loader import find_parquet_files
+    from giant.geometry import GeometryOracle
+    from giant.rollout import L1DistCollector, RolloutSummary, rollout as run_rollout
+
+    _t_setup_start = time.perf_counter()
+
     if seed is not None:
         torch.manual_seed(seed)
         np.random.seed(seed)
@@ -1416,8 +1516,11 @@ def rollout(
     _device = torch.device(device) if device else gconfig.auto_device()
     typer.echo(f"device: {_device}")
 
+    config_overrides = _parse_set_flags(set_)
     try:
-        ctx = load_for_inference(checkpoint, _device, "rollout", weights=weights.value)
+        ctx = load_for_inference(
+            checkpoint, _device, "rollout", weights=weights.value, config_overrides=config_overrides
+        )
     except CheckpointCompatibilityError as exc:
         typer.echo(f"error: {exc}", err=True)
         raise typer.Exit(1)
@@ -1451,9 +1554,11 @@ def rollout(
     # avg_tracks_per_event) — mirrors the row-group streaming `giant predict`
     # already does on its input side.
     writer: pq.ParquetWriter | None = None
+    _write_s = 0.0
 
     def _write_chunk(row: dict[str, np.ndarray]) -> None:
-        nonlocal writer
+        nonlocal writer, _write_s
+        _t0 = time.perf_counter()
         table = pa.table(row)
         if writer is None:
             table = table.replace_schema_metadata(
@@ -1464,11 +1569,14 @@ def rollout(
             )
             writer = pq.ParquetWriter(out, table.schema)
         writer.write_table(table)
+        _write_s += time.perf_counter() - _t0
 
     # Only meaningful under particle_type.target="embedding" — a
     # no-op collector otherwise, cheaper than branching the call itself.
     l1_dist_collector = L1DistCollector()
 
+    _setup_s = time.perf_counter() - _t_setup_start
+    _t_rollout_start = time.perf_counter()
     summary = run_rollout(
         model,
         sec_decoder,
@@ -1500,6 +1608,23 @@ def rollout(
     )
     if writer is not None:
         writer.close()
+    # on_chunk=_write_chunk is always passed above, so rollout() always
+    # returns the streaming-summary shape (RolloutSummary), never the
+    # materialized dict[str, np.ndarray] alternative its return type allows.
+    summary = cast(RolloutSummary, summary)
+    _rollout_s = time.perf_counter() - _t_rollout_start
+    timing = _build_rollout_timing(
+        setup_s=_setup_s,
+        rollout_s=_rollout_s,
+        write_s=_write_s,
+        n_rows=summary["n_rows"],
+        termination_reason_counts=summary["termination_reason_counts"],
+        n_seed_events=len(seeds["event_id"]),
+        device=str(_device),
+        torch_threads=torch.get_num_threads(),
+    )
+    _sample_s = timing["sample_s"]
+    n_physical_rows = timing["n_physical_rows"]
 
     l1_summary = l1_dist_collector.summary()
 
@@ -1522,6 +1647,10 @@ def rollout(
             "rollout_seed": seed,
             "n_rows": summary["n_rows"],
             "termination_reason_counts": summary["termination_reason_counts"],
+            # Wall-clock cost of this run, normalized per physical step (the
+            # comparable unit against giant.analysis.geant4_reference) — see
+            # eval_cost_per_step in giant/analysis/catalog.py.
+            "timing": timing,
             # Diagnostic — only present under
             # stage2_model.particle_type.target="embedding"; omitted (not
             # written as null) otherwise, so giant.analysis can tell "not
@@ -1532,6 +1661,7 @@ def rollout(
             # model knob (router type/n_experts, noise_dim, vocab sizes, ...)
             # is available downstream without touching this command again.
             "model_config": dict(model_cfg),
+            "config_overrides": dict(ctx.config_overrides),
             "training_epoch": ctx.epoch,
             "best_val_loss": ctx.best_val_loss,
             # [train]/[meta] from the sibling config.toml (giant.config.save_config)
@@ -1544,6 +1674,11 @@ def rollout(
 
     typer.echo(f"wrote {summary['n_rows']:,} step rows → {out}")
     typer.echo(f"terminations: {summary['termination_reason_counts']}")
+    if timing["us_per_step"] is not None:
+        typer.echo(
+            f"timing: {_rollout_s:.1f}s total ({_sample_s:.1f}s sample + {_write_s:.1f}s write), "
+            f"{timing['us_per_step']:.1f} us/step over {n_physical_rows:,} physical steps"
+        )
     typer.echo(f"reference: {ref_path}")
 
 

giant/config.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import copy
 import difflib
 import hashlib
@@ -10,12 +12,12 @@ from dataclasses import dataclass, field
 from datetime import datetime, timezone
 from enum import Enum
 from pathlib import Path
-
+from typing import TYPE_CHECKING
-import numpy as np
-import torch
 
 from giant._migration import V02_FIXED_FACTS, V02_MODEL_KEY_TO_STAGES, reject_legacy_router_expert_sizing
-from giant.model.history import HISTORY_REGISTRY
+
+if TYPE_CHECKING:
+    import torch
 
 
 class Conditioning(str, Enum):
@@ -404,6 +406,16 @@ class Stage2RouterConfig(RouterConfig):
         return {"tie_to_stage1": self.tie_to_stage1, **super().to_dict()}
 
 
+# stage2_model.n_sec.sampling choices — single source of truth for both
+# validate_config's train-time check and INFERENCE_OVERRIDES below.
+STOP_SAMPLING_CHOICES = ("greedy", "sample")
+
+# stage2_model.particle_type.other_policy choices — see ParticleTypeConfig's
+# docstring for what each means; only documented there until now, since
+# nothing validated it at train time.
+OTHER_POLICY_CHOICES = ("sample", "modal", "drop")
+
+
 @dataclass(frozen=True)
 class NSecConfig:
     # "head": a classifier over {0..k_max} on the condition encoding alone
@@ -424,11 +436,15 @@ class NSecConfig:
     # n_sec head was trained against Stage 1's own ConditionEncoder output and so has
     # to stay attached there, not just be labeled as such).
     owner: str = "stage2"
-    # mode="stop_token" only: how sample_secondaries_ar turns a slot's stop logit into a
+    # How resolve_n_sec/sample_secondaries_ar turn a count-bearing head's output into an
-    # stop/continue decision. "greedy": sigmoid(logit) >= 0.5 (deterministic). "sample":
+    # actual n_sec decision. mode="head": "greedy" is argmax over the classifier logits
-    # a Bernoulli draw at sigmoid(logit) (a real sample from the learned length
+    # (deterministic — the conditional mode, not a sample); "sample" is a categorical draw
-    # distribution, at the cost of an extra RNG draw per slot).
+    # from softmax(logits) (a real sample from the learned count distribution). mode=
-    stop_sampling: str = "greedy"
+    # "stop_token": "greedy" is sigmoid(stop_logit) >= 0.5 per slot (deterministic);
+    # "sample" is a Bernoulli draw at sigmoid(stop_logit) per slot. Renamed from
+    # "stop_sampling" (gitea #86), which is still accepted as a deprecated alias since it
+    # appears in existing checkpoints' model_config.
+    sampling: str = "greedy"
 
     @classmethod
     def from_dict(cls, d: dict | None) -> "NSecConfig":
@@ -437,7 +453,7 @@ class NSecConfig:
             mode=d.get("mode", "head"),
             lambda_weight=d.get("lambda", 0.1),
             owner=d.get("owner", "stage2"),
-            stop_sampling=d.get("stop_sampling", "greedy"),
+            sampling=d.get("sampling", d.get("stop_sampling", "greedy")),
         )
 
     def to_dict(self) -> dict:
@@ -445,7 +461,7 @@ class NSecConfig:
             "mode": self.mode,
             "lambda": self.lambda_weight,
             "owner": self.owner,
-            "stop_sampling": self.stop_sampling,
+            "sampling": self.sampling,
         }
 
 
@@ -927,6 +943,8 @@ def git_hash() -> str:
 
 
 def auto_device() -> torch.device:
+    import torch
+
     if torch.cuda.is_available():
         return torch.device("cuda")
     if torch.backends.mps.is_available():
@@ -972,6 +990,8 @@ def estimate_batch_size(
     inference (e.g. `predict`), which uses a much lower per-sample memory
     calibration since there's no backward graph or optimizer state.
     """
+    import torch
+
     if device.type != "cuda":
         raise ValueError(f"--batch-size auto is only supported on cuda devices, got {device.type!r}")
     device_index = device.index if device.index is not None else torch.cuda.current_device()
@@ -1075,6 +1095,27 @@ def _set_path(d: dict, dotted: str, value) -> None:
     cur[parts[-1]] = value
 
 
+def _pop_path(d: dict, dotted: str) -> None:
+    """Remove a dotted path from a nested dict, if present. No-op if any
+    component along the path is missing."""
+    parts = dotted.split(".")
+    cur = d
+    for part in parts[:-1]:
+        if not isinstance(cur, dict) or part not in cur:
+            return
+        cur = cur[part]
+    if isinstance(cur, dict):
+        cur.pop(parts[-1], None)
+
+
+# Config keys renamed within v0.3 itself (not part of the v0.2->v0.3 migration
+# above) — normalized by migrate_config so a config.toml still using an older
+# v0.3 key name keeps passing validate_config_keys.
+_RENAMED_KEYS = {
+    "stage2_model.n_sec.stop_sampling": "stage2_model.n_sec.sampling",  # gitea #86
+}
+
+
 def _deep_merge(base: dict, override: dict) -> dict:
     """Recursively merge `override` onto a copy of `base`.
 
@@ -1093,6 +1134,72 @@ def _deep_merge(base: dict, override: dict) -> dict:
     return result
 
 
+@dataclass(frozen=True)
+class InferenceOverride:
+    """One dotted `model_config` path that `giant.checkpoint_io.load_for_inference`
+    is allowed to change on an already-trained checkpoint, without retraining.
+
+    A path only belongs here if it affects neither module construction/tensor
+    shapes nor the data preprocessing the normalizers/vocab maps were fit
+    under — see the module docstring on `giant.model.summary` for the class
+    of key this targets (`_fingerprint`'s "plain scalar attribute" leaves),
+    and `giant.checkpoint_io.apply_config_overrides` for where this is used.
+    """
+
+    why: str
+    choices: tuple[str, ...] | None = None
+    minimum: float | None = None
+    numeric: bool = False  # int/float leaf (vs. str, the default)
+
+    def check(self, path: str, value: object) -> None:
+        if self.choices is not None:
+            if value not in self.choices:
+                raise ValueError(f"{path} = {value!r} — must be one of {self.choices}")
+            return
+        if self.numeric:
+            if isinstance(value, bool) or not isinstance(value, (int, float)):
+                raise ValueError(f"{path} = {value!r} — must be a number")
+            if self.minimum is not None and value < self.minimum:
+                raise ValueError(f"{path} = {value!r} — must be >= {self.minimum}")
+
+
+# Inference-safe dotted `model_config` paths — the allowlist gitea #87 asked
+# for, so a typo or a shape-bearing key (e.g. "stage1_model.hidden_dim")
+# raises a clear CheckpointCompatibilityError instead of surfacing as an
+# opaque load_state_dict size mismatch later. Extend this table, not a
+# per-call bypass, when a new inference-only key needs the capability.
+INFERENCE_OVERRIDES: dict[str, InferenceOverride] = {
+    "stage2_model.n_sec.sampling": InferenceOverride(
+        why="giant.sample's n_sec head/stop-token sampling reads this at sample time only (gitea #86)",
+        choices=STOP_SAMPLING_CHOICES,
+    ),
+    "stage1_model.ddpm.n_steps": InferenceOverride(
+        why="giant.model.schedule.CosineSchedule's step count, resolved at sample time",
+        numeric=True,
+        minimum=1,
+    ),
+    "stage2_model.ddpm.n_steps": InferenceOverride(
+        why="giant.model.schedule.CosineSchedule's step count, resolved at sample time",
+        numeric=True,
+        minimum=1,
+    ),
+    "stage2_model.particle_type.other_policy": InferenceOverride(
+        why="giant.rollout resolves an 'other'-bucket secondary's PDG code with this at rollout time",
+        choices=OTHER_POLICY_CHOICES,
+    ),
+    "stage1_model.router.temperature": InferenceOverride(
+        why="giant.model.routers.EnergyRouter.temperature, a plain constructor attribute",
+        numeric=True,
+        minimum=1e-6,
+    ),
+    "stage2_model.router.temperature": InferenceOverride(
+        why="giant.model.routers.EnergyRouter.temperature, a plain constructor attribute",
+        numeric=True,
+        minimum=1e-6,
+    ),
+}
+
+
 @dataclass(frozen=True)
 class FlagSpec:
     """One CLI flag's mapping into the config-overrides tree.
@@ -1271,11 +1378,21 @@ def migrate_config(cfg: dict) -> dict:
     (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.
-    """
-    if _get_path(cfg, "meta.config_version") == CONFIG_VERSION:
-        return copy.deepcopy(cfg)
 
+    Independently of the v0.2/v0.3 branch below, `_RENAMED_KEYS` normalizes
+    keys renamed within v0.3 itself (e.g. `stop_sampling` -> `sampling`,
+    gitea #86) so a config.toml written against an older v0.3 key name still
+    passes `validate_config_keys`.
+    """
     cfg = copy.deepcopy(cfg)
+    for old_path, new_path in _RENAMED_KEYS.items():
+        if _get_path(cfg, old_path) is not None and _get_path(cfg, new_path) is None:
+            _set_path(cfg, new_path, _get_path(cfg, old_path))
+        _pop_path(cfg, old_path)
+
+    if _get_path(cfg, "meta.config_version") == CONFIG_VERSION:
+        return cfg
+
     old_train = cfg.pop("train", {})
     old_model = cfg.pop("model", {})
     old_router = dict(old_model.pop("router", {}))
@@ -1512,9 +1629,9 @@ def validate_config(cfg: dict, *, resume: bool = False) -> None:
                 "conditioning to hang an EOS decision off"
             )
 
-    stop_sampling = _get_path(cfg, "stage2_model.n_sec.stop_sampling")
+    n_sec_sampling = _get_path(cfg, "stage2_model.n_sec.sampling")
-    if stop_sampling not in ("greedy", "sample"):
+    if n_sec_sampling not in STOP_SAMPLING_CHOICES:
-        raise ValueError(f"stage2_model.n_sec.stop_sampling = {stop_sampling!r} — must be 'greedy' or 'sample'")
+        raise ValueError(f"stage2_model.n_sec.sampling = {n_sec_sampling!r} — must be 'greedy' or 'sample'")
 
     precision = _get_path(cfg, "train.precision")
     if precision not in ("fp32", "bf16"):
@@ -1569,6 +1686,8 @@ def validate_config(cfg: dict, *, resume: bool = False) -> None:
                 "'energy_desc' (the only implemented ordering; see "
                 "AutoregressiveConfig.order's docstring)"
             )
+        from giant.model.history import HISTORY_REGISTRY
+
         history = _get_path(cfg, "stage2_model.autoregressive.history")
         if history not in HISTORY_REGISTRY:
             raise ValueError(
@@ -1770,6 +1889,9 @@ def resolve_default_out_dir(cfg: dict, base: Path = Path("checkpoints")) -> Path
 
 
 def seed_everything(seed: int) -> None:
+    import numpy as np
+    import torch
+
     random.seed(seed)
     np.random.seed(seed)
     torch.manual_seed(seed)
@@ -1823,6 +1945,8 @@ def build_run_meta(
     n_val_events: int,
     n_train_steps: int,
 ) -> dict:
+    import torch
+
     return {
         "config_version": CONFIG_VERSION,
         "git_hash": git_hash(),

giant/model/builders.py

@@ -138,7 +138,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
                 n_sec_head_cfg=s2_spec.heads.n_sec.to_dict(),
                 type_head_cfg=s2_spec.heads.type.to_dict(),
                 build_stop_head=stop_token,
-                stop_sampling=s2_spec.n_sec.stop_sampling,
+                n_sec_sampling=s2_spec.n_sec.sampling,
                 stop_head_cfg=s2_spec.heads.n_sec.to_dict(),
             )
         else:
@@ -168,6 +168,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
                 cond_enc=shared_cond_enc,
                 n_sec_head_cfg=s2_spec.heads.n_sec.to_dict(),
                 type_head_cfg=s2_spec.heads.type.to_dict(),
+                n_sec_sampling=s2_spec.n_sec.sampling,
             )
 
     return result

giant/model/models.py

@@ -370,6 +370,7 @@ class Stage2OneShot(StageModel):
         cond_enc: ConditionEncoder | None = None,
         n_sec_head_cfg: dict | None = None,
         type_head_cfg: dict | None = None,
+        n_sec_sampling: str = "greedy",
     ) -> None:
         super().__init__(
             pdg_vocab,
@@ -383,6 +384,7 @@ class Stage2OneShot(StageModel):
             particle_type_cfg=particle_type_cfg,
             cond_enc=cond_enc,
         )
+        self.n_sec_sampling = n_sec_sampling
         self._build_context_fusion(x_dim, context_dim, cond_out_dim)
         target = self.particle_type_cfg.target
         type_head_out_dim = None if target == "physical" else k_max * self.type_dim
@@ -498,7 +500,7 @@ class Stage2Autoregressive(StageModel):
         n_sec_head_cfg: dict | None = None,
         type_head_cfg: dict | None = None,
         build_stop_head: bool = False,
-        stop_sampling: str = "greedy",
+        n_sec_sampling: str = "greedy",
         stop_head_cfg: dict | None = None,
     ) -> None:
         super().__init__(
@@ -514,7 +516,7 @@ class Stage2Autoregressive(StageModel):
             cond_enc=cond_enc,
         )
         self.history_kind = history
-        self.stop_sampling = stop_sampling
+        self.n_sec_sampling = n_sec_sampling
         self.context_adapter = ContextAdapter(x_dim, context_dim)
         self.base_fuse = nn.Sequential(
             nn.Linear(cond_out_dim + context_dim, cond_out_dim),

giant/model/summary.py

@@ -37,7 +37,7 @@ from dataclasses import dataclass, field
 
 import torch.nn as nn
 
-from giant.config import _get_path, _set_path, leaf_paths
+from giant.config import INFERENCE_OVERRIDES, _get_path, _set_path, leaf_paths
 from giant.model.builders import build_critics, build_models
 from giant.model.trunks import RoutedTrunk
 
@@ -111,6 +111,7 @@ class ModelSummary:
     pdg_vocab: int
     mat_vocab: int
     vocab_caveats: list[str] = field(default_factory=list)
+    overridable: list[str] = field(default_factory=list)
 
 
 def _build_model_config(cfg: dict, pdg_vocab: int, mat_vocab: int) -> dict:
@@ -140,7 +141,7 @@ def _fingerprint(modules: dict[str, nn.Module]) -> list:
     exist, every parameter's/buffer's shape+dtype (never values — those are
     randomly initialized and irrelevant to *structure*), and every plain
     scalar attribute any module stores on itself (e.g. `Stage2Autoregressive
-    .stop_sampling`, `EnergyRouter.temperature`) — this is what makes a
+    .n_sec_sampling`, `EnergyRouter.temperature`) — this is what makes a
     non-parametric key's effect on construction observable."""
     sig = []
     for stage_name, module in modules.items():
@@ -234,6 +235,8 @@ def summarize_model(cfg: dict, pdg_vocab: int, mat_vocab: int) -> ModelSummary:
         else:
             inert.append(path)
 
+    overridable = sorted(p for p in in_scope if p in INFERENCE_OVERRIDES)
+
     return ModelSummary(
         modules=modules,
         consumed=sorted(consumed),
@@ -242,6 +245,7 @@ def summarize_model(cfg: dict, pdg_vocab: int, mat_vocab: int) -> ModelSummary:
         pdg_vocab=pdg_vocab,
         mat_vocab=mat_vocab,
         vocab_caveats=_vocab_caveats(cfg),
+        overridable=overridable,
     )
 
 
@@ -309,6 +313,12 @@ def render_summary(summary: ModelSummary) -> str:
     else:
         lines.append("  (none)")
 
+    if summary.overridable:
+        lines.append("")
+        lines.append("inference-overridable without retraining (giant predict/rollout --set):")
+        for path in summary.overridable:
+            lines.append(f"  {path}  ({INFERENCE_OVERRIDES[path].why})")
+
     if summary.vocab_caveats:
         lines.append("")
         lines.append("vocab placeholder caveats:")

giant/sample.py

@@ -261,7 +261,7 @@ def sample_secondaries_ar(
     slot's own stop logit (`predict_stop`, evaluated on the same prefix
     conditioning as the token itself — see `predict_type`'s docstring for
     why this needs no extra state) decides whether generation should have
-    already stopped, per `sec_decoder.stop_sampling` ("greedy": threshold at
+    already stopped, per `sec_decoder.n_sec_sampling` ("greedy": threshold at
     0; "sample": a Bernoulli draw at `sigmoid(logit)`). A row's own
     `n_sec_pred` is the first slot index where this fires; once every row in
     the batch has fired, the loop breaks before spending a model call on the
@@ -345,7 +345,7 @@ def sample_secondaries_ar(
                 slot_idx,
                 hist=hist,
             ).squeeze(1)
-            if sec_decoder.stop_sampling == "sample":
+            if sec_decoder.n_sec_sampling == "sample":
                 stop_now = torch.rand(B, device=device) < torch.sigmoid(stop_logit)
             else:
                 stop_now = stop_logit >= 0.0
@@ -496,7 +496,12 @@ def resolve_n_sec(
 
     Raises if neither stage owns any n_sec mechanism at all — the only way
     that happens is `stage2_model.n_sec.mode = "truth"`, which is not a valid
-    rollout-/predict-capable checkpoint."""
+    rollout-/predict-capable checkpoint.
+
+    `n_sec.mode = "head"` resolves the classifier logits per
+    `sec_decoder.n_sec_sampling`: "greedy" (default) takes the conditional
+    mode via argmax; "sample" draws a real sample from the learned count
+    distribution via `torch.multinomial` on the softmax — see gitea #86."""
     if n_sec_pred is not None:
         return n_sec_pred
     if getattr(sec_decoder, "stop_head", None) is not None:
@@ -508,4 +513,6 @@ def resolve_n_sec(
             "'truth' is standalone-evaluation-only"
         )
     logits = sec_decoder.predict_n_sec(cond_cont, cond_cat, stage1_out)
+    if sec_decoder.n_sec_sampling == "sample":
+        return torch.multinomial(logits.softmax(dim=-1), 1).squeeze(-1)
     return logits.argmax(dim=-1)

giant/tools/dwarf.py

@@ -5,6 +5,8 @@ simulation-fanout tools into one Typer app so there's a single command name
 (and `--help`) to remember instead of five differently-hyphenated ones.
 """
 
+from __future__ import annotations
+
 import os
 from enum import Enum
 from pathlib import Path
@@ -14,20 +16,15 @@ import typer
 from typing_extensions import Annotated
 
 from giant.config import Conditioning
-from giant.tools.bump_dataset_version import (
+
-    run_bump_gen,
+# DATA_DEFAULT/SCAN_DIR_DEFAULT are Typer option defaults (evaluated at
-    run_bump_schema,
+# decoration time below), so that one name has to stay eager — the module
-    run_create_manifest,
+# itself is stdlib-only, so it costs nothing. Every other giant.tools.*
-    run_status,
+# import here is deferred into the one command body that uses it, since
-    run_update_manifest,
+# several (steps_to_parquet: uproot/awkward/polars; warm_setup_cache:
-)
+# giant.pipeline -> torch; geometry_oracle: pandas) are expensive and
-from giant.tools.create_root_files import run_make_root
+# `dwarf --help`/tab-completion shouldn't pay for all of them upfront.
-from giant.tools.geometry_oracle import run_build_geometry_oracle
+from giant.tools.hparam_scan import DATA_DEFAULT, SCAN_DIR_DEFAULT
-from giant.tools.hparam_scan import DATA_DEFAULT, SCAN_DIR_DEFAULT, run_hparam_scan
-from giant.tools.migrate_geant_steps import run_migration
-from giant.tools.steps_to_parquet import convert_steps_to_parquet
-from giant.tools.steps_to_parquet_parallel import run_parallel_job
-from giant.tools.warm_setup_cache import run_warm_setup_cache
 
 app = typer.Typer(no_args_is_help=True)
 
@@ -121,6 +118,9 @@ def convert(
     ] = None,
 ) -> None:
     """Convert ROOT Steps tree(s) to Parquet."""
+    from giant.tools.steps_to_parquet import convert_steps_to_parquet
+    from giant.tools.steps_to_parquet_parallel import run_parallel_job
+
     if jobs < 1:
         typer.echo("error: --jobs must be >= 1", err=True)
         raise typer.Exit(1)
@@ -183,6 +183,8 @@ def migrate(
     ] = False,
 ) -> None:
     """One-time migration into the versioned raw/processed/pools/derived layout."""
+    from giant.tools.migrate_geant_steps import run_migration
+
     run_migration(str(root), execute=execute, copy=copy)
 
 
@@ -207,6 +209,8 @@ def bump_gen(
     root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
 ) -> None:
     """Cut a new raw generation."""
+    from giant.tools.bump_dataset_version import run_bump_gen
+
     run_bump_gen(
         kind=kind,
         reason=reason,
@@ -240,6 +244,8 @@ def bump_schema(
     root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
 ) -> None:
     """Cut a new schema within a gen."""
+    from giant.tools.bump_dataset_version import run_bump_schema
+
     run_bump_schema(
         kind=kind,
         gen=gen,
@@ -257,6 +263,8 @@ def status(
     root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
 ) -> None:
     """List existing gens/schemas per kind."""
+    from giant.tools.bump_dataset_version import run_status
+
     run_status(str(root))
 
 
@@ -281,6 +289,8 @@ def update_manifest(
     ] = False,
 ) -> None:
     """Repoint manifest(s) to a new gen and/or schema, verifying all target files exist."""
+    from giant.tools.bump_dataset_version import run_update_manifest
+
     run_update_manifest([str(m) for m in manifests], schema=schema, execute=execute, gen=gen)
 
 
@@ -311,6 +321,8 @@ def create_manifest(
     ] = False,
 ) -> None:
     """Create a new manifest from a list of parquet files."""
+    from giant.tools.bump_dataset_version import run_create_manifest
+
     run_create_manifest(
         [str(f) for f in files],
         execute=execute,
@@ -358,6 +370,8 @@ def make_root(
     ] = False,
 ) -> None:
     """Generate new ROOT shards via a minicalosim executable."""
+    from giant.tools.create_root_files import run_make_root
+
     _warn_if_exceeds_shared_quota(jobs, "--jobs")
     run_make_root(
         executable=executable,
@@ -423,6 +437,8 @@ def build_geometry_oracle(
     ] = 2000,
 ) -> None:
     """Fit a position -> (material, layer_id) oracle for `giant rollout`."""
+    from giant.tools.geometry_oracle import run_build_geometry_oracle
+
     run_build_geometry_oracle(
         data=data,
         out=out,
@@ -515,6 +531,8 @@ def warm_cache(
     such entry across every run) skips straight to training. See
     giant/data/setup_cache.py.
     """
+    from giant.tools.warm_setup_cache import run_warm_setup_cache
+
     flag_overrides = {
         "--val-fraction": val_fraction,
         "--seed": seed,
@@ -557,6 +575,8 @@ def hparam_scan(
     dry_run: Annotated[bool, typer.Option("--dry-run")] = False,
 ) -> None:
     """Grid-scan dropout x n_blocks x hidden_dim via sequential `giant train` runs."""
+    from giant.tools.hparam_scan import run_hparam_scan
+
     run_hparam_scan(data=data, scan_dir=scan_dir, seed=seed, dry_run=dry_run)
 
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "giant"
-version = "0.3.11"
+version = "0.3.16"
 description = "Geant4 step-function surrogate via conditional flow matching"
 readme = "README.md"
 requires-python = ">=3.12"

tests/test_catalog.py

@@ -261,3 +261,32 @@ def test_sec_count_per_step_by_species_zero_row_is_per_species(bundle):
     for j, _ in enumerate(cols):
         if j != g:
             assert ref[0][j] == 3 and sum(row[j] for row in ref[1:]) == 0
+
+
+# ---------------------------------------------------------------------------
+# eval_cost_per_step
+# ---------------------------------------------------------------------------
+
+
+def test_eval_cost_per_step_unavailable_without_timing(bundle: Bundle):
+    # `bundle`'s RolloutSpec carries no `timing` -> no rollout to compare.
+    spec = get_spec("eval_cost_per_step")
+    r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx)
+    assert r.kind == "unavailable"
+    assert r.payload["note"]
+
+
+def test_eval_cost_per_step_bar_with_timing(ctx: Context):
+    spec = get_spec("eval_cost_per_step")
+    rs = RolloutSpec(
+        "rollout",
+        _rollout_frame(),
+        timing={"us_per_step": 12.5, "write_us_per_step": 2.5},
+    )
+    b = Bundle.open([rs], _reference_frame(), ctx)
+    r = spec.finalize([spec.compute_partial(b)], ctx)
+    assert r.kind == "bar"
+    assert r.payload["series"]["rollout"] == [12.5, 2.5, 15.0]
+    assert len(r.payload["reference"]) == 3
+    assert r.payload["log_y"] is True
+    assert "rollout" in r.meta["speedup_vs_geant4_total"]

tests/test_checkpoint_io.py

@@ -14,6 +14,7 @@ from giant import config as gconfig
 from giant.checkpoint_io import (
     CheckpointCompatibilityError,
     InferenceContext,
+    apply_config_overrides,
     conditioning_axes,
     load_for_inference,
     stage_cfg,
@@ -278,3 +279,122 @@ def test_stage_cfg_new_shape_returns_subdict():
 def test_stage_cfg_v02_flat_shape_returns_empty_dict():
     model_cfg = {"hidden_dim": 32, "n_blocks": 4}
     assert stage_cfg(model_cfg, "stage2") == {}
+
+
+# ---------------------------------------------------------------------------
+# config_overrides (gitea #87)
+# ---------------------------------------------------------------------------
+
+
+def _router_model_cfg() -> dict:
+    cfg = _model_cfg()
+    cfg["stage1_model"]["router"] = {"enabled": True, "type": "energy", "n_experts": 2}
+    return cfg
+
+
+def test_config_override_n_sec_sampling_changes_stage2_attribute(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ctx = load_for_inference(
+        checkpoint,
+        torch.device("cpu"),
+        "predict",
+        config_overrides={"stage2_model.n_sec.sampling": "sample"},
+    )
+    assert ctx.stage2 is not None
+    assert ctx.stage2.n_sec_sampling == "sample"
+    assert ctx.config_overrides == {"stage2_model.n_sec.sampling": "sample"}
+
+
+def test_config_override_ddpm_n_steps_changes_context_fields(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ctx = load_for_inference(
+        checkpoint,
+        torch.device("cpu"),
+        "predict",
+        config_overrides={"stage1_model.ddpm.n_steps": 42, "stage2_model.ddpm.n_steps": 7},
+    )
+    assert ctx.stage1_ddpm_steps == 42
+    assert ctx.stage2_ddpm_steps == 7
+
+
+def test_config_override_other_policy_changes_context_field(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ctx = load_for_inference(
+        checkpoint,
+        torch.device("cpu"),
+        "predict",
+        config_overrides={"stage2_model.particle_type.other_policy": "modal"},
+    )
+    assert ctx.other_policy == "modal"
+
+
+def test_config_override_router_temperature_changes_router_attribute(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path, model_cfg=_router_model_cfg())
+    ctx = load_for_inference(
+        checkpoint,
+        torch.device("cpu"),
+        "predict",
+        config_overrides={"stage1_model.router.temperature": 1.5},
+    )
+    assert ctx.stage1 is not None
+    assert ctx.stage1.trunk.router.temperature == pytest.approx(1.5)
+
+
+def test_config_override_no_overrides_defaults_to_empty_dict(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")
+    assert ctx.config_overrides == {}
+
+
+def test_config_override_unknown_path_raises(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    with pytest.raises(CheckpointCompatibilityError, match="not an inference-safe override"):
+        load_for_inference(
+            checkpoint,
+            torch.device("cpu"),
+            "predict",
+            config_overrides={"stage2_model.n_sec.typo": "sample"},
+        )
+
+
+def test_config_override_shape_bearing_key_raises_up_front(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    with pytest.raises(CheckpointCompatibilityError, match="not an inference-safe override"):
+        load_for_inference(
+            checkpoint,
+            torch.device("cpu"),
+            "predict",
+            config_overrides={"stage1_model.hidden_dim": 999},
+        )
+
+
+def test_config_override_bad_value_raises(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    with pytest.raises(CheckpointCompatibilityError, match="must be one of"):
+        load_for_inference(
+            checkpoint,
+            torch.device("cpu"),
+            "predict",
+            config_overrides={"stage2_model.n_sec.sampling": "maybe"},
+        )
+
+
+def test_apply_config_overrides_no_overrides_returns_same_object():
+    cfg = _model_cfg()
+    assert apply_config_overrides(cfg, None) is cfg
+    assert apply_config_overrides(cfg, {}) is cfg
+
+
+def test_apply_config_overrides_migrates_legacy_flat_model_config_first():
+    legacy_cfg = {
+        "pdg_vocab": len(PDG_MAP),
+        "mat_vocab": len(MAT_MAP),
+        "hidden_dim": 32,
+        "n_blocks": 4,
+        "emb_dim": 8,
+        "dropout": 0.1,
+        "k_max": 5,
+    }
+    merged = apply_config_overrides(legacy_cfg, {"stage1_model.ddpm.n_steps": 10})
+    assert merged["stage1_model"]["ddpm"]["n_steps"] == 10
+    assert merged["stage1_model"]["hidden_dim"] == 32

tests/test_cli_predict.py

@@ -172,3 +172,43 @@ def test_predict_exits_1_on_checkpoint_missing_model_config(tmp_path):
 
     assert result.exit_code == 1
     assert "checkpoint has no model_config" in result.output
+
+
+# ---------------------------------------------------------------------------
+# --set (gitea #87)
+# ---------------------------------------------------------------------------
+
+
+def test_predict_set_flag_without_equals_exits_1(tmp_path):
+    checkpoint = tmp_path / "missing.pt"
+
+    result = runner.invoke(
+        app,
+        ["predict", "dummy.parquet", "--checkpoint", str(checkpoint), "--set", "sampling"],
+    )
+
+    assert result.exit_code == 1
+    assert "must be 'dotted.path=value'" in result.output
+
+
+def test_predict_set_flag_disallowed_path_surfaces_compat_error(tmp_path):
+    checkpoint = tmp_path / "ckpt.pt"
+    torch.save(
+        {"model_config": {"stage1_model": {}, "stage2_model": {}}, "sec_decoder": {}, "normalizer": {"sec_phys": {}}},
+        checkpoint,
+    )
+
+    result = runner.invoke(
+        app,
+        [
+            "predict",
+            "dummy.parquet",
+            "--checkpoint",
+            str(checkpoint),
+            "--set",
+            "stage1_model.hidden_dim=999",
+        ],
+    )
+
+    assert result.exit_code == 1
+    assert "not an inference-safe override" in result.output

tests/test_cli_rollout.py

@@ -8,11 +8,52 @@ from __future__ import annotations
 import torch
 from typer.testing import CliRunner
 
-from giant.cli import app
+from giant.cli import _build_rollout_timing, app
 
 runner = CliRunner()
 
 
+def test_build_rollout_timing_excludes_synthetic_rows_from_per_step_cost():
+    # 100 rows total, 30 of them synthetic termination markers (escape) ->
+    # us_per_step should be normalized over the 70 physical rows only, the
+    # same unit giant.analysis.geant4_reference measures Geant4 in.
+    timing = _build_rollout_timing(
+        setup_s=1.0,
+        rollout_s=10.0,
+        write_s=2.0,
+        n_rows=100,
+        termination_reason_counts={"escaped": 30, "natural_end": 70},
+        n_seed_events=5,
+        device="cpu",
+        torch_threads=4,
+    )
+    assert timing["n_rows"] == 100
+    assert timing["n_physical_rows"] == 70
+    assert timing["n_physical_rows"] < timing["n_rows"]
+    assert timing["sample_s"] == 8.0  # rollout_s - write_s
+    assert timing["us_per_step"] == 8.0 / 70 * 1e6
+    assert timing["write_us_per_step"] == 2.0 / 70 * 1e6
+    assert timing["ms_per_event"] == 10.0 / 5 * 1e3
+    assert timing["device"] == "cpu" and timing["torch_threads"] == 4
+
+
+def test_build_rollout_timing_handles_zero_physical_rows_and_events():
+    timing = _build_rollout_timing(
+        setup_s=1.0,
+        rollout_s=1.0,
+        write_s=0.0,
+        n_rows=5,
+        termination_reason_counts={"escaped": 5},
+        n_seed_events=0,
+        device="cpu",
+        torch_threads=1,
+    )
+    assert timing["n_physical_rows"] == 0
+    assert timing["us_per_step"] is None
+    assert timing["write_us_per_step"] is None
+    assert timing["ms_per_event"] is None
+
+
 def test_rollout_exits_1_on_checkpoint_missing_model_config(tmp_path):
     checkpoint = tmp_path / "bad.pt"
     torch.save({"sec_decoder": {}, "normalizer": {"sec_phys": {}}}, checkpoint)
@@ -31,3 +72,28 @@ def test_rollout_exits_1_on_checkpoint_missing_model_config(tmp_path):
 
     assert result.exit_code == 1
     assert "checkpoint has no model_config" in result.output
+
+
+def test_rollout_set_flag_disallowed_path_surfaces_compat_error(tmp_path):
+    checkpoint = tmp_path / "ckpt.pt"
+    torch.save(
+        {"model_config": {"stage1_model": {}, "stage2_model": {}}, "sec_decoder": {}, "normalizer": {"sec_phys": {}}},
+        checkpoint,
+    )
+
+    result = runner.invoke(
+        app,
+        [
+            "rollout",
+            "dummy.parquet",
+            "--checkpoint",
+            str(checkpoint),
+            "--geometry",
+            "dummy_geometry.pkl",
+            "--set",
+            "stage2_model.n_sec.typo=sample",
+        ],
+    )
+
+    assert result.exit_code == 1
+    assert "not an inference-safe override" in result.output

tests/test_cli_train_overrides.py

@@ -20,7 +20,7 @@ def _invoke_and_capture_cfg(monkeypatch, tmp_path: Path, args: list[str]) -> dic
     def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
         captured["cfg"] = cfg
 
-    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+    monkeypatch.setattr("giant.pipeline.run_train_job", _fake_run_train_job)
 
     result = runner.invoke(
         cli.app,
@@ -125,7 +125,7 @@ def test_stage2_init_from_and_freeze_flags_land_in_cfg_and_dont_touch_stage1(mon
 
 
 def test_batch_size_invalid_string_errors(monkeypatch, tmp_path):
-    monkeypatch.setattr(cli, "run_train_job", lambda *a, **kw: None)
+    monkeypatch.setattr("giant.pipeline.run_train_job", lambda *a, **kw: None)
     result = runner.invoke(
         cli.app,
         ["train", "dummy.parquet", "--out", str(tmp_path / "run"), "--batch-size", "not-a-number"],
@@ -140,7 +140,7 @@ def test_out_dir_resolution_prefers_explicit_out_over_resume(monkeypatch, tmp_pa
     def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
         captured["out_dir"] = out_dir
 
-    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+    monkeypatch.setattr("giant.pipeline.run_train_job", _fake_run_train_job)
 
     resume_dir = tmp_path / "resumed_run"
     resume_dir.mkdir()
@@ -161,7 +161,7 @@ def test_out_dir_resolution_falls_back_to_resume_parent(monkeypatch, tmp_path):
     def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
         captured["out_dir"] = out_dir
 
-    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+    monkeypatch.setattr("giant.pipeline.run_train_job", _fake_run_train_job)
 
     resume_dir = tmp_path / "resumed_run"
     resume_dir.mkdir()
@@ -178,7 +178,7 @@ def test_out_dir_resolution_defaults_when_neither_out_nor_resume_given(monkeypat
     def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
         captured["out_dir"] = out_dir
 
-    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+    monkeypatch.setattr("giant.pipeline.run_train_job", _fake_run_train_job)
     monkeypatch.chdir(tmp_path)
 
     result = runner.invoke(cli.app, ["train", "dummy.parquet"])
@@ -192,7 +192,7 @@ def test_batch_size_auto_estimates_and_echoes(monkeypatch, tmp_path):
     def _fake_run_train_job(*, data, cfg, out_dir, num_workers, **kwargs):
         captured["batch_size"] = cfg["train"]["batch_size"]
 
-    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+    monkeypatch.setattr("giant.pipeline.run_train_job", _fake_run_train_job)
     monkeypatch.setattr(cli.gconfig, "estimate_batch_size", lambda hidden_dim, n_blocks, device: 123)
 
     result = runner.invoke(

tests/test_condor.py

@@ -252,6 +252,22 @@ def test_compute_one_from_run_dir(tmp_path: Path):
     assert list(partial.data["r"]) == ["rollout"]
 
 
+def test_timing_survives_plot_meta_to_compute_one(tmp_path: Path):
+    yaml_path = _write_inputs(tmp_path)
+    d = yaml.safe_load(yaml_path.read_text())
+    d["timing"] = {"us_per_step": 7.0, "write_us_per_step": 1.0}
+    yaml_path.write_text(yaml.safe_dump(d))
+
+    run_dir = _prep([yaml_path])
+    meta = RunMeta.load(run_dir / "run_meta.json")
+    assert meta.rollouts[0]["plot_meta"]["timing"] == {"us_per_step": 7.0, "write_us_per_step": 1.0}
+
+    out = compute_one("eval_cost_per_step", run_dir)
+    reduced = Reduced(**Partial.load(out).data["reduced"])
+    assert reduced.kind == "bar"
+    assert reduced.payload["series"]["rollout"] == [7.0, 1.0, 8.0]
+
+
 def test_compute_reduced_explicit_paths(tmp_path: Path):
     run_dir = _prep([_write_inputs(tmp_path)])
     meta = RunMeta.load(run_dir / "run_meta.json")

tests/test_config.py

@@ -171,17 +171,40 @@ def test_n_sec_config_owner_defaults_to_stage2():
 def test_n_sec_config_owner_round_trips():
     n_sec = gconfig.NSecConfig.from_dict({"mode": "head", "owner": "stage1"})
     assert n_sec.owner == "stage1"
-    assert n_sec.to_dict() == {"mode": "head", "lambda": 0.1, "owner": "stage1", "stop_sampling": "greedy"}
+    assert n_sec.to_dict() == {"mode": "head", "lambda": 0.1, "owner": "stage1", "sampling": "greedy"}
 
 
-def test_n_sec_config_stop_sampling_defaults_to_greedy():
+def test_n_sec_config_sampling_defaults_to_greedy():
-    assert gconfig.NSecConfig().stop_sampling == "greedy"
+    assert gconfig.NSecConfig().sampling == "greedy"
 
 
-def test_n_sec_config_stop_sampling_round_trips():
+def test_n_sec_config_sampling_round_trips():
+    n_sec = gconfig.NSecConfig.from_dict({"mode": "stop_token", "sampling": "sample"})
+    assert n_sec.sampling == "sample"
+    assert n_sec.to_dict()["sampling"] == "sample"
+
+
+def test_n_sec_config_stop_sampling_alias_still_honored():
+    """gitea #86: stop_sampling was renamed to sampling; old checkpoints'
+    model_config still carries the old key and must keep working."""
     n_sec = gconfig.NSecConfig.from_dict({"mode": "stop_token", "stop_sampling": "sample"})
-    assert n_sec.stop_sampling == "sample"
+    assert n_sec.sampling == "sample"
-    assert n_sec.to_dict()["stop_sampling"] == "sample"
+    assert "stop_sampling" not in n_sec.to_dict()
+
+
+def test_n_sec_config_sampling_key_wins_over_stop_sampling_alias():
+    n_sec = gconfig.NSecConfig.from_dict({"sampling": "sample", "stop_sampling": "greedy"})
+    assert n_sec.sampling == "sample"
+
+
+def test_migrate_config_renames_stop_sampling_key():
+    cfg = {
+        "meta": {"config_version": gconfig.CONFIG_VERSION},
+        "stage2_model": {"n_sec": {"stop_sampling": "sample"}},
+    }
+    migrated = gconfig.migrate_config(cfg)
+    assert gconfig._get_path(migrated, "stage2_model.n_sec.sampling") == "sample"
+    assert gconfig._get_path(migrated, "stage2_model.n_sec.stop_sampling") is None
 
 
 # ---------------------------------------------------------------------------
@@ -851,13 +874,13 @@ def test_validate_config_stop_token_rejected_for_stage1_owner():
         assert "stop_token" in str(e) and "owner" in str(e)
 
 
-def test_validate_config_bad_stop_sampling_rejected():
+def test_validate_config_bad_n_sec_sampling_rejected():
-    cfg = _cfg_with(**{"stage2_model.n_sec.stop_sampling": "bogus"})
+    cfg = _cfg_with(**{"stage2_model.n_sec.sampling": "bogus"})
     try:
         gconfig.validate_config(cfg)
         assert False, "expected ValueError"
     except ValueError as e:
-        assert "stop_sampling" in str(e)
+        assert "sampling" in str(e)
 
 
 def test_validate_config_default_precision_is_fp32():

tests/test_model_summary.py

@@ -12,6 +12,8 @@ from giant.cli import app
 from giant.materials import MATERIAL_PROPERTIES
 from giant.model.summary import _NOT_BUILD_TIME, _built_modules, _vocab_caveats, summarize_model
 
+INFERENCE_OVERRIDES = gconfig.INFERENCE_OVERRIDES
+
 runner = CliRunner()
 
 _PDG_VOCAB = 300
@@ -53,6 +55,16 @@ def test_not_build_time_allow_list_has_no_stale_entries():
     assert not stale, f"_NOT_BUILD_TIME entries no longer in DEFAULT_CONFIG: {sorted(stale)}"
 
 
+def test_inference_overrides_allow_list_has_no_stale_entries():
+    in_scope = set(gconfig.leaf_paths(gconfig.DEFAULT_CONFIG))
+    stale = set(INFERENCE_OVERRIDES) - in_scope
+    assert not stale, f"INFERENCE_OVERRIDES entries no longer in DEFAULT_CONFIG: {sorted(stale)}"
+
+
+def test_default_config_overridable_lists_every_allowlisted_path(default_summary):
+    assert set(default_summary.overridable) == set(INFERENCE_OVERRIDES)
+
+
 def test_router_disabled_by_default_so_its_fields_are_inert(default_summary):
     assert "stage1_model.router.n_experts" in default_summary.inert
     assert "stage1_model.router.temperature" in default_summary.inert
@@ -135,3 +147,5 @@ def test_cli_default_smoke():
     assert "parameters" in result.output
     assert "trunk" in result.output
     assert "inert under this config" in result.output
+    assert "inference-overridable without retraining" in result.output
+    assert "stage2_model.n_sec.sampling" in result.output

tests/test_render.py

@@ -292,6 +292,20 @@ def test_render_one_of_each_kind(tmp_path: Path):
                 "ylabel": "frac",
             },
         ),
+        Reduced(
+            "cost",
+            "cost",
+            "bar",
+            "Cost",
+            "phase",
+            {
+                "labels": ["sample", "write", "total"],
+                "series": {"flow": [10.0, 1.0, 11.0]},
+                "reference": [5.0, 0.5, 5.5],
+                "ylabel": "us/step",
+                "log_y": True,
+            },
+        ),
         Reduced(
             "s",
             "species",

tests/test_sample.py

@@ -14,6 +14,7 @@ from giant.model.network import (
     stage2_trunk_sec_dim,
 )
 from giant.sample import (
+    resolve_n_sec,
     sample_flow,
     sample_secondaries,
     sample_secondaries_ar,
@@ -72,6 +73,7 @@ def _stage2_ar(
     mat: int = 2,
     k_max: int = 5,
     history: str = "markov",
+    n_sec_sampling: str = "greedy",
 ) -> Stage2Autoregressive:
     particle_cfg, material_cfg = _particle_material_cfg(_conditioning_for(target), emb_dim)
     return Stage2Autoregressive(
@@ -89,6 +91,7 @@ def _stage2_ar(
         history=history,
         attn_n_heads=2,
         attn_n_layers=1,
+        n_sec_sampling=n_sec_sampling,
     ).eval()
 
 
@@ -99,7 +102,7 @@ def _expected_type_dim(target: str, emb_dim: int) -> int:
 def _stage2_ar_stop_token(
     target: str,
     generator: str,
-    stop_sampling: str = "greedy",
+    n_sec_sampling: str = "greedy",
     emb_dim: int = 6,
     pdg: int = 3,
     mat: int = 2,
@@ -120,7 +123,7 @@ def _stage2_ar_stop_token(
         particle_type_cfg=ParticleTypeConfig(target=target),
         build_n_sec_head=False,
         build_stop_head=True,
-        stop_sampling=stop_sampling,
+        n_sec_sampling=n_sec_sampling,
     ).eval()
 
 
@@ -267,14 +270,14 @@ def test_sample_secondaries_ar_first_slot_has_no_history():
 # ── Stage2Autoregressive: n_sec.mode = "stop_token" ─────────────────────────
 
 
-@pytest.mark.parametrize("stop_sampling", ["greedy", "sample"])
+@pytest.mark.parametrize("n_sec_sampling", ["greedy", "sample"])
-def test_sample_secondaries_ar_stop_token_forced_stop_gives_zero_secondaries(stop_sampling):
+def test_sample_secondaries_ar_stop_token_forced_stop_gives_zero_secondaries(n_sec_sampling):
     """A stop_head pinned to a large positive logit fires at slot 0 for
     every row under both policies (greedy: sigmoid(logit) >= 0.5; sample:
     a Bernoulli draw at sigmoid(logit) ~= 1) — the loop should break before
     generating any token."""
     B, k_max = 4, 5
-    decoder = _stage2_ar_stop_token("physical", "flow", stop_sampling=stop_sampling, k_max=k_max)
+    decoder = _stage2_ar_stop_token("physical", "flow", n_sec_sampling=n_sec_sampling, k_max=k_max)
     _force_stop_head_logit(decoder, 50.0)
     cond_cont, cond_cat = _cond(B)
     stage1_out = torch.randn(B, X_DIM)
@@ -283,13 +286,13 @@ def test_sample_secondaries_ar_stop_token_forced_stop_gives_zero_secondaries(sto
     assert not sec_valid.any()
 
 
-@pytest.mark.parametrize("stop_sampling", ["greedy", "sample"])
+@pytest.mark.parametrize("n_sec_sampling", ["greedy", "sample"])
-def test_sample_secondaries_ar_stop_token_forced_never_stop_runs_to_k_max(stop_sampling):
+def test_sample_secondaries_ar_stop_token_forced_never_stop_runs_to_k_max(n_sec_sampling):
     """A stop_head pinned to a large negative logit never fires under either
     policy, so every row is capped at k_max (the safety cap, not a modeling
     ceiling)."""
     B, k_max = 4, 5
-    decoder = _stage2_ar_stop_token("physical", "flow", stop_sampling=stop_sampling, k_max=k_max)
+    decoder = _stage2_ar_stop_token("physical", "flow", n_sec_sampling=n_sec_sampling, k_max=k_max)
     _force_stop_head_logit(decoder, -50.0)
     cond_cont, cond_cat = _cond(B)
     stage1_out = torch.randn(B, X_DIM)
@@ -336,3 +339,60 @@ def test_sample_secondaries_ar_none_n_sec_pred_without_stop_head_raises():
     stage1_out = torch.randn(3, X_DIM)
     with pytest.raises(AssertionError):
         sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
+
+
+# ── resolve_n_sec: n_sec.mode = "head" sampling policy (gitea #86) ──────────
+
+
+def _force_n_sec_head_bias(decoder: Stage2Autoregressive, bias: torch.Tensor) -> None:
+    """Zeroes n_sec_head's weights and pins its bias, so predict_n_sec
+    returns `bias` (broadcast over the batch) as logits regardless of
+    conditioning — mirrors `_force_stop_head_logit`."""
+    assert decoder.n_sec_head is not None
+    last_linear = decoder.n_sec_head[-1]
+    with torch.no_grad():
+        last_linear.weight.zero_()
+        last_linear.bias.copy_(bias)
+
+
+@pytest.mark.parametrize("n_sec_sampling", ["greedy", "sample"])
+def test_resolve_n_sec_head_mode_sharply_peaked_logits_pick_dominant_class(n_sec_sampling):
+    """A logit vector overwhelmingly favoring one class gives the same
+    answer under both policies — greedy because it's the argmax, sample
+    because softmax puts ~all mass on it."""
+    B, k_max = 8, 5
+    decoder = _stage2_ar("physical", "flow", k_max=k_max, n_sec_sampling=n_sec_sampling)
+    bias = torch.full((k_max + 1,), -50.0)
+    bias[2] = 50.0
+    _force_n_sec_head_bias(decoder, bias)
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec = resolve_n_sec(decoder, decoder, cond_cont, cond_cat, stage1_out, None)
+    assert n_sec is not None
+    assert torch.equal(n_sec, torch.full((B,), 2, dtype=torch.long))
+
+
+def test_resolve_n_sec_head_mode_greedy_is_deterministic_under_flat_logits():
+    B, k_max = 32, 5
+    decoder = _stage2_ar("physical", "flow", k_max=k_max, n_sec_sampling="greedy")
+    _force_n_sec_head_bias(decoder, torch.zeros(k_max + 1))
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec = resolve_n_sec(decoder, decoder, cond_cont, cond_cat, stage1_out, None)
+    assert n_sec is not None
+    assert n_sec.unique().numel() == 1
+
+
+def test_resolve_n_sec_head_mode_sample_varies_under_flat_logits():
+    """Under a flat logit vector, a categorical draw across a large batch
+    should hit more than one class — the whole point of gitea #86: greedy
+    always collapses to one, sample should not."""
+    torch.manual_seed(0)
+    B, k_max = 256, 5
+    decoder = _stage2_ar("physical", "flow", k_max=k_max, n_sec_sampling="sample")
+    _force_n_sec_head_bias(decoder, torch.zeros(k_max + 1))
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec = resolve_n_sec(decoder, decoder, cond_cont, cond_cat, stage1_out, None)
+    assert n_sec is not None
+    assert n_sec.unique().numel() > 1

uv.lock

@@ -675,7 +675,7 @@ wheels = [
 
 [[package]]
 name = "giant"
-version = "0.3.11"
+version = "0.3.16"
 source = { editable = "." }
 dependencies = [
     { name = "numpy" },