diff --git a/giant/config.py b/giant/config.py index 5ee2eaa..9fb05a0 100644 --- a/giant/config.py +++ b/giant/config.py @@ -424,11 +424,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 - # stop/continue decision. "greedy": sigmoid(logit) >= 0.5 (deterministic). "sample": - # a Bernoulli draw at sigmoid(logit) (a real sample from the learned length - # distribution, at the cost of an extra RNG draw per slot). - stop_sampling: str = "greedy" + # How resolve_n_sec/sample_secondaries_ar turn a count-bearing head's output into an + # actual n_sec decision. mode="head": "greedy" is argmax over the classifier logits + # (deterministic — the conditional mode, not a sample); "sample" is a categorical draw + # from softmax(logits) (a real sample from the learned count distribution). mode= + # "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 +441,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 +449,7 @@ class NSecConfig: "mode": self.mode, "lambda": self.lambda_weight, "owner": self.owner, - "stop_sampling": self.stop_sampling, + "sampling": self.sampling, } @@ -1075,6 +1079,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`. @@ -1271,11 +1296,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 +1547,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") - if stop_sampling not in ("greedy", "sample"): - raise ValueError(f"stage2_model.n_sec.stop_sampling = {stop_sampling!r} — must be 'greedy' or 'sample'") + n_sec_sampling = _get_path(cfg, "stage2_model.n_sec.sampling") + if n_sec_sampling not in ("greedy", "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"): diff --git a/giant/model/builders.py b/giant/model/builders.py index 0849db2..6c1a8b3 100644 --- a/giant/model/builders.py +++ b/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 diff --git a/giant/model/models.py b/giant/model/models.py index 8fb91c3..0e26323 100644 --- a/giant/model/models.py +++ b/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), diff --git a/giant/model/summary.py b/giant/model/summary.py index 36a34d3..82bffee 100644 --- a/giant/model/summary.py +++ b/giant/model/summary.py @@ -140,7 +140,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(): diff --git a/giant/sample.py b/giant/sample.py index c1d3f15..c0ac973 100644 --- a/giant/sample.py +++ b/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) diff --git a/tests/test_config.py b/tests/test_config.py index dabfbd0..acc8fe1 100644 --- a/tests/test_config.py +++ b/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(): - assert gconfig.NSecConfig().stop_sampling == "greedy" +def test_n_sec_config_sampling_defaults_to_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.to_dict()["stop_sampling"] == "sample" + assert n_sec.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(): - cfg = _cfg_with(**{"stage2_model.n_sec.stop_sampling": "bogus"}) +def test_validate_config_bad_n_sec_sampling_rejected(): + 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(): diff --git a/tests/test_sample.py b/tests/test_sample.py index eb31f89..345a656 100644 --- a/tests/test_sample.py +++ b/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"]) -def test_sample_secondaries_ar_stop_token_forced_stop_gives_zero_secondaries(stop_sampling): +@pytest.mark.parametrize("n_sec_sampling", ["greedy", "sample"]) +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"]) -def test_sample_secondaries_ar_stop_token_forced_never_stop_runs_to_k_max(stop_sampling): +@pytest.mark.parametrize("n_sec_sampling", ["greedy", "sample"]) +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