2D Rope into develop #1540

kctezcan · 2025-12-30T14:26:09Z

Description

This is @csjfwang 's implementation of 2D rope, coming from this orignal PR: #1445

Issue Number

Closes #1445
Closes #1109

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

- Add copyright attribution for rotate_half() and apply_rotary_pos_emb() functions - Rename apply_rotary_pos_emb_2d() to rotary_pos_emb_2d() for consistency - Rename config parameter use_2D_rope to rope_2D for better extensibility when supporting different RoPE variants in the future

clessig

We can simplify the code if we change the syntax of the forward fct of the attention blocks slightly. This should probably be done consistently for all heads, even those without 2D rope.

clessig · 2026-01-05T12:06:59Z

src/weathergen/datasets/multi_stream_data_sampler.py

            )

-        return iter_start, iter_end
+        return iter_start, iter_end


Can we revert this

clessig · 2026-01-05T12:11:26Z

src/weathergen/model/engines.py

+    def forward(self, tokens, coords=None, use_reentrant=False):
        for block in self.ae_aggregation_blocks:
-            tokens = checkpoint(block, tokens, use_reentrant=use_reentrant)
+            if isinstance(block, MultiSelfAttentionHead | MultiSelfAttentionHeadLocal):


We could avoid the branches if we would follow the approach in the MLP:

WeatherGenerator/src/weathergen/model/layers.py

Line 82 in 29e2fe2

def forward(self, *args):

Or some compromise with *args and **kwargs

Personal preference is the if statement

clessig · 2026-01-05T12:11:56Z

src/weathergen/model/engines.py

+    def forward(self, tokens, coords=None, use_reentrant=False):
        for block in self.ae_global_blocks:
-            tokens = checkpoint(block, tokens, use_reentrant=use_reentrant)
+            if isinstance(block, MultiSelfAttentionHead | MultiSelfAttentionHeadLocal):


Same as above: if-statement is not necessary

but cleaner to read imho

clessig · 2026-01-05T12:12:24Z

src/weathergen/model/engines.py

        for _b_idx, block in enumerate(self.fe_blocks):
            if isinstance(block, torch.nn.modules.normalization.LayerNorm):
                tokens = block(tokens)
+            elif isinstance(block, MultiSelfAttentionHead | MultiSelfAttentionHeadLocal):


Same as above

clessig · 2026-01-05T12:12:50Z

src/weathergen/model/model.py

        )
        self.pe_global = torch.nn.Parameter(pe, requires_grad=False)

+        ### ROPE COORDS ###


standard comment:

# rope coords

clessig · 2026-01-05T12:22:44Z

src/weathergen/model/attention.py

        ks = self.lnorm_k(self.proj_heads_k(x).reshape(s)).to(self.dtype)
        vs = self.proj_heads_v(x).reshape(s).to(self.dtype)

+        if coords is not None:


MultiSelfAttenionHead is a vanilla self attention block. The 2D rope positional encoding makes strong assumption of the structure of the qs and ks. At the very least, we need some documentation on what the option means and when it makes sense. Not sure if there's another good way to ensure proper use of this option (one could have a separate class but then one should make sure there's no code duplication).

sophie-xhonneux

Did a first review of the code and left some comments that should be addressed.

Will check the correctness next

sophie-xhonneux · 2026-01-06T10:55:01Z

src/weathergen/model/model.py

+
+        if self.rope_2D:
+            # Precompute per-cell center coordinates (lat, lon in radians) for 2D RoPE.
+            # Shape: (num_healpix_cells, ae_local_num_queries, 2)


Why isn't the shape (num_healpix_cells, 2)?

sophie-xhonneux · 2026-01-06T10:55:57Z

src/weathergen/model/model.py

+        ### ROPE COORDS ###
+        self.rope_2D = cf.get("rope_2D", False)
+        if self.rope_2D:
+            self.rope_coords = torch.nn.Parameter(


I think this should be a registered buffer, not parameters

https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_buffer

sophie-xhonneux · 2026-01-06T10:58:48Z

src/weathergen/model/positional_encoding.py

+
+
+####################################################################################################
+def rotary_embedding_2d(coords, dim, base=10000.0):


There should be an option to learn the frequencies, this proposed in the original RoPE (https://arxiv.org/pdf/2403.13298v1 section 3.2 at the end) or we should have a reference explaining why not

sophie-xhonneux · 2026-01-06T11:02:03Z

src/weathergen/model/positional_encoding.py

+        k: Key tensor.
+        cos: Cosine embedding tensor.
+        sin: Sine embedding tensor.
+        position_ids: Deprecated and unused; present for API compatibility.


Let's delete, we have no requirement to be backwards compatible

sophie-xhonneux · 2026-01-06T11:04:28Z

src/weathergen/model/positional_encoding.py

+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):


This looks like 1D-rope to me?

Can we please label the function as such?

sophie-xhonneux · 2026-01-06T11:05:21Z

src/weathergen/model/positional_encoding.py

+    """
+
+    if coords.shape[-1] != 2:
+        raise ValueError(f"coords last dimension must be 2 (lat, lon); got {coords.shape[-1]}")


can we just make these asserts? it is cleaner to read

sophie-xhonneux · 2026-01-06T11:08:18Z

src/weathergen/model/attention.py

        self.flex_attention = torch.compile(flex_attention, dynamic=False)

-    def forward(self, x, ada_ln_aux=None):
+    def forward(self, x, coords=None, ada_ln_aux=None):


From a code perspective I prefer if we decide at init time if this attention module has 2D rope, because a) we may wish to learn the frequencies theta_t and b) we should not use an attention module trained with 2D-rope without it (or at least be very intentional about it)

wang85 and others added 30 commits July 16, 2025 10:07

Replace cf.rank==0 with utils.distributed.is_root

317501e

replace cf.rank==0 with weathergen.utils.distributed.is_root

77de417

Merge branch 'ecmwf:develop' into develop

6439618

Merge branch 'ecmwf:develop' into develop

8993875

Merge branch 'ecmwf:develop' into develop

f4a9d85

Merge branch 'ecmwf:develop' into develop

f8fdef4

Merge branch 'ecmwf:develop' into develop

ca89e7b

Merge branch 'ecmwf:develop' into develop

49d7a4d

Merge branch 'ecmwf:develop' into develop

f39f094

Merge branch 'ecmwf:develop' into develop

ebb03ea

Merge branch 'ecmwf:develop' into develop

f40737d

Merge branch 'ecmwf:develop' into develop

87fa078

Merge branch 'ecmwf:develop' into develop

5dfe275

Merge branch 'ecmwf:develop' into develop

b7244d9

Merge branch 'ecmwf:develop' into develop

5be41f5

Merge branch 'ecmwf:develop' into develop

39d3965

Merge branch 'ecmwf:develop' into develop

015ec88

Merge branch 'ecmwf:develop' into develop

cb1b7cc

Merge branch 'ecmwf:develop' into develop

90da4cf

Merge branch 'ecmwf:develop' into develop

f04891b

Merge branch 'ecmwf:develop' into develop

105d992

Merge branch 'ecmwf:develop' into develop

5f56073

Merge branch 'ecmwf:develop' into develop

95ee18a

Merge branch 'ecmwf:develop' into develop

3c702d3

Merge branch 'ecmwf:develop' into develop

6f14a30

Merge branch 'ecmwf:develop' into develop

5e87881

Merge branch 'ecmwf:develop' into develop

0c7d305

Merge branch 'ecmwf:develop' into develop

e43ac94

Merge branch 'ecmwf:develop' into develop

5f63bcc

Merge branch 'ecmwf:develop' into develop

c51eb94

csjfwang and others added 26 commits December 9, 2025 10:41

Merge branch 'ecmwf:develop' into develop

765276a

add 2d rope to develop

734a96b

simplify assimilate global, forecast mode config

e4be7c3

add 2d rope to forecast eigine only once

b193655

only keep global & forecast engine add 2d rope

77d95ed

simplify the code

3928e59

fix lint

b542067

small fix

8538291

fix annotation

81a52b5

fix lint

a8afd35

add annotation

5a48898

Merge branch 'ecmwf:develop' into develop

f3eb78a

Merge branch 'develop' into develop_2d_rope_final_gf

805293f

default config

dc914ea

fix default use_reentrant

d9e0504

use_2d_rope false as defaut

6e7f1ed

fix lint

7e0aff2

fix lint

2bc76dd

add 2d rope to all forecast steps

11761bd

merge confs

0d35319

more confs

c6938fb

add missing enumerate

cca2c23

def forecast config

34378c6

aux_info=None in Forecast Eng forward

e39f56b

lint

3abf188

github-project-automation bot added this to WeatherGen-dev Dec 30, 2025

clessig reviewed Jan 5, 2026

View reviewed changes

clessig mentioned this pull request Jan 5, 2026

Add 2D RoPE #1445

Closed

4 tasks

sophie-xhonneux reviewed Jan 6, 2026

View reviewed changes



		####################################################################################################
		def rotary_embedding_2d(coords, dim, base=10000.0):

		return torch.cat((-x2, x1), dim=-1)


		def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):

2D Rope into develop #1540

Are you sure you want to change the base?

2D Rope into develop #1540

Uh oh!

Conversation

kctezcan commented Dec 30, 2025

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!

sophie-xhonneux 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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants