Skip to content

use the target random number for also source random number - #2733

Merged
Jubeku merged 1 commit into
ecmwf:jk/develop-ssl-diffusion-v1/fix-streamsfrom
MeteoSwiss:jk/develop-ssl-diffusion-v1/fix-streams
Aug 7, 2026
Merged

use the target random number for also source random number#2733
Jubeku merged 1 commit into
ecmwf:jk/develop-ssl-diffusion-v1/fix-streamsfrom
MeteoSwiss:jk/develop-ssl-diffusion-v1/fix-streams

Conversation

@kctezcan

@kctezcan kctezcan commented Aug 5, 2026

Copy link
Copy Markdown
Contributor

Description

Problem: the random number used to make the noise signal for the diffusion forward process is coming from the source random draw. the random number for calculating the loss weight for diffusion is coming from the target random draw.

Fix: take the arget random number also for the source.

Remark: for this to work, the target_input in the config has to be set:

  target_input: {
    "forecasting" : {
      ...
      masking_strategy_config: {
        diffusion_rn: True,
        noise_distribution: *noise_distribution,
        sigma_min: *sigma_min,
        sigma_max: *sigma_max
        },
      ...

Issue Number

No issue

Is this PR a draft? Mark it as draft.

Checklist before asking for review

  • I have performed a self-review of my code
  • My changes comply with basic sanity checks:
    • I have fixed formatting issues with ./scripts/actions.sh lint
    • I have run unit tests with ./scripts/actions.sh unit-test
    • I have documented my code and I have updated the docstrings.
    • I have added unit tests, if relevant
  • I have tried my changes with data and code:
    • I have run the integration tests with ./scripts/actions.sh integration-test
    • (bigger changes) I have run a full training and I have written in the comment the run_id(s): launch-slurm.py --time 60
    • (bigger changes and experiments) I have shared a hegdedoc in the github issue with all the configurations and runs for this experiments
  • I have informed and aligned with people impacted by my change:
    • for config changes: the MatterMost channels and/or a design doc
    • for changes of dependencies: the MatterMost software development channel

@clessig

clessig commented Aug 5, 2026

Copy link
Copy Markdown
Collaborator

@kctezcan : if this problem also exists on develop-ssl-diffusion, could we also fix it there?

@kctezcan

kctezcan commented Aug 5, 2026

Copy link
Copy Markdown
Contributor Author

@kctezcan : if this problem also exists on develop-ssl-diffusion, could we also fix it there?

do you mean in develop-ssl-diffusion-v1?

@clessig

clessig commented Aug 5, 2026

Copy link
Copy Markdown
Collaborator

@kctezcan : if this problem also exists on develop-ssl-diffusion, could we also fix it there?

do you mean in develop-ssl-diffusion-v1?

Yes

@Jubeku

Jubeku commented Aug 6, 2026

Copy link
Copy Markdown
Contributor

Do we have a problem here for forcing streams which don't have a target mask and hence no random number? Maybe in this case my PR 2737 is complementary because a target mask is built even if target channels are empty.

@Jubeku
Jubeku merged commit c2b5078 into ecmwf:jk/develop-ssl-diffusion-v1/fix-streams Aug 7, 2026
1 check passed
Jubeku added a commit that referenced this pull request Aug 7, 2026
* fix fsdp error from unused register token module

* default to ddp_static_graph: False and fix num_params_q_aux

* get random noise from either ERA5 or ERA5_in

* fix fsdp and forcing stream target masks

* fix checkpoint loading for unsharded target encoder

* use the target random number for also source random number (#2733)

---------

Co-authored-by: kctezcan <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Status: Done

Development

Successfully merging this pull request may close these issues.

3 participants