Develop ssl - #2668
Draft
shmh40 wants to merge 13 commits into
Draft
Conversation
* optimized ema update * changes in collapse monitor and ssl teacher
Matches PR #2586 to develop
Added calculations for sine and cosine of absolute coordinates, just overwriting some of the existing encoding of relative position.
Instead of requiring "lat" and "lon", we now can specify the coords column names in stream config. This is already in develop.
…orkers (#2594) The per-worker DataLoader slicing used ceiling division, so when the epoch length was not a multiple of batch_size * num_workers the last worker yielded a partial extra batch, producing a small "foot" (blip) in the LR schedule at the end of each mini-epoch. Compute self.len by flooring against batch_size * max(1, num_workers) and re-multiplying, so the reported length is an exact multiple and no partial trailing batch is produced. Cherry-picked from 90e180a. Co-authored-by: Claude Opus 4.8 <[email protected]>
…coder) (#2592) Add is_stream_reconstructed(stream_cfg): returns False for forcing streams, else honours a per-stream `reconstruct` flag (default True). A stream with `reconstruct: false` keeps a normal (non-empty) target mask so the EMA teacher still encodes it, but gets no physical decoder / TargetPredictionEngine. The decoder-construction loops in model.py now skip a stream when `not is_stream_reconstructed(si)` instead of only when `is_stream_forcing(si)`, and the assimilation-token/target-token pairing loop skips streams absent from embed_target_coords (they have no decoder), preventing a lookup failure. Cherry-picked from 5d709ed; experiment-specific config changes excluded. Co-authored-by: Claude Opus 4.8 <[email protected]>
- Only run predict_decoders when 'masking' is in training_mode, and predict_latent when 'student_teacher' is in training_mode, so forecast-only fine-tuning no longer calls predict_latent (cherry-picked from ac53521). - Skip fully_shard() on submodules with no trainable params, so FSDP2 does not raise "_chunk_cat expects a non-empty tensor" in backward when a shard group is fully frozen during fine-tuning (from ac53521). - Skip restoring optimizer momentum buffers when istep == 0, so a checkpoint can be used as an initialization for fine-tuning without loading stale optimizer state (from 1a0b721). Config changes from the source commits are intentionally excluded. Co-authored-by: Claude Opus 4.8 <[email protected]>
* Implemented reset of weights decoder and added reset_modules config flag * Stops reset_modules after a run continues from another config * Fixed decoder reset being frozen; added overlap guard and trainable params logging * Broadcast reset modules from rank 0 to fix DDP rank divergence * Lint
github-actions
Bot
added
infra
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
4 participants
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 60