"""Tests for giant/training/.""" import copy import csv import math import tempfile from pathlib import Path from unittest.mock import MagicMock, patch import numpy as np import pytest import torch from giant.config import ParticleTypeConfig from giant.constants import ( COND_DIM, CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_SLOT_DIM, X_DIM, ) from giant.checkpoint_io import load_for_inference from giant.data.dataset import StepBatch from giant.data.transforms import Normalizer 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, WGANStageTrainer, build_checkpoint, build_stage_trainers, init_stages_from_checkpoints, train, ) from giant.training.metrics import _wandb_run_config from giant.training.stage2_inputs import ( _ar_has_prev, _assemble_stage2_ar_inputs, _assemble_stage2_ar_target, _assemble_stage2_real, _gumbel_tau, _relax_onehot_type_slice, _remaining_energy_fraction, _shift_prev, _stage2_tf_prob, _stick_fraction, _stop_target_and_mask, _type_repr, ) PDG_VOCAB = 6 MAT_VOCAB = 3 def test_gumbel_tau_at_step_zero_is_start(): assert _gumbel_tau(0, 1000, 1.0, 0.1) == 1.0 def test_gumbel_tau_at_total_steps_is_end(): assert abs(_gumbel_tau(1000, 1000, 1.0, 0.1) - 0.1) < 1e-9 def test_gumbel_tau_interpolates_linearly_midway(): assert abs(_gumbel_tau(500, 1000, 1.0, 0.1) - 0.55) < 1e-9 def test_gumbel_tau_clamps_beyond_total_steps(): assert _gumbel_tau(5000, 1000, 1.0, 0.1) == _gumbel_tau(1000, 1000, 1.0, 0.1) def test_gumbel_tau_handles_zero_total_steps(): # total_steps=0 is guarded to 1 internally: step=0 gives zero progress # (still tau_start), any step>=1 immediately clamps to full progress. assert _gumbel_tau(0, 0, 1.0, 0.1) == 1.0 assert abs(_gumbel_tau(1, 0, 1.0, 0.1) - 0.1) < 1e-9 def test_wandb_run_config_includes_full_cfg_and_param_counts(): cfg = { "train": {"lr": 3e-4}, "conditioning": {"out_dim": 128}, "stage1_model": {"generator": "flow"}, "stage2_model": {"generator": "wgan"}, } wcfg = _wandb_run_config(cfg, model_config={"pdg_vocab": 3}, param_counts={"stage1": 100}) assert wcfg["train"] == {"lr": 3e-4} assert wcfg["stage1_model"] == {"generator": "flow"} assert wcfg["stage2_model"] == {"generator": "wgan"} assert wcfg["model_config"] == {"pdg_vocab": 3} assert wcfg["param_counts"] == {"stage1": 100} def test_wandb_run_config_handles_missing_model_config(): cfg = {"train": {}, "conditioning": {}, "stage1_model": {}, "stage2_model": {}} wcfg = _wandb_run_config(cfg, model_config=None, param_counts={}) assert wcfg["model_config"] == {} # --- AR helper functions (v0.3.0 step 5) --------- def test_stick_fraction_matches_sigmoid_of_logit(): sec_cont = torch.zeros(2, 3, SEC_SLOT_DIM) sec_cont[..., 0] = torch.tensor([[0.0, 2.0, -2.0], [1.0, -1.0, 0.0]]) frac = _stick_fraction(sec_cont) assert torch.allclose(frac, torch.sigmoid(sec_cont[..., 0])) def test_remaining_energy_fraction_hand_computed(): fraction = torch.tensor([[0.5, 0.5, 1.0]]) remaining = _remaining_energy_fraction(fraction) assert torch.allclose(remaining, torch.tensor([[1.0, 0.5, 0.25]])) def test_shift_prev_shifts_and_zero_pads_slot0(): x = torch.arange(2 * 4 * 3).reshape(2, 4, 3).float() shifted = _shift_prev(x) assert torch.all(shifted[:, 0] == 0) assert torch.equal(shifted[:, 1:], x[:, :-1]) def test_ar_has_prev_false_only_at_slot_zero(): has_prev = _ar_has_prev(5, torch.device("cpu")) assert has_prev.shape == (1, 5) 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) ----------- def test_stage2_tf_prob_always_is_constant_one(): assert _stage2_tf_prob("always", 1.0, 0.0, 0, 10) == 1.0 assert _stage2_tf_prob("always", 1.0, 0.0, 9, 10) == 1.0 def test_stage2_tf_prob_never_is_constant_zero(): assert _stage2_tf_prob("never", 1.0, 1.0, 0, 10) == 0.0 assert _stage2_tf_prob("never", 1.0, 1.0, 9, 10) == 0.0 def test_stage2_tf_prob_scheduled_interpolates_linearly(): assert _stage2_tf_prob("scheduled", 1.0, 0.0, 0, 11) == 1.0 assert abs(_stage2_tf_prob("scheduled", 1.0, 0.0, 5, 11) - 0.5) < 1e-9 assert _stage2_tf_prob("scheduled", 1.0, 0.0, 10, 11) == 0.0 def test_stage2_tf_prob_scheduled_clamps_beyond_total_epochs(): end = _stage2_tf_prob("scheduled", 1.0, 0.0, 10, 11) beyond = _stage2_tf_prob("scheduled", 1.0, 0.0, 50, 11) assert beyond == end def test_stage2_tf_prob_scheduled_handles_single_epoch(): # total_epochs=1 is guarded to a denominator of 1 internally (like # _gumbel_tau's total_steps=0 guard) — epoch=0 gives zero progress. assert _stage2_tf_prob("scheduled", 1.0, 0.0, 0, 1) == 1.0 @pytest.mark.parametrize("target", ["physical", "onehot", "embedding"]) def test_type_repr_shapes_and_values(target): B, K, emb_dim = 3, 4, 6 sec_cont = torch.randn(B, K, SEC_SLOT_DIM) sec_type_idx = torch.randint(0, emb_dim, (B, K)) cond_enc = torch.nn.Module() if target == "embedding": cond_enc.pdg_emb = torch.nn.Embedding(emb_dim, emb_dim) repr_ = _type_repr(sec_type_idx, sec_cont, ParticleTypeConfig(target=target), cond_enc, emb_dim) expected_width = PARTICLE_PHYS_DIM if target == "physical" else emb_dim assert repr_.shape == (B, K, expected_width) if target == "physical": assert torch.equal(repr_, sec_cont[..., CONT_SLOT_DIM : CONT_SLOT_DIM + PARTICLE_PHYS_DIM]) if target == "onehot": assert torch.all(repr_.sum(-1) == 1.0) @pytest.mark.parametrize( "target,generator", [ ("physical", "flow"), ("physical", "wgan"), ("onehot", "flow"), ("onehot", "wgan"), ("embedding", "flow"), ("embedding", "wgan"), ], ) def test_assemble_stage2_ar_target_matches_assemble_stage2_real_flattened(target, generator): """Regression test tying the refactor together: _assemble_stage2_real is now defined as _assemble_stage2_ar_target(...).flatten(1).""" B, emb_dim = 4, 6 sec_cont = torch.randn(B, K_MAX, SEC_SLOT_DIM) sec_type_idx = torch.randint(0, emb_dim, (B, K_MAX)) cond_enc = torch.nn.Module() if target == "embedding": cond_enc.pdg_emb = torch.nn.Embedding(emb_dim, emb_dim) particle_type_cfg = ParticleTypeConfig(target=target) flat = _assemble_stage2_real(sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim) unflat = _assemble_stage2_ar_target(sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim) assert torch.equal(unflat.flatten(1), flat) def test_assemble_stage2_ar_inputs_shapes_and_history_feat_width(): B, emb_dim = 3, 6 sec_cont = torch.randn(B, K_MAX, SEC_SLOT_DIM) sec_type_idx = torch.randint(0, emb_dim, (B, K_MAX)) cond_enc = torch.nn.Module() out = _assemble_stage2_ar_inputs(sec_cont, sec_type_idx, ParticleTypeConfig(target="physical"), cond_enc, emb_dim) assert out["history_feat"].shape == (B, K_MAX, CONT_SLOT_DIM + PARTICLE_PHYS_DIM) assert out["has_prev"].shape == (B, K_MAX) assert out["remaining_frac"].shape == (B, K_MAX) assert out["slot_idx"].shape == (B, K_MAX) assert torch.all(out["slot_idx"][:, 0] == 0.0) assert torch.all(out["slot_idx"][:, -1] == 1.0) def test_relax_onehot_type_slice_grad_probe_populates_both_norms(): B, k_max, cont_dim, type_dim = 4, K_MAX, CONT_SLOT_DIM, 6 x_flat = torch.randn(B, k_max * (cont_dim + type_dim), requires_grad=True) grad_probe: dict[str, float] = {} out = _relax_onehot_type_slice(x_flat, k_max, cont_dim, type_dim, tau=0.5, grad_probe=grad_probe) out.sum().backward() assert grad_probe["cont"] >= 0.0 assert grad_probe["type"] >= 0.0 def test_relax_onehot_type_slice_grad_probe_none_is_backward_compatible(): B, k_max, cont_dim, type_dim = 4, K_MAX, CONT_SLOT_DIM, 6 x_flat = torch.randn(B, k_max * (cont_dim + type_dim), requires_grad=True) out = _relax_onehot_type_slice(x_flat, k_max, cont_dim, type_dim, tau=0.5) out.sum().backward() assert x_flat.grad is not None # --- end-to-end train() integration tests ----------------------------------- PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1} MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1} def _base_cfg(): return { "conditioning": { "out_dim": 32, "share_stages": False, "particle": dict(PARTICLE_CFG), "material": dict(MATERIAL_CFG), }, "stage1_model": { "active": True, "generator": "flow", "hidden_dim": 24, "n_res_blocks": 2, "dropout": 0.0, "lambda": 1.0, "flow": {"time_dim": 16}, "ddpm": {"time_dim": 16, "n_steps": 50}, "wgan": { "noise_dim": 16, "n_critic": 2, "gp_weight": 10.0, "critic_lr": 0.0, }, "router": {"enabled": False}, }, "stage2_model": { "active": True, "decoder": "one_shot", "generator": "wgan", "hidden_dim": 24, "n_res_blocks": 2, "dropout": 0.0, "lambda": 1.0, "k_max": K_MAX, "context_dim": 16, "n_sec": {"mode": "head", "lambda": 0.1}, # Explicit, not relying on the fallback default (which is # "onehot", matching DEFAULT_CONFIG — see issues.md Issue 1): # the "physical"-labelled cases below (and this fixture's own # comment history) intend this as the base "physical" case, # with "*_onehot"/"*_embedding" cases opting in explicitly. "particle_type": {"target": "physical", "lambda": 1.0}, "flow": {"time_dim": 16}, "ddpm": {"time_dim": 16, "n_steps": 50}, "wgan": { "noise_dim": 16, "n_critic": 2, "gp_weight": 10.0, "critic_lr": 0.0, }, "router": {"enabled": False, "tie_to_stage1": False}, }, "train": { "epochs": 2, "batch_size": 8, "lr": 3e-4, "weight_decay": 0.01, "ema_decay": 0.999, "warmup_epochs": 0, "val_fraction": 0.1, "max_val_batches": 0, "num_workers": 0, "seed": 0, "validate_every": 0, "validate_steps": 2, "wandb": False, }, } def _fake_batches(n_batches, batch_size, seed=0): g = torch.Generator().manual_seed(seed) batches = [] for _ in range(n_batches): cond_cont = torch.randn(batch_size, COND_DIM, generator=g) cond_cat = torch.stack( [ torch.randint(0, PDG_VOCAB, (batch_size,), generator=g), torch.randint(0, MAT_VOCAB, (batch_size,), generator=g), ], dim=1, ) x1 = torch.randn(batch_size, X_DIM, generator=g) n_sec = torch.randint(0, K_MAX, (batch_size,), generator=g) sec_cont = torch.randn(batch_size, K_MAX, SEC_SLOT_DIM, generator=g) proc_idx = torch.zeros(batch_size, dtype=torch.long) sec_type_idx = torch.zeros(batch_size, K_MAX, dtype=torch.long) batches.append(StepBatch(cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx)) return batches def _model_config(cfg): return { "pdg_vocab": PDG_VOCAB, "mat_vocab": MAT_VOCAB, "conditioning": cfg["conditioning"], "stage1_model": cfg["stage1_model"], "stage2_model": cfg["stage2_model"], } def _run_train(cfg, out_dir, resume_path=None, normalizer_dict=None): model_config = _model_config(cfg) models = build_models(model_config) critics = build_critics(model_config) train_loader = _fake_batches(4, cfg["train"]["batch_size"]) val_loader = _fake_batches(2, cfg["train"]["batch_size"], seed=1) train( cfg=cfg, models=models, critics=critics, train_loader=train_loader, val_loader=val_loader, device=torch.device("cpu"), out_dir=out_dir, normalizer_dict=normalizer_dict or {"cond": {}, "target": {}, "sec_phys": {}}, pdg_map={"22": 0}, mat_map={"G4_AIR": 0}, proc_map=None, model_config=model_config, total_train_batches=4, resume_path=resume_path, ) @pytest.mark.parametrize( "label,mutate", [ ("both_flow", lambda cfg: None), ("both_wgan", lambda cfg: cfg["stage1_model"].__setitem__("generator", "wgan")), ( "mixed_stage1_flow_stage2_wgan", lambda cfg: None, # already the default ), ( "mixed_stage1_wgan_stage2_flow", lambda cfg: ( cfg["stage1_model"].__setitem__("generator", "wgan"), cfg["stage2_model"].__setitem__("generator", "flow"), ), ), ("stage1_only", lambda cfg: cfg["stage2_model"].__setitem__("active", False)), ("stage2_only", lambda cfg: cfg["stage1_model"].__setitem__("active", False)), ( "both_ddpm_stage1_flow_stage2", lambda cfg: ( cfg["stage1_model"].__setitem__("generator", "ddpm"), cfg["stage2_model"].__setitem__("generator", "flow"), ), ), ( "routed_stage1_energy_gumbel", lambda cfg: cfg["stage1_model"].__setitem__( "router", { "enabled": True, "type": "energy", "n_experts": 3, "temperature": 0.5, "learn_centers": True, "lambda_balance": 0.1, "lambda_entropy": 0.01, "gumbel": True, "gumbel_tau_start": 1.0, "gumbel_tau_end": 0.1, }, ), ), ( "stage2_onehot_target_wgan", lambda cfg: cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}), ), ( "stage2_onehot_target_flow", lambda cfg: ( cfg["stage2_model"].__setitem__("generator", "flow"), cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}), ), ), ( "stage2_embedding_target_wgan", lambda cfg: ( cfg["conditioning"]["particle"].__setitem__("type", "embedding"), cfg["conditioning"]["material"].__setitem__("type", "embedding"), cfg["stage2_model"].__setitem__("particle_type", {"target": "embedding", "lambda": 1.0}), ), ), ( "stage2_embedding_target_flow", lambda cfg: ( cfg["conditioning"]["particle"].__setitem__("type", "embedding"), cfg["conditioning"]["material"].__setitem__("type", "embedding"), cfg["stage2_model"].__setitem__("generator", "flow"), cfg["stage2_model"].__setitem__("particle_type", {"target": "embedding", "lambda": 1.0}), ), ), ( "ar_wgan_onehot", lambda cfg: ( cfg["stage2_model"].__setitem__("decoder", "autoregressive"), cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}), ), ), ( "ar_wgan_physical", lambda cfg: cfg["stage2_model"].__setitem__("decoder", "autoregressive"), ), ( "ar_flow_onehot", lambda cfg: ( cfg["stage2_model"].__setitem__("decoder", "autoregressive"), cfg["stage2_model"].__setitem__("generator", "flow"), cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}), ), ), ( "ar_flow_embedding", lambda cfg: ( cfg["conditioning"]["particle"].__setitem__("type", "embedding"), cfg["conditioning"]["material"].__setitem__("type", "embedding"), cfg["stage2_model"].__setitem__("decoder", "autoregressive"), cfg["stage2_model"].__setitem__("generator", "flow"), cfg["stage2_model"].__setitem__("particle_type", {"target": "embedding", "lambda": 1.0}), ), ), ( "ar_stage2_only", lambda cfg: ( cfg["stage1_model"].__setitem__("active", False), cfg["stage2_model"].__setitem__("decoder", "autoregressive"), ), ), ( "ar_mixed_stage1_wgan_stage2_flow_onehot", lambda cfg: ( cfg["stage1_model"].__setitem__("generator", "wgan"), cfg["stage2_model"].__setitem__("generator", "flow"), cfg["stage2_model"].__setitem__("decoder", "autoregressive"), cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}), ), ), ], ) def test_train_end_to_end(label, mutate): cfg = _base_cfg() mutate(cfg) with tempfile.TemporaryDirectory() as tmp: out_dir = Path(tmp) / "run" _run_train(cfg, out_dir) assert (out_dir / "last.pt").exists() assert (out_dir / "metrics.csv").exists() ckpt = torch.load(out_dir / "last.pt", weights_only=False) if cfg["stage1_model"]["active"]: assert "model" in ckpt else: assert "model" not in ckpt if cfg["stage2_model"]["active"]: assert "sec_decoder" in ckpt else: assert "sec_decoder" not in ckpt def test_train_resume_continues_from_checkpoint(): cfg = _base_cfg() with tempfile.TemporaryDirectory() as tmp: out_dir = Path(tmp) / "run" _run_train(cfg, out_dir) ckpt_before = torch.load(out_dir / "last.pt", weights_only=False) assert ckpt_before["epoch"] == 2 cfg2 = copy.deepcopy(cfg) cfg2["train"]["epochs"] = 3 _run_train(cfg2, out_dir, resume_path=out_dir / "last.pt") ckpt_after = torch.load(out_dir / "last.pt", weights_only=False) assert ckpt_after["epoch"] == 3 assert ckpt_after["global_step"] > ckpt_before["global_step"] def test_train_raises_when_no_active_stage(): cfg = _base_cfg() cfg["stage1_model"]["active"] = False cfg["stage2_model"]["active"] = False model_config = _model_config(cfg) models = build_models(model_config) critics = build_critics(model_config) with tempfile.TemporaryDirectory() as tmp: with pytest.raises(ValueError, match="no active stage"): train( cfg=cfg, models=models, critics=critics, train_loader=_fake_batches(1, 8), val_loader=_fake_batches(1, 8), device=torch.device("cpu"), out_dir=Path(tmp) / "run", model_config=model_config, total_train_batches=1, ) def test_metrics_csv_columns_are_stage_prefixed(): cfg = _base_cfg() with tempfile.TemporaryDirectory() as tmp: out_dir = Path(tmp) / "run" _run_train(cfg, out_dir) header = (out_dir / "metrics.csv").read_text().splitlines()[0].split(",") assert "stage1/train/loss" in header assert "stage2/train/d_loss" in header assert "val/loss" in header assert "epoch" in header def test_wgan_stage_trainer_skips_generator_step_when_no_grad_this_batch(): """Regression test: on a non-generator-step batch, if this stage's model has no n_sec_head (n_sec defaults to stage 2), g_loss is a graph-less zero — .backward() must not be called on it.""" cfg = _base_cfg() cfg["stage1_model"]["generator"] = "wgan" model_config = _model_config(cfg) models = build_models(model_config) critics = build_critics(model_config) assert models["stage1"] is not None and critics["stage1"] is not None spec = StageSpec( name="stage1", is_stage2=False, generator="wgan", n_critic=1000, # never a generator step in this test ema_decay=0.0, steps_per_epoch=4, ) trainer = WGANStageTrainer(spec, models["stage1"], critics["stage1"], torch.device("cpu")) assert trainer.model.n_sec_head is None batch = _fake_batches(1, 8)[0] stats = trainer.step(batch, torch.device("cpu"), global_step=1) assert stats["did_g_step"] is False def test_flow_stage_trainer_ddpm_not_implemented_for_stage2(): spec = StageSpec(name="stage2", is_stage2=True, generator="ddpm", ddpm_n_steps=50, ema_decay=0.0) with pytest.raises(NotImplementedError): FlowDDPMStageTrainer(spec, torch.nn.Linear(1, 1), torch.device("cpu")) # --- gitea #42: freeze / init_from ------------------------------------------- def _state_dict_clone(module): return {k: v.clone() for k, v in module.state_dict().items()} def _assert_state_dicts_equal(before, after, label): for key, value in before.items(): assert torch.equal(value, after[key]), f"{label}: {key} changed while frozen" def test_frozen_flow_stage_trainer_step_does_not_update_model_or_ema(): cfg = _base_cfg() model_config = _model_config(cfg) model = build_models(model_config)["stage1"] assert model is not None spec = StageSpec(name="stage1", is_stage2=False, generator="flow", freeze=True, ema_decay=0.999, steps_per_epoch=4) trainer = FlowDDPMStageTrainer(spec, model, torch.device("cpu")) assert trainer.ema_model is not None model_before = _state_dict_clone(trainer.model) ema_before = _state_dict_clone(trainer.ema_model) for batch in _fake_batches(4, 8): trainer.step(batch, torch.device("cpu"), global_step=1) _assert_state_dicts_equal(model_before, trainer.model.state_dict(), "frozen flow model") _assert_state_dicts_equal(ema_before, trainer.ema_model.state_dict(), "frozen flow ema") def test_frozen_wgan_stage_trainer_step_does_not_update_generator_or_critic(): cfg = _base_cfg() cfg["stage1_model"]["generator"] = "wgan" model_config = _model_config(cfg) models = build_models(model_config) critics = build_critics(model_config) assert models["stage1"] is not None and critics["stage1"] is not None spec = StageSpec( name="stage1", is_stage2=False, generator="wgan", freeze=True, n_critic=1, # a generator step every batch, so a bug would surface immediately ema_decay=0.999, steps_per_epoch=4, ) trainer = WGANStageTrainer(spec, models["stage1"], critics["stage1"], torch.device("cpu")) assert trainer.ema_model is not None model_before = _state_dict_clone(trainer.model) critic_before = _state_dict_clone(trainer.critic) ema_before = _state_dict_clone(trainer.ema_model) for global_step, batch in enumerate(_fake_batches(4, 8)): trainer.step(batch, torch.device("cpu"), global_step=global_step) _assert_state_dicts_equal(model_before, trainer.model.state_dict(), "frozen wgan generator") _assert_state_dicts_equal(critic_before, trainer.critic.state_dict(), "frozen wgan critic") _assert_state_dicts_equal(ema_before, trainer.ema_model.state_dict(), "frozen wgan ema") @pytest.mark.parametrize("stage1_generator", ["flow", "wgan"]) def test_train_end_to_end_frozen_stage1_unchanged_while_stage2_trains(stage1_generator): cfg = _base_cfg() cfg["stage1_model"]["generator"] = stage1_generator cfg["stage1_model"]["freeze"] = True model_config = _model_config(cfg) models = build_models(model_config) critics = build_critics(model_config) assert models["stage1"] is not None and models["stage2"] is not None stage1_before = _state_dict_clone(models["stage1"]) stage2_before = _state_dict_clone(models["stage2"]) with tempfile.TemporaryDirectory() as tmp: train( cfg=cfg, models=models, critics=critics, train_loader=_fake_batches(4, cfg["train"]["batch_size"]), val_loader=_fake_batches(2, cfg["train"]["batch_size"], seed=1), device=torch.device("cpu"), out_dir=Path(tmp) / "run", normalizer_dict={"cond": {}, "target": {}, "sec_phys": {}}, pdg_map={"22": 0}, mat_map={"G4_AIR": 0}, proc_map=None, model_config=model_config, total_train_batches=4, ) _assert_state_dicts_equal(stage1_before, models["stage1"].state_dict(), "frozen stage1") stage2_after = models["stage2"].state_dict() assert any(not torch.equal(v, stage2_after[k]) for k, v in stage2_before.items()), ( "unfrozen stage2 should have trained" ) def test_init_stages_from_checkpoints_loads_matching_stage_and_ema_weights(tmp_path): cfg = _base_cfg() model_config = _model_config(cfg) source_models = build_models(model_config) source_critics = build_critics(model_config) source_trainers = build_stage_trainers(cfg, source_models, source_critics, torch.device("cpu"), 4) source_stage1_ema = source_trainers["stage1"].ema_model assert source_stage1_ema is not None # Diverge the source's EMA from its raw weights so a same-vs-different # check below actually distinguishes the two copy paths. for p in source_stage1_ema.parameters(): p.data.add_(1.0) ckpt_path = tmp_path / "source.pt" ckpt = build_checkpoint(source_trainers, epoch=1, global_step=1, best_val_loss=0.0, extras={}) torch.save(ckpt, ckpt_path) cfg2 = copy.deepcopy(cfg) cfg2["stage1_model"]["init_from"] = str(ckpt_path) dest_models = build_models(_model_config(cfg2)) dest_critics = build_critics(_model_config(cfg2)) dest_trainers = build_stage_trainers(cfg2, dest_models, dest_critics, torch.device("cpu"), 4) loaded = init_stages_from_checkpoints(dest_trainers) assert len(loaded) == 1 and "stage1" in loaded[0] dest_stage1_ema = dest_trainers["stage1"].ema_model assert dest_stage1_ema is not None _assert_state_dicts_equal( source_trainers["stage1"].model.state_dict(), dest_trainers["stage1"].model.state_dict(), "init_from raw" ) _assert_state_dicts_equal( source_stage1_ema.state_dict(), dest_stage1_ema.state_dict(), "init_from ema", ) # stage2 has no init_from set -- untouched fresh init, not the source's. stage2_matches_source = all( torch.equal(v, dest_trainers["stage2"].model.state_dict()[k]) for k, v in source_trainers["stage2"].model.state_dict().items() ) assert not stage2_matches_source def test_run_train_job_stage1_init_from_freeze_produces_rollout_capable_checkpoint(tmp_path): """The exact scenario gitea #42 exists for: retrain stage 2 alone against a fixed, known-good stage 1, and still get a checkpoint giant rollout can load (checkpoint_io.load_for_inference with require_stage2=True).""" normalizer_dict = { "cond": Normalizer().fit(np.zeros((1, COND_DIM), dtype=np.float32)).to_dict(), "target": Normalizer().fit(np.zeros((1, X_DIM), dtype=np.float32)).to_dict(), "sec_phys": Normalizer().fit(np.zeros((1, 2), dtype=np.float32)).to_dict(), } cfg = _base_cfg() source_out = tmp_path / "source" _run_train(cfg, source_out, normalizer_dict=normalizer_dict) source_ckpt = torch.load(source_out / "best.pt", weights_only=False) cfg2 = copy.deepcopy(cfg) cfg2["stage1_model"]["init_from"] = str(source_out / "best.pt") cfg2["stage1_model"]["freeze"] = True retrain_out = tmp_path / "retrain" _run_train(cfg2, retrain_out, normalizer_dict=normalizer_dict) ctx = load_for_inference(retrain_out / "best.pt", torch.device("cpu"), "rollout", require_stage2=True) assert ctx.stage1 is not None and ctx.stage2 is not None retrain_ckpt = torch.load(retrain_out / "best.pt", weights_only=False) for key, value in source_ckpt["model"].items(): assert torch.equal(value, retrain_ckpt["model"][key]), f"frozen stage1 {key} drifted across the retrain" def test_stage_spec_from_config_omitted_decoder_and_particle_type_match_default_config(): """Regression for issues.md Issue 1: StageSpec.from_config's own fallback defaults for stage2_model.decoder/particle_type must equal DEFAULT_CONFIG's ("autoregressive" / "onehot"), not the old, now-wrong ("one_shot" / "physical") literals a .get(key, default) call used to supply when a hand-built cfg omitted these keys.""" cfg = _base_cfg() del cfg["stage2_model"]["decoder"] del cfg["stage2_model"]["particle_type"] spec = StageSpec.from_config(cfg, "stage2", is_stage2=True, steps_per_epoch=1) assert spec.decoder == "autoregressive" assert spec.particle_type.target == "onehot" def _routed_stage1_trainer(lambda_balance, lambda_proc, lambda_entropy): cfg = _base_cfg() cfg["stage1_model"]["router"] = { "enabled": True, "type": "energy", "n_experts": 3, "temperature": 0.5, "learn_centers": True, "lambda_balance": lambda_balance, "lambda_proc": lambda_proc, "lambda_entropy": lambda_entropy, } 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) return trainers["stage1"] def test_router_aux_losses_skipped_when_lambda_zero_but_run_when_positive(): """Gitea #31: FlowDDPMStageTrainer._compute must not call router.balance_loss/classify_loss/entropy_loss when the corresponding lambda is 0 (the default) -- those calls do their own router.gate(...) forward pass that is wasted once the term is masked out of the total loss anyway. Checked both ways: zero lambdas must skip all three calls, positive lambdas must still make them (the guard must not accidentally suppress the real path).""" batch = _fake_batches(1, 4)[0] device = torch.device("cpu") trainer_zero = _routed_stage1_trainer(0.0, 0.0, 0.0) router_zero = trainer_zero.router router_zero.balance_loss = MagicMock(wraps=router_zero.balance_loss) router_zero.classify_loss = MagicMock(wraps=router_zero.classify_loss) router_zero.entropy_loss = MagicMock(wraps=router_zero.entropy_loss) stats_zero = trainer_zero.step(batch, device, global_step=1) assert router_zero.balance_loss.call_count == 0 assert router_zero.classify_loss.call_count == 0 assert router_zero.entropy_loss.call_count == 0 assert stats_zero["loss_balance"] == 0.0 assert stats_zero["loss_proc"] == 0.0 assert stats_zero["loss_entropy"] == 0.0 trainer_pos = _routed_stage1_trainer(0.1, 0.1, 0.01) router_pos = trainer_pos.router router_pos.balance_loss = MagicMock(wraps=router_pos.balance_loss) router_pos.classify_loss = MagicMock(wraps=router_pos.classify_loss) router_pos.entropy_loss = MagicMock(wraps=router_pos.entropy_loss) trainer_pos.step(batch, device, global_step=1) assert router_pos.balance_loss.call_count == 1 assert router_pos.classify_loss.call_count == 1 assert router_pos.entropy_loss.call_count == 1 # --- AR trainer wiring (v0.3.0 step 5) -------------------------------------- @pytest.mark.parametrize("teacher_forcing", ["always", "scheduled", "never"]) @pytest.mark.parametrize("history", ["markov", "attention"]) @pytest.mark.parametrize("stage2_generator", ["wgan", "flow"]) def test_build_stage_trainers_ar_scheduled_and_attention_step_runs(teacher_forcing, history, stage2_generator): """v0.3.0 step 7: history='attention' and teacher_forcing in {'scheduled', 'never'} must actually train — a stage-2 AR trainer.step() must run and produce a finite loss, for every {history} x {teacher_forcing} x {generator} combination.""" cfg = _base_cfg() cfg["stage2_model"]["decoder"] = "autoregressive" cfg["stage2_model"]["generator"] = stage2_generator cfg["stage2_model"]["autoregressive"] = { "history": history, "teacher_forcing": teacher_forcing, "tf_p_start": 1.0, "tf_p_end": 0.0, "attn_n_heads": 2, "attn_n_layers": 1, } 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) loss_key = "g_loss" if stage2_generator == "wgan" else "loss" assert math.isfinite(stats[loss_key]) @pytest.mark.parametrize("stage2_generator", ["wgan", "flow"]) def test_train_end_to_end_ar_attention_history_scheduled_teacher_forcing( stage2_generator, ): """Full `train()` run (not just one `trainer.step()` call) with history='attention' AND teacher_forcing='scheduled' together — the combination v0.3.0 step 7 exists to land — must complete and write a checkpoint + metrics.csv with finite losses throughout.""" cfg = _base_cfg() cfg["stage2_model"]["decoder"] = "autoregressive" cfg["stage2_model"]["generator"] = stage2_generator cfg["stage2_model"]["autoregressive"] = { "history": "attention", "teacher_forcing": "scheduled", "tf_p_start": 1.0, "tf_p_end": 0.0, "attn_n_heads": 2, "attn_n_layers": 1, } 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"] loss_col = "stage2/train/g_loss" if stage2_generator == "wgan" else "stage2/train/loss" assert all(math.isfinite(float(r[loss_col])) for r in rows) def test_ar_wgan_onehot_grad_norm_instrumentation_populates_metrics(): """Differentiability validation-obligation instrumentation: the trunk-gradient-norm-by-slice columns must appear and actually fire for generator='wgan' + particle_type.target='onehot' under decoder= 'autoregressive' (added at v0.3.0 step 5 to accrue evidence during the architecture comparison).""" cfg = _base_cfg() cfg["stage2_model"]["decoder"] = "autoregressive" cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0} with tempfile.TemporaryDirectory() as tmp: out_dir = Path(tmp) / "run" _run_train(cfg, out_dir) with open(out_dir / "metrics.csv", newline="") as f: rows = list(csv.DictReader(f)) assert "stage2/train/grad_norm_type_slice" in rows[0] assert "stage2/train/grad_norm_cont_slice" in rows[0] assert any(float(r["stage2/train/grad_norm_type_slice"]) > 0 for r in rows) assert any(float(r["stage2/train/grad_norm_cont_slice"]) > 0 for r in rows) def test_wgan_onehot_one_shot_also_gets_grad_norm_instrumentation(): """The instrumentation is decoder-agnostic — one_shot + wgan + onehot must populate the same columns.""" cfg = _base_cfg() cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0} with tempfile.TemporaryDirectory() as tmp: out_dir = Path(tmp) / "run" _run_train(cfg, out_dir) with open(out_dir / "metrics.csv", newline="") as f: rows = list(csv.DictReader(f)) assert any(float(r["stage2/train/grad_norm_type_slice"]) > 0 for r in rows) 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: out_dir = Path(tmp) / "run" _run_train(cfg, out_dir) 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)