lars
ad0341a9d4
- Graceful shutdown (SIGINT/SIGTERM) now actually saves a checkpoint of in-progress weights before exiting mid-epoch — it previously broke out of the epoch loop before reaching the checkpoint-save block, contradicting its own printed "saving a checkpoint" message and losing all progress since the last completed epoch. Checkpoint-dict construction is factored into a shared _build_checkpoint() helper used by both the mid-epoch and end-of-epoch save paths. - WGAN LR-schedule steps_per_epoch used the wrong denominator (n_critic + 1 instead of n_critic), causing the schedule to exhaust early and LR to floor to 0 before training completed. - --critic-lr override was silently dropped on WGAN --resume (only the generator optimizer's LR was made authoritative again after load_state_dict; optimizer_d's was not). - WGAN secondary gradient-penalty forced x_hat/grad to zero for fully-masked rows (n_sec == 0, common in a shower), adding a constant ~1.0 bias into the batch-mean GP term; such rows are now excluded from the mean. - run_train_job warns (never blocks) when --num-workers exceeds ~1/4 of the machine's CPUs, per this repo's shared-portal-machine etiquette (see CLAUDE.md's Compute environment section). Each fix has a regression test. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
461 lines
16 KiB
Python
461 lines
16 KiB
Python
import os
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
from torch.utils.data import DataLoader
|
|
|
|
from giant import config
|
|
from giant.constants import (
|
|
COND_DIM,
|
|
EMB_DIM,
|
|
K_MAX,
|
|
PARTICLE_PHYS_DIM,
|
|
SEC_SLOT_DIM,
|
|
X_DIM,
|
|
)
|
|
from giant.data import setup_cache
|
|
from giant.data.loader import (
|
|
event_id_offset,
|
|
find_parquet_files,
|
|
iter_file_chunks,
|
|
build_index_maps_from_files,
|
|
build_process_map_from_files,
|
|
)
|
|
from giant.data.transforms import (
|
|
Normalizer,
|
|
build_features,
|
|
_WelfordAccumulator,
|
|
_ReservoirSampler,
|
|
sorted_membership,
|
|
)
|
|
from giant.data.dataset import make_event_split, StreamingStepsDataset
|
|
from giant.model.network import build_models, build_critics
|
|
from giant.train import train as run_training
|
|
|
|
|
|
@dataclass
|
|
class SetupStageResult:
|
|
"""Everything `run_train_job`'s pre-epoch setup stage derives from `data`.
|
|
|
|
Also returned standalone by `run_setup_stage` for callers (e.g. `dwarf
|
|
warm-cache`) that only want to populate/refresh the setup cache sidecar
|
|
without actually training a model.
|
|
"""
|
|
|
|
files: list[Path]
|
|
pdg_map: dict[int, int]
|
|
mat_map: dict[str, int]
|
|
proc_map: dict[str, int] | None
|
|
cond_norm: Normalizer
|
|
tgt_norm: Normalizer
|
|
sec_phys_norm: Normalizer
|
|
train_events: set
|
|
val_events: set
|
|
n_train_steps: int
|
|
|
|
|
|
def run_setup_stage(
|
|
data: str | Path,
|
|
val_fraction: float,
|
|
seed: int,
|
|
conditioning: str,
|
|
router_cfg: dict,
|
|
cache_setup: bool = True,
|
|
rebuild_setup_cache: bool = False,
|
|
echo=print,
|
|
) -> SetupStageResult:
|
|
"""Scan `data` for everything training needs before the epoch loop: the
|
|
train/val event split, pdg/material vocab maps, an optional process map
|
|
(`router_cfg.type == "process"`), and the Stage-1/Stage-2 normalizers.
|
|
|
|
Reads from and writes to the `giant.data.setup_cache` sidecar when
|
|
`cache_setup` is set (`rebuild_setup_cache` ignores — but still
|
|
refreshes — any existing sidecar content). `router_cfg` may be mutated
|
|
in place: an active `EnergyRouter` (`router_cfg["type"] == "energy"`)
|
|
gets its `centers_init` seeded from real data quantiles here.
|
|
"""
|
|
files = find_parquet_files(data)
|
|
echo(f"found {len(files)} parquet file(s)")
|
|
|
|
cache: setup_cache.SetupCache | None = None
|
|
if cache_setup:
|
|
if rebuild_setup_cache:
|
|
echo("setup cache: --rebuild-setup-cache given, recomputing all sections")
|
|
loaded = None
|
|
else:
|
|
loaded = setup_cache.load(data, files, echo=echo)
|
|
cache = loaded if loaded is not None else setup_cache.SetupCache.empty(files)
|
|
|
|
if cache is not None and cache.event_index is not None:
|
|
unique_ids, counts = cache.event_index
|
|
echo(f"event index: cache hit ({len(unique_ids):,} unique events)")
|
|
else:
|
|
echo("scanning event IDs …")
|
|
unique_ids, counts = setup_cache.compute_event_index_from_files(files)
|
|
if cache is not None:
|
|
cache.event_index = (unique_ids, counts)
|
|
|
|
train_events, val_events = make_event_split(
|
|
unique_ids, val_fraction=val_fraction, seed=seed
|
|
)
|
|
events_arr = np.array(sorted(train_events))
|
|
n_train_steps = setup_cache.n_train_steps_for_split(unique_ids, counts, events_arr)
|
|
echo(
|
|
f" {int(counts.sum()):,} steps | "
|
|
f"{len(train_events)} train events | "
|
|
f"{len(val_events)} val events"
|
|
)
|
|
|
|
if cache is not None and cache.vocab is not None:
|
|
pdg_map, mat_map = cache.vocab
|
|
echo(
|
|
f"vocabulary maps: cache hit ({len(pdg_map)} PDG codes, "
|
|
f"{len(mat_map)} materials)"
|
|
)
|
|
else:
|
|
echo("building vocabulary maps …")
|
|
pdg_map, mat_map = build_index_maps_from_files(files)
|
|
echo(f" {len(pdg_map)} PDG codes | {len(mat_map)} materials")
|
|
if cache is not None:
|
|
cache.vocab = (pdg_map, mat_map)
|
|
|
|
proc_map: dict[str, int] | None = None
|
|
if router_cfg.get("enabled") and router_cfg.get("type") == "process":
|
|
n_experts = router_cfg["n_experts"]
|
|
cached_proc_map = cache.proc_maps.get(n_experts) if cache is not None else None
|
|
if cached_proc_map is not None:
|
|
proc_map = cached_proc_map
|
|
echo(
|
|
f"process vocabulary: cache hit ({len(proc_map)} labels, "
|
|
f"{n_experts} experts)"
|
|
)
|
|
else:
|
|
echo("building process vocabulary …")
|
|
proc_map = build_process_map_from_files(files, n_experts=n_experts)
|
|
echo(f" {len(proc_map)} process labels mapped to {n_experts} experts")
|
|
if cache is not None:
|
|
cache.proc_maps[n_experts] = proc_map
|
|
|
|
energy_router_active = (
|
|
router_cfg.get("enabled") and router_cfg.get("type") == "energy"
|
|
)
|
|
energy_idx = router_cfg.get("energy_idx", 3)
|
|
norm_key = setup_cache.normalizer_key(val_fraction, seed, conditioning)
|
|
entry = cache.normalizers.get(norm_key) if cache is not None else None
|
|
|
|
if entry is not None:
|
|
echo(f"normalizer: cache hit (key={norm_key!r})")
|
|
cond_norm = entry.cond_norm
|
|
tgt_norm = entry.tgt_norm
|
|
sec_phys_norm = entry.sec_phys_norm
|
|
energy_quantiles = entry.energy_quantiles
|
|
else:
|
|
echo("fitting normalizer (streaming) …")
|
|
cond_acc = _WelfordAccumulator(COND_DIM)
|
|
tgt_acc = _WelfordAccumulator(X_DIM)
|
|
sec_phys_acc = _WelfordAccumulator(PARTICLE_PHYS_DIM)
|
|
# EnergyRouter's default center spread (linspace over [-2, 2]) assumes
|
|
# the z-normalized energy column is roughly uniform, which real energy
|
|
# spectra rarely are — collect a reservoir sample here (reusing this
|
|
# same pass, not a second scan), then collapse it to a fixed quantile
|
|
# grid (setup_cache.energy_quantiles_from_sample) so centers can
|
|
# instead be seeded from actual data quantiles below. Collected
|
|
# whenever the setup cache is being populated, not only when *this*
|
|
# run's router is energy-typed, so a later run enabling
|
|
# --router-type energy against this same (val_fraction, seed,
|
|
# conditioning) key never needs to rescan just to seed centers.
|
|
collect_energy_sample = energy_router_active or cache is not None
|
|
energy_sampler = (
|
|
_ReservoirSampler(capacity=100_000) if collect_energy_sample else None
|
|
)
|
|
for i, path in enumerate(files):
|
|
for chunk in iter_file_chunks(path, offset=event_id_offset(i)):
|
|
mask = sorted_membership(chunk["event_id"], events_arr)
|
|
if not mask.any():
|
|
continue
|
|
chunk_tr = {k: v[mask] for k, v in chunk.items()}
|
|
cond_cont, _, target_s1, n_sec, sec_cont, _proc, _, _ = build_features(
|
|
chunk_tr,
|
|
pdg_map,
|
|
mat_map,
|
|
proc_map=proc_map,
|
|
require_secondaries=True,
|
|
conditioning=conditioning,
|
|
sec_phys_only=True,
|
|
)
|
|
cond_acc.update(cond_cont)
|
|
tgt_acc.update(target_s1)
|
|
if energy_sampler is not None:
|
|
energy_sampler.update(cond_cont[:, energy_idx])
|
|
sec_valid = np.arange(K_MAX)[None, :] < n_sec[:, None]
|
|
sec_phys = sec_cont[:, :, 4:6][sec_valid]
|
|
if len(sec_phys) > 0:
|
|
sec_phys_acc.update(sec_phys)
|
|
cond_norm = cond_acc.to_normalizer()
|
|
tgt_norm = tgt_acc.to_normalizer()
|
|
sec_phys_norm = sec_phys_acc.to_normalizer()
|
|
energy_quantiles = (
|
|
setup_cache.energy_quantiles_from_sample(energy_sampler.sample)
|
|
if energy_sampler is not None
|
|
else np.empty(0, dtype=np.float32)
|
|
)
|
|
if cache is not None:
|
|
cache.normalizers[norm_key] = setup_cache.NormalizerEntry(
|
|
cond_norm, tgt_norm, sec_phys_norm, n_train_steps, energy_quantiles
|
|
)
|
|
|
|
if energy_router_active and energy_quantiles.size > 0:
|
|
assert cond_norm.mean is not None and cond_norm.std is not None
|
|
levels = np.linspace(0.0, 1.0, router_cfg["n_experts"])
|
|
raw_centers = setup_cache.energy_quantile_at(energy_quantiles, levels)
|
|
centers_init = (raw_centers - cond_norm.mean[energy_idx]) / cond_norm.std[
|
|
energy_idx
|
|
]
|
|
router_cfg["centers_init"] = centers_init.astype(np.float32).tolist()
|
|
echo(
|
|
f" seeded EnergyRouter centers from data quantiles: {router_cfg['centers_init']}"
|
|
)
|
|
elif energy_router_active:
|
|
echo(
|
|
" warning: no energy samples collected — EnergyRouter falls back to "
|
|
"default centers"
|
|
)
|
|
|
|
if cache is not None:
|
|
setup_cache.save(data, files, cache, echo=echo)
|
|
|
|
return SetupStageResult(
|
|
files=files,
|
|
pdg_map=pdg_map,
|
|
mat_map=mat_map,
|
|
proc_map=proc_map,
|
|
cond_norm=cond_norm,
|
|
tgt_norm=tgt_norm,
|
|
sec_phys_norm=sec_phys_norm,
|
|
train_events=train_events,
|
|
val_events=val_events,
|
|
n_train_steps=n_train_steps,
|
|
)
|
|
|
|
|
|
def run_train_job(
|
|
data: Path,
|
|
cfg: dict,
|
|
out_dir: Path,
|
|
device: torch.device,
|
|
shuffle_buffer: int,
|
|
num_workers: int,
|
|
resume: Path | None = None,
|
|
cache_setup: bool = True,
|
|
rebuild_setup_cache: bool = False,
|
|
echo=print,
|
|
) -> None:
|
|
t, m = cfg["train"], cfg["model"]
|
|
config.seed_everything(t["seed"])
|
|
|
|
out_dir = Path(out_dir)
|
|
|
|
# Soft warning (never blocks) — CLAUDE.md's Compute environment section
|
|
# asks that shared portal machines (portal1/deepthought{,2}/bms{1..3})
|
|
# stay within ~1/4 of CPU/RAM so as not to disturb other users' jobs;
|
|
# DataLoader's num_workers has no awareness of that on its own.
|
|
cpu_count = os.cpu_count() or 1
|
|
quota = max(1, cpu_count // 4)
|
|
if num_workers > quota:
|
|
echo(
|
|
f"warning: --num-workers={num_workers} exceeds ~1/4 of this "
|
|
f"machine's {cpu_count} CPU(s) ({quota}) — portal machines are "
|
|
"shared with other users (see CLAUDE.md's Compute environment "
|
|
"section)"
|
|
)
|
|
|
|
router_cfg = m["router"]
|
|
if t["mode"] == "wgan" and router_cfg.get("enabled"):
|
|
raise ValueError(
|
|
"--mode wgan does not support --router (no routed WGAN generator/"
|
|
"critic exists) — disable one or the other"
|
|
)
|
|
|
|
conditioning = m["conditioning"]
|
|
setup = run_setup_stage(
|
|
data,
|
|
val_fraction=t["val_fraction"],
|
|
seed=t["seed"],
|
|
conditioning=conditioning,
|
|
router_cfg=router_cfg,
|
|
cache_setup=cache_setup,
|
|
rebuild_setup_cache=rebuild_setup_cache,
|
|
echo=echo,
|
|
)
|
|
files = setup.files
|
|
pdg_map, mat_map, proc_map = setup.pdg_map, setup.mat_map, setup.proc_map
|
|
cond_norm, tgt_norm, sec_phys_norm = (
|
|
setup.cond_norm,
|
|
setup.tgt_norm,
|
|
setup.sec_phys_norm,
|
|
)
|
|
train_events, val_events, n_train_steps = (
|
|
setup.train_events,
|
|
setup.val_events,
|
|
setup.n_train_steps,
|
|
)
|
|
|
|
total_train_batches = n_train_steps // t["batch_size"]
|
|
echo(f" ~{n_train_steps:,} train steps, ~{total_train_batches:,} batches")
|
|
|
|
train_ds = StreamingStepsDataset(
|
|
files=files,
|
|
split_events=train_events,
|
|
pdg_map=pdg_map,
|
|
mat_map=mat_map,
|
|
cond_normalizer=cond_norm,
|
|
target_normalizer=tgt_norm,
|
|
batch_size=t["batch_size"],
|
|
shuffle_buffer=shuffle_buffer,
|
|
shuffle=True,
|
|
proc_map=proc_map,
|
|
conditioning=conditioning,
|
|
sec_phys_normalizer=sec_phys_norm,
|
|
)
|
|
val_ds = StreamingStepsDataset(
|
|
files=files,
|
|
split_events=val_events,
|
|
pdg_map=pdg_map,
|
|
mat_map=mat_map,
|
|
cond_normalizer=cond_norm,
|
|
target_normalizer=tgt_norm,
|
|
batch_size=t["batch_size"],
|
|
shuffle=False,
|
|
proc_map=proc_map,
|
|
conditioning=conditioning,
|
|
sec_phys_normalizer=sec_phys_norm,
|
|
)
|
|
|
|
pin = device.type == "cuda"
|
|
train_loader = DataLoader(
|
|
train_ds,
|
|
batch_size=None,
|
|
num_workers=num_workers,
|
|
pin_memory=pin,
|
|
)
|
|
val_loader = DataLoader(
|
|
val_ds,
|
|
batch_size=None,
|
|
num_workers=num_workers,
|
|
pin_memory=pin,
|
|
)
|
|
|
|
emb_dim = m.get("emb_dim", EMB_DIM)
|
|
expert_hidden_dim, expert_n_blocks = config.resolve_expert_dims(
|
|
router_cfg, m["hidden_dim"], m["n_blocks"]
|
|
)
|
|
if router_cfg.get("enabled") and (expert_hidden_dim, expert_n_blocks) != (
|
|
m["hidden_dim"],
|
|
m["n_blocks"],
|
|
):
|
|
# Only reachable via an explicit router.expert_hidden_dim/n_blocks
|
|
# override (the 0/"unset" sentinel always resolves to m["hidden_dim"]/
|
|
# ["n_blocks"] — see resolve_expert_dims), so this is never a false
|
|
# positive from inheritance, only a deliberate narrow/wide-experts
|
|
# config the checkpoint dir name (_h{hidden_dim}_b{n_blocks}) won't
|
|
# reflect.
|
|
echo(
|
|
f" warning: experts are {expert_hidden_dim}x{expert_n_blocks}, "
|
|
f"different from model.hidden_dim/n_blocks ({m['hidden_dim']}x"
|
|
f"{m['n_blocks']}) — the checkpoint dir name reflects the latter, "
|
|
"not the experts actually being trained"
|
|
)
|
|
|
|
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,
|
|
"conditioning": conditioning,
|
|
"router": dict(router_cfg),
|
|
"expert_hidden_dim": expert_hidden_dim,
|
|
"expert_n_blocks": expert_n_blocks,
|
|
# Read by `predict`/`rollout` (which never receive their own --mode
|
|
# flag) to auto-detect which sampler a checkpoint needs.
|
|
"mode": t["mode"],
|
|
"noise_dim": m.get("noise_dim", 64),
|
|
}
|
|
|
|
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"
|
|
)
|
|
|
|
critic = None
|
|
sec_critic = None
|
|
if t["mode"] == "wgan":
|
|
critic, sec_critic = build_critics(model_config)
|
|
echo(
|
|
f"critic: {sum(p.numel() for p in critic.parameters()):,} parameters | "
|
|
f"sec_critic: {sum(p.numel() for p in sec_critic.parameters()):,} parameters"
|
|
)
|
|
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
meta = config.build_run_meta(
|
|
data=data,
|
|
seed=t["seed"],
|
|
n_pdg_codes=len(pdg_map),
|
|
n_materials=len(mat_map),
|
|
n_train_events=len(train_events),
|
|
n_val_events=len(val_events),
|
|
n_train_steps=n_train_steps,
|
|
)
|
|
config.save_config(cfg, out_dir, meta)
|
|
|
|
run_training(
|
|
stage1_model=stage1_model,
|
|
sec_decoder=sec_decoder,
|
|
train_loader=train_loader,
|
|
val_loader=val_loader,
|
|
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,
|
|
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),
|
|
lambda_entropy=router_cfg.get("lambda_entropy", 0.0),
|
|
gumbel_tau_start=router_cfg.get("gumbel_tau_start", 1.0),
|
|
gumbel_tau_end=router_cfg.get("gumbel_tau_end", 0.1),
|
|
normalizer_dict={
|
|
"cond": cond_norm.to_dict(),
|
|
"target": tgt_norm.to_dict(),
|
|
"sec_phys": sec_phys_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"],
|
|
validate_steps=t["validate_steps"],
|
|
max_val_batches=t["max_val_batches"],
|
|
total_train_batches=total_train_batches,
|
|
critic=critic,
|
|
sec_critic=sec_critic,
|
|
n_critic=t.get("n_critic", 5),
|
|
gp_weight=t.get("gp_weight", 10.0),
|
|
critic_lr=t.get("critic_lr") or None,
|
|
use_wandb=t.get("wandb", False),
|
|
wandb_project=t.get("wandb_project", "giant"),
|
|
wandb_run_name=t.get("wandb_run_name", ""),
|
|
wandb_log_every=t.get("wandb_log_every", 50),
|
|
)
|