diff --git a/giant/cli.py b/giant/cli.py index 2b440a8..0cd80c6 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -1043,8 +1043,12 @@ def predict( # A fresh v0.3.0 Stage1Model owns no n_sec_head — # sample_stage1 returns n_sec_pred=None then, so ask stage 2. n_sec_pred = resolve_n_sec(model, sec_decoder, cc, ck, stage1_norm, n_sec_pred) - sec_cont, sec_type, _sec_valid_pred = sample_stage2(sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps) - n_sec_pred_np = n_sec_pred.cpu().numpy() + sec_cont, sec_type, sec_valid_pred = sample_stage2(sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps) + # A stop-token decoder resolves n_sec_pred=None above — read the + # real count back off sec_valid_pred instead (a no-op round trip + # under every other n_sec.mode, where sec_valid_pred was built + # FROM n_sec_pred in the first place). + n_sec_pred_np = sec_valid_pred.sum(dim=-1).cpu().numpy() pred = stage1_norm.cpu().numpy() # normalised diff --git a/giant/config.py b/giant/config.py index b84d648..2c83541 100644 --- a/giant/config.py +++ b/giant/config.py @@ -408,18 +408,27 @@ class Stage2RouterConfig(RouterConfig): class NSecConfig: # "head": a classifier over {0..k_max} on the condition encoding alone # (no diffusion noise), callable independently at inference. - # "stop_token": an EOS-style implicit stop — accepted by the schema but - # not implemented in v0.3.0 (see validate_config). + # "stop_token": an EOS-style per-slot stop head on the autoregressive + # secondary decoder (Stage2Autoregressive only — see validate_config), + # evaluated against the generated prefix instead of conditioning alone. + # Replaces n_sec_head entirely: the two are mutually exclusive, so this + # mode builds no n_sec_head and stage2_model.n_sec.lambda instead weights + # the stop head's BCE term. # "truth": take n_sec from ground truth — standalone stage-2 evaluation # only, never for rollout. mode: str = "head" - lambda_weight: float = 0.1 # dict key "lambda" — cross-entropy weight for the head + lambda_weight: float = 0.1 # dict key "lambda" — cross-entropy/BCE weight for the head # Which stage's module physically owns the n_sec_head weights: "stage2" (default, # fresh v0.3.0 runs — Stage2OneShot/Stage2Autoregressive builds it) or "stage1" # (a migrated v0.2 checkpoint — see network._migrate_legacy_model_config, whose # 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" @classmethod def from_dict(cls, d: dict | None) -> "NSecConfig": @@ -428,10 +437,16 @@ 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"), ) def to_dict(self) -> dict: - return {"mode": self.mode, "lambda": self.lambda_weight, "owner": self.owner} + return { + "mode": self.mode, + "lambda": self.lambda_weight, + "owner": self.owner, + "stop_sampling": self.stop_sampling, + } @dataclass(frozen=True) @@ -1368,11 +1383,23 @@ def validate_config(cfg: dict) -> None: ) if _get_path(cfg, "stage2_model.n_sec.mode") == "stop_token": - raise ValueError( - "stage2_model.n_sec.mode = 'stop_token' is accepted by the schema " - "but not implemented in v0.3.0 — use 'head' (default) or 'truth' " - "(standalone stage-2 evaluation only, never for rollout)" - ) + if _get_path(cfg, "stage2_model.decoder") != "autoregressive": + raise ValueError( + "stage2_model.n_sec.mode = 'stop_token' requires " + "stage2_model.decoder = 'autoregressive' — there is no " + "per-token loop to stop under 'one_shot'" + ) + if _get_path(cfg, "stage2_model.n_sec.owner") != "stage2": + raise ValueError( + "stage2_model.n_sec.mode = 'stop_token' requires " + "stage2_model.n_sec.owner = 'stage2' — a migrated v0.2 " + "checkpoint's stage-1 n_sec_head has no per-token " + "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'") if _get_path(cfg, "stage2_model.stage1_context") == "sampled": raise ValueError( diff --git a/giant/model/builders.py b/giant/model/builders.py index de11da7..56c20bc 100644 --- a/giant/model/builders.py +++ b/giant/model/builders.py @@ -106,6 +106,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: # above. time_dim = getattr(s2_spec, generator).time_dim if objective.needs_time else 64 n_sec_owner = s2_spec.n_sec.owner + stop_token = s2_spec.n_sec.mode == "stop_token" k_max = s2_spec.k_max particle_type_cfg = s2_spec.particle_type @@ -128,7 +129,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: router=stage2_router, trunk_type=s2_spec.trunk.type, block_conditioning=s2_spec.trunk.block_conditioning, - build_n_sec_head=n_sec_owner != "stage1", + build_n_sec_head=n_sec_owner != "stage1" and not stop_token, particle_type_cfg=particle_type_cfg, history=ar_cfg.history, attn_n_heads=ar_cfg.attn_n_heads, @@ -136,6 +137,9 @@ 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(), + build_stop_head=stop_token, + stop_sampling=s2_spec.n_sec.stop_sampling, + stop_head_cfg=s2_spec.heads.n_sec.to_dict(), ) else: sec_dim = stage2_trunk_sec_dim( diff --git a/giant/model/models.py b/giant/model/models.py index 4f80fdf..4284e5e 100644 --- a/giant/model/models.py +++ b/giant/model/models.py @@ -132,13 +132,15 @@ class StageModel(nn.Module): n_sec_head_cfg: dict | None, type_head_out_dim: int | None, type_head_cfg: dict | None, + build_stop_head: bool = False, + stop_head_cfg: dict | None = None, ) -> None: """Builds `self.time_emb`, `self.trunk`, `self.n_sec_head`, - `self.type_head`. Called by a subclass's `__init__` after it has set - up its own conditioning-assembly modules — `merged_cond_dim` below - must match the width that assembly (`_cond_embed`/`_base_cond`/ - `_token_cond`, or plain `cond_enc` for `Stage1Model`) actually - produces. + `self.type_head`, `self.stop_head`. Called by a subclass's `__init__` + after it has set up its own conditioning-assembly modules — + `merged_cond_dim` below must match the width that assembly + (`_cond_embed`/`_base_cond`/`_token_cond`, or plain `cond_enc` for + `Stage1Model`) actually produces. `n_sec_head` is built iff `n_sec_head_k_max is not None` (output width `n_sec_head_k_max + 1`) — `Stage1Model` passes this only for a @@ -148,7 +150,11 @@ class StageModel(nn.Module): classes — passes `None` exactly when `particle_type_cfg.target == "physical"`) *and* the objective doesn't fold the type slice into its own trunk output (checked here, since `objective` is already needed - for the trunk itself). + for the trunk itself). `stop_head` is built iff `build_stop_head` — + only `Stage2Autoregressive` ever passes `True` (`n_sec.mode == + "stop_token"`, mutually exclusive with `n_sec_head`), a single + `cond_out_dim -> 1` logit per call, same `HeadConfig` shape rules as + the other two heads. """ objective = build_objective(self.generator_kind) has_time = objective.needs_time @@ -176,6 +182,11 @@ class StageModel(nn.Module): head_cfg = HeadConfig.from_dict(type_head_cfg) hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio)) self.type_head = build_mlp_head(cond_out_dim, type_head_out_dim, hidden, head_cfg.depth) + self.stop_head = None + if build_stop_head: + head_cfg = HeadConfig.from_dict(stop_head_cfg) + hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio)) + self.stop_head = build_mlp_head(cond_out_dim, 1, hidden, head_cfg.depth) def _require_n_sec_head(self) -> None: if self.n_sec_head is None: @@ -195,6 +206,14 @@ class StageModel(nn.Module): "directly instead)" ) + def _require_stop_head(self) -> None: + if self.stop_head is None: + raise RuntimeError( + f"this {type(self).__name__} has no stop_head — only a " + "Stage2Autoregressive built with stage2_model.n_sec.mode = " + "'stop_token' owns one" + ) + class Stage1Model(StageModel): """Predicts the 9D primary post-step vector. No `n_sec_head` — fresh runs @@ -431,7 +450,13 @@ class Stage2Autoregressive(StageModel): `context_adapter` only) feeds `predict_n_sec`, since n_sec doesn't depend on token position; `_token_cond` additionally fuses in the history encoding and two running scalars (remaining energy-budget fraction, - normalized slot index), and feeds `forward`/`predict_type`/the trunk. + normalized slot index), and feeds `forward`/`predict_type`/`predict_stop`/ + the trunk. + + `n_sec.mode = "stop_token"` (`build_stop_head=True`) replaces + `predict_n_sec`'s one-shot classifier with `predict_stop`'s per-token EOS + logit instead — the two heads are mutually exclusive (`build_n_sec_head` + is `False` whenever this is `True`, see `giant.model.builders`). """ def __init__( @@ -461,6 +486,9 @@ class Stage2Autoregressive(StageModel): cond_enc: ConditionEncoder | None = None, n_sec_head_cfg: dict | None = None, type_head_cfg: dict | None = None, + build_stop_head: bool = False, + stop_sampling: str = "greedy", + stop_head_cfg: dict | None = None, ) -> None: super().__init__( pdg_vocab, @@ -475,6 +503,7 @@ class Stage2Autoregressive(StageModel): cond_enc=cond_enc, ) self.history_kind = history + self.stop_sampling = stop_sampling self.context_adapter = ContextAdapter(x_dim, context_dim) self.base_fuse = nn.Sequential( nn.Linear(cond_out_dim + context_dim, cond_out_dim), @@ -516,6 +545,8 @@ class Stage2Autoregressive(StageModel): n_sec_head_cfg=n_sec_head_cfg, type_head_out_dim=type_head_out_dim, type_head_cfg=type_head_cfg, + build_stop_head=build_stop_head, + stop_head_cfg=stop_head_cfg, ) def _base_cond(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor: @@ -638,6 +669,38 @@ class Stage2Autoregressive(StageModel): B, K, _ = c_emb.shape return self.type_head(c_emb.reshape(B * K, -1)).view(B, K, self.type_dim) + def predict_stop( + self, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + stage1_out: torch.Tensor, + history_feat: torch.Tensor, + has_prev: torch.Tensor, + remaining_frac: torch.Tensor, + slot_idx: torch.Tensor, + hist: torch.Tensor | None = None, + ) -> torch.Tensor: + """`(B, K)` raw stop logits — `n_sec.mode = "stop_token"` only. + Evaluated on slot `k`'s own conditioning (which carries slot `k-1`'s + history, same as `predict_type`), so this is `P(n_sec == k | + prefix)`: a high logit at slot `k` means "stop before generating a + token here" — the caller (`giant.sample.sample_secondaries_ar`) + checks it before spending a model call on that slot's token.""" + self._require_stop_head() + assert self.stop_head is not None + c_emb = self._token_cond( + cond_cont, + cond_cat, + stage1_out, + history_feat, + has_prev, + remaining_frac, + slot_idx, + hist=hist, + ) + B, K, _ = c_emb.shape + return self.stop_head(c_emb.reshape(B * K, -1)).view(B, K) + class CriticModel(nn.Module): """Generator-agnostic WGAN-GP critic body: a scalar realism score, for diff --git a/giant/rollout.py b/giant/rollout.py index 03ed04d..36d5361 100644 --- a/giant/rollout.py +++ b/giant/rollout.py @@ -681,7 +681,6 @@ def _step_chunk( post_pos = reconstruct_post_pos(tr["pre_pos"], tr["pre_dir"], step_length, travel_dir_local) n_sec_pred = resolve_n_sec(stage1_model, sec_decoder, cc, ck, stage1_norm, n_sec_pred_stage1) - n_sec_np = n_sec_pred.cpu().numpy().astype(np.int64) # --- Secondaries --- # No snapping for "physical"/history-facing state elsewhere in the @@ -691,7 +690,13 @@ def _step_chunk( # decode_secondary_identity's docstring for how each # particle_type.target differs on whether PDG resolution is a real # identity decision or just a reporting label. - sec_cont, sec_type, _valid = sample_stage2(sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps) + sec_cont, sec_type, sec_valid = sample_stage2(sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps) + # A stop-token decoder resolves n_sec_pred=None above — the real count + # only exists once sample_stage2 has actually generated (or stopped + # generating) tokens, so read it back off sec_valid here. Under every + # other n_sec.mode sec_valid was built FROM n_sec_pred, so this is a + # no-op round trip in those cases. + n_sec_np = sec_valid.sum(dim=-1).cpu().numpy().astype(np.int64) sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg_code, sec_type_l1_dist = decode_secondary_identity( sec_decoder, sec_cont, diff --git a/giant/sample.py b/giant/sample.py index 374049b..c1d3f15 100644 --- a/giant/sample.py +++ b/giant/sample.py @@ -229,14 +229,14 @@ def sample_secondaries_ar( cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor, - n_sec_pred: torch.Tensor, + n_sec_pred: torch.Tensor | None, steps: int = 10, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """`Stage2Autoregressive` inference loop: one token at a time, in - descending-energy slot order, `k_max` sequential calls. Unlike training - (teacher forcing — a single parallel pass over ground-truth tokens, see - `giant.training.stage2_inputs._assemble_stage2_ar_inputs`), there is no - ground truth at inference: each token's conditioning is built + descending-energy slot order, up to `k_max` sequential calls. Unlike + training (teacher forcing — a single parallel pass over ground-truth + tokens, see `giant.training.stage2_inputs._assemble_stage2_ar_inputs`), + there is no ground truth at inference: each token's conditioning is built free-running, from the PREVIOUS TOKEN'S OWN just-generated output — the train/inference gap that is the cost of markov history's expressiveness. @@ -244,7 +244,30 @@ def sample_secondaries_ar( A `{flow,ddpm}` token costs `steps` ODE substeps; `wgan` costs one pass — the "K sequential forwards" cost applies per-token here, not once, so a flow/ddpm AR run costs ~`k_max * steps` model calls per - physics step. + physics step (or ~`n_sec * steps` under `n_sec_pred=None` below, once + every row in the batch has stopped). + + `n_sec_pred`, if given, fixes each row's secondary count up front (as + resolved by `resolve_n_sec` — `n_sec.mode` in `("head", "truth")`, or a + stop-token decoder driven by `_assemble_stage2_ar_inputs_scheduled`'s + ground-truth `n_sec`, which must run the *full* `k_max`-length free- + running self-sample regardless of the decoder's own stop head — the + scheduled-sampling training contract does not truncate). This always + runs the full `k_max`-iteration loop, masking by the given count at the + end exactly as before. + + `n_sec_pred=None` is only valid when `sec_decoder.stop_head` is set + (`n_sec.mode = "stop_token"`): before generating each slot's token, that + 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 + 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 + next slot's token — the average-case cost win the docstring above + describes. A row that never fires within `k_max` is capped there + (`K_MAX` stays a safety cap, not a modeling ceiling). Under `history="attention"` the history encoding is computed once per slot via `Stage2Autoregressive.history_step` (a KV-cache append) @@ -294,6 +317,16 @@ def sample_secondaries_ar( remaining = torch.ones(B, device=device) history_cache = sec_decoder.init_history_cache() + use_stop_token = n_sec_pred is None + if use_stop_token: + assert getattr(sec_decoder, "stop_head", None) is not None, ( + "sample_secondaries_ar called with n_sec_pred=None on a decoder " + "with no stop_head — only valid under stage2_model.n_sec.mode = " + "'stop_token'" + ) + finished = torch.zeros(B, dtype=torch.bool, device=device) + derived_n_sec = torch.full((B,), k_max, dtype=torch.long, device=device) + for k in range(k_max): has_prev = torch.full((B, 1), k >= 1, dtype=torch.bool, device=device) history_feat = prev_repr.unsqueeze(1) # (B, 1, CONT_SLOT_DIM + type_dim) @@ -301,6 +334,26 @@ def sample_secondaries_ar( slot_idx = torch.full((B, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32) hist, history_cache = sec_decoder.history_step(history_feat, has_prev, history_cache) + if use_stop_token: + stop_logit = sec_decoder.predict_stop( + cond_cont, + cond_cat, + stage1_out, + history_feat, + has_prev, + remaining_frac, + slot_idx, + hist=hist, + ).squeeze(1) + if sec_decoder.stop_sampling == "sample": + stop_now = torch.rand(B, device=device) < torch.sigmoid(stop_logit) + else: + stop_now = stop_logit >= 0.0 + derived_n_sec[stop_now & ~finished] = k + finished = finished | stop_now + if finished.all(): + break + if objective.is_adversarial: z = torch.randn(B, 1, sec_decoder.noise_dim, device=device) token = sec_decoder( @@ -362,7 +415,8 @@ def sample_secondaries_ar( prev_repr = torch.cat([stick_fraction.unsqueeze(-1), cont_k[:, 1:4], type_for_history], dim=-1) remaining = torch.clamp(remaining * (1.0 - stick_fraction), min=0.0) - sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(1) + resolved_n_sec = derived_n_sec if use_stop_token else n_sec_pred + sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < resolved_n_sec.unsqueeze(1) return sec_cont, sec_type, sec_valid @@ -397,7 +451,7 @@ def sample_stage2( cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor, - n_sec_pred: torch.Tensor, + n_sec_pred: torch.Tensor | None, steps: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Dispatches on `decoder` (one-shot vs autoregressive — the class @@ -406,9 +460,15 @@ def sample_stage2( built with `generator="ddpm"` in practice and `flow_matching_loss_secondary*` is the only stage-2 training path that exists for the non-adversarial case, so there's nothing to dispatch to here. + + `n_sec_pred=None` (from `resolve_n_sec` on a stop-token decoder) is only + meaningful for the autoregressive path — see `sample_secondaries_ar`'s + docstring; the one-shot samplers have no per-token stop mechanism to + derive a count from, so `n_sec_pred` must already be resolved for them. """ if isinstance(sec_decoder, Stage2Autoregressive): return sample_secondaries_ar(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps) + assert n_sec_pred is not None, "one-shot stage-2 decoders need a resolved n_sec_pred" if build_objective(sec_decoder.generator_kind).is_adversarial: return sample_secondaries_wgan(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred) return sample_secondaries(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps) @@ -421,20 +481,31 @@ def resolve_n_sec( cond_cat: torch.Tensor, stage1_out: torch.Tensor, n_sec_pred: torch.Tensor | None, -) -> torch.Tensor: +) -> torch.Tensor | None: """`n_sec_pred` is already populated when `stage1_model` owns a legacy `n_sec_head` (a migrated v0.2 checkpoint — see `Stage1Model`'s - docstring); otherwise ask stage 2, which owns it by default. Raises if - neither stage owns a head at all — the only way that happens is - `stage2_model.n_sec.mode` other than `"head"` (`"truth"`/`"stop_token"`), - neither of which is a valid rollout-/predict-capable checkpoint.""" + docstring); otherwise ask stage 2, which owns it by default. + + Returns `None` when `sec_decoder` owns a `stop_head` (`n_sec.mode = + "stop_token"`) instead of an `n_sec_head` — there is nothing to resolve + up front in that case, since the count only exists once + `sample_secondaries_ar` has actually generated (or stopped generating) + tokens; the caller passes this `None` straight through to `sample_stage2` + and reads the real count back off its returned `sec_valid` + (`sec_valid.sum(-1)`) afterwards. + + 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.""" if n_sec_pred is not None: return n_sec_pred + if getattr(sec_decoder, "stop_head", None) is not None: + return None if getattr(sec_decoder, "n_sec_head", None) is None: raise RuntimeError( - "checkpoint has no n_sec_head on either stage — needs " - "stage2_model.n_sec.mode = 'head' (the default); 'truth' is " - "standalone-evaluation-only and 'stop_token' isn't implemented" + "checkpoint has no n_sec_head/stop_head on either stage — needs " + "stage2_model.n_sec.mode = 'head' (the default) or 'stop_token'; " + "'truth' is standalone-evaluation-only" ) logits = sec_decoder.predict_n_sec(cond_cont, cond_cat, stage1_out) return logits.argmax(dim=-1) diff --git a/giant/training/stage2_inputs.py b/giant/training/stage2_inputs.py index 8bd3730..5277587 100644 --- a/giant/training/stage2_inputs.py +++ b/giant/training/stage2_inputs.py @@ -140,6 +140,27 @@ def _ar_has_prev(k_max: int, device: torch.device) -> torch.Tensor: return (torch.arange(k_max, device=device) >= 1).unsqueeze(0) +def _stop_target_and_mask(n_sec: torch.Tensor, k_max: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: + """`(target, mask)`, both `(B, K_MAX)`, for `n_sec.mode = "stop_token"`'s + per-slot EOS head (`Stage2Autoregressive.predict_stop`). + + `predict_stop` is evaluated on slot `k`'s own (pre-token) conditioning — + "should generation have already stopped by here" — so `target[k] = 1` + exactly at `k == n_sec` (the first invalid slot: `sample_secondaries_ar` + checks this before spending a model call generating that slot's token), + `0` elsewhere. `mask` is `k <= n_sec` — one slot *wider* than + `StageTrainer._sec_mask`'s `k < n_sec` token-content mask, since the stop + slot itself (`k == n_sec`) must be supervised even though there is no + real secondary there. A row with `n_sec == k_max` has no in-range stop + slot at all: `mask` covers the full `k_max` range (every generated token + is real) and `target` is all-zero — `sample_secondaries_ar` correctly + never breaks early for it, running into the `k_max` safety cap instead.""" + idx = torch.arange(k_max, device=device).unsqueeze(0) + target = (idx == n_sec.unsqueeze(1)).float() + mask = idx <= n_sec.unsqueeze(1) + return target, mask + + def _ar_meta(k_max: int, batch: int, device: torch.device, fraction: torch.Tensor) -> dict[str, torch.Tensor]: """`has_prev`/`remaining_frac`/`slot_idx` — the three per-token AR conditioning tensors that don't depend on *which* history representation diff --git a/giant/training/trainers.py b/giant/training/trainers.py index 70b3a05..9b80cd6 100644 --- a/giant/training/trainers.py +++ b/giant/training/trainers.py @@ -34,6 +34,7 @@ from giant.training.stage2_inputs import ( _gumbel_tau, _relax_onehot_type_slice, _stage2_tf_prob, + _stop_target_and_mask, ) @@ -89,6 +90,7 @@ class StageSpec: # loss weights lambda_weight: float = 1.0 n_sec_lambda: float = 0.1 + n_sec_mode: str = "head" # particle-type target (stage 2 only) particle_type: ParticleTypeConfig = field(default_factory=ParticleTypeConfig) @@ -144,6 +146,7 @@ class StageSpec: decoder=s2_spec.decoder if is_stage2 else "one_shot", lambda_weight=stage_spec.lambda_weight, n_sec_lambda=s2_spec.n_sec.lambda_weight, + n_sec_mode=s2_spec.n_sec.mode, particle_type=s2_spec.particle_type, particle_type_n_classes=resolve_type_n_classes( s2_spec.particle_type, cfg["conditioning"]["particle"]["emb_dim"] @@ -374,10 +377,9 @@ class StageTrainer: stage1-vs-stage2 `predict_n_sec` signature split, shared by the non-adversarial and WGAN trainers. - Gated on `n_sec_head is None`, not on `n_sec.mode`: a future - `mode="stop_token"` model (currently rejected in - `validate_config`) carries no head and would train its EOS signal in - the generator/AR loss path instead, so this correctly stays zero. + Gated on `n_sec_head is None`, not on `n_sec.mode`: a `mode = + "stop_token"` model carries no head at all (see `_stop_loss` for its + EOS signal instead), so this correctly stays zero for it. """ if self.model.n_sec_head is None: zero = torch.zeros((), device=device) @@ -391,6 +393,47 @@ class StageTrainer: nsec_acc = (logits.argmax(dim=-1) == n_sec).float().mean() return l_nsec, nsec_acc + def _stop_loss( + self, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + stage1_ctx: torch.Tensor, + n_sec: torch.Tensor, + device: torch.device, + ar_inputs: dict[str, torch.Tensor] | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """`(l_stop, stop_acc)` for `n_sec.mode = "stop_token"`'s per-slot EOS + head (`Stage2Autoregressive.predict_stop`) — zeros when this stage + owns no `stop_head` (every other `n_sec.mode`), the same gating + convention `_n_sec_loss` uses for `n_sec_head`. The two heads are + mutually exclusive (`giant.model.builders`), so exactly one of + `_n_sec_loss`/`_stop_loss` is ever non-zero for a given stage. + + Masked BCE against `_stop_target_and_mask`'s per-slot target — one + slot wider than `sec_mask` (the stop slot itself, `k == n_sec`, needs + supervision even though it holds no real secondary).""" + stop_head = getattr(self.model, "stop_head", None) + if stop_head is None: + zero = torch.zeros((), device=device) + return zero, zero + assert ar_inputs is not None + logits = self.model.predict_stop( + cond_cont, + cond_cat, + stage1_ctx, + ar_inputs["history_feat"], + ar_inputs["has_prev"], + ar_inputs["remaining_frac"], + ar_inputs["slot_idx"], + ) + target, mask = _stop_target_and_mask(n_sec, logits.size(1), device) + mask_f = mask.float() + denom = mask_f.sum().clamp(min=1) + bce = F.binary_cross_entropy_with_logits(logits, target, reduction="none") + l_stop = (bce * mask_f).sum() / denom + stop_acc = (((logits >= 0).float() == target).float() * mask_f).sum() / denom + return l_stop, stop_acc + @staticmethod def _step_optimizer(optimizer: optim.Optimizer, loss: torch.Tensor, params: list) -> float: """`zero_grad -> backward -> clip_grad_norm_(1.0) -> step`, returning @@ -477,10 +520,12 @@ class FlowDDPMStageTrainer(StageTrainer): "loss", "loss_gen", "loss_nsec", + "loss_stop", "loss_balance", "loss_proc", "loss_entropy", "nsec_acc", + "stop_acc", "loss_type", "type_acc", "grad_norm", @@ -493,6 +538,8 @@ class FlowDDPMStageTrainer(StageTrainer): "loss_gen", "loss_nsec", "nsec_acc", + "loss_stop", + "stop_acc", "loss_type", "type_acc", ) @@ -586,6 +633,7 @@ class FlowDDPMStageTrainer(StageTrainer): l_gen = self._generator_loss(cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx, ar_inputs=ar_inputs) l_nsec, nsec_acc = self._n_sec_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device) + l_stop, stop_acc = self._stop_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device, ar_inputs) l_type, type_acc = self._type_loss( cond_cont, @@ -606,7 +654,11 @@ class FlowDDPMStageTrainer(StageTrainer): if self.spec.lambda_entropy > 0: l_entropy = self.router.entropy_loss(cond_cont, cond_cat) - total = self.spec.lambda_weight * l_gen + self.spec.n_sec_lambda * l_nsec + self.particle_type_lambda * l_type + total = ( + self.spec.lambda_weight * l_gen + + self.spec.n_sec_lambda * (l_nsec + l_stop) + + self.particle_type_lambda * l_type + ) if self.spec.lambda_balance > 0: total = total + self.spec.lambda_balance * l_balance if self.spec.lambda_proc > 0: @@ -618,12 +670,14 @@ class FlowDDPMStageTrainer(StageTrainer): "loss": total, "loss_gen": l_gen, "loss_nsec": l_nsec, + "loss_stop": l_stop, "loss_type": l_type, "type_acc": type_acc, "loss_balance": l_balance, "loss_proc": l_proc, "loss_entropy": l_entropy, "nsec_acc": nsec_acc, + "stop_acc": stop_acc, } def step(self, batch: StepBatch, device: torch.device, global_step: int) -> dict: @@ -723,6 +777,8 @@ class WGANStageTrainer(StageTrainer): "gp_loss", "loss_nsec", "nsec_acc", + "loss_stop", + "stop_acc", "grad_norm_d", "grad_norm_g", ] @@ -735,11 +791,15 @@ class WGANStageTrainer(StageTrainer): self.stage_metrics = [stage_metric("lr"), stage_metric("critic_lr")] def _stage2_real_and_fake(self, batch_tensors: _Stage2RealFakeBatch, stage1_ctx, global_step, device): - """Build `(real, fake_raw, mask, critic_fn)` for stage 2, covering - both decoders and all three particle-type targets. `fake_raw` still - needs the caller's straight-through relaxation under + """Build `(real, fake_raw, mask, critic_fn, ar_inputs)` for stage 2, + covering both decoders and all three particle-type targets. `fake_raw` + still needs the caller's straight-through relaxation under `particle_type.target = "onehot"`, and neither tensor is masked-and- - multiplied on the fake side yet.""" + multiplied on the fake side yet. `ar_inputs` is `None` under + `decoder = "one_shot"`; under `"autoregressive"` it's the same dict + `_ar_inputs` built to condition `self.model` above — returned so the + caller's `_stop_loss` reuses it instead of paying for a second + (possibly self-sampling) `_ar_inputs` call.""" cond_cont, cond_cat, n_sec, sec_cont, sec_type_idx = batch_tensors B = cond_cont.size(0) type_dim = stage2_type_dim(self.particle_type_cfg, self.particle_type_n_classes) @@ -752,9 +812,10 @@ class WGANStageTrainer(StageTrainer): def critic_fn(x): return self.critic(x, cond_cont, cond_cat, stage1_ctx) + ar_inputs = None if self.decoder == "autoregressive": epoch = global_step // self.spec.steps_per_epoch - ar = self._ar_inputs(cond_cont, cond_cat, stage1_ctx, sec_cont, sec_type_idx, n_sec, epoch) + ar_inputs = self._ar_inputs(cond_cont, cond_cat, stage1_ctx, sec_cont, sec_type_idx, n_sec, epoch) real = self._sec_target(sec_cont, sec_type_idx, self.generator, flatten=False).reshape(B, -1) * mask z = torch.randn(B, k_max, self.model.noise_dim, device=device) fake_raw = self.model( @@ -762,17 +823,17 @@ class WGANStageTrainer(StageTrainer): cond_cont, cond_cat, stage1_ctx, - ar["history_feat"], - ar["has_prev"], - ar["remaining_frac"], - ar["slot_idx"], + ar_inputs["history_feat"], + ar_inputs["has_prev"], + ar_inputs["remaining_frac"], + ar_inputs["slot_idx"], ).reshape(B, -1) else: real = self._sec_target(sec_cont, sec_type_idx, self.generator, flatten=True) * mask z = torch.randn(B, self.model.noise_dim, device=device) fake_raw = self.model(z, cond_cont, cond_cat, stage1_ctx) - return real, fake_raw, mask, critic_fn + return real, fake_raw, mask, critic_fn, ar_inputs def step(self, batch: StepBatch, device: torch.device, global_step: int) -> dict: ( @@ -788,6 +849,7 @@ class WGANStageTrainer(StageTrainer): stage1_ctx = x1_s1.detach() grad_probe: dict[str, float] = {} + ar_inputs = None if not self.is_stage2: real = x1_s1 @@ -798,7 +860,7 @@ class WGANStageTrainer(StageTrainer): fake = self.model(z, cond_cont, cond_cat) mask = None else: - real, fake_raw, mask, critic_fn = self._stage2_real_and_fake( + real, fake_raw, mask, critic_fn, ar_inputs = self._stage2_real_and_fake( _Stage2RealFakeBatch(cond_cont, cond_cat, n_sec, sec_cont, sec_type_idx), stage1_ctx, global_step, @@ -840,18 +902,19 @@ class WGANStageTrainer(StageTrainer): # --- generator (+ n_sec) step --- did_g_step = global_step % self.n_critic == 0 l_nsec, nsec_acc = self._n_sec_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device) + l_stop, stop_acc = self._stop_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device, ar_inputs) - # On a non-generator-step batch with no n_sec_head on this stage - # (n_sec now defaults to stage 2), there's nothing for + # On a non-generator-step batch with no n_sec_head/stop_head on this + # stage (n_sec now defaults to stage 2), there's nothing for # the generator optimizer to do this batch — g_loss would otherwise # be a graph-less zero tensor, which .backward() rejects outright. - skip_g_step = not did_g_step and self.model.n_sec_head is None + skip_g_step = not did_g_step and self.model.n_sec_head is None and self.model.stop_head is None if did_g_step: g_loss_adv = generator_loss(critic_fn, fake) - g_loss = self.spec.lambda_weight * g_loss_adv + self.spec.n_sec_lambda * l_nsec + g_loss = self.spec.lambda_weight * g_loss_adv + self.spec.n_sec_lambda * (l_nsec + l_stop) else: g_loss_adv = torch.zeros((), device=device) - g_loss = self.spec.n_sec_lambda * l_nsec + g_loss = self.spec.n_sec_lambda * (l_nsec + l_stop) if skip_g_step: grad_norm_g = 0.0 else: @@ -869,6 +932,8 @@ class WGANStageTrainer(StageTrainer): "gp_loss": gp.item(), "loss_nsec": l_nsec.item(), "nsec_acc": nsec_acc.item(), + "loss_stop": l_stop.item(), + "stop_acc": stop_acc.item(), "did_g_step": did_g_step, "grad_norm": grad_norm_d + grad_norm_g, "grad_norm_d": grad_norm_d, diff --git a/giant/validate.py b/giant/validate.py index 1e68816..51e9a66 100644 --- a/giant/validate.py +++ b/giant/validate.py @@ -129,10 +129,7 @@ def validate_marginals( continue n_sec_pred = resolve_n_sec(stage1_model, sec_decoder, cond_cont, cond_cat, gen, n_sec_pred) - n_sec_pred_np = n_sec_pred.cpu().numpy() n_sec_np = n_sec.numpy() - all_n_sec_real.append(n_sec_np) - all_n_sec_pred.append(n_sec_pred_np) real_valid = np.arange(k_max)[None, :] < n_sec_np[:, None] # (B, k_max) real_frac = 1.0 / (1.0 + np.exp(-sec_cont[:, :, 0].numpy().astype(np.float64))) @@ -140,6 +137,13 @@ def validate_marginals( sec_cont_pred, sec_type_pred, sec_valid_pred = sample_stage2( sec_decoder, cond_cont, cond_cat, gen, n_sec_pred, steps=steps ) + # A stop-token decoder resolves n_sec_pred=None above — read the real + # count back off sec_valid_pred instead (a no-op round trip under + # every other n_sec.mode, where sec_valid_pred was built FROM + # n_sec_pred in the first place). + n_sec_pred_np = sec_valid_pred.sum(dim=-1).cpu().numpy() + all_n_sec_real.append(n_sec_np) + all_n_sec_pred.append(n_sec_pred_np) gen_frac = 1.0 / (1.0 + np.exp(-sec_cont_pred[:, :, 0].cpu().numpy().astype(np.float64))) gen_valid = sec_valid_pred.cpu().numpy() diff --git a/pyproject.toml b/pyproject.toml index 3789cdc..9d4eb10 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "giant" -version = "0.3.1" +version = "0.3.2" description = "Geant4 step-function surrogate via conditional flow matching" readme = "README.md" requires-python = ">=3.12" diff --git a/tests/test_config.py b/tests/test_config.py index 2945df6..4e71c96 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -150,7 +150,17 @@ 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"} + assert n_sec.to_dict() == {"mode": "head", "lambda": 0.1, "owner": "stage1", "stop_sampling": "greedy"} + + +def test_n_sec_config_stop_sampling_defaults_to_greedy(): + assert gconfig.NSecConfig().stop_sampling == "greedy" + + +def test_n_sec_config_stop_sampling_round_trips(): + 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" # --------------------------------------------------------------------------- @@ -725,13 +735,49 @@ def test_validate_config_tie_to_stage1_requires_stage1_active(): assert "tie_to_stage1" in str(e) -def test_validate_config_stop_token_not_implemented(): +def test_validate_config_stop_token_accepted_under_autoregressive(): + """DEFAULT_CONFIG's stage2_model.decoder is already "autoregressive" + (see test_stage2_model_config_defaults_match_documented_v030_intent), so + mode="stop_token" alone must not raise.""" cfg = _cfg_with(**{"stage2_model.n_sec.mode": "stop_token"}) + gconfig.validate_config(cfg) # must not raise + + +def test_validate_config_stop_token_rejected_under_one_shot(): + cfg = _cfg_with( + **{ + "stage2_model.n_sec.mode": "stop_token", + "stage2_model.decoder": "one_shot", + } + ) try: gconfig.validate_config(cfg) assert False, "expected ValueError" except ValueError as e: - assert "stop_token" in str(e) + assert "stop_token" in str(e) and "autoregressive" in str(e) + + +def test_validate_config_stop_token_rejected_for_stage1_owner(): + cfg = _cfg_with( + **{ + "stage2_model.n_sec.mode": "stop_token", + "stage2_model.n_sec.owner": "stage1", + } + ) + try: + gconfig.validate_config(cfg) + assert False, "expected ValueError" + except ValueError as e: + 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"}) + try: + gconfig.validate_config(cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "stop_sampling" in str(e) def test_validate_config_stage1_context_sampled_not_implemented(): diff --git a/tests/test_rollout.py b/tests/test_rollout.py index 36aa83b..91f4113 100644 --- a/tests/test_rollout.py +++ b/tests/test_rollout.py @@ -366,6 +366,7 @@ def _models_v3( k_max=6, emb_dim=4, stage2_has_n_sec_head=True, + stop_token=False, ): particle_cfg = ConditioningAxisConfig(type=conditioning, emb_dim=emb_dim, n_layers=1) material_cfg = ConditioningAxisConfig(type=conditioning, emb_dim=emb_dim, n_layers=1) @@ -417,7 +418,8 @@ def _models_v3( noise_dim=8, k_max=k_max, particle_type_cfg=particle_type_cfg, - build_n_sec_head=stage2_has_n_sec_head, + build_n_sec_head=stage2_has_n_sec_head and not stop_token, + build_stop_head=stop_token, ) return s1.eval(), s2.eval() @@ -481,6 +483,23 @@ def test_rollout_stage2_owns_n_sec_when_stage1_has_no_head(fake_material_props): assert dep + leak == pytest.approx(seeds["pre_E"][i], rel=1e-4) +@pytest.mark.parametrize("generator2", ["flow", "wgan"]) +def test_rollout_stop_token_end_to_end(fake_material_props, generator2): + """gitea #40: n_sec.mode='stop_token' (no n_sec_head on either stage — + resolve_n_sec must return None and let sample_stage2's AR loop derive + the count from its own stop head) must still run to completion, produce + secondaries, and conserve energy exactly like the 'head' mode.""" + s1, s2 = _models_v3(decoder="autoregressive", generator2=generator2, stop_token=True) + rec = _run_v3(s1, s2) + assert len(rec["event_id"]) > 0 + seeds = _seeds() + for i, ev in enumerate(seeds["event_id"]): + m = rec["event_id"] == ev + dep = rec["edep"][m].sum() + leak = rec["pre_E"][m & (rec["termination_reason"] == TERM_ESCAPED)].sum() + assert dep + leak == pytest.approx(seeds["pre_E"][i], rel=1e-4) + + def test_resolve_n_sec_raises_when_neither_stage_owns_head(fake_material_props): """Neither stage owning n_sec_head only happens for a stage2_model.n_sec.mode other than "head" — not a valid rollout-capable diff --git a/tests/test_sample.py b/tests/test_sample.py index 9ea8aa0..eb31f89 100644 --- a/tests/test_sample.py +++ b/tests/test_sample.py @@ -96,6 +96,45 @@ def _expected_type_dim(target: str, emb_dim: int) -> int: return PARTICLE_PHYS_DIM if target == "physical" else emb_dim +def _stage2_ar_stop_token( + target: str, + generator: str, + stop_sampling: str = "greedy", + emb_dim: int = 6, + pdg: int = 3, + mat: int = 2, + k_max: int = 5, +) -> Stage2Autoregressive: + particle_cfg, material_cfg = _particle_material_cfg(_conditioning_for(target), emb_dim) + return Stage2Autoregressive( + pdg_vocab=pdg, + mat_vocab=mat, + particle_cfg=particle_cfg, + material_cfg=material_cfg, + hidden_dim=32, + n_res_blocks=2, + generator=generator, + time_dim=16, + noise_dim=8, + k_max=k_max, + particle_type_cfg=ParticleTypeConfig(target=target), + build_n_sec_head=False, + build_stop_head=True, + stop_sampling=stop_sampling, + ).eval() + + +def _force_stop_head_logit(decoder: Stage2Autoregressive, logit: float) -> None: + """Zeroes stop_head's weights and pins its bias, so predict_stop returns + `logit` for every row/slot regardless of conditioning — makes the AR + loop's stop decision deterministic for testing.""" + assert decoder.stop_head is not None + last_linear = decoder.stop_head[-1] + with torch.no_grad(): + last_linear.weight.zero_() + last_linear.bias.fill_(logit) + + # ── Stage-1 n_sec ownership ────────────────────────────────────────────────── @@ -223,3 +262,77 @@ def test_sample_secondaries_ar_first_slot_has_no_history(): sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2) assert sec_cont.shape == (B, 1, CONT_SLOT_DIM) assert sec_valid.tolist() == [[False], [True], [True]] + + +# ── 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): + """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) + _force_stop_head_logit(decoder, 50.0) + cond_cont, cond_cat = _cond(B) + stage1_out = torch.randn(B, X_DIM) + sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2) + assert sec_valid.shape == (B, k_max) + 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): + """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) + _force_stop_head_logit(decoder, -50.0) + cond_cont, cond_cat = _cond(B) + stage1_out = torch.randn(B, X_DIM) + sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2) + assert sec_valid.all() + + +@pytest.mark.parametrize("generator", ["flow", "wgan"]) +def test_sample_secondaries_ar_stop_token_valid_mask_is_always_a_prefix(generator): + """Without forcing the stop head, per-row stop timing varies — but + sec_valid must always be a contiguous prefix (slot k valid implies every + slot < k is also valid), matching the "head"/"truth" contract.""" + B, k_max = 6, 5 + decoder = _stage2_ar_stop_token("physical", generator, k_max=k_max) + cond_cont, cond_cat = _cond(B) + stage1_out = torch.randn(B, X_DIM) + _, _, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2) + n = sec_valid.sum(dim=-1) + expected = torch.arange(k_max).unsqueeze(0) < n.unsqueeze(1) + assert torch.equal(sec_valid, expected) + + +def test_sample_secondaries_ar_stop_token_explicit_n_sec_pred_ignores_stop_head(): + """The scheduled-sampling training contract: passing n_sec_pred + explicitly (as _assemble_stage2_ar_inputs_scheduled's self-sample call + does, with ground-truth n_sec) must run the full k_max loop and mask by + the given count, even though the decoder owns a stop_head that would + otherwise stop early.""" + B, k_max = 3, 5 + decoder = _stage2_ar_stop_token("physical", "flow", k_max=k_max) + _force_stop_head_logit(decoder, 50.0) # would stop immediately if consulted + cond_cont, cond_cat = _cond(B) + stage1_out = torch.randn(B, X_DIM) + n_sec_pred = torch.tensor([0, 2, k_max]) + _, _, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2) + for i, n in enumerate(n_sec_pred.tolist()): + assert sec_valid[i, :n].all() + assert not sec_valid[i, n:].any() + + +def test_sample_secondaries_ar_none_n_sec_pred_without_stop_head_raises(): + decoder = _stage2_ar("physical", "flow", k_max=5) # head mode: no stop_head + cond_cont, cond_cat = _cond(3) + stage1_out = torch.randn(3, X_DIM) + with pytest.raises(AssertionError): + sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2) diff --git a/tests/test_train.py b/tests/test_train.py index 98f7f0b..b5da165 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -20,7 +20,7 @@ from giant.constants import ( X_DIM, ) from giant.data.dataset import StepBatch -from giant.model.network import build_critics, build_models +from giant.model.network import Stage2Autoregressive, build_critics, build_models from giant.training import ( FlowDDPMStageTrainer, StageSpec, @@ -40,6 +40,7 @@ from giant.training.stage2_inputs import ( _shift_prev, _stage2_tf_prob, _stick_fraction, + _stop_target_and_mask, _type_repr, ) @@ -120,6 +121,23 @@ def test_ar_has_prev_false_only_at_slot_zero(): assert has_prev.tolist() == [[False, True, True, True, True]] +def test_stop_target_and_mask_hand_computed(): + # k_max=5; n_sec=0 (no real secondaries, stop slot is 0), n_sec=2 + # (stop slot is 2), n_sec=5 (== k_max: no in-range stop slot at all). + n_sec = torch.tensor([0, 2, 5]) + target, mask = _stop_target_and_mask(n_sec, 5, torch.device("cpu")) + assert target.tolist() == [ + [1, 0, 0, 0, 0], + [0, 0, 1, 0, 0], + [0, 0, 0, 0, 0], + ] + assert mask.tolist() == [ + [True, False, False, False, False], + [True, True, True, False, False], + [True, True, True, True, True], + ] + + # --- _stage2_tf_prob (v0.3.0 step 7) ----------- @@ -740,6 +758,71 @@ def test_wgan_onehot_one_shot_also_gets_grad_norm_instrumentation(): assert any(float(r["stage2/train/grad_norm_cont_slice"]) > 0 for r in rows) +# --- n_sec.mode = "stop_token" (gitea #40) ---------------------------------- + + +def _stop_token_cfg(): + cfg = _base_cfg() + cfg["stage2_model"]["decoder"] = "autoregressive" + cfg["stage2_model"]["n_sec"] = {"mode": "stop_token", "lambda": 0.1} + return cfg + + +@pytest.mark.parametrize("stage2_generator", ["wgan", "flow"]) +def test_build_stage_trainers_stop_token_step_runs(stage2_generator): + """A stop_token AR stage-2 trainer.step() must run and emit a finite + loss_stop for both non-adversarial (flow) and WGAN generators — the two + trainer subclasses wire the stop head's BCE term in independently.""" + cfg = _stop_token_cfg() + cfg["stage2_model"]["generator"] = stage2_generator + model_config = _model_config(cfg) + models = build_models(model_config) + critics = build_critics(model_config) + trainers = build_stage_trainers(cfg, models, critics, torch.device("cpu"), total_train_batches=4) + trainer = trainers["stage2"] + batch = _fake_batches(1, 4)[0] + stats = trainer.step(batch, torch.device("cpu"), global_step=1) + assert math.isfinite(stats["loss_stop"]) + assert math.isfinite(stats["stop_acc"]) + + +def test_stop_token_model_has_stop_head_not_n_sec_head(): + cfg = _stop_token_cfg() + model_config = _model_config(cfg) + models = build_models(model_config) + stage2 = models["stage2"] + assert isinstance(stage2, Stage2Autoregressive) + assert stage2.n_sec_head is None + assert stage2.stop_head is not None + + +def test_head_mode_model_has_n_sec_head_not_stop_head(): + """Sanity check on the other side of the gate — the default 'head' mode + must be unaffected by the stop_head plumbing.""" + cfg = _base_cfg() + cfg["stage2_model"]["decoder"] = "autoregressive" + model_config = _model_config(cfg) + models = build_models(model_config) + stage2 = models["stage2"] + assert isinstance(stage2, Stage2Autoregressive) + assert stage2.n_sec_head is not None + assert stage2.stop_head is None + + +def test_train_end_to_end_stop_token(): + """Full train() run with n_sec.mode='stop_token' must complete and write + a checkpoint + metrics.csv with finite losses throughout.""" + cfg = _stop_token_cfg() + with tempfile.TemporaryDirectory() as tmp: + out_dir = Path(tmp) / "run" + _run_train(cfg, out_dir) + assert (out_dir / "last.pt").exists() + with open(out_dir / "metrics.csv", newline="") as f: + rows = list(csv.DictReader(f)) + assert len(rows) == cfg["train"]["epochs"] + assert all(math.isfinite(float(r["stage2/train/loss_stop"])) for r in rows) + + def test_wgan_physical_omits_grad_norm_slice_columns(): cfg = _base_cfg() # _base_cfg's stage2_model.particle_type.target is "physical" with tempfile.TemporaryDirectory() as tmp: diff --git a/uv.lock b/uv.lock index c4e2c8e..f56201e 100644 --- a/uv.lock +++ b/uv.lock @@ -633,7 +633,7 @@ wheels = [ [[package]] name = "giant" -version = "0.3.1" +version = "0.3.2" source = { editable = "." } dependencies = [ { name = "numpy" },