Give Stage1Model/Stage2OneShot/Stage2Autoregressive a shared StageModel base (gitea #39) #56

Merged

lars merged 3 commits from fix/issue-39 into master 2026-08-14 15:16:03 +02:00

Conversation 0 Commits 3 Files Changed 22

+620 -311

lars commented 2026-08-14 14:39:40 +02:00

Owner

Copy Link

Copy Source

Stage1Model, Stage2OneShot and Stage2Autoregressive each independently
implemented ~90 near-identical lines of init scaffolding:
build-or-share cond_enc, particle_type_cfg normalisation, objective ->
time_emb -> merged_cond_dim -> build_trunk, and the n_sec_head/type_head
classifier heads (plus their identical RuntimeError guards). Now unblocked
by #33 (trunk registry), #34 (block-conditioning registry) and #36
(build_mlp_head), which settled what belongs in the shared base.

Adds StageModel(nn.Module) owning all of that: init builds/shares
cond_enc and normalises particle_type_cfg; _build_trunk_and_heads,
called by each subclass after it sets up its own conditioning-assembly
modules (cond_enc alone for Stage1Model, a context-fusion path for the
two Stage2 classes), builds the objective/time embedding/trunk and the
n_sec_head/type_head guarded by the shared _require_n_sec_head/
require_type_head (Stage1Model overrides the n_sec guard since its
message points at stage 2, not stage 1). Public init signatures,
attribute names, and forward/predict* behaviour are unchanged.

Verified with a pre/post state_dict-key-set diff against the
pre-refactor classes (bit-identical) before writing this commit, plus
new parametrized tests pinning each class's state_dict key set and the
generator -> time_emb contract the base now owns. tests/test_migration_
v02_v03.py's existing bit-identical old-vs-new forward comparison and
the rest of tests/test_network.py's per-class coverage pass unchanged.

Co-Authored-By: Claude Opus 5 noreply@anthropic.com

Stage1Model, Stage2OneShot and Stage2Autoregressive each independently implemented ~90 near-identical lines of __init__ scaffolding: build-or-share cond_enc, particle_type_cfg normalisation, objective -> time_emb -> merged_cond_dim -> build_trunk, and the n_sec_head/type_head classifier heads (plus their identical RuntimeError guards). Now unblocked by #33 (trunk registry), #34 (block-conditioning registry) and #36 (build_mlp_head), which settled what belongs in the shared base. Adds StageModel(nn.Module) owning all of that: __init__ builds/shares cond_enc and normalises particle_type_cfg; _build_trunk_and_heads, called by each subclass after it sets up its own conditioning-assembly modules (cond_enc alone for Stage1Model, a context-fusion path for the two Stage2 classes), builds the objective/time embedding/trunk and the n_sec_head/type_head guarded by the shared _require_n_sec_head/ _require_type_head (Stage1Model overrides the n_sec guard since its message points at stage 2, not stage 1). Public __init__ signatures, attribute names, and forward/predict_* behaviour are unchanged. Verified with a pre/post state_dict-key-set diff against the pre-refactor classes (bit-identical) before writing this commit, plus new parametrized tests pinning each class's state_dict key set and the generator -> time_emb contract the base now owns. tests/test_migration_ v02_v03.py's existing bit-identical old-vs-new forward comparison and the rest of tests/test_network.py's per-class coverage pass unchanged. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

lars added 1 commit 2026-08-14 14:39:41 +02:00

Give Stage1Model/Stage2OneShot/Stage2Autoregressive a shared StageModel base (gitea #39) ... 1a3c907571

Stage1Model, Stage2OneShot and Stage2Autoregressive each independently
implemented ~90 near-identical lines of __init__ scaffolding:
build-or-share cond_enc, particle_type_cfg normalisation, objective ->
time_emb -> merged_cond_dim -> build_trunk, and the n_sec_head/type_head
classifier heads (plus their identical RuntimeError guards). Now unblocked
by #33 (trunk registry), #34 (block-conditioning registry) and #36
(build_mlp_head), which settled what belongs in the shared base.

Adds StageModel(nn.Module) owning all of that: __init__ builds/shares
cond_enc and normalises particle_type_cfg; _build_trunk_and_heads,
called by each subclass after it sets up its own conditioning-assembly
modules (cond_enc alone for Stage1Model, a context-fusion path for the
two Stage2 classes), builds the objective/time embedding/trunk and the
n_sec_head/type_head guarded by the shared _require_n_sec_head/
_require_type_head (Stage1Model overrides the n_sec guard since its
message points at stage 2, not stage 1). Public __init__ signatures,
attribute names, and forward/predict_* behaviour are unchanged.

Verified with a pre/post state_dict-key-set diff against the
pre-refactor classes (bit-identical) before writing this commit, plus
new parametrized tests pinning each class's state_dict key set and the
generator -> time_emb contract the base now owns. tests/test_migration_
v02_v03.py's existing bit-identical old-vs-new forward comparison and
the rest of tests/test_network.py's per-class coverage pass unchanged.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

lars added 1 commit 2026-08-14 15:04:46 +02:00

Pass ConditioningAxisConfig/ParticleTypeConfig themselves instead of raw dicts (gitea #38) ... c4b12b5e7a

build_models/build_critics parsed model_config into frozen dataclasses
(ConditioningConfig, Stage2ModelConfig, ...) but then threw the parsed
sub-objects away and passed the original raw dicts (conditioning["particle"],
s2_spec.particle_type.to_dict()) down into ConditionEncoder/StageModel/etc,
which re-read them with their own hardcoded .get(key, default) fallbacks —
each an independent copy of a fact the dataclass already stated once. Worst
instance: giant/training/trainers.py:236 converted an already-parsed
ParticleTypeConfig back into a dict for no reason.

Threads ConditioningAxisConfig (particle_cfg/material_cfg) and
ParticleTypeConfig (particle_type_cfg) as the actual dataclass instances
through every signature that used to type them dict: ConditionEncoder,
StageModel/CriticModel, resolve_type_n_classes/stage2_type_dim/
stage2_trunk_sec_dim, giant/model/builders.py, giant/sample.py,
giant/training/stage2_inputs.py, giant/training/trainers.py (StageSpec/
StageTrainer), giant/pipeline.py, giant/rollout.py, giant/validate.py — so ty
now catches a misspelled field instead of it silently falling back. No
config-schema change: config.toml/checkpoint model_config keep the same
nested-dict shape; only what happens after the existing X.from_dict(...)
parse changes.

User-confirmed scope decision: both axes (particle_cfg/material_cfg and
particle_type_cfg), not just the more heavily-duplicated particle_type_cfg
axis, and not stopping at the two most literal parse-then-discard round
trips — matching the issue's own proposal.

Preserved-default decision: StageModel's particle_type_cfg=None sentinel
(hit only by direct/test construction — build_models always passes an
explicit particle_type) still resolves to ParticleTypeConfig(target=
"physical"), not ParticleTypeConfig()'s own target="onehot" config-file
default — switching it would have silently grown an unused, gradient-less
type_head on every test that constructs Stage2OneShot/Stage2Autoregressive
without particle_type_cfg=, breaking their "every param has a grad" checks.

New tests in tests/test_network.py: ConditionEncoder/StageModel store the
exact ConditioningAxisConfig/ParticleTypeConfig instance passed in (identity,
not just equality) — no internal dict round-trip — and build_models's output
carries real dataclass instances end to end, not the plain dicts it produced
before this fix.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

lars added 1 commit 2026-08-14 15:05:57 +02:00

Bump patch version to 0.3.1 ... cc37a55183

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

lars merged commit 4f092c4528 into master 2026-08-14 15:16:03 +02:00

lars deleted branch fix/issue-39 2026-08-14 15:16:03 +02:00

Sign in to join this conversation.

Reviewers

No Reviewers

Labels

Clear labels

architecture

bug

Silent/wrong behavior, no crash

experiment

feature

New feature matching the v0.3 redesign

investigation

modularity

Making the network genuinely pluggable

performance

phase2

validation

No Label

Milestone

No items

No Milestone

Projects

Clear projects

No project

Assignees

Clear assignees

lars (Lars Bogner)

No Assignees

1 Participants

Notifications

Subscribe

Due Date

No due date set.

Dependencies

No dependencies set.

Reference: lars/giant#56