lars
87e37ebe14
CI / Format (ruff format) (push) Successful in 29s
CI / Lint (ruff check) (push) Successful in 30s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 38s
CI / Lint (ruff check) (pull_request) Successful in 31s
CI / Format (ruff format) (pull_request) Successful in 44s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 47s
CI / Tests (push) Successful in 4m32s
CI / Tests (pull_request) Successful in 4m25s
stage{1,2}_model.active = false already trains one stage alone, but the
checkpoint it writes holds only that stage, so giant rollout refuses it --
the "retrain stage 2 alone against a fixed, known-good stage 1" experiment
the 2026-08-03 species failure calls for wasn't runnable end to end.
Adds stage{1,2}_model.init_from (a checkpoint .pt to load this stage's
weights from before training) and .freeze (never update them), symmetric
across both stages. Both stages stay active = true, so both get built and
both land in the output checkpoint -- the frozen stage is merely
initialized from disk instead of from scratch.
Decisions made during planning:
- Soft freeze: forward/backward still run every batch (loss/grad_norm stay
meaningful, no autograd special-casing), only optimizer.step() (and, for
the frozen stage, lr_sched.step()/EMA update) is skipped -- weights are
byte-identical for the whole run. This is StageTrainer._step_optimizer,
shared by the non-adversarial path and both halves (generator + critic)
of the WGAN path, so a frozen WGAN stage's critic freezes too.
- validate_config requires init_from whenever freeze = true, unless the run
is a --resume (a resumed frozen stage's weights come from the resume
checkpoint instead) -- freezing a randomly-initialized model is almost
certainly a mistake.
- CLI flags on both `giant train` and `giant new-run`
(--stage{1,2}-init-from/--stage{1,2}-freeze), matching every other
per-stage model knob's existing treatment.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
1126 lines
44 KiB
Python
1126 lines
44 KiB
Python
"""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)
|