From bac541240f580b3cf92c5f49b3f924d7d0596b8d Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Wed, 8 Jul 2026 15:17:19 +0200 Subject: [PATCH 01/11] Add mixture-of-experts routing prototype for Stage 1 and Stage 2 Both stages can now route through a pluggable Router (EnergyRouter as the first implementation, a soft turn-on gate over pre-step log-energy) into several small ExpertTrunks instead of one monolithic trunk. Trains as a differentiable soft mixture and dispatches to a single expert per row at eval time, which is the source of the per-call speedup this prototype is after (issue #5's ~10x native-Geant4 budget). Disabled by default, so existing configs/checkpoints are unaffected; build_models() centralizes routed-vs-monolith construction across train/predict/rollout. Co-Authored-By: Claude Sonnet 5 --- giant/cli.py | 87 ++++++---- giant/config.py | 55 ++++++- giant/model/network.py | 366 ++++++++++++++++++++++++++++++++++++++++- giant/pipeline.py | 47 ++---- giant/train.py | 38 ++++- tests/test_router.py | 296 +++++++++++++++++++++++++++++++++ 6 files changed, 815 insertions(+), 74 deletions(-) create mode 100644 tests/test_router.py diff --git a/giant/cli.py b/giant/cli.py index 4b61b10..42f8ad0 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -39,31 +39,30 @@ from giant.data.transforms import ( Normalizer, ) from giant.geometry import GeometryOracle -from giant.model.network import DenoisingMLP, SecondaryDecoder +from giant.model.network import build_models from giant.pipeline import run_train_job from giant.rollout import rollout as run_rollout from giant.sample import sample_flow, sample_secondaries, snap_type_to_pdg_idx -_STAGE1_MODEL_KEYS = { - "pdg_vocab", - "mat_vocab", - "hidden_dim", - "n_blocks", - "emb_dim", - "dropout", - "k_max", -} -_SEC_DECODER_MODEL_KEYS = { - "pdg_vocab", - "mat_vocab", - "hidden_dim", - "n_blocks", - "emb_dim", - "dropout", -} - app = typer.Typer(no_args_is_help=True) + +def _batch_size_estimate_dims(model_cfg: dict) -> tuple[int, int]: + """Pick the (hidden_dim, n_blocks) that dominate per-call activation memory. + + Routed models spend their FLOPs in the (smaller) expert trunks, not the + monolith's hidden_dim/n_blocks, so estimate_batch_size needs the expert + dims instead when routing is enabled. + """ + router_cfg = model_cfg.get("router") + if router_cfg and router_cfg.get("enabled"): + return ( + model_cfg.get("expert_hidden_dim", 128), + model_cfg.get("expert_n_blocks", 3), + ) + return model_cfg["hidden_dim"], model_cfg["n_blocks"] + + _CEPH_PREDICTIONS = Path("/ceph/lbogner/geant_steps/predictions") @@ -159,6 +158,23 @@ def train( "--dropout", "-d", help="Dropout probability in ResBlocks (default: 0.1)" ), ] = None, + router: Annotated[ + Optional[bool], + typer.Option( + "--router/--no-router", + help="Route both stages through a mixture of small experts " + "instead of one monolithic trunk (see model.router in config.toml)", + ), + ] = None, + router_type: Annotated[ + Optional[str], + typer.Option( + "--router-type", help="Router implementation name (see ROUTER_REGISTRY)" + ), + ] = None, + n_experts: Annotated[ + Optional[int], typer.Option("--n-experts", help="Number of routed experts") + ] = None, val_fraction: Annotated[ Optional[float], typer.Option("--val-fraction", "-f") ] = None, @@ -238,7 +254,7 @@ def train( }.items() if v is not None } - cli_model = { + cli_model: dict[str, object] = { k: v for k, v in { "hidden_dim": hidden_dim, @@ -248,6 +264,17 @@ def train( }.items() if v is not None } + cli_router = { + k: v + for k, v in { + "enabled": router, + "type": router_type, + "n_experts": n_experts, + }.items() + if v is not None + } + if cli_router: + cli_model["router"] = cli_router cfg = gconfig.merge_cli_overrides( gconfig.DEFAULT_CONFIG, config, cli_train, cli_model ) @@ -256,9 +283,10 @@ def train( _device = torch.device(device) if device else gconfig.auto_device() if batch_size_auto: + est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(m) try: t["batch_size"] = gconfig.estimate_batch_size( - m["hidden_dim"], m["n_blocks"], _device + est_hidden_dim, est_n_blocks, _device ) except ValueError as exc: typer.echo(f"error: {exc}", err=True) @@ -386,10 +414,11 @@ def predict( model_cfg = ckpt["model_config"] if batch_size_auto: + est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(model_cfg) try: batch_size_value = gconfig.estimate_batch_size( - model_cfg["hidden_dim"], - model_cfg["n_blocks"], + est_hidden_dim, + est_n_blocks, _device, training=False, ) @@ -408,15 +437,10 @@ def predict( cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"]) tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"]) - model = DenoisingMLP( - **{k: v for k, v in model_cfg.items() if k in _STAGE1_MODEL_KEYS} - ) + model, sec_decoder = build_models(model_cfg) model.load_state_dict(ckpt["model"]) model.to(_device).eval() - sec_decoder = SecondaryDecoder( - **{k: v for k, v in model_cfg.items() if k in _SEC_DECODER_MODEL_KEYS} - ) sec_decoder.load_state_dict(ckpt["sec_decoder"]) sec_decoder.to(_device).eval() @@ -759,12 +783,9 @@ def rollout( cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"]) tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"]) - model = DenoisingMLP(**{k: v for k, v in model_cfg.items() if k in _STAGE1_MODEL_KEYS}) + model, sec_decoder = build_models(model_cfg) model.load_state_dict(ckpt["model"]) model.to(_device).eval() - sec_decoder = SecondaryDecoder( - **{k: v for k, v in model_cfg.items() if k in _SEC_DECODER_MODEL_KEYS} - ) sec_decoder.load_state_dict(ckpt["sec_decoder"]) sec_decoder.to(_device).eval() typer.echo(f"loaded checkpoint: {checkpoint}") diff --git a/giant/config.py b/giant/config.py index 7001151..b5008a6 100644 --- a/giant/config.py +++ b/giant/config.py @@ -28,6 +28,16 @@ DEFAULT_CONFIG: dict = { "n_blocks": 6, "emb_dim": 16, "dropout": 0.1, + "router": { + "enabled": False, + "type": "energy", # selects the Router impl from ROUTER_REGISTRY + "n_experts": 4, + "expert_hidden_dim": 128, + "expert_n_blocks": 3, + "temperature": 0.5, # energy-router kwarg + "learn_centers": True, # energy-router kwarg + "lambda_balance": 0.0, # optional load-balance aux loss weight + }, }, } @@ -164,15 +174,29 @@ def merge_cli_overrides( train_overrides: dict, model_overrides: dict, ) -> dict: - """Resolve config as defaults -> TOML file -> explicit CLI flags.""" + """Resolve config as defaults -> TOML file -> explicit CLI flags. + + `model.router` is deep-merged one level (rather than replaced wholesale) + at each stage, so a TOML file or CLI flag only overriding e.g. + `router.enabled` doesn't drop the rest of the router defaults. + """ cfg = {"train": dict(defaults["train"]), "model": dict(defaults["model"])} + cfg["model"]["router"] = dict(defaults["model"]["router"]) if config_path is not None: file_cfg = load_toml(config_path) - for section in ("train", "model"): - cfg[section].update(file_cfg.get(section, {})) + cfg["train"].update(file_cfg.get("train", {})) + file_model = dict(file_cfg.get("model", {})) + file_router = file_model.pop("router", None) + cfg["model"].update(file_model) + if file_router: + cfg["model"]["router"].update(file_router) warn_if_git_hash_mismatch(file_cfg, config_path) + model_overrides = dict(model_overrides) + router_overrides = model_overrides.pop("router", None) cfg["train"].update(train_overrides) cfg["model"].update(model_overrides) + if router_overrides: + cfg["model"]["router"].update(router_overrides) return cfg @@ -184,17 +208,38 @@ def seed_everything(seed: int) -> None: torch.cuda.manual_seed_all(seed) +def _toml_value(v) -> str: + if isinstance(v, bool): + return "true" if v else "false" + if isinstance(v, str): + return repr(v) + return str(v) + + def save_config(cfg: dict, out_dir: Path, meta: dict) -> None: lines = [] + # One-level-nested dict values (e.g. model.router) are rendered as their + # own [section.subsection] table after the parent section, since TOML + # doesn't accept a bare dict as a `key = value` scalar line. + nested_sections: list[tuple[str, dict]] = [] for section, values in cfg.items(): lines.append(f"[{section}]") for k, v in values.items(): - lines.append(f"{k:<14} = {repr(v) if isinstance(v, str) else v}") + if isinstance(v, dict): + nested_sections.append((f"{section}.{k}", v)) + continue + lines.append(f"{k:<14} = {_toml_value(v)}") + lines.append("") + + for name, values in nested_sections: + lines.append(f"[{name}]") + for k, v in values.items(): + lines.append(f"{k:<14} = {_toml_value(v)}") lines.append("") lines.append("[meta]") for k, v in meta.items(): - lines.append(f"{k:<14} = {repr(v) if isinstance(v, str) else v}") + lines.append(f"{k:<14} = {_toml_value(v)}") (out_dir / "config.toml").write_text("\n".join(lines)) diff --git a/giant/model/network.py b/giant/model/network.py index b64b48e..91c06a1 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -1,9 +1,10 @@ +import inspect import math import torch import torch.nn as nn -from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM +from giant.constants import COND_DIM, EMB_DIM, K_MAX, SEC_DIM, X_DIM class SinusoidalEmbedding(nn.Module): @@ -236,3 +237,366 @@ class SecondaryDecoder(nn.Module): for block in self.blocks: x = block(x, cond) return self.out_proj(x) + + +class Router(nn.Module): + """Contract for a pluggable mixture-of-experts routing axis. + + Subclasses implement `gate` (soft partition-of-unity weights over + experts, used in train mode for a fully differentiable mixture); + `top1` and `balance_loss` have working defaults so a new routing axis + is usually a one-method add. See `ROUTER_REGISTRY` / `build_router`. + """ + + def __init__(self, n_experts: int) -> None: + super().__init__() + self.n_experts = n_experts + + def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + """(B, n_experts) soft weights, rows summing to 1.""" + raise NotImplementedError + + def top1(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + """(B,) hard expert index, used for eval-time grouped dispatch.""" + return self.gate(cond_cont, cond_cat).argmax(dim=-1) + + def balance_loss( + self, cond_cont: torch.Tensor, cond_cat: torch.Tensor + ) -> torch.Tensor: + """Importance CV^2 load-balancing auxiliary loss (Shazeer et al. 2017).""" + importance = self.gate(cond_cont, cond_cat).sum(dim=0) # (n_experts,) + return (importance.std() / (importance.mean() + 1e-8)) ** 2 + + +ROUTER_REGISTRY: dict[str, type[Router]] = {} + + +def register_router(name: str): + def decorator(cls: type[Router]) -> type[Router]: + ROUTER_REGISTRY[name] = cls + return cls + + return decorator + + +def build_router(name: str, n_experts: int, **kwargs) -> Router: + """Factory: look up a `Router` subclass by name from the registry. + + Every registered router type is fed the same `model.router` config + dict; kwargs not declared by that type's constructor are silently + dropped, so per-type hyperparameters (e.g. EnergyRouter's + `temperature`) can coexist in one config without special-casing. + """ + if name not in ROUTER_REGISTRY: + raise ValueError( + f"unknown router type {name!r}; available: {sorted(ROUTER_REGISTRY)}" + ) + cls = ROUTER_REGISTRY[name] + accepted = set(inspect.signature(cls.__init__).parameters) - {"self", "n_experts"} + filtered = {k: v for k, v in kwargs.items() if k in accepted} + return cls(n_experts=n_experts, **filtered) + + +@register_router("energy") +class EnergyRouter(Router): + """Soft turn-on gate over normalized pre-step log-energy. + + Reads `cond_cont[:, energy_idx]` (ignores cond_cat). Learnable (or + fixed) 1-D centers, initialized spread across [-2, 2] — roughly the + z-normalized energy range. `gate(e) = softmax_i(-(e - c_i)^2 / tau)`, + differentiable in e; as tau -> 0 this hardens to nearest-center + (Voronoi) selection, which is exactly what `top1` uses at eval. + """ + + def __init__( + self, + n_experts: int = 4, + temperature: float = 0.5, + learn_centers: bool = True, + energy_idx: int = 3, + ) -> None: + super().__init__(n_experts) + self.temperature = temperature + self.energy_idx = energy_idx + centers = torch.linspace(-2.0, 2.0, n_experts) + if learn_centers: + self.centers = nn.Parameter(centers) + else: + self.register_buffer("centers", centers) + + def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + e = cond_cont[:, self.energy_idx].unsqueeze(-1) # (B, 1) + d2 = (e - self.centers.unsqueeze(0)) ** 2 # (B, n_experts) + return torch.softmax(-d2 / self.temperature, dim=-1) + + +class ExpertTrunk(nn.Module): + """One small expert: `input_proj -> ResBlock stack -> out_proj`. + + Same shape as the monolithic DenoisingMLP/SecondaryDecoder trunk, but + intended to be narrower/shallower (per-call cost is the whole point). + """ + + def __init__( + self, + in_dim: int, + hidden_dim: int, + n_blocks: int, + merged_cond_dim: int, + dropout: float = 0.1, + ) -> None: + super().__init__() + self.input_proj = nn.Linear(in_dim, hidden_dim) + self.blocks = nn.ModuleList( + [ + ResBlock(hidden_dim, merged_cond_dim, dropout=dropout) + for _ in range(n_blocks) + ] + ) + self.out_proj = nn.Linear(hidden_dim, in_dim) + + def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor: + x = self.input_proj(x) + for block in self.blocks: + x = block(x, cond) + return self.out_proj(x) + + +def _route_forward( + experts: nn.ModuleList, + router: Router, + x: torch.Tensor, + cond: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + training: bool, +) -> torch.Tensor: + """Shared dispatch for both Routed* trunks. + + Train mode: full soft mixture `sum_i gate_i * expert_i(x)` — fully + differentiable, N-expert compute. Eval mode: grouped top-1 dispatch — + each row runs exactly one (small) expert, which is the actual source + of the per-call speedup this architecture is for. + """ + if training: + weights = router.gate(cond_cont, cond_cat) # (B, n_experts) + out = torch.zeros_like(x) + for i, expert in enumerate(experts): + out = out + weights[:, i : i + 1] * expert(x, cond) + return out + + idx = router.top1(cond_cont, cond_cat) # (B,) + out = torch.zeros_like(x) + for i, expert in enumerate(experts): + mask = idx == i + if mask.any(): + out[mask] = expert(x[mask], cond[mask]) + return out + + +class RoutedDenoisingMLP(nn.Module): + """Routed drop-in for `DenoisingMLP`. + + Shares the time embedding, `ConditionEncoder`, and `n_sec_head` (all + tiny) across experts and routes only the trunk (where the FLOPs are). + Same `forward`/`predict_n_sec`/`pdg_embedding_weight` signatures as + `DenoisingMLP`, so sample.py/rollout.py/validate.py need no changes. + """ + + def __init__( + self, + pdg_vocab: int, + mat_vocab: int, + router: Router, + expert_hidden_dim: int = 128, + expert_n_blocks: int = 3, + emb_dim: int = EMB_DIM, + time_dim: int = 64, + cond_out_dim: int = 128, + x_dim: int = X_DIM, + dropout: float = 0.1, + k_max: int = K_MAX, + ) -> None: + super().__init__() + self.router = router + self.time_emb = SinusoidalEmbedding(time_dim) + self.cond_enc = ConditionEncoder( + pdg_vocab=pdg_vocab, + mat_vocab=mat_vocab, + emb_dim=emb_dim, + out_dim=cond_out_dim, + ) + merged_cond_dim = time_dim + cond_out_dim + self.experts = nn.ModuleList( + [ + ExpertTrunk( + x_dim, + expert_hidden_dim, + expert_n_blocks, + merged_cond_dim, + dropout=dropout, + ) + for _ in range(router.n_experts) + ] + ) + self.n_sec_head = nn.Sequential( + nn.Linear(cond_out_dim, cond_out_dim), + nn.SiLU(), + nn.Linear(cond_out_dim, k_max + 1), + ) + + def forward( + self, + x_t: torch.Tensor, + t: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + ) -> torch.Tensor: + t_emb = self.time_emb(t) + c_emb = self.cond_enc(cond_cont, cond_cat) + cond = torch.cat([t_emb, c_emb], dim=-1) + return _route_forward( + self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training + ) + + def predict_n_sec( + self, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + ) -> torch.Tensor: + """Return n_sec logits (B, K_MAX+1) from conditioning alone.""" + c_emb = self.cond_enc(cond_cont, cond_cat) + return self.n_sec_head(c_emb) + + def pdg_embedding_weight(self) -> torch.Tensor: + """Return the PDG embedding table weights for secondary type targets.""" + return self.cond_enc.pdg_emb.weight + + +class RoutedSecondaryDecoder(nn.Module): + """Routed drop-in for `SecondaryDecoder`. + + Shares the time embedding and `SecondaryConditionEncoder` across + experts and routes only the trunk. Same `forward` signature as + `SecondaryDecoder`. + """ + + def __init__( + self, + pdg_vocab: int, + mat_vocab: int, + router: Router, + expert_hidden_dim: int = 128, + expert_n_blocks: int = 3, + emb_dim: int = EMB_DIM, + time_dim: int = 64, + cond_out_dim: int = 128, + stage1_proj_dim: int = 64, + sec_dim: int = SEC_DIM, + dropout: float = 0.1, + ) -> None: + super().__init__() + self.router = router + self.time_emb = SinusoidalEmbedding(time_dim) + self.cond_enc = SecondaryConditionEncoder( + pdg_vocab=pdg_vocab, + mat_vocab=mat_vocab, + emb_dim=emb_dim, + cond_out_dim=cond_out_dim, + stage1_proj_dim=stage1_proj_dim, + out_dim=cond_out_dim, + ) + merged_cond_dim = time_dim + cond_out_dim + self.experts = nn.ModuleList( + [ + ExpertTrunk( + sec_dim, + expert_hidden_dim, + expert_n_blocks, + merged_cond_dim, + dropout=dropout, + ) + for _ in range(router.n_experts) + ] + ) + + def forward( + self, + x_t: torch.Tensor, + t: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + stage1_out: torch.Tensor, + ) -> torch.Tensor: + t_emb = self.time_emb(t) + c_emb = self.cond_enc(cond_cont, cond_cat, stage1_out) + cond = torch.cat([t_emb, c_emb], dim=-1) + return _route_forward( + self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training + ) + + +_STAGE1_MODEL_KEYS = { + "pdg_vocab", + "mat_vocab", + "hidden_dim", + "n_blocks", + "emb_dim", + "dropout", + "k_max", +} +_SEC_DECODER_MODEL_KEYS = { + "pdg_vocab", + "mat_vocab", + "hidden_dim", + "n_blocks", + "emb_dim", + "dropout", +} + + +def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]: + """Construct (stage1, sec_decoder) from a persisted/CLI model_config dict. + + Dispatches to the routed pair when `model_config["router"]["enabled"]` + is truthy; a missing/absent "router" key (pre-routing checkpoints) + falls back to the monolithic pair unchanged, so this is a drop-in + replacement for the ad-hoc constructions it replaces. + """ + router_cfg = model_config.get("router") + if router_cfg and router_cfg.get("enabled"): + router_kwargs = { + k: v + for k, v in router_cfg.items() + if k not in ("enabled", "type", "n_experts") + } + shared = dict( + pdg_vocab=model_config["pdg_vocab"], + mat_vocab=model_config["mat_vocab"], + expert_hidden_dim=model_config.get("expert_hidden_dim", 128), + expert_n_blocks=model_config.get("expert_n_blocks", 3), + emb_dim=model_config.get("emb_dim", EMB_DIM), + dropout=model_config.get("dropout", 0.1), + ) + stage1 = RoutedDenoisingMLP( + router=build_router( + router_cfg["type"], router_cfg["n_experts"], **router_kwargs + ), + k_max=model_config.get("k_max", K_MAX), + **shared, + ) + sec_decoder = RoutedSecondaryDecoder( + router=build_router( + router_cfg["type"], router_cfg["n_experts"], **router_kwargs + ), + **shared, + ) + return stage1, sec_decoder + + stage1 = DenoisingMLP( + **{k: v for k, v in model_config.items() if k in _STAGE1_MODEL_KEYS} + ) + sec_decoder = SecondaryDecoder( + **{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS} + ) + return stage1, sec_decoder diff --git a/giant/pipeline.py b/giant/pipeline.py index 593455c..16d655f 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -14,7 +14,7 @@ from giant.data.loader import ( ) from giant.data.transforms import build_features, _WelfordAccumulator from giant.data.dataset import make_event_split, StreamingStepsDataset -from giant.model.network import DenoisingMLP, SecondaryDecoder +from giant.model.network import build_models from giant.train import train as run_training @@ -114,23 +114,22 @@ def run_train_job( "update giant/constants.py if emb_dim changed" ) - stage1_model = DenoisingMLP( - pdg_vocab=len(pdg_map), - mat_vocab=len(mat_map), - hidden_dim=m["hidden_dim"], - n_blocks=m["n_blocks"], - emb_dim=emb_dim, - dropout=m["dropout"], - k_max=K_MAX, - ) - sec_decoder = SecondaryDecoder( - pdg_vocab=len(pdg_map), - mat_vocab=len(mat_map), - hidden_dim=m["hidden_dim"], - n_blocks=m["n_blocks"], - emb_dim=emb_dim, - dropout=m["dropout"], - ) + router_cfg = m["router"] + model_config = { + "pdg_vocab": len(pdg_map), + "mat_vocab": len(mat_map), + "hidden_dim": m["hidden_dim"], + "n_blocks": m["n_blocks"], + "emb_dim": emb_dim, + "dropout": m["dropout"], + "k_max": K_MAX, + "sec_slot_dim": SEC_SLOT_DIM, + "router": dict(router_cfg), + "expert_hidden_dim": router_cfg["expert_hidden_dim"], + "expert_n_blocks": router_cfg["expert_n_blocks"], + } + + stage1_model, sec_decoder = build_models(model_config) echo( f"stage1: {sum(p.numel() for p in stage1_model.parameters()):,} parameters | " f"sec_decoder: {sum(p.numel() for p in sec_decoder.parameters()):,} parameters" @@ -148,17 +147,6 @@ def run_train_job( ) config.save_config(cfg, out_dir, meta) - model_config = { - "pdg_vocab": len(pdg_map), - "mat_vocab": len(mat_map), - "hidden_dim": m["hidden_dim"], - "n_blocks": m["n_blocks"], - "emb_dim": emb_dim, - "dropout": m["dropout"], - "k_max": K_MAX, - "sec_slot_dim": SEC_SLOT_DIM, - } - run_training( stage1_model=stage1_model, sec_decoder=sec_decoder, @@ -172,6 +160,7 @@ def run_train_job( out_dir=out_dir, lambda_nsec=t.get("lambda_nsec", 0.1), lambda_s2=t.get("lambda_s2", 1.0), + lambda_balance=router_cfg.get("lambda_balance", 0.0), normalizer_dict={"cond": cond_norm.to_dict(), "target": tgt_norm.to_dict()}, pdg_map={str(k): v for k, v in pdg_map.items()}, mat_map={str(k): v for k, v in mat_map.items()}, diff --git a/giant/train.py b/giant/train.py index c21deb7..b30a616 100644 --- a/giant/train.py +++ b/giant/train.py @@ -26,10 +26,12 @@ _METRICS_FIELDS = [ "train_loss_s1", "train_loss_nsec", "train_loss_s2", + "train_loss_balance", "val_loss", "val_loss_s1", "val_loss_nsec", "val_loss_s2", + "val_loss_balance", "lr", "epoch_time_s", ] @@ -108,8 +110,9 @@ def _compute_losses( device: torch.device, lambda_nsec: float, lambda_s2: float, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Compute (total_loss, L_s1, L_nsec, L_s2) for one batch.""" + lambda_balance: float = 0.0, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute (total_loss, L_s1, L_nsec, L_s2, L_balance) for one batch.""" cond_cont, cond_cat, x1_s1, n_sec, sec_cont, sec_pdg_idx = batch cond_cont = cond_cont.to(device) cond_cat = cond_cat.to(device) @@ -150,8 +153,21 @@ def _compute_losses( sec_mask, ) + # Optional MoE load-balance auxiliary loss: only present when both stages + # are routed (RoutedDenoisingMLP/RoutedSecondaryDecoder carry `.router`, + # the monolith models don't), computed on cond_cont alone (cheap — no + # trunk compute) so it's reported even when lambda_balance == 0. + if hasattr(stage1_model, "router") and hasattr(sec_decoder, "router"): + l_balance = stage1_model.router.balance_loss( + cond_cont, cond_cat + ) + sec_decoder.router.balance_loss(cond_cont, cond_cat) + else: + l_balance = torch.zeros((), device=device) + total = l_s1 + lambda_nsec * l_nsec + lambda_s2 * l_s2 - return total, l_s1, l_nsec, l_s2 + if lambda_balance > 0: + total = total + lambda_balance * l_balance + return total, l_s1, l_nsec, l_s2, l_balance def train( @@ -167,6 +183,7 @@ def train( out_dir: str | Path, lambda_nsec: float = 0.1, lambda_s2: float = 1.0, + lambda_balance: float = 0.0, normalizer_dict: dict | None = None, pdg_map: dict | None = None, mat_map: dict | None = None, @@ -245,6 +262,7 @@ def train( train_s1_sum = 0.0 train_nsec_sum = 0.0 train_s2_sum = 0.0 + train_balance_sum = 0.0 train_n = 0 ema_loss = 0.0 bar = tqdm( @@ -256,7 +274,7 @@ def train( dynamic_ncols=True, ) for batch in bar: - loss, l_s1, l_nsec, l_s2 = _compute_losses( + loss, l_s1, l_nsec, l_s2, l_balance = _compute_losses( stage1_model, sec_decoder, batch, @@ -265,6 +283,7 @@ def train( device, lambda_nsec, lambda_s2, + lambda_balance, ) optimizer.zero_grad() loss.backward() @@ -277,6 +296,7 @@ def train( train_s1_sum += l_s1.item() * B train_nsec_sum += l_nsec.item() * B train_s2_sum += l_s2.item() * B + train_balance_sum += l_balance.item() * B train_n += B ema_loss = ( batch_loss if train_n == B else 0.95 * ema_loss + 0.05 * batch_loss @@ -299,10 +319,11 @@ def train( val_s1_sum = 0.0 val_nsec_sum = 0.0 val_s2_sum = 0.0 + val_balance_sum = 0.0 val_n = 0 with torch.no_grad(): for batch in val_loader: - loss, l_s1, l_nsec, l_s2 = _compute_losses( + loss, l_s1, l_nsec, l_s2, l_balance = _compute_losses( stage1_model, sec_decoder, batch, @@ -311,12 +332,14 @@ def train( device, lambda_nsec, lambda_s2, + lambda_balance, ) B = batch[0].size(0) val_loss_sum += loss.item() * B val_s1_sum += l_s1.item() * B val_nsec_sum += l_nsec.item() * B val_s2_sum += l_s2.item() * B + val_balance_sum += l_balance.item() * B val_n += B val_loss = val_loss_sum / max(val_n, 1) epoch_time = time.monotonic() - epoch_start @@ -328,7 +351,8 @@ def train( f" train {train_loss:.4f}" f" (s1={train_s1_sum / max(train_n, 1):.3f}" f" nsec={train_nsec_sum / max(train_n, 1):.3f}" - f" s2={train_s2_sum / max(train_n, 1):.3f})" + f" s2={train_s2_sum / max(train_n, 1):.3f}" + f" bal={train_balance_sum / max(train_n, 1):.3f})" f" val {val_loss:.4f}" f" lr {current_lr:.2e} {epoch_time:.1f}s{marker}" ) @@ -339,10 +363,12 @@ def train( "train_loss_s1": train_s1_sum / max(train_n, 1), "train_loss_nsec": train_nsec_sum / max(train_n, 1), "train_loss_s2": train_s2_sum / max(train_n, 1), + "train_loss_balance": train_balance_sum / max(train_n, 1), "val_loss": val_loss, "val_loss_s1": val_s1_sum / max(val_n, 1), "val_loss_nsec": val_nsec_sum / max(val_n, 1), "val_loss_s2": val_s2_sum / max(val_n, 1), + "val_loss_balance": val_balance_sum / max(val_n, 1), "lr": current_lr, "epoch_time_s": epoch_time, } diff --git a/tests/test_router.py b/tests/test_router.py new file mode 100644 index 0000000..36b7e73 --- /dev/null +++ b/tests/test_router.py @@ -0,0 +1,296 @@ +"""Tests for the mixture-of-experts routing prototype (giant/model/network.py).""" + +import torch + +from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM +from giant.model.network import ( + DenoisingMLP, + EnergyRouter, + ROUTER_REGISTRY, + RoutedDenoisingMLP, + RoutedSecondaryDecoder, + SecondaryDecoder, + build_models, + build_router, +) + + +def _cond(B=8, pdg=3, mat=2): + cond_cont = torch.randn(B, COND_DIM) + cond_cat = torch.stack( + [torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1 + ) + return cond_cont, cond_cat + + +def _routed_stage1(n_experts=4, pdg=3, mat=2, **router_kwargs): + router = build_router("energy", n_experts, **router_kwargs) + return RoutedDenoisingMLP( + pdg_vocab=pdg, + mat_vocab=mat, + router=router, + expert_hidden_dim=16, + expert_n_blocks=2, + ) + + +def _routed_sec_decoder(n_experts=4, pdg=3, mat=2, **router_kwargs): + router = build_router("energy", n_experts, **router_kwargs) + return RoutedSecondaryDecoder( + pdg_vocab=pdg, + mat_vocab=mat, + router=router, + expert_hidden_dim=16, + expert_n_blocks=2, + ) + + +# ── Router / EnergyRouter contract ────────────────────────────────────────── + + +def test_energy_router_registered(): + assert ROUTER_REGISTRY["energy"] is EnergyRouter + + +def test_energy_router_gate_partition_of_unity(): + router = EnergyRouter(n_experts=4) + cond_cont, cond_cat = _cond(16) + g = router.gate(cond_cont, cond_cat) + assert g.shape == (16, 4) + torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0) + + +def test_energy_router_top1_matches_gate_argmax(): + router = EnergyRouter(n_experts=4) + cond_cont, cond_cat = _cond(16) + assert torch.equal( + router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1) + ) + + +def test_energy_router_hardens_as_temperature_shrinks(): + """As tau -> 0 the soft gate should converge to a one-hot at the argmax.""" + router = EnergyRouter(n_experts=4, temperature=1e-4) + cond_cont, cond_cat = _cond(16) + g = router.gate(cond_cont, cond_cat) + top1 = router.top1(cond_cont, cond_cat) + onehot = torch.nn.functional.one_hot(top1, num_classes=4).float() + torch.testing.assert_close(g, onehot, atol=1e-3, rtol=0) + + +def test_energy_router_balance_loss_is_nonnegative_scalar(): + router = EnergyRouter(n_experts=4) + cond_cont, cond_cat = _cond(16) + loss = router.balance_loss(cond_cont, cond_cat) + assert loss.shape == () + assert loss.item() >= 0.0 + + +def test_build_router_ignores_unrecognized_kwargs(): + # lambda_balance is a model_config.router key but not an EnergyRouter kwarg + router = build_router("energy", 4, temperature=0.3, lambda_balance=0.5) + assert isinstance(router, EnergyRouter) + assert router.temperature == 0.3 + + +def test_build_router_unknown_type_raises(): + try: + build_router("nonexistent", 4) + except ValueError: + return + raise AssertionError("expected ValueError for unknown router type") + + +# ── RoutedDenoisingMLP ─────────────────────────────────────────────────────── + + +def test_routed_denoising_mlp_output_shape_train_and_eval(): + B = 8 + model = _routed_stage1() + x_t = torch.randn(B, X_DIM) + t = torch.rand(B) + cond_cont, cond_cat = _cond(B) + + model.train() + out_train = model(x_t, t, cond_cont, cond_cat) + assert out_train.shape == (B, X_DIM) + + model.eval() + with torch.no_grad(): + out_eval = model(x_t, t, cond_cont, cond_cat) + assert out_eval.shape == (B, X_DIM) + + +def test_routed_denoising_mlp_gradients_flow_in_train_mode(): + """Soft mixture in train mode should touch every expert's parameters.""" + B = 8 + model = _routed_stage1(n_experts=3) + x_t = torch.randn(B, X_DIM) + t = torch.rand(B) + cond_cont, cond_cat = _cond(B) + model.train() + flow_loss = model(x_t, t, cond_cont, cond_cat).sum() + nsec_loss = model.predict_n_sec(cond_cont, cond_cat).sum() + (flow_loss + nsec_loss).backward() + for name, p in model.named_parameters(): + assert p.grad is not None, f"no grad for {name}" + + +def test_routed_denoising_mlp_eval_dispatch_matches_manual_grouping(): + """Eval-mode grouped top-1 dispatch must equal running each row through + its assigned expert individually (batch order shouldn't matter).""" + B = 12 + model = _routed_stage1(n_experts=4) + model.eval() + x_t = torch.randn(B, X_DIM) + t = torch.rand(B) + cond_cont, cond_cat = _cond(B) + + with torch.no_grad(): + batched = model(x_t, t, cond_cont, cond_cat) + + t_emb = model.time_emb(t) + c_emb = model.cond_enc(cond_cont, cond_cat) + cond = torch.cat([t_emb, c_emb], dim=-1) + idx = model.router.top1(cond_cont, cond_cat) + manual = torch.zeros_like(x_t) + for i in range(B): + manual[i] = model.experts[int(idx[i])](x_t[i : i + 1], cond[i : i + 1])[0] + + torch.testing.assert_close(batched, manual, atol=1e-5, rtol=1e-4) + + +def test_routed_denoising_mlp_predict_n_sec_shape(): + B = 6 + model = _routed_stage1() + cond_cont, cond_cat = _cond(B) + logits = model.predict_n_sec(cond_cont, cond_cat) + assert logits.shape == (B, K_MAX + 1) + + +def test_routed_denoising_mlp_pdg_embedding_weight_shape(): + model = _routed_stage1(pdg=5, mat=2) + from giant.constants import EMB_DIM + + assert model.pdg_embedding_weight().shape == (5, EMB_DIM) + + +# ── RoutedSecondaryDecoder ─────────────────────────────────────────────────── + + +def test_routed_secondary_decoder_output_shape_train_and_eval(): + B = 8 + decoder = _routed_sec_decoder() + x_t = torch.randn(B, SEC_DIM) + t = torch.rand(B) + cond_cont, cond_cat = _cond(B) + stage1_out = torch.randn(B, X_DIM) + + decoder.train() + out_train = decoder(x_t, t, cond_cont, cond_cat, stage1_out) + assert out_train.shape == (B, SEC_DIM) + + decoder.eval() + with torch.no_grad(): + out_eval = decoder(x_t, t, cond_cont, cond_cat, stage1_out) + assert out_eval.shape == (B, SEC_DIM) + + +def test_routed_secondary_decoder_gradients_flow(): + B = 4 + decoder = _routed_sec_decoder(n_experts=3) + x_t = torch.randn(B, SEC_DIM) + t = torch.rand(B) + cond_cont, cond_cat = _cond(B) + stage1_out = torch.randn(B, X_DIM) + decoder.train() + decoder(x_t, t, cond_cont, cond_cat, stage1_out).sum().backward() + for name, p in decoder.named_parameters(): + assert p.grad is not None, f"no grad for {name}" + + +# ── build_models dispatch ──────────────────────────────────────────────────── + + +def test_build_models_monolith_when_router_absent(): + model_config = dict( + pdg_vocab=4, + mat_vocab=2, + hidden_dim=32, + n_blocks=2, + emb_dim=16, + dropout=0.1, + k_max=K_MAX, + ) + stage1, sec_decoder = build_models(model_config) + assert isinstance(stage1, DenoisingMLP) + assert isinstance(sec_decoder, SecondaryDecoder) + + +def test_build_models_monolith_when_router_disabled(): + model_config = dict( + pdg_vocab=4, + mat_vocab=2, + hidden_dim=32, + n_blocks=2, + emb_dim=16, + dropout=0.1, + k_max=K_MAX, + router={"enabled": False, "type": "energy", "n_experts": 4}, + ) + stage1, sec_decoder = build_models(model_config) + assert isinstance(stage1, DenoisingMLP) + assert isinstance(sec_decoder, SecondaryDecoder) + + +def test_build_models_routed_when_enabled(): + model_config = dict( + pdg_vocab=4, + mat_vocab=2, + emb_dim=16, + dropout=0.1, + k_max=K_MAX, + expert_hidden_dim=16, + expert_n_blocks=2, + router={ + "enabled": True, + "type": "energy", + "n_experts": 4, + "temperature": 0.5, + "learn_centers": True, + "lambda_balance": 0.0, + }, + ) + stage1, sec_decoder = build_models(model_config) + assert isinstance(stage1, RoutedDenoisingMLP) + assert isinstance(sec_decoder, RoutedSecondaryDecoder) + assert len(stage1.experts) == 4 + assert len(sec_decoder.experts) == 4 + + +def test_build_models_routed_pair_is_drop_in_for_sample_flow(): + """Exercise the exact calling convention giant/sample.py uses.""" + from giant.sample import sample_flow, sample_secondaries + + model_config = dict( + pdg_vocab=3, + mat_vocab=2, + emb_dim=16, + dropout=0.1, + k_max=K_MAX, + expert_hidden_dim=8, + expert_n_blocks=1, + router={"enabled": True, "type": "energy", "n_experts": 2}, + ) + stage1, sec_decoder = build_models(model_config) + B = 5 + cond_cont, cond_cat = _cond(B, pdg=3, mat=2) + stage1_norm, n_sec_pred = sample_flow(stage1, cond_cont, cond_cat, steps=2) + assert stage1_norm.shape == (B, X_DIM) + assert n_sec_pred.shape == (B,) + + sec_cont, sec_type_emb, sec_valid = sample_secondaries( + sec_decoder, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2 + ) + assert sec_cont.shape == (B, K_MAX, 4) + assert sec_valid.shape == (B, K_MAX) -- 2.39.5 From 0b3ece52ed612f8e6762dc8681bdc997fa321edc Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Wed, 8 Jul 2026 16:15:09 +0200 Subject: [PATCH 02/11] Add ProcessRouter for physics-process-based expert gating Routes on the physics process (Compton, phot, brems, ...) that ends a step, supervised by a small classifier since process is a post-step outcome unobservable at gate time. Threads a process label end-to-end through the data pipeline (loader, build_features, dataset batches, training loss/checkpointing) alongside the existing EnergyRouter. --- giant/cli.py | 2 +- giant/config.py | 6 +++ giant/data/dataset.py | 45 ++++++++++++------- giant/data/loader.py | 37 ++++++++++++++++ giant/data/transforms.py | 14 +++++- giant/model/network.py | 78 +++++++++++++++++++++++++++++++++ giant/pipeline.py | 17 ++++++-- giant/train.py | 41 +++++++++++++++--- giant/validate.py | 4 +- tests/test_loader.py | 36 +++++++++++++++- tests/test_router.py | 93 ++++++++++++++++++++++++++++++++++++++++ tests/test_transforms.py | 43 ++++++++++++++++++- 12 files changed, 385 insertions(+), 31 deletions(-) diff --git a/giant/cli.py b/giant/cli.py index 42f8ad0..d38cdc0 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -472,7 +472,7 @@ def predict( nonlocal writer, total if coord == Coord.local: - cond_cont, cond_cat, target_raw, _, _, _, _, _ = build_features( + cond_cont, cond_cat, target_raw, _, _, _, _, _, _ = build_features( piece, pdg_map, mat_map ) cond_cont = cond_norm.transform(cond_cont) diff --git a/giant/config.py b/giant/config.py index b5008a6..d0ba445 100644 --- a/giant/config.py +++ b/giant/config.py @@ -37,6 +37,12 @@ DEFAULT_CONFIG: dict = { "temperature": 0.5, # energy-router kwarg "learn_centers": True, # energy-router kwarg "lambda_balance": 0.0, # optional load-balance aux loss weight + "emb_dim": 8, # process-router kwarg: its own pdg/mat embedding width + "hidden_dim": 64, # process-router kwarg: its classifier's hidden width + "lambda_proc": 0.0, # process-router kwarg: supervised process-CE weight + # (0.0 still trains a working router — the gate gets gradient + # through the downstream flow loss like EnergyRouter's centers — + # but only lambda_proc > 0 grounds it in the true `process` label) }, }, } diff --git a/giant/data/dataset.py b/giant/data/dataset.py index 433c5f3..ce05905 100644 --- a/giant/data/dataset.py +++ b/giant/data/dataset.py @@ -36,7 +36,7 @@ class StreamingStepsDataset(IterableDataset): numpy slicing instead of a per-row Python loop in the default collate. Each batch is a tuple: - (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx) + (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx, proc_idx) where: cond_cont: (B, COND_DIM) float32 cond_cat: (B, 2) int64 @@ -44,6 +44,8 @@ class StreamingStepsDataset(IterableDataset): n_sec: (B,) int64 — true secondary count per step sec_cont: (B, K_MAX, 4) float32 — [stick_logit, local_dir] per slot sec_pdg_idx: (B, K_MAX) int64 — PDG model-index per secondary slot + proc_idx: (B,) int64 — process-class label (ProcessRouter supervision + only; zeros when `proc_map` is None) """ def __init__( @@ -57,6 +59,7 @@ class StreamingStepsDataset(IterableDataset): batch_size: int, shuffle_buffer: int = 65536, shuffle: bool = True, + proc_map: dict[str, int] | None = None, ) -> None: self.files = list(files) self.split_events = split_events @@ -68,6 +71,7 @@ class StreamingStepsDataset(IterableDataset): self.batch_size = batch_size self.shuffle_buffer = max(shuffle_buffer, batch_size) self.shuffle = shuffle + self.proc_map = proc_map def __iter__(self): worker_info = torch.utils.data.get_worker_info() @@ -85,6 +89,7 @@ class StreamingStepsDataset(IterableDataset): buf_nsec: list[np.ndarray] = [] buf_sec: list[np.ndarray] = [] buf_spdg: list[np.ndarray] = [] + buf_proc: list[np.ndarray] = [] buf_n = 0 for path in files: @@ -94,14 +99,16 @@ class StreamingStepsDataset(IterableDataset): continue chunk = {k: v[mask] for k, v in chunk.items()} - cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx, _, _ = ( - build_features( - chunk, - self.pdg_map, - self.mat_map, - cond_normalizer=self.cond_normalizer, - target_normalizer=self.target_normalizer, - ) + ( + cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx, + proc_idx, _, _, + ) = build_features( + chunk, + self.pdg_map, + self.mat_map, + cond_normalizer=self.cond_normalizer, + target_normalizer=self.target_normalizer, + proc_map=self.proc_map, ) buf_cont.append(cond_cont) buf_cat.append(cond_cat) @@ -109,19 +116,24 @@ class StreamingStepsDataset(IterableDataset): buf_nsec.append(n_sec) buf_sec.append(sec_cont) buf_spdg.append(sec_pdg_idx) + buf_proc.append(proc_idx) buf_n += len(cond_cont) if buf_n >= self.shuffle_buffer: - buf_cont, buf_cat, buf_tgt, buf_nsec, buf_sec, buf_spdg, buf_n = ( + ( + buf_cont, buf_cat, buf_tgt, buf_nsec, buf_sec, buf_spdg, + buf_proc, buf_n, + ) = ( yield from self._flush( buf_cont, buf_cat, buf_tgt, buf_nsec, buf_sec, buf_spdg, - final=False, + buf_proc, final=False, ) ) if buf_n > 0: yield from self._flush( - buf_cont, buf_cat, buf_tgt, buf_nsec, buf_sec, buf_spdg, final=True + buf_cont, buf_cat, buf_tgt, buf_nsec, buf_sec, buf_spdg, buf_proc, + final=True, ) def _flush( @@ -132,6 +144,7 @@ class StreamingStepsDataset(IterableDataset): buf_nsec: list[np.ndarray], buf_sec: list[np.ndarray], buf_spdg: list[np.ndarray], + buf_proc: list[np.ndarray], final: bool, ): cont = np.concatenate(buf_cont) @@ -140,11 +153,12 @@ class StreamingStepsDataset(IterableDataset): nsec = np.concatenate(buf_nsec) sec = np.concatenate(buf_sec) spdg = np.concatenate(buf_spdg) + proc = np.concatenate(buf_proc) if self.shuffle: idx = np.random.permutation(len(cont)) cont, cat, tgt = cont[idx], cat[idx], tgt[idx] - nsec, sec, spdg = nsec[idx], sec[idx], spdg[idx] + nsec, sec, spdg, proc = nsec[idx], sec[idx], spdg[idx], proc[idx] bs = self.batch_size n = len(cont) @@ -158,13 +172,14 @@ class StreamingStepsDataset(IterableDataset): torch.from_numpy(nsec[start:end]).long(), torch.from_numpy(sec[start:end]).float(), torch.from_numpy(spdg[start:end]).long(), + torch.from_numpy(proc[start:end]).long(), ) if final: - return [], [], [], [], [], [], 0 + return [], [], [], [], [], [], [], 0 rem = n_full * bs return ( [cont[rem:]], [cat[rem:]], [tgt[rem:]], - [nsec[rem:]], [sec[rem:]], [spdg[rem:]], + [nsec[rem:]], [sec[rem:]], [spdg[rem:]], [proc[rem:]], n - rem, ) diff --git a/giant/data/loader.py b/giant/data/loader.py index fa0ba79..a215157 100644 --- a/giant/data/loader.py +++ b/giant/data/loader.py @@ -95,6 +95,16 @@ def _df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]: "layer_id": df["layer_id"].to_numpy(dtype=np.int32), "n_sec": df["child_track_ids"].apply(len).to_numpy(dtype=np.int32), "e_sec": df["e_sec"].to_numpy(dtype=np.float32), + # The physics process that ended the step (e.g. "compt", "phot", + # "eBrem") — a post-step outcome, so it's a router/classifier + # supervision label only, never conditioning (see build_process_map* + # / ProcessRouter). Guarded like has_sec_lists: older parquet + # conversions predating this column still load fine. + "process": ( + df["process"].to_numpy(dtype=object) + if "process" in df.columns + else np.full(len(df), "", dtype=object) + ), "step_length": df["step_length"].to_numpy(dtype=np.float32), "post_E": df["post_E"].to_numpy(dtype=np.float32), "delta_e": (df["pre_E"] - df["post_E"]).to_numpy(dtype=np.float32), @@ -192,3 +202,30 @@ def build_index_maps_from_files( {v: i for i, v in enumerate(sorted(pdg_vals))}, {v: i for i, v in enumerate(sorted(mat_vals))}, ) + + +def build_process_map_from_files( + files: list[Path], n_experts: int +) -> dict[str, int]: + """Scan the `process` column and build a frequency-capped process->index map. + + Physics processes have a long tail (rare nuclear captures, decays, ...) + while `ProcessRouter` needs a fixed number of expert slots, so only the + `n_experts - 1` most frequent processes get their own index; every rarer + process is bucketed into a shared "other" index (`n_experts - 1`). This + mirrors how `build_features` clamps the n_sec label to K_MAX for the + fixed-width n_sec_head classifier. + """ + counts: dict[str, int] = {} + for path in files: + df = pd.read_parquet(path, columns=["process"]) + for name, count in df["process"].value_counts().items(): + name = str(name) + counts[name] = counts.get(name, 0) + int(count) + ranked = sorted(counts, key=lambda name: counts[name], reverse=True) + keep = ranked[: max(n_experts - 1, 0)] + proc_map = {name: i for i, name in enumerate(keep)} + other_idx = n_experts - 1 + for name in ranked[len(keep) :]: + proc_map[name] = other_idx + return proc_map diff --git a/giant/data/transforms.py b/giant/data/transforms.py index f5df8db..e8df166 100644 --- a/giant/data/transforms.py +++ b/giant/data/transforms.py @@ -433,6 +433,7 @@ def build_features( cond_normalizer: Normalizer | None = None, target_normalizer: Normalizer | None = None, fit: bool = False, + proc_map: dict[str, int] | None = None, ) -> tuple[ np.ndarray, np.ndarray, @@ -440,16 +441,20 @@ def build_features( np.ndarray, np.ndarray, np.ndarray, + np.ndarray, Normalizer | None, Normalizer | None, ]: - """Assemble (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx) arrays. + """Assemble (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx, proc_idx) arrays. target_s1: (N, 9) Stage-1 primary post-step target (unchanged from Phase 1) n_sec: (N,) integer secondary counts (target for n_sec head) sec_cont: (N, K_MAX, 4) continuous secondary targets [stick_logit, dir_local] sec_pdg_idx: (N, K_MAX) integer PDG model-indices; used to look up embedding targets in the training loop + proc_idx: (N,) integer process-class label (ProcessRouter supervision only — + never conditioning). Zeros when `proc_map` is None or the loaded + data has no "process" column (e.g. pre-conversion parquet files). """ from giant.constants import K_MAX @@ -523,6 +528,12 @@ def build_features( if target_normalizer is not None: target_s1 = target_normalizer.transform(target_s1) + process = data.get("process") + if proc_map is not None and process is not None: + proc_idx = np.array([proc_map[str(p)] for p in process], dtype=np.int64) + else: + proc_idx = np.zeros(len(cond_cat), dtype=np.int64) + return ( cond_cont, cond_cat, @@ -530,6 +541,7 @@ def build_features( n_sec, sec_cont, sec_pdg_idx, + proc_idx, cond_normalizer, target_normalizer, ) diff --git a/giant/model/network.py b/giant/model/network.py index 91c06a1..749e950 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -3,6 +3,7 @@ import math import torch import torch.nn as nn +import torch.nn.functional as F from giant.constants import COND_DIM, EMB_DIM, K_MAX, SEC_DIM, X_DIM @@ -267,6 +268,19 @@ class Router(nn.Module): importance = self.gate(cond_cont, cond_cat).sum(dim=0) # (n_experts,) return (importance.std() / (importance.mean() + 1e-8)) ** 2 + def classify_loss( + self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor + ) -> torch.Tensor: + """Optional supervised auxiliary loss shaping the router's own belief. + + Default: none (a scalar 0), for routers like EnergyRouter that read a + quantity directly off cond_cont/cond_cat and need no label. Routers + gating on an unobservable pre-step quantity (e.g. ProcessRouter, + which predicts the physics process that will end the step) override + this to supervise their internal classifier against the true label. + """ + return torch.zeros((), device=cond_cont.device) + ROUTER_REGISTRY: dict[str, type[Router]] = {} @@ -330,6 +344,64 @@ class EnergyRouter(Router): return torch.softmax(-d2 / self.temperature, dim=-1) +@register_router("process") +class ProcessRouter(Router): + """Routes on the physics process expected to end the step. + + Unlike EnergyRouter (which reads a quantity that's already known at + pre-step time), the process — Compton, photoelectric, brems, ... — is a + *post-step outcome*: it can't be read off cond_cont/cond_cat directly. + Instead this router runs a small classifier over pre-step conditioning + (its own pdg/material embeddings, kept separate from the trunk's + ConditionEncoder) that predicts it, one class per expert slot + (`n_experts` doubles as the number of process classes — see + `build_process_map_from_files`, which caps the process vocabulary to + exactly this many classes, bucketing rare processes into a shared + "other" slot). + + The classifier is supervised by `classify_loss` against the true + `process` label (see `giant/train.py`) — a *training-time* signal only; + `gate`/`top1` never see it, so eval-time dispatch (rollout, predict) + needs no ground truth, same as every other Router. This sidesteps the + gradient/differentiability problem that sank the earlier + process-conditioned-flow proposal (see the archived decision doc): the + hard categorical choice only ever feeds a non-differentiable expert + *dispatch*, never the flow's own conditioning path. + """ + + def __init__( + self, + n_experts: int, + pdg_vocab: int, + mat_vocab: int, + emb_dim: int = 8, + hidden_dim: int = 64, + ) -> None: + super().__init__(n_experts) + self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim) + self.mat_emb = nn.Embedding(mat_vocab, emb_dim) + self.classifier = nn.Sequential( + nn.Linear(COND_DIM + 2 * emb_dim, hidden_dim), + nn.SiLU(), + nn.Linear(hidden_dim, n_experts), + ) + + def logits(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + """(B, n_experts) raw process-classifier logits, one class per expert.""" + pdg_e = self.pdg_emb(cond_cat[:, 0]) + mat_e = self.mat_emb(cond_cat[:, 1]) + h = torch.cat([cond_cont, pdg_e, mat_e], dim=-1) + return self.classifier(h) + + def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + return torch.softmax(self.logits(cond_cont, cond_cat), dim=-1) + + def classify_loss( + self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor + ) -> torch.Tensor: + return F.cross_entropy(self.logits(cond_cont, cond_cat), labels) + + class ExpertTrunk(nn.Module): """One small expert: `input_proj -> ResBlock stack -> out_proj`. @@ -570,6 +642,12 @@ def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]: for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts") } + # Not every router needs these (EnergyRouter doesn't declare them, so + # build_router's kwarg filtering drops them silently) but + # ProcessRouter needs its own pdg/material embeddings sized to match + # the checkpoint's vocab, same as the trunk's ConditionEncoder. + router_kwargs.setdefault("pdg_vocab", model_config["pdg_vocab"]) + router_kwargs.setdefault("mat_vocab", model_config["mat_vocab"]) shared = dict( pdg_vocab=model_config["pdg_vocab"], mat_vocab=model_config["mat_vocab"], diff --git a/giant/pipeline.py b/giant/pipeline.py index 16d655f..931c5fd 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -11,6 +11,7 @@ from giant.data.loader import ( load_event_ids, iter_file_chunks, build_index_maps_from_files, + build_process_map_from_files, ) from giant.data.transforms import build_features, _WelfordAccumulator from giant.data.dataset import make_event_split, StreamingStepsDataset @@ -54,6 +55,13 @@ def run_train_job( pdg_map, mat_map = build_index_maps_from_files(files) echo(f" {len(pdg_map)} PDG codes | {len(mat_map)} materials") + router_cfg = m["router"] + proc_map: dict[str, int] | None = None + if router_cfg.get("enabled") and router_cfg.get("type") == "process": + echo("building process vocabulary …") + proc_map = build_process_map_from_files(files, n_experts=router_cfg["n_experts"]) + echo(f" {len(proc_map)} process labels mapped to {router_cfg['n_experts']} experts") + echo("fitting normalizer (streaming) …") cond_acc = _WelfordAccumulator(COND_DIM) tgt_acc = _WelfordAccumulator(X_DIM) @@ -63,8 +71,8 @@ def run_train_job( if not mask.any(): continue chunk_tr = {k: v[mask] for k, v in chunk.items()} - cond_cont, _, target_s1, _n_sec, _sec_cont, _sec_pdg, _, _ = build_features( - chunk_tr, pdg_map, mat_map + cond_cont, _, target_s1, _n_sec, _sec_cont, _sec_pdg, _proc, _, _ = ( + build_features(chunk_tr, pdg_map, mat_map, proc_map=proc_map) ) cond_acc.update(cond_cont) tgt_acc.update(target_s1) @@ -81,6 +89,7 @@ def run_train_job( batch_size=t["batch_size"], shuffle_buffer=shuffle_buffer, shuffle=True, + proc_map=proc_map, ) val_ds = StreamingStepsDataset( files=files, @@ -91,6 +100,7 @@ def run_train_job( target_normalizer=tgt_norm, batch_size=t["batch_size"], shuffle=False, + proc_map=proc_map, ) pin = device.type == "cuda" @@ -114,7 +124,6 @@ def run_train_job( "update giant/constants.py if emb_dim changed" ) - router_cfg = m["router"] model_config = { "pdg_vocab": len(pdg_map), "mat_vocab": len(mat_map), @@ -161,9 +170,11 @@ def run_train_job( lambda_nsec=t.get("lambda_nsec", 0.1), lambda_s2=t.get("lambda_s2", 1.0), lambda_balance=router_cfg.get("lambda_balance", 0.0), + lambda_proc=router_cfg.get("lambda_proc", 0.0), normalizer_dict={"cond": cond_norm.to_dict(), "target": tgt_norm.to_dict()}, pdg_map={str(k): v for k, v in pdg_map.items()}, mat_map={str(k): v for k, v in mat_map.items()}, + proc_map=proc_map, model_config=model_config, resume_path=resume, validate_every=t["validate_every"], diff --git a/giant/train.py b/giant/train.py index b30a616..35f3d8c 100644 --- a/giant/train.py +++ b/giant/train.py @@ -27,11 +27,13 @@ _METRICS_FIELDS = [ "train_loss_nsec", "train_loss_s2", "train_loss_balance", + "train_loss_proc", "val_loss", "val_loss_s1", "val_loss_nsec", "val_loss_s2", "val_loss_balance", + "val_loss_proc", "lr", "epoch_time_s", ] @@ -111,15 +113,19 @@ def _compute_losses( lambda_nsec: float, lambda_s2: float, lambda_balance: float = 0.0, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Compute (total_loss, L_s1, L_nsec, L_s2, L_balance) for one batch.""" - cond_cont, cond_cat, x1_s1, n_sec, sec_cont, sec_pdg_idx = batch + lambda_proc: float = 0.0, +) -> tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor +]: + """Compute (total_loss, L_s1, L_nsec, L_s2, L_balance, L_proc) for one batch.""" + cond_cont, cond_cat, x1_s1, n_sec, sec_cont, sec_pdg_idx, proc_idx = batch cond_cont = cond_cont.to(device) cond_cat = cond_cat.to(device) x1_s1 = x1_s1.to(device) n_sec = n_sec.to(device) sec_cont = sec_cont.to(device) sec_pdg_idx = sec_pdg_idx.to(device) + proc_idx = proc_idx.to(device) # Stage-1 flow loss if mode == "flow": @@ -161,13 +167,21 @@ def _compute_losses( l_balance = stage1_model.router.balance_loss( cond_cont, cond_cat ) + sec_decoder.router.balance_loss(cond_cont, cond_cat) + # Supervised router auxiliary loss (e.g. ProcessRouter's process + # classifier); a scalar 0 for routers with no such loss (EnergyRouter). + l_proc = stage1_model.router.classify_loss( + cond_cont, cond_cat, proc_idx + ) + sec_decoder.router.classify_loss(cond_cont, cond_cat, proc_idx) else: l_balance = torch.zeros((), device=device) + l_proc = torch.zeros((), device=device) total = l_s1 + lambda_nsec * l_nsec + lambda_s2 * l_s2 if lambda_balance > 0: total = total + lambda_balance * l_balance - return total, l_s1, l_nsec, l_s2, l_balance + if lambda_proc > 0: + total = total + lambda_proc * l_proc + return total, l_s1, l_nsec, l_s2, l_balance, l_proc def train( @@ -184,9 +198,11 @@ def train( lambda_nsec: float = 0.1, lambda_s2: float = 1.0, lambda_balance: float = 0.0, + lambda_proc: float = 0.0, normalizer_dict: dict | None = None, pdg_map: dict | None = None, mat_map: dict | None = None, + proc_map: dict | None = None, model_config: dict | None = None, resume_path: str | Path | None = None, validate_every: int = 0, @@ -263,6 +279,7 @@ def train( train_nsec_sum = 0.0 train_s2_sum = 0.0 train_balance_sum = 0.0 + train_proc_sum = 0.0 train_n = 0 ema_loss = 0.0 bar = tqdm( @@ -274,7 +291,7 @@ def train( dynamic_ncols=True, ) for batch in bar: - loss, l_s1, l_nsec, l_s2, l_balance = _compute_losses( + loss, l_s1, l_nsec, l_s2, l_balance, l_proc = _compute_losses( stage1_model, sec_decoder, batch, @@ -284,6 +301,7 @@ def train( lambda_nsec, lambda_s2, lambda_balance, + lambda_proc, ) optimizer.zero_grad() loss.backward() @@ -297,6 +315,7 @@ def train( train_nsec_sum += l_nsec.item() * B train_s2_sum += l_s2.item() * B train_balance_sum += l_balance.item() * B + train_proc_sum += l_proc.item() * B train_n += B ema_loss = ( batch_loss if train_n == B else 0.95 * ema_loss + 0.05 * batch_loss @@ -320,10 +339,11 @@ def train( val_nsec_sum = 0.0 val_s2_sum = 0.0 val_balance_sum = 0.0 + val_proc_sum = 0.0 val_n = 0 with torch.no_grad(): for batch in val_loader: - loss, l_s1, l_nsec, l_s2, l_balance = _compute_losses( + loss, l_s1, l_nsec, l_s2, l_balance, l_proc = _compute_losses( stage1_model, sec_decoder, batch, @@ -333,6 +353,7 @@ def train( lambda_nsec, lambda_s2, lambda_balance, + lambda_proc, ) B = batch[0].size(0) val_loss_sum += loss.item() * B @@ -340,6 +361,7 @@ def train( val_nsec_sum += l_nsec.item() * B val_s2_sum += l_s2.item() * B val_balance_sum += l_balance.item() * B + val_proc_sum += l_proc.item() * B val_n += B val_loss = val_loss_sum / max(val_n, 1) epoch_time = time.monotonic() - epoch_start @@ -352,7 +374,8 @@ def train( f" (s1={train_s1_sum / max(train_n, 1):.3f}" f" nsec={train_nsec_sum / max(train_n, 1):.3f}" f" s2={train_s2_sum / max(train_n, 1):.3f}" - f" bal={train_balance_sum / max(train_n, 1):.3f})" + f" bal={train_balance_sum / max(train_n, 1):.3f}" + f" proc={train_proc_sum / max(train_n, 1):.3f})" f" val {val_loss:.4f}" f" lr {current_lr:.2e} {epoch_time:.1f}s{marker}" ) @@ -364,11 +387,13 @@ def train( "train_loss_nsec": train_nsec_sum / max(train_n, 1), "train_loss_s2": train_s2_sum / max(train_n, 1), "train_loss_balance": train_balance_sum / max(train_n, 1), + "train_loss_proc": train_proc_sum / max(train_n, 1), "val_loss": val_loss, "val_loss_s1": val_s1_sum / max(val_n, 1), "val_loss_nsec": val_nsec_sum / max(val_n, 1), "val_loss_s2": val_s2_sum / max(val_n, 1), "val_loss_balance": val_balance_sum / max(val_n, 1), + "val_loss_proc": val_proc_sum / max(val_n, 1), "lr": current_lr, "epoch_time_s": epoch_time, } @@ -401,6 +426,8 @@ def train( ckpt["pdg_map"] = pdg_map if mat_map is not None: ckpt["mat_map"] = mat_map + if proc_map is not None: + ckpt["proc_map"] = proc_map if model_config is not None: ckpt["model_config"] = model_config diff --git a/giant/validate.py b/giant/validate.py index 959e375..3372a06 100644 --- a/giant/validate.py +++ b/giant/validate.py @@ -85,8 +85,8 @@ def validate_marginals( for i, batch in enumerate(val_loader): if n_batches is not None and i >= n_batches: break - # Batch is (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx). - cond_cont, cond_cat, x1, n_sec, sec_cont, sec_pdg_idx = batch + # Batch is (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx, proc_idx). + cond_cont, cond_cat, x1, n_sec, sec_cont, sec_pdg_idx, _proc_idx = batch cond_cont = cond_cont.to(device) cond_cat = cond_cat.to(device) diff --git a/tests/test_loader.py b/tests/test_loader.py index 6c562f3..703d9d3 100644 --- a/tests/test_loader.py +++ b/tests/test_loader.py @@ -1,6 +1,7 @@ +import pandas as pd import pytest -from giant.data.loader import find_parquet_files +from giant.data.loader import build_process_map_from_files, find_parquet_files def _touch(path): @@ -58,3 +59,36 @@ def test_manifest_with_no_entries_raises(tmp_path): with pytest.raises(FileNotFoundError): find_parquet_files(manifest) + + +def test_build_process_map_from_files_keeps_most_frequent(tmp_path): + """process counts: eIoni=5, phot=3, compt=2, Rayl=1 — with n_experts=3, only + the top 2 (eIoni, phot) get their own index; compt/Rayl share the "other" + (last) index.""" + process = ( + ["eIoni"] * 5 + ["phot"] * 3 + ["compt"] * 2 + ["Rayl"] * 1 + ) + path = tmp_path / "shard-000.parquet" + pd.DataFrame({"process": process}).to_parquet(path) + + proc_map = build_process_map_from_files([path], n_experts=3) + + assert proc_map["eIoni"] == 0 + assert proc_map["phot"] == 1 + assert proc_map["compt"] == 2 + assert proc_map["Rayl"] == 2 + assert set(proc_map.values()) <= {0, 1, 2} + + +def test_build_process_map_from_files_spans_multiple_files(tmp_path): + path_a = tmp_path / "a.parquet" + path_b = tmp_path / "b.parquet" + pd.DataFrame({"process": ["eIoni"] * 3 + ["phot"] * 1}).to_parquet(path_a) + pd.DataFrame({"process": ["phot"] * 4 + ["compt"] * 1}).to_parquet(path_b) + + # phot: 1+4=5 total > eIoni: 3 > compt: 1 + proc_map = build_process_map_from_files([path_a, path_b], n_experts=3) + + assert proc_map["phot"] == 0 + assert proc_map["eIoni"] == 1 + assert proc_map["compt"] == 2 diff --git a/tests/test_router.py b/tests/test_router.py index 36b7e73..0755ebb 100644 --- a/tests/test_router.py +++ b/tests/test_router.py @@ -6,6 +6,7 @@ from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM from giant.model.network import ( DenoisingMLP, EnergyRouter, + ProcessRouter, ROUTER_REGISTRY, RoutedDenoisingMLP, RoutedSecondaryDecoder, @@ -101,6 +102,98 @@ def test_build_router_unknown_type_raises(): raise AssertionError("expected ValueError for unknown router type") +# ── ProcessRouter ──────────────────────────────────────────────────────────── + + +def test_process_router_registered(): + assert ROUTER_REGISTRY["process"] is ProcessRouter + + +def test_process_router_gate_partition_of_unity(): + router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2) + cond_cont, cond_cat = _cond(16) + g = router.gate(cond_cont, cond_cat) + assert g.shape == (16, 4) + torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0) + + +def test_process_router_top1_matches_gate_argmax(): + router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2) + cond_cont, cond_cat = _cond(16) + assert torch.equal( + router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1) + ) + + +def test_process_router_balance_loss_is_nonnegative_scalar(): + router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2) + cond_cont, cond_cat = _cond(16) + loss = router.balance_loss(cond_cont, cond_cat) + assert loss.shape == () + assert loss.item() >= 0.0 + + +def test_process_router_classify_loss_decreases_with_training(): + """The classifier should be able to fit an arbitrary label assignment — + a sanity check that gradients actually flow to the process classifier.""" + torch.manual_seed(0) + router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2) + cond_cont, cond_cat = _cond(32) + labels = torch.randint(0, 4, (32,)) + + opt = torch.optim.Adam(router.parameters(), lr=0.05) + first = router.classify_loss(cond_cont, cond_cat, labels).item() + for _ in range(50): + opt.zero_grad() + loss = router.classify_loss(cond_cont, cond_cat, labels) + loss.backward() + opt.step() + last = loss.item() + assert last < first + + +def test_energy_router_classify_loss_defaults_to_zero(): + """Routers with no supervised signal (EnergyRouter) fall back to the + Router base class's zero-loss default.""" + router = EnergyRouter(n_experts=4) + cond_cont, cond_cat = _cond(16) + labels = torch.randint(0, 4, (16,)) + loss = router.classify_loss(cond_cont, cond_cat, labels) + assert loss.shape == () + assert loss.item() == 0.0 + + +def test_build_router_process_type_uses_pdg_mat_vocab(): + router = build_router("process", 4, pdg_vocab=5, mat_vocab=3, emb_dim=8) + assert isinstance(router, ProcessRouter) + assert router.pdg_emb.num_embeddings == 5 + assert router.mat_emb.num_embeddings == 3 + + +def test_build_models_routed_with_process_router(): + model_config = dict( + pdg_vocab=4, + mat_vocab=2, + emb_dim=16, + dropout=0.1, + k_max=K_MAX, + expert_hidden_dim=16, + expert_n_blocks=2, + router={ + "enabled": True, + "type": "process", + "n_experts": 3, + "lambda_proc": 1.0, + }, + ) + stage1, sec_decoder = build_models(model_config) + assert isinstance(stage1, RoutedDenoisingMLP) + assert isinstance(stage1.router, ProcessRouter) + assert len(stage1.experts) == 3 + assert stage1.router.pdg_emb.num_embeddings == 4 + assert stage1.router.mat_emb.num_embeddings == 2 + + # ── RoutedDenoisingMLP ─────────────────────────────────────────────────────── diff --git a/tests/test_transforms.py b/tests/test_transforms.py index d92f347..ad8ddfb 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -235,7 +235,48 @@ def test_build_features_clamps_n_sec_label_to_k_max(): pdg_map = {11: 0} mat_map = {"PbWO4": 0} - _, _, _, n_sec, _, _, _, _ = build_features(data, pdg_map, mat_map) + _, _, _, n_sec, _, _, _, _, _ = build_features(data, pdg_map, mat_map) assert n_sec.max() <= K_MAX np.testing.assert_array_equal(n_sec, [0, 5, K_MAX]) + + +def _minimal_step_data(N: int, process: np.ndarray | None = None) -> dict: + rng = np.random.default_rng(0) + data = { + "pdg": np.full(N, 11, dtype=np.int32), + "material": np.full(N, "PbWO4", dtype=object), + "pre_pos": rng.standard_normal((N, 3)).astype(np.float32), + "pre_E": np.full(N, 10.0, dtype=np.float32), + "pre_dir": np.tile(np.array([0, 0, 1], dtype=np.float32), (N, 1)), + "layer_id": np.zeros(N, dtype=np.int32), + "n_sec": np.zeros(N, dtype=np.int32), + "e_sec": np.full(N, 1.0, dtype=np.float32), + "step_length": np.full(N, 1.0, dtype=np.float32), + "post_E": np.full(N, 9.0, dtype=np.float32), + "edep": np.full(N, 1.0, dtype=np.float32), + "post_dir": np.tile(np.array([0, 0, 1], dtype=np.float32), (N, 1)), + "post_pos": rng.standard_normal((N, 3)).astype(np.float32), + } + if process is not None: + data["process"] = process + return data + + +def test_build_features_proc_idx_zero_without_proc_map(): + data = _minimal_step_data(3, process=np.array(["compt", "phot", "eIoni"], dtype=object)) + pdg_map, mat_map = {11: 0}, {"PbWO4": 0} + + *_, proc_idx, _, _ = build_features(data, pdg_map, mat_map) + + np.testing.assert_array_equal(proc_idx, [0, 0, 0]) + + +def test_build_features_proc_idx_looks_up_proc_map(): + data = _minimal_step_data(3, process=np.array(["compt", "phot", "eIoni"], dtype=object)) + pdg_map, mat_map = {11: 0}, {"PbWO4": 0} + proc_map = {"compt": 0, "phot": 1, "eIoni": 2} + + *_, proc_idx, _, _ = build_features(data, pdg_map, mat_map, proc_map=proc_map) + + np.testing.assert_array_equal(proc_idx, [0, 1, 2]) -- 2.39.5 From f387178dbf97e555d47cedad91f3702fab556e44 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Thu, 9 Jul 2026 09:02:51 +0200 Subject: [PATCH 03/11] Error on missing secondary lists instead of silently zeroing Stage-2 targets MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A parquet that carries child_track_ids/e_sec but was never run through the parent->child join lacks the per-secondary columns (sec_E_list/sec_pdg_list/ sec_dir_list). build_features would fall back to all-zero sec_cont/sec_pdg_idx, collapsing every secondary to PDG index 0 and a constant energy fraction — a broken Stage 2 that trained with no error (single-species validation tables). Add an opt-in require_secondaries flag that raises when n_sec > 0 but the lists are absent, and enable it on the training paths (StreamingStepsDataset and the normalizer-fit pass). giant predict keeps the default False for Stage-1-only use. Co-Authored-By: Claude Opus 4.8 --- giant/data/dataset.py | 1 + giant/data/transforms.py | 25 +++++++++++++++++++++++++ giant/pipeline.py | 5 ++++- tests/test_transforms.py | 27 +++++++++++++++++++++++++++ 4 files changed, 57 insertions(+), 1 deletion(-) diff --git a/giant/data/dataset.py b/giant/data/dataset.py index ce05905..afdb051 100644 --- a/giant/data/dataset.py +++ b/giant/data/dataset.py @@ -109,6 +109,7 @@ class StreamingStepsDataset(IterableDataset): cond_normalizer=self.cond_normalizer, target_normalizer=self.target_normalizer, proc_map=self.proc_map, + require_secondaries=True, ) buf_cont.append(cond_cont) buf_cat.append(cond_cat) diff --git a/giant/data/transforms.py b/giant/data/transforms.py index e8df166..4567fff 100644 --- a/giant/data/transforms.py +++ b/giant/data/transforms.py @@ -434,6 +434,7 @@ def build_features( target_normalizer: Normalizer | None = None, fit: bool = False, proc_map: dict[str, int] | None = None, + require_secondaries: bool = False, ) -> tuple[ np.ndarray, np.ndarray, @@ -455,6 +456,11 @@ def build_features( proc_idx: (N,) integer process-class label (ProcessRouter supervision only — never conditioning). Zeros when `proc_map` is None or the loaded data has no "process" column (e.g. pre-conversion parquet files). + + require_secondaries: when True, raise if any step has n_sec > 0 but the + per-secondary list columns are absent (a mis-converted file that would + otherwise silently zero all Stage-2 targets). Training paths set this; + Stage-1-only callers (e.g. `giant predict`) leave it False. """ from giant.constants import K_MAX @@ -515,6 +521,25 @@ def build_features( sec_pdg_list ).astype(np.int64) else: + # Guard against silently training Stage 2 on zeroed targets: if any step + # actually spawned secondaries (n_sec > 0, from child_track_ids) but the + # per-secondary columns are absent, the file was never run through the + # parent->child join (steps_to_parquet._add_secondary_attributes / + # `dwarf convert`). Zero-filling here would collapse every secondary to + # PDG index 0 and a constant energy fraction — a broken Stage 2 with no + # error. Callers that only need Stage-1 (e.g. `giant predict`) keep the + # default require_secondaries=False. + if require_secondaries and n_sec_raw.max(initial=0) > 0: + n_with_sec = int((n_sec_raw > 0).sum()) + raise ValueError( + f"{n_with_sec} step(s) have secondaries (n_sec > 0) but the " + "per-secondary columns (sec_E_list / sec_pdg_list / sec_dx_list " + "…) are missing. This parquet was not run through the " + "parent->child join (steps_to_parquet._add_secondary_attributes " + "/ `dwarf convert`); training on it would silently zero all " + "Stage-2 targets. Re-convert the file, or pass " + "require_secondaries=False for Stage-1-only use." + ) N = len(n_sec) sec_cont = np.zeros((N, K_MAX, 4), dtype=np.float32) sec_pdg_idx = np.zeros((N, K_MAX), dtype=np.int64) diff --git a/giant/pipeline.py b/giant/pipeline.py index 931c5fd..c06d49f 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -72,7 +72,10 @@ def run_train_job( continue chunk_tr = {k: v[mask] for k, v in chunk.items()} cond_cont, _, target_s1, _n_sec, _sec_cont, _sec_pdg, _proc, _, _ = ( - build_features(chunk_tr, pdg_map, mat_map, proc_map=proc_map) + build_features( + chunk_tr, pdg_map, mat_map, proc_map=proc_map, + require_secondaries=True, + ) ) cond_acc.update(cond_cont) tgt_acc.update(target_s1) diff --git a/tests/test_transforms.py b/tests/test_transforms.py index ad8ddfb..7c5eed8 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -280,3 +280,30 @@ def test_build_features_proc_idx_looks_up_proc_map(): *_, proc_idx, _, _ = build_features(data, pdg_map, mat_map, proc_map=proc_map) np.testing.assert_array_equal(proc_idx, [0, 1, 2]) + + +def test_build_features_require_secondaries_raises_when_lists_missing(): + """A parquet with n_sec > 0 but no per-secondary list columns was never run + through the parent->child join; require_secondaries must catch it instead of + silently zeroing every Stage-2 target (regression: this collapsed the + secondary species to a single PDG index during training).""" + data = _minimal_step_data(3) + data["n_sec"] = np.array([0, 2, 1], dtype=np.int32) # secondaries, but no lists + pdg_map, mat_map = {11: 0}, {"PbWO4": 0} + + with pytest.raises(ValueError, match="per-secondary columns"): + build_features(data, pdg_map, mat_map, require_secondaries=True) + + +def test_build_features_require_secondaries_ok_when_no_secondaries(): + """require_secondaries only fires when secondaries actually exist; a file + with n_sec == 0 everywhere (e.g. Stage-1-only) must still load.""" + data = _minimal_step_data(3) # n_sec all zero + pdg_map, mat_map = {11: 0}, {"PbWO4": 0} + + _, _, _, _, sec_cont, sec_pdg_idx, *_ = build_features( + data, pdg_map, mat_map, require_secondaries=True + ) + + assert not sec_cont.any() + assert not sec_pdg_idx.any() -- 2.39.5 From 5dc086e35aca8ac2de157c977af3df5d734a3a85 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Thu, 9 Jul 2026 10:08:04 +0200 Subject: [PATCH 04/11] Derive a unique per-job seed for minicalosim shard generation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Concurrent job launches in create_root_files.py can start within the same wall-clock second, and minicalosim's default seed falls back to time(NULL) in that case — so two "independent" shards could silently get identical RNG state and produce byte-identical physics. Requires the companion MINICALOSIM_SEED env-var support in the minicalosim repo. Co-Authored-By: Claude Sonnet 5 --- scripts/create_root_files.py | 21 +++++++++++++-- tests/test_create_root_files.py | 45 ++++++++++++++++++++++++++++++++- 2 files changed, 63 insertions(+), 3 deletions(-) diff --git a/scripts/create_root_files.py b/scripts/create_root_files.py index 495ae8a..1b1c8f2 100644 --- a/scripts/create_root_files.py +++ b/scripts/create_root_files.py @@ -25,6 +25,7 @@ import shutil import subprocess import sys import uuid +import zlib from dataclasses import dataclass from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path @@ -104,6 +105,20 @@ def plan_jobs( return jobs +def job_seed(kind: str, gen: str, job: SimJob) -> int: + """Deterministic RNG seed for one sim job, unique per (kind, gen, detector, config, shard). + + Jobs run concurrently (ThreadPoolExecutor below) and can start within the + same wall-clock second; minicalosim's default seed falls back to + time(NULL) in that case, so two concurrently-launched jobs can silently + get identical RNG state and produce byte-identical physics despite + landing in separate shard files. Deriving the seed from the full job + identity instead keeps it both unique and reproducible. + """ + key = f"{kind}|{gen}|{job.detector}|{job.config or ''}|{job.shard_index}" + return zlib.crc32(key.encode()) & 0x7FFFFFFF + + def run_job( job: SimJob, executable: Path, @@ -124,7 +139,8 @@ def run_job( cmd.append(job.config) cmd.append(str(events_per_file)) - result = subprocess.run(cmd, cwd=workdir, capture_output=True, text=True) + env = dict(os.environ, MINICALOSIM_SEED=str(job_seed(kind, gen, job))) + result = subprocess.run(cmd, cwd=workdir, capture_output=True, text=True, env=env) if result.returncode != 0: return JobResult( @@ -254,7 +270,8 @@ def run_make_root( / job.detector / f"shard-{job.shard_index:03d}.root" ) - print(f" {' '.join(cmd)} -> {dest}") + seed = job_seed(kind, gen, job) + print(f" MINICALOSIM_SEED={seed} {' '.join(cmd)} -> {dest}") if not execute: print("\nDry run only — pass --execute to apply.") diff --git a/tests/test_create_root_files.py b/tests/test_create_root_files.py index 72d8955..4d21f5e 100644 --- a/tests/test_create_root_files.py +++ b/tests/test_create_root_files.py @@ -9,6 +9,7 @@ from scripts import create_root_files parse_detector_spec = create_root_files.parse_detector_spec next_shard_index = create_root_files.next_shard_index plan_jobs = create_root_files.plan_jobs +job_seed = create_root_files.job_seed run_job = create_root_files.run_job run_all = create_root_files.run_all SimJob = create_root_files.SimJob @@ -35,7 +36,13 @@ start = time.time() time.sleep({sleep}) end = time.time() payload = json.dumps( - {{"argv": sys.argv[1:], "cwd": os.getcwd(), "start": start, "end": end}} + {{ + "argv": sys.argv[1:], + "cwd": os.getcwd(), + "start": start, + "end": end, + "seed": os.environ.get("MINICALOSIM_SEED"), + }} ) for i in range({output_count}): with open(f"out_{{i}}.root", "w") as f: @@ -130,6 +137,42 @@ def test_plan_jobs_multiple_detectors_each_start_independently(tmp_path): assert by_detector["sampling_fe_scint"] == [0, 1] +def test_job_seed_deterministic(): + job = SimJob(detector="pbwo4", config=None, shard_index=3) + assert job_seed("steps", "gen1", job) == job_seed("steps", "gen1", job) + + +def test_job_seed_varies_by_shard_index(): + a = SimJob(detector="pbwo4", config=None, shard_index=0) + b = SimJob(detector="pbwo4", config=None, shard_index=1) + assert job_seed("steps", "gen1", a) != job_seed("steps", "gen1", b) + + +def test_job_seed_varies_by_detector(): + a = SimJob(detector="pbwo4", config=None, shard_index=0) + b = SimJob(detector="sampling_pb_scint", config="pb_scint", shard_index=0) + assert job_seed("steps", "gen1", a) != job_seed("steps", "gen1", b) + + +def test_job_seed_varies_by_gen(): + job = SimJob(detector="pbwo4", config=None, shard_index=0) + assert job_seed("steps", "gen1", job) != job_seed("steps", "gen2", job) + + +def test_run_job_passes_deterministic_seed_env_var(tmp_path): + fake = _write_fake_executable(tmp_path / "fake_exe.py") + (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True) + tmp_root = tmp_path / ".sim-tmp" + tmp_root.mkdir() + + job = SimJob(detector="pbwo4", config=None, shard_index=5) + result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root) + + assert result.dest is not None + payload = json.loads(result.dest.read_text()) + assert payload["seed"] == str(job_seed("steps", "gen1", job)) + + def test_run_job_moves_output_to_correct_shard_path(tmp_path): fake = _write_fake_executable(tmp_path / "fake_exe.py") gen_dir = tmp_path / "raw" / "steps" / "gen1" -- 2.39.5 From 028fa13b7bf347324ea571d34d8926f1b7072450 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Thu, 9 Jul 2026 14:36:06 +0200 Subject: [PATCH 05/11] Drop orphaned child tracks instead of nulling secondary targets MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A listed child_track_id can fail to match any first-step row (e.g. a secondary absorbed below the tracking threshold at birth). The parent->child left join in _add_secondary_attributes left these as nulls, which silently became NaN once the parquet round-tripped through the loader's float32 padding — poisoning every later secondary slot in that step via the cumulative "remaining budget" in encode_secondaries, while e_sec quietly undercounted and n_sec (from len(child_track_ids)) overcounted relative to the actual lists. Drop orphans from both the per-secondary lists and child_track_ids itself so downstream counts stay consistent, and thread the per-file orphaned count back through convert_steps_to_parquet so both the sequential and --jobs>1 batch paths in `dwarf convert` can report an aggregate total instead of relying on grepping printed output. Also floors encode_secondaries' slot-0 budget to _EPS (matching the i>0 branch), fixing a harmless but noisy 0/0 divide warning on zero-secondary steps. Co-Authored-By: Claude Sonnet 5 --- giant/data/transforms.py | 2 +- scripts/dwarf.py | 9 +++++- scripts/steps_to_parquet.py | 47 +++++++++++++++++++++++----- scripts/steps_to_parquet_parallel.py | 11 +++++++ tests/test_steps_to_parquet.py | 39 +++++++++++++++++++++-- 5 files changed, 96 insertions(+), 12 deletions(-) diff --git a/giant/data/transforms.py b/giant/data/transforms.py index 4567fff..fadceea 100644 --- a/giant/data/transforms.py +++ b/giant/data/transforms.py @@ -310,7 +310,7 @@ def encode_secondaries( stick_logits = np.zeros((N, K), dtype=np.float32) for i in range(K): if i == 0: - remaining = e_sec + remaining = np.maximum(e_sec, _EPS) else: remaining = np.maximum(e_sec - sec_E_list[:, :i].sum(axis=1), _EPS) f = np.clip(sec_E_list[:, i].astype(np.float64) / remaining, _EPS, 1.0 - _EPS) diff --git a/scripts/dwarf.py b/scripts/dwarf.py index f728bcf..e096834 100644 --- a/scripts/dwarf.py +++ b/scripts/dwarf.py @@ -116,14 +116,21 @@ def convert( "error: --output can only be used with a single input file", err=True ) raise typer.Exit(1) + total_orphaned = 0 for root_file in root_files: - convert_steps_to_parquet( + _, n_orphaned = convert_steps_to_parquet( root_file, output_path=output, batch_size=batch_size, tree_name=tree, compression=compression_value, ) + total_orphaned += n_orphaned + if total_orphaned: + typer.echo( + f"\n{total_orphaned} orphaned child track(s) dropped across " + f"{len(root_files)} file(s)." + ) return if output is not None: diff --git a/scripts/steps_to_parquet.py b/scripts/steps_to_parquet.py index cb74fe5..880c965 100644 --- a/scripts/steps_to_parquet.py +++ b/scripts/steps_to_parquet.py @@ -4,7 +4,7 @@ See `uv run dwarf convert --help` for the CLI. """ from pathlib import Path -from typing import Literal +from typing import Literal, cast import awkward as ak import polars as pl @@ -13,7 +13,7 @@ import uproot ParquetCompression = Literal["lz4", "uncompressed", "snappy", "gzip", "brotli", "zstd"] -def _add_secondary_attributes(df: pl.DataFrame) -> pl.DataFrame: +def _add_secondary_attributes(df: pl.DataFrame) -> tuple[pl.DataFrame, int]: """Add per-step secondary attributes via the parent→child track join. For each step that spawns secondaries, collects each child track's birth @@ -27,6 +27,18 @@ def _add_secondary_attributes(df: pl.DataFrame) -> pl.DataFrame: Steps with no children get 0.0 / empty lists. The full event must be present in `df` (it is — the writer concatenates before calling this). + + A listed child_track_id can fail to match any row in `first_step` — the + child track never took a recorded step (e.g. absorbed below the tracking + threshold at birth). Such orphans carry no physical secondary data, so + they're dropped from child_track_ids/sec_*_list rather than left as nulls: + a null in a float32 list silently becomes NaN once the parquet round-trips + through the loader (`giant/data/loader.py:_pad_list_col`), and that NaN + poisons every later secondary slot in the same step via the cumulative-sum + "remaining budget" in `encode_secondaries`. + + Returns (df, n_orphaned) — the caller uses the count to report/aggregate + across files rather than relying solely on the printed message here. """ first_step = ( df.sort("step_no") @@ -41,6 +53,8 @@ def _add_secondary_attributes(df: pl.DataFrame) -> pl.DataFrame: .rename({"track_id": "child_track_id"}) ) + child_track_id_dtype = cast(pl.List, df.schema["child_track_ids"]).inner + exploded = ( df.select(["event_id", "child_track_ids"]) .with_row_index("_step_row") @@ -51,11 +65,20 @@ def _add_secondary_attributes(df: pl.DataFrame) -> pl.DataFrame: joined = exploded.join(first_step, on=["event_id", "child_track_id"], how="left") + n_orphaned = joined["child_E"].null_count() + if n_orphaned: + print( + f" dropping {n_orphaned} orphaned child_track_id(s) with no " + "recorded first step (absorbed below tracking threshold?)" + ) + joined = joined.drop_nulls("child_E") + # Sort each step's secondaries by descending energy, then aggregate into lists per_step = ( joined.sort("child_E", descending=True) .group_by("_step_row") .agg( + pl.col("child_track_id").alias("child_track_ids"), pl.col("child_E").sum().alias("e_sec"), pl.col("child_E").alias("sec_E_list"), pl.col("child_pdg").alias("sec_pdg_list"), @@ -67,11 +90,14 @@ def _add_secondary_attributes(df: pl.DataFrame) -> pl.DataFrame: empty_list_f64 = pl.Series("x", [[]], dtype=pl.List(pl.Float64)) empty_list_i32 = pl.Series("x", [[]], dtype=pl.List(pl.Int32)) + empty_list_child_id = pl.Series("x", [[]], dtype=pl.List(child_track_id_dtype)) - return ( - df.with_row_index("_step_row") + out = ( + df.drop("child_track_ids") + .with_row_index("_step_row") .join(per_step, on="_step_row", how="left") .with_columns( + pl.col("child_track_ids").fill_null(empty_list_child_id), pl.col("e_sec").fill_null(0.0).cast(pl.Float64), pl.col("sec_E_list").fill_null(empty_list_f64), pl.col("sec_pdg_list").fill_null(empty_list_i32), @@ -81,6 +107,7 @@ def _add_secondary_attributes(df: pl.DataFrame) -> pl.DataFrame: ) .drop("_step_row") ) + return out, n_orphaned def _batch_to_polars(batch: ak.Array) -> pl.DataFrame: @@ -106,7 +133,7 @@ def convert_steps_to_parquet( batch_size: str = "100 MB", tree_name: str = "Steps", compression: ParquetCompression = "snappy", -) -> Path: +) -> tuple[Path, int]: """Read *tree_name* from *root_path* and write it to a Parquet file. Reads in batches of *batch_size* so that peak ROOT-deserialization memory @@ -122,6 +149,11 @@ def convert_steps_to_parquet( integer row count (500_000). tree_name: Name of the TTree inside the ROOT file. compression: Parquet compression codec (snappy | lz4 | zstd | gzip | none). + + Returns (output_path, n_orphaned) — n_orphaned is the count of dropped + orphaned child_track_ids (see `_add_secondary_attributes`), 0 if the tree + has no child_track_ids column at all. Callers converting many files use + it to aggregate a total instead of grepping the printed per-file message. """ root_path = Path(root_path) if output_path is None: @@ -144,11 +176,12 @@ def convert_steps_to_parquet( df = pl.concat(batches) # Steps tree carries the parent→child links needed to derive secondary energy; # other trees (e.g. Hits) don't, so only augment when the column is present. + n_orphaned = 0 if "child_track_ids" in df.columns: print("\nComputing per-step secondary attributes …", end=" ", flush=True) - df = _add_secondary_attributes(df) + df, n_orphaned = _add_secondary_attributes(df) print(f"\nWriting {output_path} …", end=" ", flush=True) df.write_parquet(output_path, compression=compression) print(f"done ({output_path.stat().st_size / 1e6:.1f} MB)") - return output_path + return output_path, n_orphaned diff --git a/scripts/steps_to_parquet_parallel.py b/scripts/steps_to_parquet_parallel.py index d89a607..dc1c5fe 100644 --- a/scripts/steps_to_parquet_parallel.py +++ b/scripts/steps_to_parquet_parallel.py @@ -26,6 +26,12 @@ from pathlib import Path GEN_RE = re.compile(r"^gen\d+$") SCHEMA_RE = re.compile(r"^schema(\d+)$") +# Matches the per-file orphan-drop message printed by +# steps_to_parquet._add_secondary_attributes — each subprocess's count is +# parsed back out of its captured stdout since there's no in-process return +# value across the subprocess boundary. +_ORPHAN_RE = re.compile(r"dropping (\d+) orphaned child_track_id") + class DestinationError(ValueError): pass @@ -210,4 +216,9 @@ def run_parallel_job( print(f" {root_file}", file=sys.stderr) raise SystemExit(1) + total_orphaned = sum( + int(m.group(1)) for _, _, stdout, _ in results for m in _ORPHAN_RE.finditer(stdout) + ) + if total_orphaned: + print(f"\n{total_orphaned} orphaned child track(s) dropped across {len(results)} file(s).") print(f"\nAll {len(results)} conversion(s) completed.") diff --git a/tests/test_steps_to_parquet.py b/tests/test_steps_to_parquet.py index 6ec6918..9810d6e 100644 --- a/tests/test_steps_to_parquet.py +++ b/tests/test_steps_to_parquet.py @@ -22,7 +22,8 @@ def _frame() -> pl.DataFrame: def test_e_sec_sums_child_first_step_energy(): - out = steps_to_parquet._add_secondary_attributes(_frame()) + out, n_orphaned = steps_to_parquet._add_secondary_attributes(_frame()) + assert n_orphaned == 0 e_sec = dict( zip(zip(out["track_id"], out["step_no"], out["event_id"]), out["e_sec"]) ) @@ -31,7 +32,7 @@ def test_e_sec_sums_child_first_step_energy(): def test_e_sec_zero_when_no_children(): - out = steps_to_parquet._add_secondary_attributes(_frame()) + out, _ = steps_to_parquet._add_secondary_attributes(_frame()) childless = out.filter( (pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 1) ) @@ -40,6 +41,38 @@ def test_e_sec_zero_when_no_children(): def test_e_sec_preserves_row_count_and_order(): df = _frame() - out = steps_to_parquet._add_secondary_attributes(df) + out, _ = steps_to_parquet._add_secondary_attributes(df) assert out.height == df.height assert out["pre_E"].to_list() == df["pre_E"].to_list() + + +def test_orphaned_child_track_is_dropped_not_nulled(): + """A listed child_track_id with no first step of its own (e.g. absorbed + below the tracking threshold at birth) must not leave a null in + sec_E_list/sec_pdg_list/etc: that null turns into NaN once the parquet + round-trips through the loader, poisoning every later secondary slot in + the step via encode_secondaries' cumulative "remaining budget". It must + also be dropped from child_track_ids itself, so n_sec (len(child_track_ids) + downstream) matches the actual, orphan-free secondary lists.""" + df = pl.DataFrame( + { + "event_id": [0, 0, 0], + "track_id": [1, 1, 2], + "step_no": [0, 1, 0], + "pre_E": [100.0, 80.0, 15.0], + "pdg": [11, 11, 22], + "pre_dx": [0.0, 0.0, 1.0], + "pre_dy": [0.0, 0.0, 0.0], + "pre_dz": [1.0, 1.0, 0.0], + # track 3 is listed as a child but never appears with its own step. + "child_track_ids": [[2, 3], [], []], + } + ) + out, n_orphaned = steps_to_parquet._add_secondary_attributes(df) + row = out.filter((pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 0)) + + assert n_orphaned == 1 + assert row["child_track_ids"].to_list() == [[2]] + assert row["e_sec"].item() == 15.0 + assert row["sec_E_list"].to_list() == [[15.0]] + assert None not in row["sec_E_list"].item() -- 2.39.5 From 05d5dee6066bee7061fd8eb67ed7f98c1356d311 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Wed, 15 Jul 2026 10:00:36 +0200 Subject: [PATCH 06/11] Apply ruff format after merging phase2-secondary-prediction The merged proc_idx/proc_map plumbing wasn't run through ruff format before merging; reflow only, no logic changes. --- giant/data/dataset.py | 52 +++++++++++++++++++++------- giant/data/loader.py | 4 +-- giant/pipeline.py | 13 +++++-- scripts/steps_to_parquet_parallel.py | 8 +++-- tests/test_loader.py | 4 +-- tests/test_steps_to_parquet.py | 4 ++- tests/test_transforms.py | 8 +++-- 7 files changed, 67 insertions(+), 26 deletions(-) diff --git a/giant/data/dataset.py b/giant/data/dataset.py index afdb051..f319d93 100644 --- a/giant/data/dataset.py +++ b/giant/data/dataset.py @@ -100,8 +100,15 @@ class StreamingStepsDataset(IterableDataset): chunk = {k: v[mask] for k, v in chunk.items()} ( - cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx, - proc_idx, _, _, + cond_cont, + cond_cat, + target_s1, + n_sec, + sec_cont, + sec_pdg_idx, + proc_idx, + _, + _, ) = build_features( chunk, self.pdg_map, @@ -122,18 +129,34 @@ class StreamingStepsDataset(IterableDataset): if buf_n >= self.shuffle_buffer: ( - buf_cont, buf_cat, buf_tgt, buf_nsec, buf_sec, buf_spdg, - buf_proc, buf_n, - ) = ( - yield from self._flush( - buf_cont, buf_cat, buf_tgt, buf_nsec, buf_sec, buf_spdg, - buf_proc, final=False, - ) + buf_cont, + buf_cat, + buf_tgt, + buf_nsec, + buf_sec, + buf_spdg, + buf_proc, + buf_n, + ) = yield from self._flush( + buf_cont, + buf_cat, + buf_tgt, + buf_nsec, + buf_sec, + buf_spdg, + buf_proc, + final=False, ) if buf_n > 0: yield from self._flush( - buf_cont, buf_cat, buf_tgt, buf_nsec, buf_sec, buf_spdg, buf_proc, + buf_cont, + buf_cat, + buf_tgt, + buf_nsec, + buf_sec, + buf_spdg, + buf_proc, final=True, ) @@ -180,7 +203,12 @@ class StreamingStepsDataset(IterableDataset): return [], [], [], [], [], [], [], 0 rem = n_full * bs return ( - [cont[rem:]], [cat[rem:]], [tgt[rem:]], - [nsec[rem:]], [sec[rem:]], [spdg[rem:]], [proc[rem:]], + [cont[rem:]], + [cat[rem:]], + [tgt[rem:]], + [nsec[rem:]], + [sec[rem:]], + [spdg[rem:]], + [proc[rem:]], n - rem, ) diff --git a/giant/data/loader.py b/giant/data/loader.py index 86dd989..4ddbf7f 100644 --- a/giant/data/loader.py +++ b/giant/data/loader.py @@ -202,9 +202,7 @@ def build_index_maps_from_files( ) -def build_process_map_from_files( - files: list[Path], n_experts: int -) -> dict[str, int]: +def build_process_map_from_files(files: list[Path], n_experts: int) -> dict[str, int]: """Scan the `process` column and build a frequency-capped process->index map. Physics processes have a long tail (rare nuclear captures, decays, ...) diff --git a/giant/pipeline.py b/giant/pipeline.py index 2a0ba96..f160231 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -59,8 +59,12 @@ def run_train_job( proc_map: dict[str, int] | None = None if router_cfg.get("enabled") and router_cfg.get("type") == "process": echo("building process vocabulary …") - proc_map = build_process_map_from_files(files, n_experts=router_cfg["n_experts"]) - echo(f" {len(proc_map)} process labels mapped to {router_cfg['n_experts']} experts") + proc_map = build_process_map_from_files( + files, n_experts=router_cfg["n_experts"] + ) + echo( + f" {len(proc_map)} process labels mapped to {router_cfg['n_experts']} experts" + ) echo("fitting normalizer (streaming) …") cond_acc = _WelfordAccumulator(COND_DIM) @@ -73,7 +77,10 @@ def run_train_job( chunk_tr = {k: v[mask] for k, v in chunk.items()} cond_cont, _, target_s1, _n_sec, _sec_cont, _sec_pdg, _proc, _, _ = ( build_features( - chunk_tr, pdg_map, mat_map, proc_map=proc_map, + chunk_tr, + pdg_map, + mat_map, + proc_map=proc_map, require_secondaries=True, ) ) diff --git a/scripts/steps_to_parquet_parallel.py b/scripts/steps_to_parquet_parallel.py index dc1c5fe..b2ee41e 100644 --- a/scripts/steps_to_parquet_parallel.py +++ b/scripts/steps_to_parquet_parallel.py @@ -217,8 +217,12 @@ def run_parallel_job( raise SystemExit(1) total_orphaned = sum( - int(m.group(1)) for _, _, stdout, _ in results for m in _ORPHAN_RE.finditer(stdout) + int(m.group(1)) + for _, _, stdout, _ in results + for m in _ORPHAN_RE.finditer(stdout) ) if total_orphaned: - print(f"\n{total_orphaned} orphaned child track(s) dropped across {len(results)} file(s).") + print( + f"\n{total_orphaned} orphaned child track(s) dropped across {len(results)} file(s)." + ) print(f"\nAll {len(results)} conversion(s) completed.") diff --git a/tests/test_loader.py b/tests/test_loader.py index 703d9d3..af609dc 100644 --- a/tests/test_loader.py +++ b/tests/test_loader.py @@ -65,9 +65,7 @@ def test_build_process_map_from_files_keeps_most_frequent(tmp_path): """process counts: eIoni=5, phot=3, compt=2, Rayl=1 — with n_experts=3, only the top 2 (eIoni, phot) get their own index; compt/Rayl share the "other" (last) index.""" - process = ( - ["eIoni"] * 5 + ["phot"] * 3 + ["compt"] * 2 + ["Rayl"] * 1 - ) + process = ["eIoni"] * 5 + ["phot"] * 3 + ["compt"] * 2 + ["Rayl"] * 1 path = tmp_path / "shard-000.parquet" pd.DataFrame({"process": process}).to_parquet(path) diff --git a/tests/test_steps_to_parquet.py b/tests/test_steps_to_parquet.py index 9810d6e..8e839c1 100644 --- a/tests/test_steps_to_parquet.py +++ b/tests/test_steps_to_parquet.py @@ -69,7 +69,9 @@ def test_orphaned_child_track_is_dropped_not_nulled(): } ) out, n_orphaned = steps_to_parquet._add_secondary_attributes(df) - row = out.filter((pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 0)) + row = out.filter( + (pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 0) + ) assert n_orphaned == 1 assert row["child_track_ids"].to_list() == [[2]] diff --git a/tests/test_transforms.py b/tests/test_transforms.py index 86ab26d..c95a36b 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -272,7 +272,9 @@ def _step_data_no_sec_lists(n_sec: np.ndarray) -> dict: def test_build_features_proc_idx_zero_without_proc_map(): - data = _minimal_step_data(3, process=np.array(["compt", "phot", "eIoni"], dtype=object)) + data = _minimal_step_data( + 3, process=np.array(["compt", "phot", "eIoni"], dtype=object) + ) pdg_map, mat_map = {11: 0}, {"PbWO4": 0} *_, proc_idx, _, _ = build_features(data, pdg_map, mat_map) @@ -281,7 +283,9 @@ def test_build_features_proc_idx_zero_without_proc_map(): def test_build_features_proc_idx_looks_up_proc_map(): - data = _minimal_step_data(3, process=np.array(["compt", "phot", "eIoni"], dtype=object)) + data = _minimal_step_data( + 3, process=np.array(["compt", "phot", "eIoni"], dtype=object) + ) pdg_map, mat_map = {11: 0}, {"PbWO4": 0} proc_map = {"compt": 0, "phot": 1, "eIoni": 2} -- 2.39.5 From ac5fbd14b404f25398fa6a7d4e337fb7e29eed36 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Wed, 15 Jul 2026 10:13:07 +0200 Subject: [PATCH 07/11] Add PdgRouter for particle-type-based expert gating MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Routes on the pre-step PDG code, which — unlike ProcessRouter's process label — is already known at gate time (a conditioning input), so no supervision is needed and classify_loss falls back to the zero default. Generalizes EnergyRouter's soft-turn-on-then-Voronoi trick from a 1-D distance to a small learned PDG embedding space: its own embedding table maps each PDG code to a point, and n_experts learnable centers partition that space. Co-Authored-By: Claude Sonnet 5 --- giant/config.py | 6 +-- giant/model/network.py | 42 +++++++++++++++++++ tests/test_router.py | 95 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 140 insertions(+), 3 deletions(-) diff --git a/giant/config.py b/giant/config.py index d0ba445..7e4f8e7 100644 --- a/giant/config.py +++ b/giant/config.py @@ -34,10 +34,10 @@ DEFAULT_CONFIG: dict = { "n_experts": 4, "expert_hidden_dim": 128, "expert_n_blocks": 3, - "temperature": 0.5, # energy-router kwarg - "learn_centers": True, # energy-router kwarg + "temperature": 0.5, # energy/pdg-router kwarg + "learn_centers": True, # energy/pdg-router kwarg "lambda_balance": 0.0, # optional load-balance aux loss weight - "emb_dim": 8, # process-router kwarg: its own pdg/mat embedding width + "emb_dim": 8, # process/pdg-router kwarg: own pdg(/mat) embedding width "hidden_dim": 64, # process-router kwarg: its classifier's hidden width "lambda_proc": 0.0, # process-router kwarg: supervised process-CE weight # (0.0 still trains a working router — the gate gets gradient diff --git a/giant/model/network.py b/giant/model/network.py index 05487d5..7a6e8ca 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -344,6 +344,48 @@ class EnergyRouter(Router): return torch.softmax(-d2 / self.temperature, dim=-1) +@register_router("pdg") +class PdgRouter(Router): + """Soft turn-on gate over a learned PDG embedding. + + Unlike ProcessRouter's process label, PDG code is already known at + pre-step time (it's a conditioning input, `cond_cat[:, 0]`), so no + supervision is needed — `classify_loss` falls back to the Router base + class's zero-loss default, same as EnergyRouter. Because PDG is + categorical rather than a scalar, this generalizes EnergyRouter's + soft-turn-on-then-Voronoi trick from a 1-D distance to a distance in a + small embedding space: its own embedding table (kept separate from the + trunk's ConditionEncoder, same reasoning as ProcessRouter's own + pdg/mat embeddings) maps each PDG code to a point, and `n_experts` + learnable (or fixed) centers partition that space. + `gate(pdg) = softmax_i(-||emb(pdg) - c_i||^2 / tau)`. + """ + + def __init__( + self, + n_experts: int, + pdg_vocab: int, + emb_dim: int = 8, + temperature: float = 0.5, + learn_centers: bool = True, + ) -> None: + super().__init__(n_experts) + self.temperature = temperature + self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim) + centers = torch.randn(n_experts, emb_dim) * 0.1 + if learn_centers: + self.centers = nn.Parameter(centers) + else: + self.register_buffer("centers", centers) + + def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + e = self.pdg_emb(cond_cat[:, 0]) # (B, emb_dim) + d2 = ((e.unsqueeze(1) - self.centers.unsqueeze(0)) ** 2).sum( + -1 + ) # (B, n_experts) + return torch.softmax(-d2 / self.temperature, dim=-1) + + @register_router("process") class ProcessRouter(Router): """Routes on the physics process expected to end the step. diff --git a/tests/test_router.py b/tests/test_router.py index 0755ebb..eef1c3f 100644 --- a/tests/test_router.py +++ b/tests/test_router.py @@ -6,6 +6,7 @@ from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM from giant.model.network import ( DenoisingMLP, EnergyRouter, + PdgRouter, ProcessRouter, ROUTER_REGISTRY, RoutedDenoisingMLP, @@ -102,6 +103,100 @@ def test_build_router_unknown_type_raises(): raise AssertionError("expected ValueError for unknown router type") +# ── PdgRouter ──────────────────────────────────────────────────────────────── + + +def test_pdg_router_registered(): + assert ROUTER_REGISTRY["pdg"] is PdgRouter + + +def test_pdg_router_gate_partition_of_unity(): + router = PdgRouter(n_experts=4, pdg_vocab=3) + cond_cont, cond_cat = _cond(16) + g = router.gate(cond_cont, cond_cat) + assert g.shape == (16, 4) + torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0) + + +def test_pdg_router_top1_matches_gate_argmax(): + router = PdgRouter(n_experts=4, pdg_vocab=3) + cond_cont, cond_cat = _cond(16) + assert torch.equal( + router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1) + ) + + +def test_pdg_router_hardens_as_temperature_shrinks(): + """As tau -> 0 the soft gate should converge to a one-hot at the argmax.""" + router = PdgRouter(n_experts=4, pdg_vocab=3, temperature=1e-4) + cond_cont, cond_cat = _cond(16) + g = router.gate(cond_cont, cond_cat) + top1 = router.top1(cond_cont, cond_cat) + onehot = torch.nn.functional.one_hot(top1, num_classes=4).float() + torch.testing.assert_close(g, onehot, atol=1e-3, rtol=0) + + +def test_pdg_router_balance_loss_is_nonnegative_scalar(): + router = PdgRouter(n_experts=4, pdg_vocab=3) + cond_cont, cond_cat = _cond(16) + loss = router.balance_loss(cond_cont, cond_cat) + assert loss.shape == () + assert loss.item() >= 0.0 + + +def test_pdg_router_classify_loss_defaults_to_zero(): + """PDG is already known at gate time (unlike ProcessRouter's process + label), so no supervision is needed — falls back to Router's default.""" + router = PdgRouter(n_experts=4, pdg_vocab=3) + cond_cont, cond_cat = _cond(16) + labels = torch.randint(0, 4, (16,)) + loss = router.classify_loss(cond_cont, cond_cat, labels) + assert loss.shape == () + assert loss.item() == 0.0 + + +def test_pdg_router_only_reads_pdg_column(): + """Gate must depend on cond_cat[:, 0] (pdg) only, not cond_cont or material.""" + router = PdgRouter(n_experts=4, pdg_vocab=3) + cond_cont, cond_cat = _cond(16) + g_before = router.gate(cond_cont, cond_cat) + + cond_cont_perturbed = torch.randn_like(cond_cont) + cond_cat_diff_mat = cond_cat.clone() + cond_cat_diff_mat[:, 1] = (cond_cat_diff_mat[:, 1] + 1) % 2 + g_after = router.gate(cond_cont_perturbed, cond_cat_diff_mat) + + torch.testing.assert_close(g_before, g_after, atol=1e-6, rtol=0) + + +def test_build_router_pdg_type_uses_pdg_vocab(): + router = build_router("pdg", 4, pdg_vocab=5, mat_vocab=3, emb_dim=8) + assert isinstance(router, PdgRouter) + assert router.pdg_emb.num_embeddings == 5 + + +def test_build_models_routed_with_pdg_router(): + model_config = dict( + pdg_vocab=4, + mat_vocab=2, + emb_dim=16, + dropout=0.1, + k_max=K_MAX, + expert_hidden_dim=16, + expert_n_blocks=2, + router={ + "enabled": True, + "type": "pdg", + "n_experts": 3, + }, + ) + stage1, sec_decoder = build_models(model_config) + assert isinstance(stage1, RoutedDenoisingMLP) + assert isinstance(stage1.router, PdgRouter) + assert len(stage1.experts) == 3 + assert stage1.router.pdg_emb.num_embeddings == 4 + + # ── ProcessRouter ──────────────────────────────────────────────────────────── -- 2.39.5 From 550dc679c7122de37980845212611e8970c92959 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Wed, 15 Jul 2026 11:10:04 +0200 Subject: [PATCH 08/11] Stream giant rollout output instead of buffering the whole run MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _Recorder previously accumulated every generated step across all events/ tracks/steps in Python lists, materialised once at the end and written via a single pq.write_table — memory scaled with n_events * max_steps * avg_tracks_per_event. rollout() now takes an optional on_chunk callback that streams each non-empty batch immediately (fixed per-key dtypes via _RECORD_DTYPES keep every chunk's table schema identical, which pq.ParquetWriter requires across writes); giant rollout wires this to an incrementally-written ParquetWriter, mirroring the row-group streaming giant predict already does on its input side. Without on_chunk, rollout() keeps its old buffered return for existing callers/tests. Co-Authored-By: Claude Sonnet 5 --- giant/cli.py | 44 ++++++++++------ giant/rollout.py | 113 ++++++++++++++++++++++++++++++++++++------ tests/test_rollout.py | 92 ++++++++++++++++++++++++++++++++++ 3 files changed, 218 insertions(+), 31 deletions(-) diff --git a/giant/cli.py b/giant/cli.py index 39eace5..f36e074 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -804,7 +804,29 @@ def rollout( seeds = _seed_from_data(files, n_events) typer.echo(f"seeded {len(seeds['event_id']):,} shower(s)") - records = run_rollout( + out, dataset_path, pred_uuid = _resolve_prediction_output(data, out) + out.parent.mkdir(parents=True, exist_ok=True) + + # Written incrementally as each batch of steps is produced, rather than + # buffering the whole run (which scales with n_events * max_steps * + # avg_tracks_per_event) — mirrors the row-group streaming `giant predict` + # already does on its input side. + writer: pq.ParquetWriter | None = None + + def _write_chunk(row: dict[str, np.ndarray]) -> None: + nonlocal writer + table = pa.table(row) + if writer is None: + table = table.replace_schema_metadata( + { + PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE, + PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION, + } + ) + writer = pq.ParquetWriter(out, table.schema) + writer.write_table(table) + + summary = run_rollout( model, sec_decoder, oracle, @@ -820,18 +842,10 @@ def rollout( device=_device, max_tracks_per_event=max_tracks_per_event, escape_threshold=escape_threshold, + on_chunk=_write_chunk, ) - - out, dataset_path, pred_uuid = _resolve_prediction_output(data, out) - out.parent.mkdir(parents=True, exist_ok=True) - - table = pa.table(records).replace_schema_metadata( - { - PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE, - PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION, - } - ) - pq.write_table(table, out) + if writer is not None: + writer.close() ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset_path) ref = yaml.safe_load(ref_path.read_text()) @@ -848,10 +862,8 @@ def rollout( ) ref_path.write_text(yaml.dump(ref, default_flow_style=False, sort_keys=False)) - n_rows = len(records["event_id"]) - reasons = Counter(r for r in records["termination_reason"].tolist() if r) - typer.echo(f"wrote {n_rows:,} step rows → {out}") - typer.echo(f"terminations: {dict(reasons)}") + typer.echo(f"wrote {summary['n_rows']:,} step rows → {out}") + typer.echo(f"terminations: {summary['termination_reason_counts']}") typer.echo(f"reference: {ref_path}") diff --git a/giant/rollout.py b/giant/rollout.py index 109faa0..228c3c5 100644 --- a/giant/rollout.py +++ b/giant/rollout.py @@ -16,6 +16,9 @@ treated as detector leakage and not deposited. from __future__ import annotations +from collections import Counter +from typing import Callable + import numpy as np import torch @@ -89,32 +92,95 @@ def _concat_frontiers(parts: list[dict[str, np.ndarray]]) -> dict[str, np.ndarra return {k: np.concatenate([p[k] for p in parts], axis=0) for k in parts[0]} -class _Recorder: - """Accumulates per-step rows into column lists, materialised at the end.""" +# Fixed per-key dtype, so every chunk table has an identical schema — needed +# for `giant rollout --on_chunk` to stream chunks straight into one +# pq.ParquetWriter (which requires matching schemas across writes), and a +# side benefit even in the buffered path since np.concatenate would otherwise +# silently upcast any stray int32/float32 chunk to the majority dtype. +_RECORD_DTYPES: dict[str, type] = { + "event_id": np.int64, + "track_id": np.int64, + "parent_id": np.int64, + "generation": np.int64, + "step_no": np.int64, + "pdg": np.int64, + "pre_x": np.float64, + "pre_y": np.float64, + "pre_z": np.float64, + "pre_E": np.float64, + "pre_dx": np.float64, + "pre_dy": np.float64, + "pre_dz": np.float64, + "post_x": np.float64, + "post_y": np.float64, + "post_z": np.float64, + "post_E": np.float64, + "post_dx": np.float64, + "post_dy": np.float64, + "post_dz": np.float64, + "edep": np.float64, + "step_length": np.float64, + "material": object, + "layer_id": np.int64, + "n_sec_pred": np.int64, + "termination_reason": object, +} - def __init__(self) -> None: - self._cols: dict[str, list] = {k: [] for k in _RECORD_KEYS} + +class _Recorder: + """Accumulates per-step rows into column lists, materialised at the end — + or, when `sink` is given, streams each non-empty chunk to it immediately + instead, keeping only row-count / termination-reason summaries in memory. + + The streaming path is what lets `giant rollout` write output incrementally + (see `rollout`'s `on_chunk` parameter): without it, a whole run's steps — + scaling with `n_events * max_steps * avg_tracks_per_event` — would sit in + RAM until the very end. + """ + + def __init__( + self, sink: Callable[[dict[str, np.ndarray]], None] | None = None + ) -> None: + self._sink = sink + self._cols: dict[str, list] | None = ( + None if sink is not None else {k: [] for k in _RECORD_KEYS} + ) + self.n_rows = 0 + self.termination_reason_counts: Counter[str] = Counter() def add(self, **cols) -> None: n = len(cols["event_id"]) if n == 0: return - for k in _RECORD_KEYS: - v = cols[k] - self._cols[k].append(np.asarray(v).reshape(n)) + row = { + k: np.asarray(cols[k], dtype=_RECORD_DTYPES[k]).reshape(n) + for k in _RECORD_KEYS + } + self.n_rows += n + reasons = row["termination_reason"] + nonempty = reasons[reasons != ""] + if len(nonempty): + for r, c in zip(*np.unique(nonempty, return_counts=True)): + self.termination_reason_counts[str(r)] += int(c) + + if self._sink is not None: + self._sink(row) + else: + assert self._cols is not None + for k in _RECORD_KEYS: + self._cols[k].append(row[k]) def to_dict(self) -> dict[str, np.ndarray]: + assert self._cols is not None, ( + "to_dict() is unavailable when streaming to a sink — use " + "n_rows/termination_reason_counts instead" + ) out = {} for k, chunks in self._cols.items(): if chunks: out[k] = np.concatenate(chunks, axis=0) else: - out[k] = np.empty( - 0, - dtype=object - if k in ("material", "termination_reason") - else np.float64, - ) + out[k] = np.empty(0, dtype=_RECORD_DTYPES[k]) return out @@ -209,8 +275,20 @@ def rollout( device: torch.device | None = None, max_tracks_per_event: int | None = None, escape_threshold: float | None = None, + on_chunk: Callable[[dict[str, np.ndarray]], None] | None = None, ) -> dict[str, np.ndarray]: - """Run showers to completion; return a step-record dict (see _RECORD_KEYS).""" + """Run showers to completion. + + By default, returns a step-record dict (see _RECORD_KEYS) with the whole + run's rows materialised in memory. + + If `on_chunk` is given, every non-empty batch of rows is streamed to it as + soon as it's produced instead — no per-run buffering — and this returns a + small summary dict instead: `{"n_rows": int, "termination_reason_counts": + dict[str, int]}`. Use this for large `--n-events`/`--max-steps` runs, + where the full record set would otherwise scale with + `n_events * max_steps * avg_tracks_per_event`. + """ device = device or torch.device("cpu") stage1_model.eval() sec_decoder.eval() @@ -227,7 +305,7 @@ def rollout( seeds["pre_E"], seeds["pre_dir"], ) - rec = _Recorder() + rec = _Recorder(sink=on_chunk) while len(frontier["event_id"]) > 0: next_parts: list[dict[str, np.ndarray]] = [] @@ -257,6 +335,11 @@ def rollout( ) frontier = _concat_frontiers(next_parts) + if on_chunk is not None: + return { + "n_rows": rec.n_rows, + "termination_reason_counts": dict(rec.termination_reason_counts), + } return rec.to_dict() diff --git a/tests/test_rollout.py b/tests/test_rollout.py index f744755..4d7c747 100644 --- a/tests/test_rollout.py +++ b/tests/test_rollout.py @@ -1,5 +1,6 @@ """Tests for the autoregressive shower rollout driver.""" +from collections import Counter from pathlib import Path from unittest.mock import patch @@ -161,3 +162,94 @@ def test_max_tracks_cap_conserves_energy(): leak = rec["pre_E"][m & (rec["termination_reason"] == TERM_ESCAPED)].sum() assert dep + leak == pytest.approx(seeds["pre_E"][i], rel=1e-4) assert len(np.unique(rec["track_id"][m])) <= 3 + + +# ── Streaming output (on_chunk) ────────────────────────────────────────────── + + +def _run_streaming(on_chunk, **kwargs): + torch.manual_seed(0) + np.random.seed(0) + s1, s2 = _models() + cond, tgt = _norms() + seeds = kwargs.pop("seeds", None) or _seeds() + return rollout( + s1, + s2, + _oracle(), + seeds, + cond, + tgt, + PDG_MAP, + MAT_MAP, + energy_cutoff=kwargs.pop("energy_cutoff", 1.0), + max_steps=kwargs.pop("max_steps", 30), + steps=4, + batch_size=128, + max_tracks_per_event=kwargs.pop("max_tracks_per_event", 300), + escape_threshold=kwargs.pop("escape_threshold", 1e9), + on_chunk=on_chunk, + ) + + +def test_on_chunk_receives_every_row_exactly_once(): + """Concatenating the streamed chunks must reproduce the buffered result.""" + from giant.rollout import _RECORD_KEYS + + buffered = _run() + + chunks: list[dict[str, np.ndarray]] = [] + summary = _run_streaming(chunks.append) + + streamed = {k: np.concatenate([c[k] for c in chunks]) for k in _RECORD_KEYS} + assert summary["n_rows"] == len(buffered["event_id"]) + assert len(streamed["event_id"]) == len(buffered["event_id"]) + for k in _RECORD_KEYS: + np.testing.assert_array_equal(streamed[k], buffered[k]) + + +def test_on_chunk_summary_termination_reason_counts_match_buffered(): + buffered = _run() + summary = _run_streaming(lambda row: None) + + expected = Counter(r for r in buffered["termination_reason"].tolist() if r) + assert summary["termination_reason_counts"] == dict(expected) + + +def test_on_chunk_never_buffers_full_records(): + """Streaming mode must not accumulate rows for later to_dict() retrieval.""" + from giant.rollout import _Recorder + + rec = _Recorder(sink=lambda row: None) + rec.add( + event_id=np.array([0]), + track_id=np.array([0]), + parent_id=np.array([-1]), + generation=np.array([0]), + step_no=np.array([0]), + pdg=np.array([11]), + pre_x=np.array([0.0]), + pre_y=np.array([0.0]), + pre_z=np.array([0.0]), + pre_E=np.array([1.0]), + pre_dx=np.array([0.0]), + pre_dy=np.array([0.0]), + pre_dz=np.array([1.0]), + post_x=np.array([0.0]), + post_y=np.array([0.0]), + post_z=np.array([1.0]), + post_E=np.array([0.0]), + post_dx=np.array([0.0]), + post_dy=np.array([0.0]), + post_dz=np.array([1.0]), + edep=np.array([1.0]), + step_length=np.array([1.0]), + material=np.array(["G4_AIR"], dtype=object), + layer_id=np.array([0]), + n_sec_pred=np.array([0]), + termination_reason=np.array(["natural_end"], dtype=object), + ) + assert rec.n_rows == 1 + assert rec.termination_reason_counts == {"natural_end": 1} + with pytest.raises(AssertionError): + rec.to_dict() -- 2.39.5 From f3fec8bcb38f7c12717abeffc5ecc3a51927b71a Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Wed, 15 Jul 2026 13:53:21 +0200 Subject: [PATCH 09/11] Add ComposedRouter for multi-axis MoE gating Route on several independent axes at once (e.g. energy x pdg), each with its own expert count and hyperparameters. The joint gate is the outer product of per-axis softmax gates, so it stays a partition of unity and top1/balance_loss factor per-axis. Config uses flat axis{i}_{field} keys in model.router (TOML/CLI friendly), also settable via repeatable --router-axis flags. Co-Authored-By: Claude Opus 4.8 --- giant/cli.py | 51 +++++++++- giant/config.py | 7 ++ giant/model/network.py | 147 +++++++++++++++++++++++---- tests/test_router.py | 220 +++++++++++++++++++++++++++++++++++++++++ 4 files changed, 405 insertions(+), 20 deletions(-) diff --git a/giant/cli.py b/giant/cli.py index f36e074..b7a3a5f 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -63,6 +63,43 @@ def _batch_size_estimate_dims(model_cfg: dict) -> tuple[int, int]: return model_cfg["hidden_dim"], model_cfg["n_blocks"] +def _coerce_scalar(value: str) -> object: + """Best-effort str -> bool/int/float, else leave as str. + + CLI flag values always arrive as strings; router kwargs like + `n_experts` (int) or `temperature` (float) need to come out typed the + same way a TOML file's native types would, since they're merged into + the same `model.router` dict as file-sourced config. + """ + if value.lower() in ("true", "false"): + return value.lower() == "true" + try: + return int(value) + except ValueError: + pass + try: + return float(value) + except ValueError: + pass + return value + + +def _parse_router_axis_flags(specs: list[str]) -> dict[str, object]: + """Parse repeated `--router-axis "type:key=val,key=val"` flags into + `axis{i}_{field}` flat keys (see `_parse_composed_axes` in + giant.model.network), indexed by flag order — the Nth `--router-axis` + becomes axis N. + """ + out: dict[str, object] = {} + for i, spec in enumerate(specs): + axis_type, _, rest = spec.partition(":") + out[f"axis{i}_type"] = axis_type + for pair in filter(None, rest.split(",")): + key, _, val = pair.partition("=") + out[f"axis{i}_{key}"] = _coerce_scalar(val) + return out + + _CEPH_PREDICTIONS = Path("/ceph/lbogner/geant_steps/predictions") @@ -175,6 +212,16 @@ def train( n_experts: Annotated[ Optional[int], typer.Option("--n-experts", help="Number of routed experts") ] = None, + router_axis: Annotated[ + Optional[list[str]], + typer.Option( + "--router-axis", + help="Composed-router axis spec 'type:key=val,key=val' (repeatable; " + "Nth flag = axis N). Use with --router-type composed instead of " + "--n-experts, e.g. --router-axis 'energy:n_experts=4' " + "--router-axis 'pdg:n_experts=3,emb_dim=8'", + ), + ] = None, val_fraction: Annotated[ Optional[float], typer.Option("--val-fraction", "-f") ] = None, @@ -264,7 +311,7 @@ def train( }.items() if v is not None } - cli_router = { + cli_router: dict[str, object] = { k: v for k, v in { "enabled": router, @@ -273,6 +320,8 @@ def train( }.items() if v is not None } + if router_axis: + cli_router.update(_parse_router_axis_flags(router_axis)) if cli_router: cli_model["router"] = cli_router cfg = gconfig.merge_cli_overrides( diff --git a/giant/config.py b/giant/config.py index 7e4f8e7..8ce67f5 100644 --- a/giant/config.py +++ b/giant/config.py @@ -43,6 +43,13 @@ DEFAULT_CONFIG: dict = { # (0.0 still trains a working router — the gate gets gradient # through the downstream flow loss like EnergyRouter's centers — # but only lambda_proc > 0 grounds it in the true `process` label) + # type = "composed" routes on multiple axes at once (e.g. energy x + # pdg), each with its own expert count/hyperparameters. Axes are + # NOT in these defaults (there's no meaningful default axis list) + # — set them as flat axis{i}_{field} keys instead of "n_experts", + # e.g. axis0_type = "energy", axis0_n_experts = 4, axis1_type = + # "pdg", axis1_n_experts = 3, axis1_emb_dim = 8. See + # giant.model.network._parse_composed_axes / `--router-axis`. }, }, } diff --git a/giant/model/network.py b/giant/model/network.py index 7a6e8ca..57d25df 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -1,5 +1,6 @@ import inspect import math +import re import torch import torch.nn as nn @@ -444,6 +445,79 @@ class ProcessRouter(Router): return F.cross_entropy(self.logits(cond_cont, cond_cat), labels) +class ComposedRouter(Router): + """Joint router over independent axes (e.g. energy x pdg), outer-product gated. + + Wraps N already-built sub-routers, each free to have its own + `n_experts` and hyperparameters (an `EnergyRouter(n_experts=4, ...)` + composed with a `PdgRouter(n_experts=3, ...)` needs no axis to match + the other's expert count). The joint gate is the outer product of the + per-axis softmax gates, flattened to `(B, prod(n_experts_i))` — still a + partition of unity, since each factor is one. Because the axes are + routed independently, the joint argmax factors into the per-axis + argmaxes, so `top1` (inherited from `Router`) costs no more than + routing each axis alone despite the multiplicative expert count; the + same is true of `balance_loss` (inherited, computed on the flattened + joint gate — now one importance term per *joint* expert cell). + + Not registered in `ROUTER_REGISTRY` / buildable via `build_router`, + since those assume one `n_experts` int shared by a single router type; + use `build_composed_router` instead, which resolves a list of per-axis + specs (each independently typed and sized) through `build_router`. + """ + + def __init__(self, routers: list[Router]) -> None: + if not routers: + raise ValueError("ComposedRouter needs at least one sub-router") + n_experts = 1 + for r in routers: + n_experts *= r.n_experts + super().__init__(n_experts) + self.routers = nn.ModuleList(routers) + + def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + joint = self.routers[0].gate(cond_cont, cond_cat) # (B, n_0) + for router in self.routers[1:]: + g = router.gate(cond_cont, cond_cat) # (B, n_i) + joint = (joint.unsqueeze(-1) * g.unsqueeze(1)).flatten( + 1 + ) # (B, prod so far) + return joint + + def classify_loss( + self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor + ) -> torch.Tensor: + """Sum of each sub-router's own classify_loss (0 for unsupervised axes).""" + total = torch.zeros((), device=cond_cont.device) + for router in self.routers: + total = total + router.classify_loss(cond_cont, cond_cat, labels) + return total + + +def build_composed_router(specs: list[dict], **shared_kwargs) -> ComposedRouter: + """Build a `ComposedRouter` from a list of per-axis router specs. + + Each spec is a `{"type": ..., "n_experts": ..., ...per-axis kwargs}` + dict resolved through `build_router` exactly like a single-axis router + config, so axes can differ in both expert count and hyperparameters + (e.g. an energy axis's `temperature` vs a pdg axis's `emb_dim`). + `shared_kwargs` (`pdg_vocab`, `mat_vocab`, ...) are merged under each + spec, with the spec's own keys taking precedence. + """ + routers = [ + build_router( + spec["type"], + spec["n_experts"], + **{ + **shared_kwargs, + **{k: v for k, v in spec.items() if k not in ("type", "n_experts")}, + }, + ) + for spec in specs + ] + return ComposedRouter(routers) + + class ExpertTrunk(nn.Module): """One small expert: `input_proj -> ResBlock stack -> out_proj`. @@ -669,6 +743,54 @@ _SEC_DECODER_MODEL_KEYS = { } +_AXIS_KEY_RE = re.compile(r"^axis(\d+)_(.+)$") + + +def _parse_composed_axes(router_cfg: dict) -> list[dict]: + """Regroup `axis{i}_{field}` flat keys into a list of per-axis spec dicts. + + Flat keys (rather than a nested list-of-dicts) keep composed-router + config expressible in the same one-level-of-nesting TOML/CLI shape as + every other router option (`model.router` stays a flat table of + scalars) — e.g. `axis0_type = "energy"`, `axis0_n_experts = 4`, + `axis1_type = "pdg"`, `axis1_n_experts = 3`, `axis1_emb_dim = 8`. + Axis indices must be contiguous from 0; order follows the index, not + dict insertion order (TOML/CLI merging doesn't preserve it reliably). + """ + axes: dict[int, dict] = {} + for key, value in router_cfg.items(): + m = _AXIS_KEY_RE.match(key) + if m is None: + continue + idx, field = int(m.group(1)), m.group(2) + axes.setdefault(idx, {})[field] = value + missing = set(range(len(axes))) - axes.keys() + if missing: + raise ValueError(f"composed router config has gaps at axis indices {missing}") + return [axes[i] for i in range(len(axes))] + + +def _build_router_from_cfg(router_cfg: dict, pdg_vocab: int, mat_vocab: int) -> Router: + """Resolve one `model.router` config into a `Router`, single-axis or composed. + + `router_cfg["type"] == "composed"` reads `axis{i}_{field}` flat keys + (see `_parse_composed_axes`) instead of a single `type`/`n_experts` pair. + """ + shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab) + if router_cfg["type"] == "composed": + return build_composed_router(_parse_composed_axes(router_cfg), **shared_vocab) + router_kwargs = { + k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts") + } + # Not every router needs these (EnergyRouter doesn't declare them, so + # build_router's kwarg filtering drops them silently) but ProcessRouter + # needs its own pdg/material embeddings sized to match the checkpoint's + # vocab, same as the trunk's ConditionEncoder. + router_kwargs.setdefault("pdg_vocab", pdg_vocab) + router_kwargs.setdefault("mat_vocab", mat_vocab) + return build_router(router_cfg["type"], router_cfg["n_experts"], **router_kwargs) + + def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]: """Construct (stage1, sec_decoder) from a persisted/CLI model_config dict. @@ -679,36 +801,23 @@ def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]: """ router_cfg = model_config.get("router") if router_cfg and router_cfg.get("enabled"): - router_kwargs = { - k: v - for k, v in router_cfg.items() - if k not in ("enabled", "type", "n_experts") - } - # Not every router needs these (EnergyRouter doesn't declare them, so - # build_router's kwarg filtering drops them silently) but - # ProcessRouter needs its own pdg/material embeddings sized to match - # the checkpoint's vocab, same as the trunk's ConditionEncoder. - router_kwargs.setdefault("pdg_vocab", model_config["pdg_vocab"]) - router_kwargs.setdefault("mat_vocab", model_config["mat_vocab"]) + pdg_vocab = model_config["pdg_vocab"] + mat_vocab = model_config["mat_vocab"] shared = dict( - pdg_vocab=model_config["pdg_vocab"], - mat_vocab=model_config["mat_vocab"], + pdg_vocab=pdg_vocab, + mat_vocab=mat_vocab, expert_hidden_dim=model_config.get("expert_hidden_dim", 128), expert_n_blocks=model_config.get("expert_n_blocks", 3), emb_dim=model_config.get("emb_dim", EMB_DIM), dropout=model_config.get("dropout", 0.1), ) stage1 = RoutedDenoisingMLP( - router=build_router( - router_cfg["type"], router_cfg["n_experts"], **router_kwargs - ), + router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab), k_max=model_config.get("k_max", K_MAX), **shared, ) sec_decoder = RoutedSecondaryDecoder( - router=build_router( - router_cfg["type"], router_cfg["n_experts"], **router_kwargs - ), + router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab), **shared, ) return stage1, sec_decoder diff --git a/tests/test_router.py b/tests/test_router.py index eef1c3f..31f6b54 100644 --- a/tests/test_router.py +++ b/tests/test_router.py @@ -4,6 +4,7 @@ import torch from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM from giant.model.network import ( + ComposedRouter, DenoisingMLP, EnergyRouter, PdgRouter, @@ -12,6 +13,7 @@ from giant.model.network import ( RoutedDenoisingMLP, RoutedSecondaryDecoder, SecondaryDecoder, + build_composed_router, build_models, build_router, ) @@ -289,6 +291,224 @@ def test_build_models_routed_with_process_router(): assert stage1.router.mat_emb.num_embeddings == 2 +# ── ComposedRouter ─────────────────────────────────────────────────────────── + + +def test_composed_router_n_experts_is_product(): + router = ComposedRouter( + [EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)] + ) + assert router.n_experts == 12 + + +def test_composed_router_gate_partition_of_unity(): + router = ComposedRouter( + [EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)] + ) + cond_cont, cond_cat = _cond(16, pdg=5) + g = router.gate(cond_cont, cond_cat) + assert g.shape == (16, 12) + torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0) + + +def test_composed_router_gate_is_outer_product_of_sub_gates(): + energy_router = EnergyRouter(n_experts=4) + pdg_router = PdgRouter(n_experts=3, pdg_vocab=5) + router = ComposedRouter([energy_router, pdg_router]) + cond_cont, cond_cat = _cond(16, pdg=5) + + g_energy = energy_router.gate(cond_cont, cond_cat) # (16, 4) + g_pdg = pdg_router.gate(cond_cont, cond_cat) # (16, 3) + expected = (g_energy.unsqueeze(-1) * g_pdg.unsqueeze(1)).flatten(1) # (16, 12) + + torch.testing.assert_close(router.gate(cond_cont, cond_cat), expected) + + +def test_composed_router_top1_factors_into_per_axis_argmax(): + """Joint argmax over the outer product must equal the pair of per-axis + argmaxes, flattened with the same row-major index convention as gate().""" + energy_router = EnergyRouter(n_experts=4) + pdg_router = PdgRouter(n_experts=3, pdg_vocab=5) + router = ComposedRouter([energy_router, pdg_router]) + cond_cont, cond_cat = _cond(16, pdg=5) + + joint_idx = router.top1(cond_cont, cond_cat) + energy_idx = energy_router.top1(cond_cont, cond_cat) + pdg_idx = pdg_router.top1(cond_cont, cond_cat) + expected = energy_idx * pdg_router.n_experts + pdg_idx + + assert torch.equal(joint_idx, expected) + + +def test_composed_router_supports_different_expert_counts_per_axis(): + router = ComposedRouter( + [EnergyRouter(n_experts=5), PdgRouter(n_experts=2, pdg_vocab=5)] + ) + assert router.n_experts == 10 + cond_cont, cond_cat = _cond(8, pdg=5) + assert router.gate(cond_cont, cond_cat).shape == (8, 10) + + +def test_composed_router_classify_loss_sums_sub_router_losses(): + """energy/pdg both default to zero, so the composed loss should too.""" + router = ComposedRouter( + [EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)] + ) + cond_cont, cond_cat = _cond(16, pdg=5) + labels = torch.randint(0, 4, (16,)) + loss = router.classify_loss(cond_cont, cond_cat, labels) + assert loss.shape == () + assert loss.item() == 0.0 + + +def test_composed_router_rejects_empty_router_list(): + try: + ComposedRouter([]) + except ValueError: + return + raise AssertionError("expected ValueError for empty router list") + + +def test_composed_router_not_in_registry(): + assert "composed" not in ROUTER_REGISTRY + + +# ── _parse_composed_axes (axis{i}_{field} flat-key config convention) ─────── + + +def test_parse_composed_axes_groups_indexed_keys(): + from giant.model.network import _parse_composed_axes + + router_cfg = { + "enabled": True, + "type": "composed", + "axis0_type": "energy", + "axis0_n_experts": 4, + "axis0_temperature": 0.3, + "axis1_type": "pdg", + "axis1_n_experts": 3, + "axis1_emb_dim": 6, + } + axes = _parse_composed_axes(router_cfg) + assert axes == [ + {"type": "energy", "n_experts": 4, "temperature": 0.3}, + {"type": "pdg", "n_experts": 3, "emb_dim": 6}, + ] + + +def test_parse_composed_axes_ignores_unrelated_keys(): + from giant.model.network import _parse_composed_axes + + router_cfg = { + "enabled": True, + "type": "composed", + "lambda_balance": 0.0, + "axis0_type": "energy", + "axis0_n_experts": 4, + } + axes = _parse_composed_axes(router_cfg) + assert axes == [{"type": "energy", "n_experts": 4}] + + +def test_parse_composed_axes_raises_on_index_gap(): + from giant.model.network import _parse_composed_axes + + router_cfg = { + "type": "composed", + "axis0_type": "energy", + "axis0_n_experts": 4, + # axis1 missing entirely + "axis2_type": "pdg", + "axis2_n_experts": 3, + } + try: + _parse_composed_axes(router_cfg) + except ValueError: + return + raise AssertionError("expected ValueError for a gap in axis indices") + + +def test_build_composed_router_resolves_per_axis_specs(): + router = build_composed_router( + [ + {"type": "energy", "n_experts": 4, "temperature": 0.3}, + {"type": "pdg", "n_experts": 3, "emb_dim": 6}, + ], + pdg_vocab=5, + mat_vocab=2, + ) + assert isinstance(router, ComposedRouter) + assert router.n_experts == 12 + energy_router, pdg_router = router.routers + assert isinstance(energy_router, EnergyRouter) + assert energy_router.temperature == 0.3 + assert isinstance(pdg_router, PdgRouter) + assert pdg_router.pdg_emb.num_embeddings == 5 + assert pdg_router.pdg_emb.embedding_dim == 6 + + +def test_build_models_routed_with_composed_router(): + model_config = dict( + pdg_vocab=5, + mat_vocab=2, + emb_dim=16, + dropout=0.1, + k_max=K_MAX, + expert_hidden_dim=16, + expert_n_blocks=2, + router={ + "enabled": True, + "type": "composed", + "axis0_type": "energy", + "axis0_n_experts": 4, + "axis1_type": "pdg", + "axis1_n_experts": 3, + }, + ) + stage1, sec_decoder = build_models(model_config) + assert isinstance(stage1, RoutedDenoisingMLP) + assert isinstance(stage1.router, ComposedRouter) + assert len(stage1.experts) == 12 + assert len(sec_decoder.experts) == 12 + # stage1 and sec_decoder must not share router weights (same convention + # as the single-axis routers built by build_models). + assert stage1.router is not sec_decoder.router + + +def test_build_models_routed_pair_composed_router_is_drop_in_for_sample_flow(): + from giant.sample import sample_flow, sample_secondaries + + model_config = dict( + pdg_vocab=3, + mat_vocab=2, + emb_dim=16, + dropout=0.1, + k_max=K_MAX, + expert_hidden_dim=8, + expert_n_blocks=1, + router={ + "enabled": True, + "type": "composed", + "axis0_type": "energy", + "axis0_n_experts": 2, + "axis1_type": "pdg", + "axis1_n_experts": 2, + }, + ) + stage1, sec_decoder = build_models(model_config) + B = 5 + cond_cont, cond_cat = _cond(B, pdg=3, mat=2) + stage1_norm, n_sec_pred = sample_flow(stage1, cond_cont, cond_cat, steps=2) + assert stage1_norm.shape == (B, X_DIM) + assert n_sec_pred.shape == (B,) + + sec_cont, sec_type_emb, sec_valid = sample_secondaries( + sec_decoder, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2 + ) + assert sec_cont.shape == (B, K_MAX, 4) + assert sec_valid.shape == (B, K_MAX) + + # ── RoutedDenoisingMLP ─────────────────────────────────────────────────────── -- 2.39.5 From 24a83486ecbb4936b3322f45148d2f4edfb70bbb Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Wed, 15 Jul 2026 16:22:25 +0200 Subject: [PATCH 10/11] Scale auto batch-size estimate by MoE expert count during training Training runs the full soft mixture (every expert over the whole batch), so routed activation memory scales with the expert count; the old estimate used one expert's dims and would overshoot free VRAM by a factor of n_experts. Fold the expert count into n_blocks for the training path (inference's top-1 dispatch still just partitions the batch, so one expert's dims bound it). Co-Authored-By: Claude Opus 4.8 --- giant/cli.py | 44 ++++++++++++++++++++++++++++++++++++-------- 1 file changed, 36 insertions(+), 8 deletions(-) diff --git a/giant/cli.py b/giant/cli.py index b7a3a5f..82fe028 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -1,7 +1,9 @@ from collections import Counter from datetime import date, datetime, timezone from enum import Enum +import math from pathlib import Path +import re from typing import Optional import uuid as uuid_mod @@ -47,19 +49,43 @@ from giant.sample import sample_flow, sample_secondaries, snap_type_to_pdg_idx app = typer.Typer(no_args_is_help=True) -def _batch_size_estimate_dims(model_cfg: dict) -> tuple[int, int]: +def _router_total_experts(router_cfg: dict) -> int: + """Total expert count for a router config, single-axis or composed. + + A composed router runs one expert per *joint* cell, so its count is the + product of the per-axis `axis{i}_n_experts` (mirrors + `ComposedRouter.__init__` in giant.model.network); a single-axis router + just reports its own `n_experts`. + """ + if router_cfg.get("type") == "composed": + axis_counts = { + m.group(1): int(v) + for k, v in router_cfg.items() + if (m := re.match(r"^axis(\d+)_n_experts$", k)) + } + return math.prod(axis_counts.values()) if axis_counts else 1 + return int(router_cfg.get("n_experts", 1)) + + +def _batch_size_estimate_dims(model_cfg: dict, training: bool) -> tuple[int, int]: """Pick the (hidden_dim, n_blocks) that dominate per-call activation memory. Routed models spend their FLOPs in the (smaller) expert trunks, not the monolith's hidden_dim/n_blocks, so estimate_batch_size needs the expert - dims instead when routing is enabled. + dims instead when routing is enabled. Training runs the full soft mixture + (every expert on the whole batch), so its activation memory scales with + the expert count; inference does top-1 dispatch (each row hits one + expert), so the batch just partitions across experts and one expert's + dims already bound it. estimate_batch_size scales memory linearly with + hidden_dim * n_blocks, so the training multiplier folds into n_blocks. """ router_cfg = model_cfg.get("router") if router_cfg and router_cfg.get("enabled"): - return ( - model_cfg.get("expert_hidden_dim", 128), - model_cfg.get("expert_n_blocks", 3), - ) + hidden_dim = model_cfg.get("expert_hidden_dim", 128) + n_blocks = model_cfg.get("expert_n_blocks", 3) + if training: + n_blocks *= _router_total_experts(router_cfg) + return hidden_dim, n_blocks return model_cfg["hidden_dim"], model_cfg["n_blocks"] @@ -332,7 +358,7 @@ def train( _device = torch.device(device) if device else gconfig.auto_device() if batch_size_auto: - est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(m) + est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(m, training=True) try: t["batch_size"] = gconfig.estimate_batch_size( est_hidden_dim, est_n_blocks, _device @@ -463,7 +489,9 @@ def predict( model_cfg = ckpt["model_config"] if batch_size_auto: - est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(model_cfg) + est_hidden_dim, est_n_blocks = _batch_size_estimate_dims( + model_cfg, training=False + ) try: batch_size_value = gconfig.estimate_batch_size( est_hidden_dim, -- 2.39.5 From 5a0e98c5d1940da0b0025ff00522836098471780 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Fri, 17 Jul 2026 11:02:12 +0200 Subject: [PATCH 11/11] Add EMA weights, weight decay, step-based LR schedule, and grad-norm logging to training Gives flow-matching sampling a cleaner EMA shadow copy to draw from (--ema-decay, --weights raw|ema in predict/rollout), fixes the LR warmup/cosine schedule stepping once per epoch even when an epoch is tens of thousands of steps, and caps the per-epoch val-loss pass (--max-val-batches) so large val sets don't dominate epoch time. Co-Authored-By: Claude Sonnet 5 --- giant/cli.py | 85 +++++++++++++++++++++++++++++++++++++---- giant/config.py | 5 +++ giant/pipeline.py | 3 ++ giant/train.py | 96 ++++++++++++++++++++++++++++++++++++++++------- 4 files changed, 168 insertions(+), 21 deletions(-) diff --git a/giant/cli.py b/giant/cli.py index 82fe028..935bc98 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -183,6 +183,39 @@ class Coord(str, Enum): local = "local" +class Weights(str, Enum): + raw = "raw" + ema = "ema" + + +def _load_model_weights( + model: torch.nn.Module, + sec_decoder: torch.nn.Module, + ckpt: dict, + weights: "Weights", + checkpoint_path: Path, +) -> None: + """Load either the raw or EMA state dicts from a training checkpoint. + + EMA weights (giant.train's shadow copy, see --ema-decay) only exist in + checkpoints written after that feature landed, so `ema` fails loudly + rather than silently falling back to raw weights a caller didn't ask for. + """ + if weights == Weights.raw: + model_key, sec_key = "model", "sec_decoder" + else: + model_key, sec_key = "model_ema", "sec_decoder_ema" + if model_key not in ckpt or sec_key not in ckpt: + typer.echo( + f"error: {checkpoint_path} has no EMA weights (trained before " + "--ema-decay, or with --ema-decay 0) — use --weights raw", + err=True, + ) + raise typer.Exit(1) + model.load_state_dict(ckpt[model_key]) + sec_decoder.load_state_dict(ckpt[sec_key]) + + @app.command() def train( data: Annotated[ @@ -209,6 +242,18 @@ def train( ), ] = None, lr: Annotated[Optional[float], typer.Option("--lr", "-l")] = None, + weight_decay: Annotated[ + Optional[float], + typer.Option("--weight-decay", "-W", help="AdamW weight decay (default: 0.01)"), + ] = None, + ema_decay: Annotated[ + Optional[float], + typer.Option( + "--ema-decay", + help="EMA decay for a shadow copy of the model weights, saved " + "alongside the raw weights in checkpoints (0 disables; default: 0.9999)", + ), + ] = None, warmup_epochs: Annotated[ Optional[int], typer.Option("--warmup-epochs", "-w") ] = None, @@ -272,6 +317,14 @@ def train( "(ignored in ddpm mode, which always runs the full schedule)", ), ] = None, + max_val_batches: Annotated[ + Optional[int], + typer.Option( + "--max-val-batches", + help="Cap the per-epoch val-loss pass to N batches (0 = full " + "val set every epoch; default: 200)", + ), + ] = None, shuffle_buffer: Annotated[ int, typer.Option( @@ -318,12 +371,15 @@ def train( "epochs": epochs, "batch_size": batch_size_value, "lr": lr, + "weight_decay": weight_decay, + "ema_decay": ema_decay, "warmup_epochs": warmup_epochs, "val_fraction": val_fraction, "num_workers": num_workers, "seed": seed, "validate_every": validate_every, "validate_steps": validate_steps, + "max_val_batches": max_val_batches, }.items() if v is not None } @@ -439,6 +495,15 @@ def predict( steps: Annotated[ int, typer.Option("--steps", "-s", help="Flow matching ODE steps") ] = 10, + weights: Annotated[ + Weights, + typer.Option( + "--weights", + help="raw: the live training weights. ema: the EMA shadow copy " + "(see --ema-decay in `giant train`) — usually cleaner samples, " + "requires a checkpoint trained with EMA enabled.", + ), + ] = Weights.raw, device: Annotated[ Optional[str], typer.Option("--device", "-d", help="cpu | cuda | mps (default: auto)"), @@ -515,13 +580,11 @@ def predict( tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"]) model, sec_decoder = build_models(model_cfg) - model.load_state_dict(ckpt["model"]) + _load_model_weights(model, sec_decoder, ckpt, weights, checkpoint) model.to(_device).eval() - - sec_decoder.load_state_dict(ckpt["sec_decoder"]) sec_decoder.to(_device).eval() - typer.echo(f"loaded checkpoint: {checkpoint}") + typer.echo(f"loaded checkpoint: {checkpoint} (weights: {weights.value})") gconfig.warn_if_checkpoint_config_mismatch(checkpoint) # --- Output path --- @@ -810,6 +873,15 @@ def rollout( int, typer.Option("--steps", "-s", help="Flow matching ODE steps per model call"), ] = 10, + weights: Annotated[ + Weights, + typer.Option( + "--weights", + help="raw: the live training weights. ema: the EMA shadow copy " + "(see --ema-decay in `giant train`) — usually cleaner samples, " + "requires a checkpoint trained with EMA enabled.", + ), + ] = Weights.raw, batch_size: Annotated[ int, typer.Option("--batch-size", "-b", help="Tracks stepped per model forward") ] = 4096, @@ -865,11 +937,10 @@ def rollout( tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"]) model, sec_decoder = build_models(model_cfg) - model.load_state_dict(ckpt["model"]) + _load_model_weights(model, sec_decoder, ckpt, weights, checkpoint) model.to(_device).eval() - sec_decoder.load_state_dict(ckpt["sec_decoder"]) sec_decoder.to(_device).eval() - typer.echo(f"loaded checkpoint: {checkpoint}") + typer.echo(f"loaded checkpoint: {checkpoint} (weights: {weights.value})") oracle = GeometryOracle.load(geometry) typer.echo( diff --git a/giant/config.py b/giant/config.py index 8ce67f5..4834f3c 100644 --- a/giant/config.py +++ b/giant/config.py @@ -14,6 +14,11 @@ DEFAULT_CONFIG: dict = { "epochs": 100, "batch_size": 4096, "lr": 3e-4, + "weight_decay": 0.01, # AdamW default — exposed so it can be tuned + "ema_decay": 0.9999, # EMA of model weights for sampling; 0 disables + # per-epoch val loss (not the marginal/KL validate_every pass) is + # capped to this many batches; 0 = full val set every epoch + "max_val_batches": 200, "val_fraction": 0.1, "num_workers": 4, "seed": 0, diff --git a/giant/pipeline.py b/giant/pipeline.py index f160231..d6522c7 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -174,6 +174,8 @@ def run_train_job( mode=t["mode"], epochs=t["epochs"], lr=t["lr"], + weight_decay=t["weight_decay"], + ema_decay=t["ema_decay"], warmup_epochs=t["warmup_epochs"], device=device, out_dir=out_dir, @@ -189,5 +191,6 @@ def run_train_job( resume_path=resume, validate_every=t["validate_every"], validate_steps=t["validate_steps"], + max_val_batches=t["max_val_batches"], total_train_batches=total_train_batches, ) diff --git a/giant/train.py b/giant/train.py index 35f3d8c..d796956 100644 --- a/giant/train.py +++ b/giant/train.py @@ -1,3 +1,4 @@ +import copy import csv import math import os @@ -35,6 +36,7 @@ _METRICS_FIELDS = [ "val_loss_balance", "val_loss_proc", "lr", + "grad_norm", "epoch_time_s", ] @@ -103,6 +105,14 @@ def _build_sec_x1( return x1_s2.flatten(1) # (B, SEC_DIM) +@torch.no_grad() +def _update_ema( + ema_model: torch.nn.Module, model: torch.nn.Module, decay: float +) -> None: + for ema_p, p in zip(ema_model.parameters(), model.parameters()): + ema_p.mul_(decay).add_(p, alpha=1 - decay) + + def _compute_losses( stage1_model: torch.nn.Module, sec_decoder: torch.nn.Module, @@ -195,6 +205,8 @@ def train( warmup_epochs: int, device: torch.device, out_dir: str | Path, + weight_decay: float = 0.01, + ema_decay: float = 0.9999, lambda_nsec: float = 0.1, lambda_s2: float = 1.0, lambda_balance: float = 0.0, @@ -207,6 +219,7 @@ def train( resume_path: str | Path | None = None, validate_every: int = 0, validate_steps: int = 10, + max_val_batches: int = 0, total_train_batches: int = 0, ) -> None: out_dir = Path(out_dir) @@ -215,15 +228,38 @@ def train( stage1_model = stage1_model.to(device) sec_decoder = sec_decoder.to(device) - all_params = list(stage1_model.parameters()) + list(sec_decoder.parameters()) - optimizer = optim.AdamW(all_params, lr=lr) + # Flow-matching/diffusion models sample noticeably better from an EMA of + # the weights than from the raw SGD-noisy ones — buffers (e.g. the fixed + # sinusoidal-embedding freqs, or non-learned router centers) never change + # after this initial copy, so only parameters need the running average. + ema_stage1_model: torch.nn.Module | None = None + ema_sec_decoder: torch.nn.Module | None = None + if ema_decay > 0: + ema_stage1_model = copy.deepcopy(stage1_model).eval() + ema_sec_decoder = copy.deepcopy(sec_decoder).eval() + for p in ema_stage1_model.parameters(): + p.requires_grad_(False) + for p in ema_sec_decoder.parameters(): + p.requires_grad_(False) - def _lr_lambda(epoch: int) -> float: - if warmup_epochs > 0 and epoch < warmup_epochs: - return (epoch + 1) / warmup_epochs - t = epoch - warmup_epochs - T = max(epochs - warmup_epochs, 1) - return 0.5 * (1.0 + math.cos(math.pi * t / T)) + all_params = list(stage1_model.parameters()) + list(sec_decoder.parameters()) + optimizer = optim.AdamW(all_params, lr=lr, weight_decay=weight_decay) + + # Warmup/decay in units of optimizer steps rather than epochs: at large + # dataset sizes a single epoch can be tens of thousands of steps, and an + # epoch-granularity schedule would leave warmup/cosine decay unable to + # move within it. Requires an accurate `total_train_batches` (steps per + # epoch); the only caller, run_train_job, always supplies one. + steps_per_epoch = max(total_train_batches, 1) + warmup_steps = warmup_epochs * steps_per_epoch + total_steps = max(epochs * steps_per_epoch, 1) + + def _lr_lambda(step: int) -> float: + if warmup_steps > 0 and step < warmup_steps: + return (step + 1) / warmup_steps + t = step - warmup_steps + T = max(total_steps - warmup_steps, 1) + return 0.5 * (1.0 + math.cos(math.pi * min(t, T) / T)) lr_sched = optim.lr_scheduler.LambdaLR(optimizer, _lr_lambda) @@ -235,6 +271,12 @@ def train( ckpt = torch.load(resume_path, map_location=device, weights_only=False) stage1_model.load_state_dict(ckpt["model"]) sec_decoder.load_state_dict(ckpt["sec_decoder"]) + if ema_decay > 0: + assert ema_stage1_model is not None and ema_sec_decoder is not None + ema_stage1_model.load_state_dict(ckpt.get("model_ema", ckpt["model"])) + ema_sec_decoder.load_state_dict( + ckpt.get("sec_decoder_ema", ckpt["sec_decoder"]) + ) optimizer.load_state_dict(ckpt["optimizer"]) lr_sched.load_state_dict(ckpt["lr_sched"]) start_epoch = ckpt.get("epoch", 0) + 1 @@ -271,7 +313,6 @@ def train( with _GracefulShutdown() as shutdown: for epoch in range(start_epoch, epochs + 1): epoch_start = time.monotonic() - current_lr = optimizer.param_groups[0]["lr"] stage1_model.train() sec_decoder.train() train_loss_sum = 0.0 @@ -281,7 +322,10 @@ def train( train_balance_sum = 0.0 train_proc_sum = 0.0 train_n = 0 + train_batches = 0 + grad_norm_sum = 0.0 ema_loss = 0.0 + ema_grad_norm = 0.0 bar = tqdm( train_loader, desc=f" epoch {epoch:{epoch_w}d}/{epochs}", @@ -305,11 +349,17 @@ def train( ) optimizer.zero_grad() loss.backward() - torch.nn.utils.clip_grad_norm_(all_params, 1.0) + grad_norm = torch.nn.utils.clip_grad_norm_(all_params, 1.0) optimizer.step() + lr_sched.step() + if ema_decay > 0: + assert ema_stage1_model is not None and ema_sec_decoder is not None + _update_ema(ema_stage1_model, stage1_model, ema_decay) + _update_ema(ema_sec_decoder, sec_decoder, ema_decay) B = batch[0].size(0) batch_loss = loss.item() + batch_grad_norm = grad_norm.item() train_loss_sum += batch_loss * B train_s1_sum += l_s1.item() * B train_nsec_sum += l_nsec.item() * B @@ -317,10 +367,19 @@ def train( train_balance_sum += l_balance.item() * B train_proc_sum += l_proc.item() * B train_n += B + train_batches += 1 + grad_norm_sum += batch_grad_norm ema_loss = ( batch_loss if train_n == B else 0.95 * ema_loss + 0.05 * batch_loss ) - bar.set_postfix_str(f"loss={ema_loss:.4f}", refresh=False) + ema_grad_norm = ( + batch_grad_norm + if train_batches == 1 + else 0.95 * ema_grad_norm + 0.05 * batch_grad_norm + ) + bar.set_postfix_str( + f"loss={ema_loss:.4f} gnorm={ema_grad_norm:.3f}", refresh=False + ) if shutdown.requested: break @@ -330,7 +389,8 @@ def train( break train_loss = train_loss_sum / max(train_n, 1) - lr_sched.step() + train_grad_norm = grad_norm_sum / max(train_batches, 1) + current_lr = optimizer.param_groups[0]["lr"] stage1_model.eval() sec_decoder.eval() @@ -342,7 +402,9 @@ def train( val_proc_sum = 0.0 val_n = 0 with torch.no_grad(): - for batch in val_loader: + for val_batch_idx, batch in enumerate(val_loader): + if max_val_batches > 0 and val_batch_idx >= max_val_batches: + break loss, l_s1, l_nsec, l_s2, l_balance, l_proc = _compute_losses( stage1_model, sec_decoder, @@ -377,7 +439,8 @@ def train( f" bal={train_balance_sum / max(train_n, 1):.3f}" f" proc={train_proc_sum / max(train_n, 1):.3f})" f" val {val_loss:.4f}" - f" lr {current_lr:.2e} {epoch_time:.1f}s{marker}" + f" lr {current_lr:.2e} gnorm {train_grad_norm:.3f}" + f" {epoch_time:.1f}s{marker}" ) metrics_writer.writerow( { @@ -395,6 +458,7 @@ def train( "val_loss_balance": val_balance_sum / max(val_n, 1), "val_loss_proc": val_proc_sum / max(val_n, 1), "lr": current_lr, + "grad_norm": train_grad_norm, "epoch_time_s": epoch_time, } ) @@ -420,6 +484,10 @@ def train( "epoch": epoch, "best_val_loss": best_val_loss, } + if ema_decay > 0: + assert ema_stage1_model is not None and ema_sec_decoder is not None + ckpt["model_ema"] = ema_stage1_model.state_dict() + ckpt["sec_decoder_ema"] = ema_sec_decoder.state_dict() if normalizer_dict is not None: ckpt["normalizer"] = normalizer_dict if pdg_map is not None: -- 2.39.5