From 7f62141445e39345409eee1a5a59ec762ac00e1f Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Thu, 18 Jun 2026 17:52:14 +0200 Subject: [PATCH] Make sampler step count configurable for validation validate_marginals and collect_samples could already vary flow ODE steps for inference (giant predict --steps), but training-time marginal validation and DDIM evaluation were stuck at hardcoded defaults. Add a validate_steps config/CLI option and forward steps to sample_ddim consistently. Co-Authored-By: Claude Sonnet 4.6 --- giant/analysis.py | 19 +++++++++++++------ giant/cli.py | 8 ++++++++ giant/config.py | 1 + giant/pipeline.py | 1 + giant/train.py | 8 +++++++- giant/validate.py | 13 +++++++++++-- 6 files changed, 41 insertions(+), 9 deletions(-) diff --git a/giant/analysis.py b/giant/analysis.py index c2d2c10..5de1e68 100644 --- a/giant/analysis.py +++ b/giant/analysis.py @@ -39,8 +39,9 @@ Three tiers of checks, building on the aggregate marginal/KL check in generation artifact. `collect_samples` takes a `steps` argument (forwarded to the flow ODE -integrator) so a later sampler-step-count ablation can sweep it by calling this -function repeatedly without any new plumbing. +integrator or, in ddim mode, the DDIM substep count) so a later +sampler-step-count ablation can sweep it by calling this function repeatedly +without any new plumbing. """ from __future__ import annotations @@ -205,11 +206,17 @@ def collect_samples( bundle: ModelBundle, val_loader: DataLoader, n_batches: int | None = None, - steps: int = 10, + steps: int | None = None, ) -> SampleCollection: - """Run the sampler over `val_loader`, pairing generations with real targets + conditioning.""" + """Run the sampler over `val_loader`, pairing generations with real targets + conditioning. + + `steps` overrides the number of sampler steps (flow ODE steps or DDIM + substeps); `None` keeps each sampler's own default. Unused in "ddpm" + mode, which always runs the full schedule. + """ model = bundle.model device = bundle.device + steps_kw = {} if steps is None else {"steps": steps} cond_list, real_list, gen_list = [], [], [] for i, (cond_cont, cond_cat, x1) in enumerate(val_loader): @@ -219,11 +226,11 @@ def collect_samples( cond_cat = cond_cat.to(device) if bundle.mode == "flow": - gen = sample_flow(model, cond_cont, cond_cat, steps=steps) + gen = sample_flow(model, cond_cont, cond_cat, **steps_kw) elif bundle.mode == "ddpm": gen = sample_ddpm(model, cond_cont, cond_cat, bundle.schedule) else: - gen = sample_ddim(model, cond_cont, cond_cat, bundle.schedule) + gen = sample_ddim(model, cond_cont, cond_cat, bundle.schedule, **steps_kw) cond_full = torch.cat([cond_cont.cpu(), cond_cat.cpu().float()], dim=-1) cond_list.append(cond_full.numpy()) diff --git a/giant/cli.py b/giant/cli.py index a248e2e..389c3d0 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -78,6 +78,13 @@ def train( Optional[int], typer.Option(help="Run marginal+KL validation every N epochs (0 disables)"), ] = None, + validate_steps: Annotated[ + Optional[int], + typer.Option( + help="Flow matching ODE steps used during marginal validation " + "(ignored in ddpm mode, which always runs the full schedule)" + ), + ] = None, shuffle_buffer: Annotated[ int, typer.Option(help="Rows held in RAM per worker for shuffling") ] = 65536, @@ -105,6 +112,7 @@ def train( "num_workers": num_workers, "seed": seed, "validate_every": validate_every, + "validate_steps": validate_steps, }.items() if v is not None } diff --git a/giant/config.py b/giant/config.py index 0872044..9724f88 100644 --- a/giant/config.py +++ b/giant/config.py @@ -18,6 +18,7 @@ DEFAULT_CONFIG: dict = { "num_workers": 4, "seed": 0, "validate_every": 10, + "validate_steps": 10, }, "model": { "hidden_dim": 256, diff --git a/giant/pipeline.py b/giant/pipeline.py index 565ce00..895b244 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -152,4 +152,5 @@ def run_train_job( model_config=model_config, resume_path=resume, validate_every=t["validate_every"], + validate_steps=t["validate_steps"], ) diff --git a/giant/train.py b/giant/train.py index 2db15fa..f904fc5 100644 --- a/giant/train.py +++ b/giant/train.py @@ -70,6 +70,7 @@ def train( model_config: dict | None = None, resume_path: str | Path | None = None, validate_every: int = 0, + validate_steps: int = 10, ) -> None: out_dir = Path(out_dir) out_dir.mkdir(parents=True, exist_ok=True) @@ -170,7 +171,12 @@ def train( if validate_every > 0 and epoch % validate_every == 0: print(f"[epoch {epoch}] marginal validation:") validate_marginals( - model, val_loader, mode=mode, schedule=ddpm_schedule, device=device + model, + val_loader, + mode=mode, + schedule=ddpm_schedule, + device=device, + steps=validate_steps, ) ckpt: dict = { diff --git a/giant/validate.py b/giant/validate.py index 1814b91..afed080 100644 --- a/giant/validate.py +++ b/giant/validate.py @@ -6,6 +6,10 @@ from giant.constants import LOCAL_TARGET_NAMES from giant.sample import sample_flow, sample_ddpm, sample_ddim +def _kw(steps: int | None) -> dict[str, int]: + return {} if steps is None else {"steps": steps} + + def _histogram_kl( p_samples: np.ndarray, q_samples: np.ndarray, bins: int = 50, eps: float = 1e-8 ) -> float: @@ -32,12 +36,17 @@ def validate_marginals( device: torch.device | None = None, n_batches: int | None = None, kl_bins: int = 50, + steps: int | None = None, ) -> dict[str, np.ndarray]: """Compare per-dimension marginals of generated vs. real steps. Returns {"real": (N,9), "generated": (N,9), "kl_divergence": (9,)} in normalised space. `kl_divergence[j]` is KL(real || generated) for dimension j, estimated from a shared histogram over both samples. + + `steps` overrides the number of sampler steps (flow ODE steps or DDIM + substeps); `None` keeps each sampler's own default. Unused in "ddpm" + mode, which always runs the full schedule. """ if device is None: device = next(model.parameters()).device @@ -51,11 +60,11 @@ def validate_marginals( cond_cat = cond_cat.to(device) if mode == "flow": - gen = sample_flow(model, cond_cont, cond_cat) + gen = sample_flow(model, cond_cont, cond_cat, **_kw(steps)) elif mode == "ddpm": gen = sample_ddpm(model, cond_cont, cond_cat, schedule) else: - gen = sample_ddim(model, cond_cont, cond_cat, schedule) + gen = sample_ddim(model, cond_cont, cond_cat, schedule, **_kw(steps)) all_real.append(x1.numpy()) all_gen.append(gen.cpu().numpy())