6h-step stochastic bake-off subset: arms 1,2,3,6,7 + arm9 (7 + γ0.5 whitening) - #1392
Conversation
… whitening) Reruns the stochastic-ACE bake-off protocol at a 6h step to test whether the daily arms' small-scale precipitation power deficit is a timestep artifact. Each config is its daily counterpart with only the 6-hourly dataset, 6-hourly normalization stats (beaker dataset mounted at /statsdata), and x4 inference step counts at fixed lead time; arm 9 adds g0.5 whitening to arm 7 (whitening the spectral-power term alone, the cell the daily bake-off left untested).
|
Claude: independent pre-review (fresh context, no authoring history). One run-killing bug, plus two judgment calls and a doc nit. Not ready to launch as-is. 1. BLOCKER — every
|
… values Pre-review findings: the daily data_path is a wrapper directory containing the store named by file_pattern, but the 6-hourly store has no wrapper, so the substituted path pointed inside the store and every loader's glob would come up empty. Use data_path: /climate-default/ + unchanged file_pattern (the form the prior 6h stochastic runs trained with). Also revert the weather blocks' forward_steps_in_memory 5->20 (a memory/IO chunk size, not a horizon) and note the ensembles step-index edit in the README.
|
Claude: All four findings addressed in efec091:
Also noted in the description that |
|
Claude: re-review of Blocker fix verifiedAll 20 No drift
Independent invariant checkParsed all six configs and asserted, per inference block: every I did not re-run Description claims spot-checked against measurements
The README's Data and Horizons bullets now match the diff, including the |
| global_mean_removal: | ||
| kind: shared | ||
| append_as_input: true |
There was a problem hiding this comment.
Get rid of the global mean removal in all configurations, we didn't mean to use it here.
Jeremy: the daily arms' global_mean_removal (kind: shared, append_as_input: true) was not meant to be part of this recipe. Removed from all six configs; README notes it as a deliberate divergence from the daily arms alongside the timestep.
Arms 1, 7 and 9 were killed by host-RAM exhaustion in the first launch wave. Set num_data_workers 8 -> 4 on every loader (training, validation, all five inference) and the training loader's prefetch_factor 4 -> 2, then relaunch. Both are queue-depth knobs; sample order comes from the seeded distributed sampler, so they are expected to be results-neutral. Arms 2, 3 and 6 are unchanged (3 and 6 are still running).
…ed config Arm 3 died 2026-08-05 at epoch 69/80 to a 30-min NCCL allreduce watchdog timeout, four minutes after a second full-footprint arm landed on its node. Its beaker autoResume job carried the pre-fix config and starved 5.5 h unscheduled; it was cancelled in favour of this continuation. arm3-50-50-ec-resume-e69.yaml adds resume_results (existing_dir /prior-results, resume_wandb true) over result dataset 01KYS0V4WTP34MX4MERGRAY7ZD, and takes the same num_data_workers 8->4 / prefetch_factor 4->2 as arms 1/7/9. max_epochs stays 80. No fme/ code differs between the checkpoint's commit (95364df) and this base, so the checkpoint is compatible. The launch line reuses the ORIGINAL wandb run name so resume_wandb keeps the run id without renaming its display name.
The bake-off base recipe is the pre-training donor nzccs8zd, which is the main ERA5 baseline with exactly filter_num_groups: 16 and spectral_ratio: 0.125 carried over. It sets neither residual_prediction nor global_mean_removal. Both reached these arms because the daily configs adopted the 4-degree/daily v2 architecture block wholesale rather than building base + those two knobs; the arms' builder.config is byte-for-byte v2's, minus clip_latent_global_means. global_mean_removal was caught and dropped from the 6h subset on 2026-07-29; residual_prediction was not. 6h subset (all six arms): - residual_prediction: true -> false. - Loader knobs made uniform at num_data_workers 4 / prefetch_factor 2. Only arms 1, 7 and 9 carried the measured fix; arms 2, 3 and 6 still ran the pre-fix 8/4 that caused the host-RAM OOM and wedge failures. - 8 GPUs per arm (was 4), so a job holds a whole node and the co-scheduling behind those failures cannot recur. batch_size is the global batch, so the effective batch and gradient are unchanged. - Drop arm3-50-50-ec-resume-e69.yaml: it resumed an epoch-69 checkpoint trained with residual prediction, which is not loadable on this recipe. - Job names gain a "-nores" suffix to separate them from the superseded wave. Daily wave: add arm1-90-10-ec-nores.yaml, arm 1 with residual_prediction false and nothing else changed, as a one-knob A/B against the completed arm 1 (qhv9zf95) to test whether that wave's small-scale precipitation power deficit is caused by residual prediction. global_mean_removal and the 8/4 loader knobs are kept so the comparison stays one-knob. The eight completed daily arms are untouched. Also corrects both READMEs, which credited nzccs8zd with residual prediction and shared global-mean removal. It has neither. Launch targets: ai2/ace workspace, urgent priority, ai2/jupiter, with CM_PRIORITY=high for the external balancer.
Jeremy's call (2026-08-07): run the ablation on the corrected bake-off recipe rather than hold global_mean_removal to match the old arm. residual_prediction and global_mean_removal are both settings the daily wave inherited from the 4-degree/daily v2 architecture block; the donor nzccs8zd sets neither. Against the completed arm 1 (qhv9zf95) the diff is therefore two knobs, residual and GMR, so a change in the small-scale precipitation spectrum is attributable to the pair rather than to residual prediction alone. Noted in the README and the launcher comment so the result is not over-read. Job renamed to ...-nores-nogmr-rs0 so it does not collide with the cancelled GMR-retaining run (wandb do85dswp). Loader knobs unchanged at 8/4.
1 participant
Configs only — nothing is launched from this PR; launches are Jeremy's call after review.
Reruns the stochastic-ACE bake-off protocol at a 6h step, to test whether the daily arms' small-scale precipitation power deficit (reports#51) is a daily-timestep artifact. Based on the daily bake-off branch at its exact launch SHA (42ab578), so the training code — including the unmerged spectral-power-CRPS term — is identical to what produced the daily arms, and the PR diff is the new
configs/baselines/stochastic-ace-bakeoff-6h/directory alone.Arms (per the daily numbering): 1 (0.9/0.1 + EC), 2 (0.9/0.1, no EC), 3 (0.5/0.5 + EC), 6 (0.8/0.1/0.1 + EC), 7 (0.9/0/0.1 + EC), and new arm 9 = arm 7 + γ0.5 per-sample whitening — the cell reports#51 left untested (with
energy_score_weight: 0, the shared whitening operator applies to the spectral-power term alone; verified against the loss-builder code path at this SHA).Each config is its daily counterpart with three mechanical edits (generated by an assertion-checked script, not hand-edited) plus one recipe change requested in review —
global_mean_removal(kind: shared,append_as_input: true) is dropped from all six arms:2026-03-19-era5-1deg-8layer-daily-1940-2025.zarr/→data_path: /climate-default/, with the unchangedfile_patternselecting the 6-hourly store2026-03-19-era5-1deg-8layer-1940-2025.zarrdirectly — the 6-hourly store has no wrapper directory, and this is the form the prior 6h stochastic runs (4s0rnth6) trained with. Same variable set, same ACE2 train/val/inference split, same 06Z IC dates (the 6h zarr spans 1940-01-01T12–2025-12-31T18 at 00/06/12/18Z, so all ICs and horizons are in range).andrep/2026-03-19-era5-1deg-8layer-stats-1990-2019mounted at/statsdata(same stats the prior 6h stochastic pretrains used; residual scaling is timestep-dependent — e.g. PRESsfc residual std 250.4 at 6h vs 641.5 daily — so the daily stats can't be reused). The mount rides the run script's existing# arg:header mechanism.step_means/ensemblesindex 5 → 20.forward_steps_in_memory(a memory/IO chunk size, not a horizon) is unchanged everywhere.Arm 9 additionally adds the
energy_score_whiteningblock in arm 8's exact syntax.Held fixed deliberately (timestep and the global-mean-removal drop are the only differences vs the daily arms): model (fg16/sr0.125 NoiseConditionedSFNO), 1-step training, seed 0,
max_epochs: 80, batch size 8, LR 1e-4, EMA decay 0.999, inline-inference cadence (every 2 epochs). Review consequences: 80 epochs at 6h is ~4× the optimizer steps and wall time per arm — the 6h zarr has 125,646 timesteps vs the daily 31,411 (exactly 4.0×), and daily arm 1 took ~54 h on 4 GPUs, so expect ~9 days per arm on 4 GPUs (24 GPUs held ~9 days for the set of six). EMA/LR-schedule shapes also differ in step terms at fixed epoch count, and each inline inference pass is ~4×. If step-matching (20 epochs) is preferred over epoch-matching, say so and I'll adjust.Validation done: all six configs pass
fme.ace.validate_config --config_type trainat this SHA (schema only — it does not touch the filesystem, hence the path bug pre-review caught); all sixstepper_training.lossblocks were built through the realLossConfig.buildpath, confirming the sp-term/whitening wiring per arm; and the data/stats path forms match the prior finished 6h run.run-train.shis the daily script with only the launch-target block changed (guardrails byte-identical); job names followace2s-bakeoff-6h-arm*-rs0.🤖 Generated with Claude Code