Fix bugs related to uninitialised GPU memory and istep by sophie-xhonneux · Pull Request #2722 · ecmwf/WeatherGenerator

sophie-xhonneux · 2026-08-04T07:24:05Z

Description

Change torch.empty to torch.zeros in StreamEmbedder to ensure there is no uninitialised GPU memory, also skip streams
with stream embed Identity when considering memory to allocate (#2721)

Fix istep not being passed along in loss_calculator.py (#2720)

Config changes for space JEPA (no need to review)

Remark: Linting is broken, but this was true before.

Issue Number

Closes #2720 #2721

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

…ophiex/dev-ssl/ssl-and-diff-july

…/val_workers' into sophiex/dev-ssl/ssl-and-diff-july

ToDo: using ramp-up EMA in teacher seems to cause Nans, figure out why?

…diff-july

…poisoning (#2717) NaN * 0 = NaN, so a non-finite per-token loss at a mask-excluded position poisoned the whole per-sample sum in _masked_l1_loss_per_sample. Use torch.where(mask.bool(), ...) so excluded positions contribute exactly 0; NaNs inside the mask still propagate as before.

…2718) The skip decision was evaluated only on rank 0 and broadcast, so a NaN/spike loss on any other rank was never skipped and its gradients poisoned the all-reduce. Now every rank checks its local loss and the per-rank flags are combined with an all-reduce MAX (logical OR), so a spike on any rank skips the batch on all ranks. Logging remains rank-0-only.

clessig

Minor comments. This PR seems to have some unrelated changes and would break FSDP again, I think.

clessig · 2026-08-04T08:53:09Z

+        # Intermediate encoder levels are tapped before the encoder's final LayerNorm, so they
+        # arrive on very different (and much larger) scales than the final level. Normalize each
+        # level before concatenating so the fusion and the downstream SSL head see a common scale.
+        self.level_norms = nn.ModuleList([nn.LayerNorm(dim_embed) for _ in range(num_levels)])


This change seems unrelated to the problem that is addressed in the PR?

clessig · 2026-08-04T08:53:19Z


    def forward(self, levels: list[torch.Tensor]) -> torch.Tensor:
-        return self.proj(torch.cat(levels, dim=-1))
+        normed = [norm(level) for norm, level in zip(self.level_norms, levels, strict=True)]


Same as above

clessig · 2026-08-04T08:53:53Z

-            num_params_q_aux = (
-                np.prod(self.encoder.q_aux.shape) if self.encoder.q_aux.requires_grad else 0
-            )
+        # if self.encoder.q_aux:


This change was introduced to make FSDP work

see my comment below, the statement breaks in single GPU mode if q_aux exists because it is a nn.Parameter tensor so not bool or can be reduced to bool

clessig · 2026-08-04T08:59:04Z

        print(f" Learnable spatial queries: {num_params_q_cells:,}")
-        if self.encoder.q_aux:
-            print(f" Learnable auxiliary queries: {num_params_q_aux:,}")
+        # if self.encoder.q_aux:


Same as above

we are printing parameters, this should not break FSDP, it doesn't create the parameters

shmh40 · 2026-08-04T09:01:17Z

@@ -0,0 +1,86 @@
+# (C) Copyright 2024 WeatherGenerator contributors.


shmh40 · 2026-08-04T09:01:43Z

-            num_params_q_aux = (
-                np.prod(self.encoder.q_aux.shape) if self.encoder.q_aux.requires_grad else 0
-            )
+        # if self.encoder.q_aux:


Need to revert this I think

this broke in single GPU mode for me :/ can you explain the error to me?

Here is the PR:

#2686

we are in the function that prints parameters, does not seem like it will break FSDP

Yes I think you're right it is not fsdp but whether it breaks I think depends if you have register tokens or not right?

You do have register tokens (I guess), so now it works. But if we want to run with 0 register tokens with this code, we are now doing shape of None and we get an error?

So I think we should actually have:

if self.encoder.q_aux is not None:

yes! you are right, I couldn't think of the right check at midnight yesterday so I commented it out, but what you right is correct I think

shmh40 · 2026-08-04T09:02:33Z

    token_counts = mask_f.sum(dim=-1)
    valid_samples = token_counts > 0
-    per_sample_loss = (per_token_loss * mask_f).sum(dim=-1) / token_counts.clamp(min=1.0)
+    # torch.where (not multiply): NaN at mask-excluded positions must not poison the sum.


Can merge this PR in last so it is after the PR that explicitly does this

shmh40 · 2026-08-04T09:32:54Z

    if OmegaConf.is_dict(conf):
        for key in list(conf.keys()):
-            key = str(key)
+            # key = str(key) # this lines breaks target_source_correspondence please stop re introducing it


Do we need this for anything else in the config?

I get around it with e.g. target_source_correspondence: {"0": {"0": independent}} in target source I think, in case it is necessary for something else in the config

I would ask the question the other way around, is there any reason this line is needed, it breaks the type assumptions of YAML

…diff-july

clessig

Can we clean up the PR so that we don't change 20 files unrelated to the problems. In engines.py are also changes that are unrelated?

sophie-xhonneux and others added 11 commits July 24, 2026 22:04

Make the configs and pipeline

28df344

Merge remote-tracking branch 'origin/develop-ssl-diffusion-v1' into s…

15a41dd

…ophiex/dev-ssl/ssl-and-diff-july

Introduce data_loading.num_workers_validation

7eb0076

Fixing missing get() and linting

dc31fa4

Add configs, fix reconstruct flags

1554291

Fix source/target lists so we don't have student teacher NaNs

5a41e35

Merge remote-tracking branch 'origin/clessig/develop-ssl-diffusion-v1…

1c80cd9

…/val_workers' into sophiex/dev-ssl/ssl-and-diff-july

Add layernorms for deep fusion ssl

9de7ec0

Try to hack around NaN issues

cef7b96

Fix bugs

a055f47

ToDo: using ramp-up EMA in teacher seems to cause Nans, figure out why?

Fix errors when streams diagnostic

5028b60

github-project-automation Bot added this to WeatherGen-dev Aug 4, 2026

sophie-xhonneux changed the base branch from develop to develop-ssl-diffusion-v1 August 4, 2026 07:25

Merge branch 'develop-ssl-diffusion-v1' into sophiex/dev-ssl/ssl-and-…

6e7028b

…diff-july

github-actions Bot added infra Issues related to infrastructure model Related to model training or definition (not generic infra) labels Aug 4, 2026

sophie-xhonneux added 2 commits August 4, 2026 10:09

clessig reviewed Aug 4, 2026

View reviewed changes

shmh40 reviewed Aug 4, 2026

View reviewed changes

Guard against no q_aux existing

5791b30

shmh40 reviewed Aug 4, 2026

View reviewed changes

Merge branch 'develop-ssl-diffusion-v1' into sophiex/dev-ssl/ssl-and-…

468882c

…diff-july

clessig reviewed Aug 4, 2026

View reviewed changes

		@@ -0,0 +1,86 @@
		# (C) Copyright 2024 WeatherGenerator contributors.

Uh oh!

Conversation

sophie-xhonneux commented Aug 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Issue Number

Checklist before asking for review

Uh oh!

clessig left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shmh40 Aug 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

clessig left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

sophie-xhonneux commented Aug 4, 2026 •

edited

Loading

shmh40 Aug 4, 2026 •

edited

Loading