From 96aad375c884268755b972609dea7db37cb2cf71 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Wed, 26 Aug 2026 12:08:42 +0200 Subject: [PATCH] Seed each epoch's RNG from (seed, epoch) (gitea #83) The per-epoch training fan-out only makes sense if epoch k is the same epoch either way, and the shuffle fix alone wasn't enough: run_train_job calls seed_everything(train.seed) at process start, so a fresh job restarted the torch/numpy stream at epoch 1's state and drew different flow/WGAN noise than the corresponding epoch of a single long run. giant.config.epoch_seed derives a per-epoch seed, and the training loop reseeds from it at the top of every epoch. Verified on a 3-epoch toy run: the chained workflow's concatenated metrics.csv is now byte-identical to a single `giant train --epochs 3` with the same seed (it matched only on epoch 1 before). Co-Authored-By: Claude Opus 5 --- CLAUDE.md | 2 +- giant/config.py | 13 +++++++++++++ giant/training/loop.py | 5 +++++ 3 files changed, 19 insertions(+), 1 deletion(-) diff --git a/CLAUDE.md b/CLAUDE.md index 5c5805d..7487205 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -103,7 +103,7 @@ Secondary energies are a **stick-breaking partition of the `e_sec` budget** from **Input is one or more `giant rollout` YAML sidecars** (`run.py:load_rollout_yamls`): each YAML's `output`/`dataset` keys name its rollout parquet and seed file (= the reference truth); every supplied YAML must resolve to the same `dataset`, checked up front with a clear error otherwise (the premise is "N candidates vs one ground truth"). Each rollout's series name comes from a repeated `--label` CLI flag, else the YAML stem (N>1), else `"rollout"` (a single YAML). `prep` creates a **run directory** (`/analysis_runs/analysis_/` by default, `--run-dir` to override) holding `shared.json`, `run_meta.json` (`RunMeta.rollouts: list[{name,path,plot_meta}]`, insertion order = CLI order = every plot's series order), `reduced_partial/`, `reduced/`, `plots/`. **Compute/merge/render split:** `giant analyze prep a.yaml [b.yaml ...] --chunks N` records `N` in `run_meta.json`, and the workflow's `AnalysisComputeTask` runs one HTCondor job per (plot, chunk) pair (`compute-one --id --chunk --run-dir`, polars/numpy only — no LaTeX on workers), each streaming over an `event_id`-disjoint slice (`event_id % N == chunk`) of the reference **and every rollout** and writing a small `reduced_partial/__.json`; every `PlotSpec` splits into a `compute_partial`/`finalize` pair so chunks can be summed/concatenated back per rollout (`chunkable=False` specs — the checkpoint-bound diagnostics, already bounded/subsampled — always run as a single chunk). The local `giant analyze render ` first joins every plot's chunk partials into `reduced/.json` (`merge_all`, a no-op join when `N=1`; `merge-one` does a single plot for debugging), then turns those into the styled PDF/gallery tree. `giant analyze metrics ` is a separate, unrelated entry point: training-progress plots straight from a run's `metrics.csv`. -**Workflow orchestration** (`giant/workflow/`, `giant workflow run` CLI): b2luigi is the **only sanctioned way to run a multi-step pipeline**; `giant`/`dwarf` are single-step primitives the tasks invoke. One workflow TOML (`configs/workflow_example.toml`) parameterises a whole experiment — `[workflow]`/`[condor]`/`[dataset]`/`[geometry]` plus repeated `[[train]]`/`[[rollout]]`/`[[analysis]]` tables, each cross-referenced by name — and `spec.py` parses it into frozen dataclasses, rejecting unknown keys and dangling references. Every task's output directory is `//name=/spec_hash=/…`, where the 8-hex `spec_hash` covers that task's resolved sub-spec **and its transitive parents**, so an edited spec re-runs exactly the affected subtree instead of silently reusing stale outputs. The DAG (`tasks.py`): `DatasetTask` (external, fails fast if `/ceph` isn't mounted) → `WarmCacheTask` / `GeometryOracleTask` → `TrainEpochTask(name, milestone)` → `TrainTask` → `RolloutTask` → `AnalysisPrepTask` → `AnalysisComputeTask(name, plot_id, chunk)` → `AnalysisRenderTask` → `WorkflowTask`. Training is fanned out into **one short GPU job per epoch** (`epochs_per_job` trades queue waits back), chained by `--resume` on the previous job's `last.pt` — the loop already handles that unchanged — and `TrainTask` republishes `best.pt`/`last.pt`/a concatenated `metrics.csv` so nothing downstream sees the fan-out. `StreamingStepsDataset.set_epoch` (called per epoch by `training/loop.py`) seeds shuffling from `(seed, epoch, worker_id)` so epoch *k*'s batch order is the same either way. `AnalysisRenderTask` is always local (the only step importing plotstyle/LaTeX); `htcondor.py` holds the CPU/GPU submit settings, with the GPU requirement strings (`TARGET.ProvidesEtpCeph` + device/memory pins) ported from the `condor-gpu-train-rollout` branch. `run.py` is the script b2luigi re-executes on workers (`--spec` forwarded via `task_cmd_additional_args`, so a worker resolves the identical graph); `giant workflow run` is a thin exec of it. Needs `uv sync --extra cpu --extra workflow`. +**Workflow orchestration** (`giant/workflow/`, `giant workflow run` CLI): b2luigi is the **only sanctioned way to run a multi-step pipeline**; `giant`/`dwarf` are single-step primitives the tasks invoke. One workflow TOML (`configs/workflow_example.toml`) parameterises a whole experiment — `[workflow]`/`[condor]`/`[dataset]`/`[geometry]` plus repeated `[[train]]`/`[[rollout]]`/`[[analysis]]` tables, each cross-referenced by name — and `spec.py` parses it into frozen dataclasses, rejecting unknown keys and dangling references. Every task's output directory is `//name=/spec_hash=/…`, where the 8-hex `spec_hash` covers that task's resolved sub-spec **and its transitive parents**, so an edited spec re-runs exactly the affected subtree instead of silently reusing stale outputs. The DAG (`tasks.py`): `DatasetTask` (external, fails fast if `/ceph` isn't mounted) → `WarmCacheTask` / `GeometryOracleTask` → `TrainEpochTask(name, milestone)` → `TrainTask` → `RolloutTask` → `AnalysisPrepTask` → `AnalysisComputeTask(name, plot_id, chunk)` → `AnalysisRenderTask` → `WorkflowTask`. Training is fanned out into **one short GPU job per epoch** (`epochs_per_job` trades queue waits back), chained by `--resume` on the previous job's `last.pt` — the loop already handles that unchanged — and `TrainTask` republishes `best.pt`/`last.pt`/a concatenated `metrics.csv` so nothing downstream sees the fan-out. `StreamingStepsDataset.set_epoch` and `config.epoch_seed` (both applied per epoch by `training/loop.py`) derive the batch order and the global RNG state from `(seed, epoch)`, so epoch *k* is bit-identical either way — verified by diffing a chained run's `metrics.csv` against a single 3-epoch `giant train`. `AnalysisRenderTask` is always local (the only step importing plotstyle/LaTeX); `htcondor.py` holds the CPU/GPU submit settings, with the GPU requirement strings (`TARGET.ProvidesEtpCeph` + device/memory pins) ported from the `condor-gpu-train-rollout` branch. `run.py` is the script b2luigi re-executes on workers (`--spec` forwarded via `task_cmd_additional_args`, so a worker resolves the identical graph); `giant workflow run` is a thin exec of it. Needs `uv sync --extra cpu --extra workflow`. **Shower rollout** (`giant/rollout.py`, `giant rollout` CLI): autoregressively steps the two-stage model into a full shower, advancing tracks breadth-first (every sweep steps all active tracks once, in `batch_size` chunks, so many tracks share each forward pass). Each primary post-step becomes the next pre-step, secondaries are pushed as new tracks, and per-step `material`/`layer_id` come from a `GeometryOracle` (`giant/geometry.py`, built via `dwarf build-geometry-oracle`) that learns position → (material, layer_id) from data and flags detector escape by nearest-neighbour distance. Tracks terminate on one of the `TERM_*` reasons in `constants.py` (energy cutoff, max steps, escape, natural end, unknown pdg, max tracks); energy is deposited locally on every stop except escape (leakage), so showers conserve energy by construction. `giant/checkpoint_io.py` is the shared checkpoint → ready-to-run-models path used by both `predict` and `rollout`. diff --git a/giant/config.py b/giant/config.py index 5ee2eaa..a7d7fe7 100644 --- a/giant/config.py +++ b/giant/config.py @@ -1769,6 +1769,19 @@ def resolve_default_out_dir(cfg: dict, base: Path = Path("checkpoints")) -> Path return out_dir +def epoch_seed(seed: int, epoch: int) -> int: + """Per-epoch derivative of the run seed. + + Reseeding the global RNGs from this at the top of every epoch makes epoch + *k* draw the same noise whether it runs inside one long `giant train` or + as its own resumed job in a per-epoch workflow chain + (`giant/workflow/tasks.py:TrainEpochTask`) — without it, a fresh process + would restart the stream at epoch 1's state. Mirrors what + `StreamingStepsDataset.set_epoch` does for the batch order. + """ + return (int(seed) * 1_000_003 + int(epoch)) % (2**32) + + def seed_everything(seed: int) -> None: random.seed(seed) np.random.seed(seed) diff --git a/giant/training/loop.py b/giant/training/loop.py index ea6b6dd..49f9d0b 100644 --- a/giant/training/loop.py +++ b/giant/training/loop.py @@ -18,6 +18,7 @@ import torch from torch.utils.data import DataLoader from tqdm import tqdm +from giant import config from giant.data.loader import TopNMap from giant.data.setup_cache import topnmap_to_json from giant.training.checkpoint import build_checkpoint, init_stages_from_checkpoints, load_checkpoint @@ -184,6 +185,10 @@ def train( if device.type == "cuda": torch.cuda.reset_peak_memory_stats(device) collector.start_epoch(epoch) + # Epoch-aware RNG: same noise (and, below, same batch order) for + # epoch k whether the run is one process or a chain of per-epoch + # jobs. See giant.config.epoch_seed. + config.seed_everything(config.epoch_seed(t["seed"], epoch)) # Epoch-aware shuffle stream (see StreamingStepsDataset.set_epoch): # keeps epoch k's batch order identical whether it runs here or as # its own resumed per-epoch job in a b2luigi workflow.