From 48faaee79dbad35a027fe7d733fd52b86b9e088e Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Mon, 17 Aug 2026 12:27:13 +0200 Subject: [PATCH] Implement stage2_model.stage1_context = "sampled" (gitea #41) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Stage 2 was trained on ground-truth stage-1 outcomes but deployed on sampled ones, and in a rollout that gap compounds over every step of every track — the same train/inference gap teacher_forcing="scheduled" already closes within stage 2, just never applied at the stage boundary. "sampled" was declared in the schema but rejected loudly by validate_config as unimplemented; this lands the real implementation. Mirrors the existing scheduled-sampling precedent rather than a hard switch: new stage2_model.ctx_p_start/ctx_p_end (defaults 1.0 -> 0.0) linearly ramp P(condition on ground truth) from epoch 0 to the final epoch, so stage 2 doesn't chase a wildly moving stage-1 target early in training. Per the plan discussed with the user: the sample is drawn from stage 1's sampling_model() (EMA weights when present, matching what inference actually deploys), mixed per example via a Bernoulli draw (never blended within a row), and validation always uses the ground truth regardless of the schedule. Fixes a latent bug the same pattern would otherwise have hit: every sampler in giant/sample.py flips its model to .eval() with no restore, so sampling from the raw (non-EMA) stage-1 model mid-step now explicitly restores its .training flag afterward to avoid silently corrupting stage 1's own training mode for the rest of the epoch. validate_config now enforces stage1_context in {"truth", "sampled"}, requires both stages active for "sampled" (nothing to sample from otherwise), range-checks ctx_p_start/ctx_p_end, and rejects the ctx_p_start = ctx_p_end = 1.0 configuration as an unadvertised no-op identical to "truth". Co-Authored-By: Claude Opus 5 --- giant/config.py | 40 +++++++-- giant/model/summary.py | 13 +-- giant/training/stage2_inputs.py | 28 +++++-- giant/training/trainers.py | 116 +++++++++++++++++++++++--- tests/test_config.py | 70 +++++++++++++++- tests/test_config_consumed_keys.py | 7 -- tests/test_train.py | 127 ++++++++++++++++++++++++++++- 7 files changed, 360 insertions(+), 41 deletions(-) diff --git a/giant/config.py b/giant/config.py index 3873e86..b26c606 100644 --- a/giant/config.py +++ b/giant/config.py @@ -677,6 +677,13 @@ class Stage2ModelConfig: # output, closing the train/inference gap at the cost of a sampling pass # per batch and a moving target early in training. stage1_context: str = "truth" + # Ramp for "sampled": P(condition on the ground-truth stage-1 outcome + # rather than a fresh sample), linearly interpolated from ctx_p_start + # (epoch 0) to ctx_p_end (the final epoch) — the same scheduled-sampling + # shape as autoregressive.tf_p_start/tf_p_end, so stage 2 doesn't chase a + # wildly moving stage-1 target in early epochs. Unread under "truth". + ctx_p_start: float = 1.0 + ctx_p_end: float = 0.0 n_sec: NSecConfig = field(default_factory=NSecConfig) particle_type: ParticleTypeConfig = field(default_factory=ParticleTypeConfig) autoregressive: AutoregressiveConfig = field(default_factory=AutoregressiveConfig) @@ -701,6 +708,8 @@ class Stage2ModelConfig: k_max=d.get("k_max", 15), context_dim=d.get("context_dim", 64), stage1_context=d.get("stage1_context", "truth"), + ctx_p_start=d.get("ctx_p_start", 1.0), + ctx_p_end=d.get("ctx_p_end", 0.0), n_sec=NSecConfig.from_dict(d.get("n_sec")), particle_type=ParticleTypeConfig.from_dict(d.get("particle_type")), autoregressive=AutoregressiveConfig.from_dict(d.get("autoregressive")), @@ -724,6 +733,8 @@ class Stage2ModelConfig: "k_max": self.k_max, "context_dim": self.context_dim, "stage1_context": self.stage1_context, + "ctx_p_start": self.ctx_p_start, + "ctx_p_end": self.ctx_p_end, "n_sec": self.n_sec.to_dict(), "particle_type": self.particle_type.to_dict(), "autoregressive": self.autoregressive.to_dict(), @@ -1420,13 +1431,28 @@ def validate_config(cfg: dict) -> None: 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( - "stage2_model.stage1_context = 'sampled' is accepted by the schema " - "but not implemented — trainers.py always trains stage 2 against " - "the ground-truth stage-1 output; use 'truth' (default) instead " - "(see issues.md Issue 16 for the planned implementation)" - ) + stage1_context = _get_path(cfg, "stage2_model.stage1_context") + if stage1_context not in ("truth", "sampled"): + raise ValueError(f"stage2_model.stage1_context = {stage1_context!r} — must be 'truth' or 'sampled'") + if stage1_context == "sampled": + if not (_get_path(cfg, "stage1_model.active") and _get_path(cfg, "stage2_model.active")): + raise ValueError( + "stage2_model.stage1_context = 'sampled' requires both " + "stage1_model.active and stage2_model.active = true — there is " + "no stage-1 model to sample from in a stage-2-only run" + ) + ctx_p_start = _get_path(cfg, "stage2_model.ctx_p_start") + ctx_p_end = _get_path(cfg, "stage2_model.ctx_p_end") + for name, value in (("ctx_p_start", ctx_p_start), ("ctx_p_end", ctx_p_end)): + if not (0.0 <= value <= 1.0): + raise ValueError(f"stage2_model.{name} = {value} — must be in [0, 1]") + if ctx_p_start == 1.0 and ctx_p_end == 1.0: + raise ValueError( + "stage2_model.stage1_context = 'sampled' with ctx_p_start = " + "ctx_p_end = 1.0 always conditions on the ground truth — " + "identical to 'truth' but silently so; use 'truth' instead or " + "lower ctx_p_end" + ) if ( _get_path(cfg, "stage2_model.n_sec.mode") == "truth" diff --git a/giant/model/summary.py b/giant/model/summary.py index ff8fdc6..d41473b 100644 --- a/giant/model/summary.py +++ b/giant/model/summary.py @@ -23,11 +23,11 @@ baked into a static table. Keys legitimately owned by the trainer/sampler/rollout rather than by `build_models`/`build_critics` (loss weights, WGAN-GP training -hyperparameters, teacher-forcing schedules, ...) are cataloged in -`_NOT_BUILD_TIME` below so the report doesn't flag them as suspicious. A -couple of leaves are inert under every config today — -`stage2_model.autoregressive.order`, `stage2_model.stage1_context` — matching -`tests/test_config_consumed_keys.py`'s own `_KNOWN_UNUSED` entries; they are +hyperparameters, teacher-forcing and stage1-context schedules, ...) are +cataloged in `_NOT_BUILD_TIME` below so the report doesn't flag them as +suspicious. One leaf is inert under every config today — +`stage2_model.autoregressive.order` — matching +`tests/test_config_consumed_keys.py`'s own `_KNOWN_UNUSED` entry; it is deliberately *not* in `_NOT_BUILD_TIME`, since "always inert" is itself the finding those two tests independently converge on. """ @@ -73,6 +73,9 @@ _NOT_BUILD_TIME: dict[str, str] = { "stage2_model.autoregressive.teacher_forcing": "giant/training/stage2_inputs.py's training-time input assembly", "stage2_model.autoregressive.tf_p_start": "trainers.py's teacher-forcing schedule", "stage2_model.autoregressive.tf_p_end": "trainers.py's teacher-forcing schedule", + "stage2_model.stage1_context": "trainers.py's stage1/stage2 boundary — StageTrainer._stage1_context", + "stage2_model.ctx_p_start": "trainers.py's stage1-context sampling schedule", + "stage2_model.ctx_p_end": "trainers.py's stage1-context sampling schedule", "stage1_model.router.lambda_balance": "trainers.py's load-balancing auxiliary loss weight", "stage1_model.router.lambda_entropy": "trainers.py's entropy-regularization auxiliary loss weight", "stage1_model.router.lambda_proc": "trainers.py's supervised process-classification auxiliary loss weight", diff --git a/giant/training/stage2_inputs.py b/giant/training/stage2_inputs.py index 5277587..ca0b6fd 100644 --- a/giant/training/stage2_inputs.py +++ b/giant/training/stage2_inputs.py @@ -202,21 +202,37 @@ def _assemble_stage2_ar_inputs( return {"history_feat": history_feat, **_ar_meta(K, B, device, fraction)} +def _linear_schedule(p_start: float, p_end: float, epoch: int, total_epochs: int) -> float: + """Linear interpolation from `p_start` (epoch 0) to `p_end` (the final + epoch) — standard scheduled sampling (Bengio et al. 2015), shared by + every train-time schedule keyed on epoch.""" + frac = epoch / max(total_epochs - 1, 1) + frac = min(max(frac, 0.0), 1.0) + return p_start + (p_end - p_start) * frac + + def _stage2_tf_prob(mode: str, p_start: float, p_end: float, epoch: int, total_epochs: int) -> float: """P(condition slot k+1 on the TRUE token k rather than the model's own prediction), for the current epoch (`stage2_model.autoregressive.teacher_forcing`). `"always"`/`"never"` are the two degenerate constants; `"scheduled"` - linearly interpolates - `p_start` (epoch 0) to `p_end` (the final epoch) — standard scheduled - sampling (Bengio et al. 2015).""" + linearly interpolates `p_start` to `p_end` via `_linear_schedule`.""" if mode == "always": return 1.0 if mode == "never": return 0.0 - frac = epoch / max(total_epochs - 1, 1) - frac = min(max(frac, 0.0), 1.0) - return p_start + (p_end - p_start) * frac + return _linear_schedule(p_start, p_end, epoch, total_epochs) + + +def _ctx_truth_prob(mode: str, p_start: float, p_end: float, epoch: int, total_epochs: int) -> float: + """P(condition stage 2 on the TRUE stage-1 outcome rather than a fresh + stage-1 sample), for the current epoch (`stage2_model.stage1_context`). + `"truth"` is the degenerate constant 1.0; `"sampled"` linearly + interpolates `ctx_p_start` to `ctx_p_end` via `_linear_schedule` — the + stage-boundary counterpart of `_stage2_tf_prob`.""" + if mode == "truth": + return 1.0 + return _linear_schedule(p_start, p_end, epoch, total_epochs) def _history_repr_from_ar_sample( diff --git a/giant/training/trainers.py b/giant/training/trainers.py index 9b80cd6..159d8b1 100644 --- a/giant/training/trainers.py +++ b/giant/training/trainers.py @@ -27,10 +27,12 @@ from giant.constants import CONT_SLOT_DIM from giant.data.dataset import StepBatch from giant.model.network import Router, build_objective, resolve_type_n_classes, stage2_type_dim from giant.model.wgan import generator_loss, gradient_penalty +from giant.sample import sample_stage1 from giant.training.metrics import MetricSpec, stage_metric, train_metric, val_metric from giant.training.stage2_inputs import ( _assemble_stage2_ar_inputs_scheduled, _assemble_stage2_ar_target, + _ctx_truth_prob, _gumbel_tau, _relax_onehot_type_slice, _stage2_tf_prob, @@ -117,6 +119,11 @@ class StageSpec: tf_p_end: float = 1.0 ar_sample_steps: int = 10 + # stage-1/stage-2 boundary (stage 2 only) + stage1_context: str = "truth" + ctx_p_start: float = 1.0 + ctx_p_end: float = 0.0 + # generator-specific ddpm_n_steps: int = 1000 n_critic: int = 5 @@ -169,6 +176,9 @@ class StageSpec: teacher_forcing=s2_spec.autoregressive.teacher_forcing if is_stage2 else cls.teacher_forcing, tf_p_start=s2_spec.autoregressive.tf_p_start if is_stage2 else cls.tf_p_start, tf_p_end=s2_spec.autoregressive.tf_p_end if is_stage2 else cls.tf_p_end, + stage1_context=s2_spec.stage1_context if is_stage2 else cls.stage1_context, + ctx_p_start=s2_spec.ctx_p_start if is_stage2 else cls.ctx_p_start, + ctx_p_end=s2_spec.ctx_p_end if is_stage2 else cls.ctx_p_end, # AR self-sampling under scheduled/never teacher forcing reuses # train.validate_steps as its flow-matching ODE step count — no # dedicated config key for this (the autoregressive config lists @@ -186,13 +196,16 @@ class StageSpec: class StageTrainer: """One active stage's optimizer(s), EMA, and per-batch step. - Reads only the shared `StepBatch` (`giant.data.dataset`) — stage 2 always - conditions on the ground-truth `x1_s1` (`stage2_model.stage1_context = - "truth"`, stage-level teacher forcing; `"sampled"` is not implemented), - so stage trainers never need each other's output at train time. This means - "stage-2-only training is a cheap ablation, not new plumbing" falls out - for free: a trainer only exists for active stages, and inactive stages - are simply never constructed. + Reads only the shared `StepBatch` (`giant.data.dataset`) by default — stage + 2 conditions on the ground-truth `x1_s1` (`stage2_model.stage1_context = + "truth"`, stage-level teacher forcing), so "stage-2-only training is a + cheap ablation, not new plumbing" falls out for free: a trainer only + exists for active stages, and inactive stages are simply never + constructed. `stage2_model.stage1_context = "sampled"` is the one + exception — `build_stage_trainers` wires the stage-2 trainer to the + stage-1 one via `attach_stage1` so it can draw a real stage-1 sample + (`giant.sample.sample_stage1`) instead, scheduled by `ctx_p_start`/ + `ctx_p_end` (see `_stage1_context`). Grad-norm clipping is per-stage here — v0.2's single shared optimizer clipped both stages' gradients jointly; splitting per stage is a small, @@ -246,6 +259,18 @@ class StageTrainer: for p in self.ema_model.parameters(): p.requires_grad_(False) + #: Set by `attach_stage1` when `stage2_model.stage1_context = + #: "sampled"` — the stage-1 `StageTrainer` this (stage-2) trainer + #: draws its context sample from. `None` for stage 1 itself, and for + #: stage 2 under "truth". + self.stage1_source: "StageTrainer | None" = None + + def attach_stage1(self, stage1_trainer: "StageTrainer") -> None: + """Wires this (stage-2) trainer to the stage-1 trainer it should + sample from under `stage2_model.stage1_context = "sampled"` — see + `build_stage_trainers`.""" + self.stage1_source = stage1_trainer + # --- schedule ------------------------------------------------------- def _init_lr_schedule(self, optimizer: optim.Optimizer, warmup_steps: int, total_steps: int) -> None: @@ -289,6 +314,58 @@ class StageTrainer: for module in self._modules: module.eval() + # --- stage-1/stage-2 boundary (shared by both trainer subclasses) --- + + def _stage1_context( + self, + x1_s1: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + epoch: int | None, + ) -> torch.Tensor: + """The stage-1 outcome stage 2 conditions on this batch. + + `epoch=None` means "always ground truth" regardless of + `spec.stage1_context` — the same val-loss convention `_ar_inputs` + uses, so validation stays a stable, non-stochastic comparison. + Otherwise, under `stage1_context = "sampled"`, each example + independently uses the ground truth with probability `p_truth` + (`_ctx_truth_prob`, ramped by `ctx_p_start`/`ctx_p_end`) and a fresh + `giant.sample.sample_stage1` draw from `stage1_source.sampling_model()` + otherwise — a real sampling pass, not a cheap proxy, matching + `_assemble_stage2_ar_inputs_scheduled`'s precedent for the equivalent + in-stage-2 self-sample. Mixed per example (not per-dimension): a row + is either the real ground-truth 9D vector or a real sample, never an + elementwise blend of the two. + """ + x1_s1 = x1_s1.detach() + if self.stage1_source is None or epoch is None: + return x1_s1 + p_truth = _ctx_truth_prob( + self.spec.stage1_context, + self.spec.ctx_p_start, + self.spec.ctx_p_end, + epoch, + self.spec.epochs, + ) + if p_truth >= 1.0: + return x1_s1 + + stage1_model = self.stage1_source.sampling_model() + was_training = stage1_model.training + sampled, _ = sample_stage1( + stage1_model, + cond_cont, + cond_cat, + steps=self.spec.ar_sample_steps, + ddpm_steps=self.stage1_source.spec.ddpm_n_steps, + ) + if was_training: + stage1_model.train() + + use_truth = torch.rand(x1_s1.size(0), 1, device=x1_s1.device) < p_truth + return torch.where(use_truth, x1_s1, sampled).detach() + # --- stage-2 secondary assembly (shared by both trainer subclasses) --- def _ar_inputs( @@ -608,9 +685,10 @@ class FlowDDPMStageTrainer(StageTrainer): def _compute(self, batch: StepBatch, device: torch.device, epoch: int | None = None) -> dict: """`epoch=None` (the `val_loss` path) always uses full teacher - forcing (`p_tf=1.0`) regardless of `spec.teacher_forcing` — validation - should stay a stable, non-stochastic ground-truth comparison; only - the training `step` path schedules `p_tf` by epoch.""" + forcing (`p_tf=1.0`) and the ground-truth stage-1 context, regardless + of `spec.teacher_forcing`/`spec.stage1_context` — validation should + stay a stable, non-stochastic ground-truth comparison; only the + training `step` path schedules `p_tf`/`p_truth` by epoch.""" ( cond_cont, cond_cat, @@ -621,7 +699,7 @@ class FlowDDPMStageTrainer(StageTrainer): sec_type_idx, ) = _batch_to_device(batch, device) sec_mask = self._sec_mask(n_sec, sec_cont.size(1), device) - stage1_ctx = x1_s1.detach() + stage1_ctx = self._stage1_context(x1_s1, cond_cont, cond_cat, epoch) x1_s2 = None ar_inputs = None @@ -846,7 +924,8 @@ class WGANStageTrainer(StageTrainer): sec_type_idx, ) = _batch_to_device(batch, device) B = cond_cont.size(0) - stage1_ctx = x1_s1.detach() + epoch = global_step // self.spec.steps_per_epoch + stage1_ctx = self._stage1_context(x1_s1, cond_cont, cond_cat, epoch) grad_probe: dict[str, float] = {} ar_inputs = None @@ -996,7 +1075,13 @@ def build_stage_trainers( total_train_batches: int, ) -> dict[str, StageTrainer]: """One trainer per active stage — `models[name] is None` means that stage - is `active = false` and is simply never constructed.""" + is `active = false` and is simply never constructed. + + `stage2_model.stage1_context = "sampled"` additionally wires the + stage-2 trainer to the stage-1 one (`StageTrainer.attach_stage1`) so it + can draw a real stage-1 sample instead of only ever seeing the + ground-truth stage-1 outcome — `validate_config` already guarantees both + stages are active whenever that config value is set.""" trainers: dict[str, StageTrainer] = {} for name, is_stage2 in (("stage1", False), ("stage2", True)): model = models.get(name) @@ -1011,4 +1096,9 @@ def build_stage_trainers( trainers[name] = WGANStageTrainer(spec, model, critic, device) else: trainers[name] = FlowDDPMStageTrainer(spec, model, device) + + stage2 = trainers.get("stage2") + stage1 = trainers.get("stage1") + if stage2 is not None and stage1 is not None and stage2.spec.stage1_context == "sampled": + stage2.attach_stage1(stage1) return trainers diff --git a/tests/test_config.py b/tests/test_config.py index 4e71c96..0eee62a 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -780,13 +780,79 @@ def test_validate_config_bad_stop_sampling_rejected(): assert "stop_sampling" in str(e) -def test_validate_config_stage1_context_sampled_not_implemented(): +def test_validate_config_stage1_context_sampled_accepted_with_both_stages_active(): + """gitea #41: 'sampled' is now implemented, so DEFAULT_CONFIG's + stage1_model/stage2_model.active = true (both) must let it through.""" cfg = _cfg_with(**{"stage2_model.stage1_context": "sampled"}) + gconfig.validate_config(cfg) # must not raise + + +def test_validate_config_bad_stage1_context_rejected(): + cfg = _cfg_with(**{"stage2_model.stage1_context": "bogus"}) try: gconfig.validate_config(cfg) assert False, "expected ValueError" except ValueError as e: - assert "sampled" in str(e) + assert "stage1_context" in str(e) + + +def test_validate_config_stage1_context_sampled_requires_stage1_active(): + cfg = _cfg_with( + **{ + "stage2_model.stage1_context": "sampled", + "stage1_model.active": False, + } + ) + try: + gconfig.validate_config(cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "sampled" in str(e) and "stage1_model.active" in str(e) + + +def test_validate_config_stage1_context_sampled_requires_stage2_active(): + cfg = _cfg_with( + **{ + "stage2_model.stage1_context": "sampled", + "stage2_model.active": False, + } + ) + try: + gconfig.validate_config(cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "sampled" in str(e) and "stage2_model.active" in str(e) + + +@pytest.mark.parametrize("key", ["ctx_p_start", "ctx_p_end"]) +@pytest.mark.parametrize("value", [-0.1, 1.1]) +def test_validate_config_ctx_p_out_of_range_rejected(key, value): + cfg = _cfg_with( + **{ + "stage2_model.stage1_context": "sampled", + f"stage2_model.{key}": value, + } + ) + try: + gconfig.validate_config(cfg) + assert False, "expected ValueError" + except ValueError as e: + assert key in str(e) + + +def test_validate_config_stage1_context_sampled_always_truth_rejected_as_noop(): + cfg = _cfg_with( + **{ + "stage2_model.stage1_context": "sampled", + "stage2_model.ctx_p_start": 1.0, + "stage2_model.ctx_p_end": 1.0, + } + ) + try: + gconfig.validate_config(cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "ctx_p_start" in str(e) and "ctx_p_end" in str(e) def test_validate_config_n_sec_truth_rejected_for_rollout_capable_checkpoint(): diff --git a/tests/test_config_consumed_keys.py b/tests/test_config_consumed_keys.py index 9672a8c..84d1a76 100644 --- a/tests/test_config_consumed_keys.py +++ b/tests/test_config_consumed_keys.py @@ -54,13 +54,6 @@ _EXCLUDED_FILES = ("giant/model/_legacy.py",) # it. If a key here starts showing up as consumed, the fix landed and this # entry is stale — see test_known_unused_allow_list_has_no_stale_entries. _KNOWN_UNUSED = { - "stage2_model.stage1_context": ( - "issues.md Issue 1 — trainers.py hardcodes stage1_ctx to the " - "ground-truth stage-1 output; 'sampled' is now rejected loudly by " - "validate_config (not silently accepted), but the key still isn't " - "read by any build/train consumer file since only 'truth' can pass " - "validation — see Issue 16 for the real implementation" - ), "stage2_model.autoregressive.order": ( "gitea #30 — validate_config now checks order is 'energy_desc', but " "nothing in the build/train/rollout consumer whitelist reads the " diff --git a/tests/test_train.py b/tests/test_train.py index b5da165..3708bfe 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -5,7 +5,7 @@ import csv import math import tempfile from pathlib import Path -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest import torch @@ -21,6 +21,7 @@ from giant.constants import ( ) from giant.data.dataset import StepBatch from giant.model.network import Stage2Autoregressive, build_critics, build_models +from giant.sample import sample_stage1 as trainers_sample_stage1 from giant.training import ( FlowDDPMStageTrainer, StageSpec, @@ -831,3 +832,127 @@ def test_wgan_physical_omits_grad_norm_slice_columns(): header = (out_dir / "metrics.csv").read_text().splitlines()[0].split(",") assert "stage2/train/grad_norm_type_slice" not in header assert "stage2/train/grad_norm_cont_slice" not in header + + +# --- stage2_model.stage1_context = "sampled" (gitea #41) -------------------- + + +def _sampled_ctx_cfg(ema_decay=0.999): + cfg = _base_cfg() + cfg["stage1_model"]["generator"] = "flow" + cfg["stage2_model"]["generator"] = "flow" + cfg["stage2_model"]["stage1_context"] = "sampled" + cfg["stage2_model"]["ctx_p_start"] = 0.0 + cfg["stage2_model"]["ctx_p_end"] = 0.0 + cfg["train"]["ema_decay"] = ema_decay + return cfg + + +def _build_sampled_trainers(cfg): + model_config = _model_config(cfg) + models = build_models(model_config) + critics = build_critics(model_config) + return build_stage_trainers(cfg, models, critics, torch.device("cpu"), total_train_batches=4) + + +def test_build_stage_trainers_attaches_stage1_only_under_sampled(): + trainers = _build_sampled_trainers(_sampled_ctx_cfg()) + assert trainers["stage2"].stage1_source is trainers["stage1"] + assert trainers["stage1"].stage1_source is None + + +def test_build_stage_trainers_leaves_stage1_source_none_under_truth(): + """Regression guard for the old silent no-op: 'truth' (the default) must + never attach a stage1_source, so _stage1_context short-circuits without + ever calling sample_stage1.""" + cfg = _base_cfg() + trainers = _build_sampled_trainers(cfg) + assert trainers["stage2"].stage1_source is None + + +def test_stage1_context_sampled_calls_sample_stage1_and_differs_from_truth(): + cfg = _sampled_ctx_cfg() + trainers = _build_sampled_trainers(cfg) + stage1, stage2 = trainers["stage1"], trainers["stage2"] + batch = _fake_batches(1, 4)[0] + cond_cont, cond_cat, x1_s1 = batch.cond_cont, batch.cond_cat, batch.target_s1 + + with patch("giant.training.trainers.sample_stage1", wraps=trainers_sample_stage1) as spy: + ctx = stage2._stage1_context(x1_s1, cond_cont, cond_cat, epoch=0) + assert spy.call_count == 1 + assert spy.call_args.args[0] is stage1.sampling_model() + assert not torch.equal(ctx, x1_s1) + + +def test_stage1_context_truth_default_never_calls_sample_stage1(): + cfg = _base_cfg() + trainers = _build_sampled_trainers(cfg) + stage2 = trainers["stage2"] + batch = _fake_batches(1, 4)[0] + cond_cont, cond_cat, x1_s1 = batch.cond_cont, batch.cond_cat, batch.target_s1 + + with patch("giant.training.trainers.sample_stage1", wraps=trainers_sample_stage1) as spy: + ctx = stage2._stage1_context(x1_s1, cond_cont, cond_cat, epoch=0) + assert spy.call_count == 0 + assert torch.equal(ctx, x1_s1) + + +def test_stage1_context_val_epoch_none_uses_ground_truth_even_under_sampled(): + cfg = _sampled_ctx_cfg() + trainers = _build_sampled_trainers(cfg) + stage2 = trainers["stage2"] + batch = _fake_batches(1, 4)[0] + cond_cont, cond_cat, x1_s1 = batch.cond_cont, batch.cond_cat, batch.target_s1 + + with patch("giant.training.trainers.sample_stage1", wraps=trainers_sample_stage1) as spy: + ctx = stage2._stage1_context(x1_s1, cond_cont, cond_cat, epoch=None) + assert spy.call_count == 0 + assert torch.equal(ctx, x1_s1) + + +def test_stage1_context_sampled_preserves_stage1_training_mode(): + """Every sampler in giant/sample.py flips its model to .eval() as a side + effect with no restore of its own (see sample_flow). Sampling from the + RAW stage-1 model (ema_decay=0, so sampling_model() returns self.model, + the same weights the stage-1 trainer is actively training on) must not + silently leave it in eval mode for the rest of the epoch's stage-1 + updates.""" + cfg = _sampled_ctx_cfg(ema_decay=0.0) + trainers = _build_sampled_trainers(cfg) + stage1, stage2 = trainers["stage1"], trainers["stage2"] + stage1.train_mode() + assert stage1.model.training + batch = _fake_batches(1, 4)[0] + cond_cont, cond_cat, x1_s1 = batch.cond_cont, batch.cond_cat, batch.target_s1 + + stage2._stage1_context(x1_s1, cond_cont, cond_cat, epoch=0) + assert stage1.model.training + + +@pytest.mark.parametrize("stage2_generator", ["wgan", "flow"]) +def test_build_stage_trainers_sampled_step_runs(stage2_generator): + """Both trainer subclasses' call sites (FlowDDPMStageTrainer._compute, + WGANStageTrainer.step) must run end to end under 'sampled' and produce a + finite loss.""" + cfg = _sampled_ctx_cfg() + cfg["stage2_model"]["generator"] = stage2_generator + trainers = _build_sampled_trainers(cfg) + trainer = trainers["stage2"] + batch = _fake_batches(1, 4)[0] + stats = trainer.step(batch, torch.device("cpu"), global_step=1) + loss_key = "g_loss" if stage2_generator == "wgan" else "loss" + assert math.isfinite(stats[loss_key]) + + +def test_train_end_to_end_stage1_context_sampled(): + """Full train() run with stage1_context='sampled' must complete and + write a checkpoint + metrics.csv with finite losses throughout.""" + cfg = _sampled_ctx_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"])) for r in rows)