Repository navigation
Conversation
…iex/dev/warm-and-frozen-teachers
…iex/dev/warm-and-frozen-teachers
…to be (at least here) identical
…x/dev/exp-week9-local2global
We identified a bug whereby context loss is only computed on the last layer. This will be fixed in the next commits, but we need to save the current state for reproducibility.
* optimze zarr reading * add test * rm tests/test_obs_get.py --------- Co-authored-by: kasravi1 <j.kasravi@fz-jeulich.de> Co-authored-by: Florian Scheidl <36441589+florianscheidl@users.noreply.github.com>
* vectorize hpy splits * add test for vectorize hpy splits * add test_hpy_splits_known_layout_hl1 * move comments to docstring * Removed the equivalence tests --------- Co-authored-by: kasravi1 <j.kasravi@fz-jeulich.de> Co-authored-by: Florian Scheidl <36441589+florianscheidl@users.noreply.github.com>
…#2793) * vectorize normalize_target_channels * add test to vectorize normalize_target_channels * Type _normalize/_denormalize stats as arrays and drop the unused dict branch * Remove the dict-stats normalize test; production only passes arrays --------- Co-authored-by: kasravi1 <j.kasravi@fz-jeulich.de> Co-authored-by: Florian Scheidl <36441589+florianscheidl@users.noreply.github.com>
* implement stochastic churn (#2799) (cherry picked from commit dc5d610) * Fix load_decoder_chkpt silently loading no decoder under FSDP (#2800) load_model normalises the "module." state_dict prefix in both of its branches (sharded and unsharded). load_decoder_from_checkpoint was written later and copied only the unsharded half, so its sharded path -- taken whenever with_ddp and with_fsdp, i.e. every decoder finetune here -- looked each parameter up in model.state_dict() using the raw checkpoint key. For a DDP-saved backbone (keys prefixed "module.") against an unprefixed model state dict, every one of the 131 lookups missed, each was skipped individually, and load_state_dict({}) succeeded. The run then trained a randomly initialised decoder. Three things hid it: the per-parameter warning fired once per rank (thousands of lines), the unexpected-keys report was empty because nothing was loaded, and "Loaded 0 decoder tensors" was a single INFO line. Six completed decoder finetunes off cw6a4szu, plus one off nhv6tkln, trained random decoders this way; the two that worked (new73gkw, oliqngmf backbones) did so only because those checkpoints happen to be saved without the prefix. - Extract _align_module_prefix / _strip_module_prefix and use them in all four places the prefix logic was duplicated, including the decoder overlay's sharded path, which is the actual fix. - _strip_module_prefix removes only a leading prefix; the previous key.replace("module.", "") stripped the substring anywhere in the path. - Raise instead of loading nothing: matching N tensors but loading 0 is now a RuntimeError, as is requesting an overlay from a checkpoint with no decoder. The condition is identical on every rank, so all ranks raise together rather than one hanging the others. - Collapse the per-parameter skip warning into one rank-0 summary and report loaded/matched counts. Verified against all four backbone checkpoints: the old raw lookup gives 0/131 for cw6a4szu and nhv6tkln against an unprefixed model (reproducing the bug) and 99/99 for new73gkw and oliqngmf (reproducing why those worked); with the fix all four match 100% against both key conventions. Claude-Session: https://claude.ai/code/session_01NsFj2pVh4pxs2UxthxewWy Co-authored-by: Claude Opus 5 (cherry picked from commit 980f640) * Register MLP norm only in self.layers, not also as self.lnorm (#2817) * Register MLP norm only in self.layers, not also as self.lnorm MLP.__init__ bound one norm module to two attribute paths for non-DiT MLPs: self.lnorm = norm(dim_in, eps=norm_eps) if not is_dit: self.layers.append(self.lnorm) state_dict() therefore emitted the same tensor twice, as "lnorm.*" and as "layers.0.*". Consequences: - FSDP2 state_dict() reports duplicate parameter paths for the shared module, while named_modules() deduplicates and keeps only the canonical "layers.0" path, so the alias has no module to initialise. - Every checkpoint written this way carries two identical copies of each norm tensor. Loading such a checkpoint into de-aliased code reports the "lnorm.*" copies as unmatched -- harmless in itself, but it buries real mismatches in noise. Register the norm only in self.layers. forward() already applies it through the self.layers loop and special-cases AdaLayerNorm to pass ada_ln_aux, so behaviour is unchanged. The DiT branch keeps self.lnorm: it is an AdaLNZero applied explicitly in forward() and never appended to self.layers, so it was never aliased. Verified: a non-DiT MLP now exposes only layers.* in state_dict(), with and without dim_aux; forward passes unchanged; ruff clean. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NsFj2pVh4pxs2UxthxewWy * remove comments * added short comment --------- Co-authored-by: Claude Opus 5 <noreply@anthropic.com> (cherry picked from commit 455a76a) * Add load_encoder_chkpt, and do not re-apply either overlay on continuation (#2821) Two changes to the checkpoint-overlay machinery in model_interface.py. 1. load_encoder_chkpt -- the mirror of load_decoder_chkpt. Filters a checkpoint to `encoder.*` (stream embedders, ae_* engines, learned queries) and overlays it after the primary load, leaving the forecast engine and decoder from the primary checkpoint untouched. This allows pairing one run's forecast engine with a different run's encoder. The two selections are disjoint, so both overlays may be configured at once. The shared body of both loaders is factored into _load_part_from_checkpoint rather than duplicated. This is deliberate: the decoder overlay was originally written by copying only the *unsharded* half of load_model's `module.` prefix handling, so every sharded lookup missed and the run silently trained a randomly initialised decoder. One body means the two parts cannot drift apart that way again. 2. Neither overlay is re-applied when continuing a run. The primary load is an if/elif: a continuation loads from `run_id_contd` -- the run's own checkpoint, which already contains whatever training that part has had -- and never consults `load_chkpt`. Re-overlaying afterwards reverted a decoder-finetune's decoder to the borrowed original on every chained-job restart, discarding the finetuning done since the previous restart. All 12 decoder finetunes of an in-flight sweep were affected, 4-6 reverts each, and this is also the cause of their loss spikes at chain handoffs (2-3.4x, recovering over 300-1150 steps): the learning rate is continuous across restarts and optimizer momentum is fully restored, so neither explains them; the partial decoder revert does. Known limitation: `run_id_contd` is not a train_continue signal -- inference passes `--from-run-id` through it -- so both overlays are now also skipped during inference. If the overlay path is needed at inference, gate on an explicit `overlay_decoder` / `overlay_encoder` parameter instead of on `run_id_contd`; note target_and_aux_utils.py is a third call site that also passes a non-None `run_id_contd`. Claude-Session: https://claude.ai/code/session_01NsFj2pVh4pxs2UxthxewWy Co-authored-by: Claude Opus 5 <noreply@anthropic.com> (cherry picked from commit 6d358af) * shorten comments * Address PR comments 1 * remove test file * remove comment and duplicated line --------- Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
github-actions
Bot
added
data
shmh40
self-requested a review
October 14, 2026 07:34
…2962) * Introduce keyword arguments for MLP and cleanup if/else clauses * Incorporated review comments. Upgraded diffusion conditioning class to value object for simpler access
* Intermediate state of refactoring. Writing separated from validation_io.py and plotting outsourced to dedicated scripts * Some fixes and adding an explicit dimension for 'denoising_step' for diffusion dev * Change dev outputs format from zarr to zip to save inodes * Intermediate fixes * Fixed separate decoder loading * Implemented Christian's PR review suggestions and cleaned var names * Incorporated 2nd round of comments. Strict separation of train and inference loggers
* Fix data loader RNG seed collisions across ranks, workers and stages
The seed was computed as
rng_seed * (rank+1)*73 * (worker+1)*37 * (mini_epoch+13)*7
which depends only on the product (rank+1)*(worker+1)*(mini_epoch+13).
Different combinations therefore shared a seed, e.g. rank 0/worker 1 and
rank 1/worker 0 (with 8 ranks x 10 workers x 2 mini_epochs only 78 of 160
seeds were distinct). With num_workers=0 the rank did not enter the seed
at all, and train and validation samplers used the same seed.
Create the generator from np.random.SeedSequence with rng_seed, rank,
loader worker, mini_epoch, iteration count and stage as separate entropy
values, so that every combination yields an independent stream. The
iteration count covers persistent loader workers, whose mini_epoch is not
advanced.
* Update docstring
* Rename _make_rng to _reset_rng and clarify the reset counter
- Rename _make_rng() to _reset_rng() and move it next to reset()
- Rename _num_iters to _num_rng_resets: it counts _reset_rng() calls
(passes over the dataset) in this process and keeps the seed unique per
pass for persistent loader workers, whose mini_epoch is not advanced
- Shorten the _reset_rng() docstring
* rm diffusion config
* Add configurable noise embedding with log-spaced Fourier features
The DiT timestep embedding (frequencies 10000^(-k/half) in [1e-4, 1]) was designed
for integer timesteps in [0, 1000]; over the ~1-3 unit range of c_noise = ln(sigma)/4
~220/256 channels are constant and the embedding is effectively rank 2.
New `noise_embedding: {type, f_min, f_max, include_raw}` config block. `log_fourier`
uses log-spaced frequencies in [f_min, f_max] (fp32, computed in forward so meta-device
init is safe) plus the raw c_noise channel. Default `dit_legacy` is unchanged, so
existing checkpoints and configs are unaffected.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01S3h3zrLh9S8VWCEi2M4h1T
* Warn when noise_embedding.type is unset and recommend log_fourier (8pi)
The default stays dit_legacy for now so existing checkpoints keep loading.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01S3h3zrLh9S8VWCEi2M4h1T
* shorten comments + remove tests
---------
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Issue Number
Is this PR a draft? Mark it as draft.
Checklist before asking for review
./scripts/actions.sh lint./scripts/actions.sh unit-test./scripts/actions.sh integration-testlaunch-slurm.py --time 60FastEvaluation