Adding partial convolution option in Samudra by William-gregory · Pull Request #1348 · ai2cm/ace

William-gregory · 2026-07-08T09:05:52Z

This PR adds a partial convolution option to the Samudra model, motivated by the need to reduce grid-scale noise in Samudra rollouts (see comment below). The core change are to:

fme.ace.models.ocean.m2lines.samudra
fme.ace.models.ocean.m2lines.layers
fme.ace.registry.m2lines

The partial convolution option has also been included in fme.ace.models.ocean.m2lines.test_samudra.

The PR relies on in_names in the ModuleConfig and ModuleSelector build methods. This was needed so that masks can be aligned with individual channels in the input tensor. Therefore, the PR also includes updates to:

fme.ace.models.registry
fme.core.models.mlp
fme.core.registry
fme.core.step.single_module

jpdunc23

Thank you for adding this! These are very promising results, so we should definitely try to get this merged without too much delay.

For the PR description, please follow https://github.com/ai2cm/ace/blob/main/.github/pull_request_template.md. Plots are good for PR conversation, but not so good in the description which becomes the body of the commit message once the PR is squashed-and-merged. Also, I recommend a more formal style for the description, e.g. avoiding first person ("Adds xyz feature to" vs "I'm adding this feature because"). The exact content of your current description would be good to copy paste into a comment in the PR discussion before making edits.

Asking @elynnwu if she can help review.

jpdunc23 · 2026-07-08T19:49:34Z

+        return x.permute(0, 3, 1, 2).contiguous()
+
+
+class PartialConv2d(torch.nn.Module):


What is the provenance of this class? The idea seems to come from https://github.com/NVIDIA/partialconv/blob/master/models/partialconv2d.py but the implementation looks pretty different.

Assuming this is related to the NVIDIA project, please include a link to https://arxiv.org/abs/1804.07723 in the docstring.

Or if there is some other provenance, we give attribution somehow.

jpdunc23 · 2026-07-08T19:59:10Z

        n_in_channels: int,
        n_out_channels: int,
        dataset_info: DatasetInfo,
+        in_names: list[str] | None = None,


Hmm, I think we can avoid this change. We have dataset_info here which has an attribute spatial_mask_provider, so ideally we should be able to construct the mask tensor with the right network channel order. The hard part is ensuring that the order is the same as the packer order.

@mcgibbon any thoughts?

jpdunc23 · 2026-07-08T20:05:54Z

+        x = x.permute(0, 2, 3, 1)
+        x = self.norm(x)
+        return x.permute(0, 3, 1, 2).contiguous()


We should be able to avoid the permute here, which can be quite slow. See e.g.

ace/fme/core/models/conditional_sfno/layers.py

Line 95 in 437397b

class ChannelLayerNorm(nn.Module):

@mcgibbon how do you feel about models outside of core/models/conditional_sfno reusing some pieces? Or would it be better to copy-paste the implementation?

jpdunc23 · 2026-07-08T20:15:21Z

-                    norm_kwargs=self.norm_kwargs,
-                    upscale_factor=self.upscale_factor,
-                    checkpoint_strategy=self.checkpoint_strategy,
+            if i == 0:


Why is it that only the first layer uses partial convolutions? Do the masks not flow through the layers?

This seems somewhat incongruous with having MaskedAvgPool2d for every down layer, but maybe there's a good reason to do it this way.

jpdunc23 · 2026-07-08T20:33:06Z

+        output = self.skip(residual_input) + y
+
+        # Latent features are only valid over actual ocean columns.
+        return output * ocean_column_mask


If I understand correctly, ocean_column_mask is just the surface mask. The invalid area of the surface mask (i.e., land) is always a subset of the invalid area of deeper layers, so does output * ocean_column_mask actually change the output in any way?

jpdunc23 · 2026-07-08T20:36:47Z

+from fme.core.dataset_info import DatasetInfo


 class Samudra(torch.nn.Module):


With this amount of changes to an existing model, I'm tempted to create a new class (e.g., MaskAwareSamudra) rather than risk introducing potentially subtle changes to existing checkpoints.

@elynnwu thoughts?

Yes, it seems like there're a lot of changes here that might be worth splitting it off. The one concern for doing that is we'll eventually diverge too much.

I'm wondering if we should do a bit more ablations to nail down exactly what is causing this noisy behavior. I've previously also had this issue that eventually (mostly) went away due to the way masks are defined and training longer.

If possible, maybe we should try a these ablations:

Input masking only. Multiply the input by the per-channel wet mask (x * channel_mask) before the existing layer 0. No architecture change at all — layer 0 stays a plain ConvNeXtBlock, pooling stays plain AvgPool

Input masking + mask-as-feature. Same as above, plus concatenate the per-channel mask as extra input channels (so layer 0's in_channels doubles, but it's still an ordinary convolution — no partial-conv renormalization, no masked pooling)

This will tell us if the effect is coming from masking the input, mask-conditioning, or the full partial-conv in this PR.

Thanks @elynnwu. Just to clarify on 1. Is this different to what the original code is doing already? I was under the impression that inputs/outputs are already being masked before/after each forward pass?

elynnwu · 2026-07-08T21:10:25Z

-            convblock.append(
-                torch.nn.LayerNorm(in_channels * upscale_factor, **self.norm_kwargs)
-            )
+            convblock.append(ChannelLayerNorm(in_channels * upscale_factor))


This changes state-dict keys (convblock.N.weight → convblock.N.norm.weight), which breaks loading any existing norm="layer" Samudra checkpoint. This also drops **self.norm_kwargs, is this intended? Our baseline Samudra uses instance norm, so it may not be as critical.

Good point on the backwards compat. I think the code crashed when i tried layer norm with the initial torch.nn.LayerNorm(in_channels * upscale_factor, **self.norm_kwargs) strategy and it seemed I needed the permute step. Given we're just using instance norm anyway, maybe its best if I restore this to its original implementation?

elynnwu · 2026-07-08T21:26:54Z

+from fme.core.dataset_info import DatasetInfo


 class Samudra(torch.nn.Module):


Yes, it seems like there're a lot of changes here that might be worth splitting it off. The one concern for doing that is we'll eventually diverge too much.

I'm wondering if we should do a bit more ablations to nail down exactly what is causing this noisy behavior. I've previously also had this issue that eventually (mostly) went away due to the way masks are defined and training longer.

elynnwu · 2026-07-08T21:27:35Z

+from fme.core.dataset_info import DatasetInfo


 class Samudra(torch.nn.Module):


If possible, maybe we should try a these ablations:

Input masking only. Multiply the input by the per-channel wet mask (x * channel_mask) before the existing layer 0. No architecture change at all — layer 0 stays a plain ConvNeXtBlock, pooling stays plain AvgPool

Input masking + mask-as-feature. Same as above, plus concatenate the per-channel mask as extra input channels (so layer 0's in_channels doubles, but it's still an ordinary convolution — no partial-conv renormalization, no masked pooling)

This will tell us if the effect is coming from masking the input, mask-conditioning, or the full partial-conv in this PR.

William-gregory · 2026-07-09T13:37:57Z

Thanks @jpdunc23 and @elynnwu for looking through all of this! I'm pasting my original PR description with the figures here, following recommendation from @jpdunc23. I will look into your comments and ablation suggestions asap!

"In my training of uncoupled Samudra (using OM4 data), I was finding that my resultant bias maps were quite noisy. The noise was most prevalent near the surface, but could be seen at all depths. I was wondering if this noise was being introduced by convolutions and average pooling being computed across ocean and land grid cells at all depths, which then causes noise to mix across the channels. I therefore added a partial convolution operation to Samudra, which can be turned on in the config by partial_convolutions: true. I've found that this both reduces noise in the resultant bias maps and makes the model converge faster. The wandb plots below show simass rmse and sst bias during inline inference. The simass error without partial convolutions has a pretty consistent profile to what I've seen across my OM4 training and Ai2's CM4 training (flat for some number of epochs, and then begins to drop). With partial convolutions, simass error starts to drop from first epoch. The sst bias map then shows much less grid-scale noise.

To add the partial convolution functionality, I had to edit the ModuleConfig and ModuleSelector build methods (and therefore core.step.single_module) to pass in_names to the model. This is so that I can align each channel with its corresponding mask in dataset_info. Note that the partial convolution is just applied to the first layer of the network, where channels are defined on the physical domain."

William-gregory added 4 commits July 7, 2026 05:56

add partial conv option in samudra

245823b

hard code samudra layer0 norm

38716c3

fix missing in_names in module build

e62527e

Merge branch 'main' into samudra_partialConv

8e13a49

jpdunc23 self-requested a review July 8, 2026 16:41

jpdunc23 reviewed Jul 8, 2026

View reviewed changes

elynnwu reviewed Jul 8, 2026

View reviewed changes

Merge branch 'main' into samudra_partialConv

e38f739

		return x.permute(0, 3, 1, 2).contiguous()


		class PartialConv2d(torch.nn.Module):

		from fme.core.dataset_info import DatasetInfo


		class Samudra(torch.nn.Module):

Uh oh!

Conversation

William-gregory commented Jul 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jpdunc23 left a comment • 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!

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!

William-gregory commented Jul 9, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

William-gregory commented Jul 8, 2026 •

edited

Loading

jpdunc23 left a comment •

edited

Loading