lars/giant
1
0
Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity
Files
da7cde3ef9670d9fd9f25a3ec6463c0e8b0e1d45
giant/tests/test_train.py
T
lars 200c6d243b v0.3.0 step 7: AttentionHistory (KV-cached) + scheduled/never teacher forcing 
AttentionHistory (giant/model/network.py) adds causal self-attention over
the emitted-secondary prefix as the alternative to MarkovHistory, with a
parallel forward() for training and an init_cache()/step() KV-cache path
for sample.py's per-slot AR inference loop, wired into
Stage2Autoregressive via history="attention".

giant/train.py adds _stage2_tf_prob and _assemble_stage2_ar_inputs_scheduled,
mixing ground-truth history with a detached sample_secondaries_ar self-sample
per slot so teacher_forcing="scheduled"/"never" close the train/inference gap
teacher_forcing="always" always avoided; wired into both stage-2 AR trainers.

config.py's validate_config no longer rejects these two previously
unimplemented schema values.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-08-07 13:15:56 +02:00

736 lines
26 KiB
Python
Raw Blame History

"""Tests for giant/train.py."""
import copy
import csv
import math
import tempfile
from pathlib import Path
import pytest
import torch
from giant.constants import (
COND_DIM,
CONT_SLOT_DIM,
K_MAX,
PARTICLE_PHYS_DIM,
SEC_SLOT_DIM,
X_DIM,
)
from giant.model.network import build_critics, build_models
from giant.train import (
FlowDDPMStageTrainer,
WGANStageTrainer,
_ar_has_prev,
_assemble_stage2_ar_inputs,
_assemble_stage2_ar_target,
_assemble_stage2_real,
_build_stage_trainers,
_gumbel_tau,
_relax_onehot_type_slice,
_remaining_energy_fraction,
_shift_prev,
_stage2_tf_prob,
_stick_fraction,
_type_repr,
_wandb_run_config,
train,
)
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, docs/v0.3.0-design.md §6) ---------
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]]
# --- _stage2_tf_prob (docs/v0.3.0-design.md §3.3, 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, {"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 = {"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, {"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},
"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(
(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):
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={"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, decision 1), 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
trainer = WGANStageTrainer(
name="stage1",
model=models["stage1"],
critic=critics["stage1"],
is_stage2=False,
lambda_weight=1.0,
n_sec_lambda=0.1,
n_critic=1000, # never a generator step in this test
gp_weight=10.0,
lr=3e-4,
critic_lr=0.0,
ema_decay=0.0,
warmup_epochs=0,
epochs=1,
steps_per_epoch=4,
device=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():
with pytest.raises(NotImplementedError):
FlowDDPMStageTrainer(
name="stage2",
model=torch.nn.Linear(1, 1),
is_stage2=True,
generator="ddpm",
lambda_weight=1.0,
n_sec_lambda=0.1,
lambda_balance=0.0,
lambda_proc=0.0,
lambda_entropy=0.0,
gumbel_tau_start=1.0,
gumbel_tau_end=0.1,
lr=3e-4,
weight_decay=0.01,
ema_decay=0.0,
warmup_epochs=0,
epochs=1,
steps_per_epoch=1,
ddpm_n_steps=50,
device=torch.device("cpu"),
)
# --- 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():
"""§11.4 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 per the design doc's
instruction 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)
def test_wgan_physical_omits_grad_norm_slice_columns():
cfg = _base_cfg() # default stage2_model has no particle_type -> "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
Reference in New Issue View Git Blame Copy Permalink