Make ResBlock's conditioning-injection mechanism selectable (gitea #34)
CI / Format (ruff format) (push) Successful in 29s
CI / Lint (ruff check) (push) Successful in 33s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 32s
CI / Tests (push) Successful in 1m58s
CI / Format (ruff format) (pull_request) Successful in 28s
CI / Lint (ruff check) (pull_request) Successful in 31s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 27s
CI / Tests (pull_request) Successful in 2m5s

ResBlock injected conditioning exactly one way — h = linear1(h) +
cond_proj(cond), a conditional bias, the weakest standard option for a
model whose entire job is to be conditional. Adds BLOCK_REGISTRY
(giant/model/layers.py), mirroring the TRUNK_REGISTRY/ROUTER_REGISTRY
registry+factory idiom (gitea #33), with two new drop-in alternatives:
FilmResBlock (per-channel scale+shift modulating the norm output,
zero-init so conditioning has no effect at construction) and
AdaLNResBlock (DiT-style AdaLN-Zero — the norm's own affine is replaced
by a conditioning-derived scale/shift, plus a zero-init gate on the
residual branch, making the block the exact identity function at init).

Selected per stage via a new stage{1,2}_model.trunk.block_conditioning
config leaf ("add" | "film" | "adaln", default "add"), threaded through
build_trunk/build_expert_body/RoutedTrunk and the three stage model
constructors. Default stays "add" and ResBlock's body is unchanged, so
existing configs/checkpoints are bit-identical to before this change.

Decided during planning: the new field lives on the existing TrunkConfig
rather than a new top-level block/blocks config section; the WGAN
CriticModel (which builds its own ResBlock stack outside TRUNK_REGISTRY)
and the issue's mentioned blocks.norm/blocks.activation axes are both
left out of scope.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
2026-08-14 09:41:50 +02:00
parent dc4cad7d11
commit 0f95e0eaae
8 changed files with 280 additions and 14 deletions
+29
View File
@@ -85,6 +85,15 @@ def test_trunk_config_defaults_to_resmlp_for_both_stages():
assert gconfig.DEFAULT_CONFIG["stage2_model"]["trunk"]["type"] == "resmlp"
def test_trunk_config_defaults_block_conditioning_to_add_for_both_stages():
"""gitea #34: a pre-existing config with no `block_conditioning` key
must reproduce today's additive-bias behaviour exactly."""
assert gconfig.Stage1ModelConfig().trunk.block_conditioning == "add"
assert gconfig.Stage2ModelConfig().trunk.block_conditioning == "add"
assert gconfig.DEFAULT_CONFIG["stage1_model"]["trunk"]["block_conditioning"] == "add"
assert gconfig.DEFAULT_CONFIG["stage2_model"]["trunk"]["block_conditioning"] == "add"
def test_particle_type_config_n_classes_defaults_to_zero_and_round_trips():
"""gitea #29: n_classes=0 means "inherit conditioning.particle.emb_dim"
— the default must stay 0 so an existing config.toml with no
@@ -936,6 +945,26 @@ def test_validate_config_keys_rejects_unknown_trunk_key():
assert "type" in str(e)
def test_validate_config_keys_allows_block_conditioning():
cfg = _cfg_with(
**{
"stage1_model.trunk.block_conditioning": "film",
"stage2_model.trunk.block_conditioning": "adaln",
}
)
gconfig.validate_config_keys(cfg) # must not raise
def test_validate_config_keys_rejects_unknown_block_conditioning_key():
cfg = _cfg_with(**{"stage1_model.trunk.block_conditioning_o": "film"}) # typo
try:
gconfig.validate_config_keys(cfg)
assert False, "expected ValueError"
except ValueError as e:
assert "stage1_model.trunk.block_conditioning_o" in str(e)
assert "block_conditioning" in str(e)
def test_merge_cli_overrides_rejects_typo_in_toml_file(tmp_path):
path = tmp_path / "config.toml"
path.write_text("[meta]\nconfig_version = 3\n\n[stage1_model]\nn_res_block = 12\n")