Test turning all grad checkpointings off

config/default_config.yml

@@ -6,6 +6,7 @@ embed_centroids_local_coords: False
 embed_size_centroids: 0
 embed_unembed_mode: "block"
 embed_dropout_rate: 0.1
+embed_gradient_checkpoint_mode: False
 
 target_cell_local_prediction: True
 
@@ -17,11 +18,13 @@ ae_local_with_qk_lnorm: True
 
 ae_local_num_queries: 1
 ae_local_queries_per_cell: False
+ae_local_blocks_grdient_checkpoint_mode: False
 ae_adapter_num_heads: 16
 ae_adapter_embed: 128
 ae_adapter_with_qk_lnorm: True
 ae_adapter_with_residual: True
 ae_adapter_dropout_rate: 0.1
+ae_adapter_grdient_checkpoint_mode: False
 
 ae_global_dim_embed: 2048
 ae_global_num_blocks: 8
@@ -33,12 +36,14 @@ ae_global_with_qk_lnorm: True
 ae_global_att_dense_rate: 1.0
 ae_global_block_factor: 64
 ae_global_mlp_hidden_factor: 2
+assimilate_global_gradient_checkpoint_mode: False 
 
 decoder_type: PerceiverIOCoordConditioning # CrossAttentionAdaNormConditioning
 pred_adapter_kv: False
 pred_self_attention: True
 pred_dyadic_dims: False
 pred_mlp_adaln: True
+pred_gradient_checkpoint_mode: False
 
 # number of steps offset applied to first target window; if set to zero and forecast_steps=0 then
 # one is training an auto-encoder

src/weathergen/model/embeddings.py

@@ -34,6 +34,7 @@ def __init__(
         norm_type="LayerNorm",
         embed_size_centroids=64,
         unembed_mode="full",
+        embed_gradient_checkpoint_mode=True,
         stream_name="stream_embed",
     ):
         """Constructor
@@ -59,6 +60,7 @@ def __init__(
         self.num_heads = num_heads
         self.embed_size_centroids = embed_size_centroids
         self.unembed_mode = unembed_mode
+        self.embed_gradient_checkpoint_mode = embed_gradient_checkpoint_mode
 
         norm = torch.nn.LayerNorm if norm_type == "LayerNorm" else RMSNorm
 
@@ -148,23 +150,47 @@ def __init__(
     def forward_channels(self, x_in, centroids):
         peh = positional_encoding_harmonic
 
-        # embed provided input data
-        x = peh(checkpoint(self.embed, x_in.transpose(-2, -1), use_reentrant=False))
+        if self.embed_gradient_checkpoint_mode:
+            # embed provided input data
+            x = peh(checkpoint(self.embed, x_in.transpose(-2, -1), use_reentrant=False))
 
-        for layer in self.layers:
-            x = checkpoint(layer, x, use_reentrant=False)
+            for layer in self.layers:
+                x = checkpoint(layer, x, use_reentrant=False)
+
+            # read out
+            if self.unembed_mode == "full":
+                out = checkpoint(
+                    self.unembed,
+                    self.ln_final(x.flatten(-2, -1)),
+                    use_reentrant=False,
+                )
+            elif self.unembed_mode == "block":
+                out = [
+                    checkpoint(ue, ln(x[:, i]), use_reentrant=False)
+                    for i, (ue, ln) in enumerate(zip(self.unembed, self.ln_final, strict=True))
+                ]
+                out = torch.stack(out, dim=1).flatten(-2, -1)
+            else:
+                assert False
 
-        # read out
-        if self.unembed_mode == "full":
-            out = checkpoint(self.unembed, self.ln_final(x.flatten(-2, -1)), use_reentrant=False)
-        elif self.unembed_mode == "block":
-            out = [
-                checkpoint(ue, ln(x[:, i]), use_reentrant=False)
-                for i, (ue, ln) in enumerate(zip(self.unembed, self.ln_final, strict=True))
-            ]
-            out = torch.stack(out, dim=1).flatten(-2, -1)
         else:
-            assert False
+            # embed provided input data
+            x = peh(self.embed(x_in.transpose(-2, -1)))
+
+            for layer in self.layers:
+                x = layer(x)
+
+            # read out
+            if self.unembed_mode == "full":
+                out = self.unembed(self.ln_final(x.flatten(-2, -1)))
+            elif self.unembed_mode == "block":
+                out = [
+                    ue(ln(x[:, i]))
+                    for i, (ue, ln) in enumerate(zip(self.unembed, self.ln_final, strict=True))
+                ]
+                out = torch.stack(out, dim=1).flatten(-2, -1)
+            else:
+                assert False
 
         # append centroids
         if self.embed_size_centroids > 0:

src/weathergen/model/engines.py

@@ -71,6 +71,7 @@ def create(self) -> torch.nn.ModuleList:
                         norm_type=self.cf.norm_type,
                         embed_size_centroids=self.cf.embed_size_centroids,
                         unembed_mode=self.cf.embed_unembed_mode,
+                        embed_gradient_checkpoint_mode=self.cf.embed_gradient_checkpoint_mode,
                         stream_name=stream_name,
                     )
                 )

src/weathergen/model/model.py

@@ -743,8 +743,12 @@ def assimilate_local(
                 tokens_global_all += [tokens_global_c]
                 continue
 
-            for block in self.ae_local_blocks:
-                tokens_c = checkpoint(block, tokens_c, cell_lens_c, use_reentrant=False)
+            if self.cf.ae_local_blocks_grdient_checkpoint_mode:
+                for block in self.ae_local_blocks:
+                    tokens_c = checkpoint(block, tokens_c, cell_lens_c, use_reentrant=False)
+            else:
+                for block in self.ae_local_blocks:
+                    tokens_c = block(tokens_c, cell_lens_c)
 
             if self.cf.latent_noise_kl_weight > 0.0:
                 tokens_c, posteriors_c = self.interpolate_latents.interpolate_with_noise(
@@ -754,15 +758,24 @@ def assimilate_local(
             else:
                 tokens_c, posteriors = tokens_c, 0.0
 
-            for block in self.ae_adapter:
-                tokens_global_c = checkpoint(
-                    block,
-                    tokens_global_c,
-                    tokens_c,
-                    q_cells_lens_c,
-                    cell_lens_c,
-                    use_reentrant=False,
-                )
+            if self.cf.ae_adapter_grdient_checkpoint_mode:
+                for block in self.ae_adapter:
+                    tokens_global_c = checkpoint(
+                        block,
+                        tokens_global_c,
+                        tokens_c,
+                        q_cells_lens_c,
+                        cell_lens_c,
+                        use_reentrant=False,
+                    )
+            else:
+                for block in self.ae_adapter:
+                    tokens_global_c = block(
+                        tokens_global_c,
+                        tokens_c,
+                        q_cells_lens_c,
+                        cell_lens_c,
+                    )
 
             tokens_global_all += [tokens_global_c]
 
@@ -787,8 +800,12 @@ def assimilate_global(self, model_params: ModelParams, tokens: torch.Tensor) ->
         """
 
         # global assimilation engine and adapter
-        for block in self.ae_global_blocks:
-            tokens = checkpoint(block, tokens, use_reentrant=False)
+        if self.cf.assimilate_global_gradient_checkpoint_mode:
+            for block in self.ae_global_blocks:
+                tokens = checkpoint(block, tokens, use_reentrant=False)
+        else:
+            for block in self.ae_global_blocks:
+                tokens = block(tokens)
 
         return tokens
 
@@ -855,18 +872,30 @@ def predict(
             ## embed token coords, concatenating along batch dimension
             # (which is taking care of through the varlen attention)
             # arguably we should to the mixed precision policy when creating the model in FSDP
-            tc_tokens = torch.cat(
-                [
-                    checkpoint(
-                        tc_embed,
-                        streams_data[i_b][ii].target_coords[fstep],
-                        use_reentrant=False,
-                    )
-                    if len(streams_data[i_b][ii].target_coords[fstep].shape) > 1
-                    else streams_data[i_b][ii].target_coords[fstep]
-                    for i_b in range(len(streams_data))
-                ]
-            )
+            if self.cf.pred_gradient_checkpoint_mode:
+                tc_tokens = torch.cat(
+                    [
+                        checkpoint(
+                            tc_embed,
+                            streams_data[i_b][ii].target_coords[fstep],
+                            use_reentrant=False,
+                        )
+                        if len(streams_data[i_b][ii].target_coords[fstep].shape) > 1
+                        else streams_data[i_b][ii].target_coords[fstep]
+                        for i_b in range(len(streams_data))
+                    ]
+                )
+            else:
+                tc_tokens = torch.cat(
+                    [
+                        tc_embed(
+                            streams_data[i_b][ii].target_coords[fstep],
+                        )
+                        if len(streams_data[i_b][ii].target_coords[fstep].shape) > 1
+                        else streams_data[i_b][ii].target_coords[fstep]
+                        for i_b in range(len(streams_data))
+                    ]
+                )
 
             # skip when coordinate embeddings yields nan (i.e. the coord embedding network diverged)
             if torch.isnan(tc_tokens).any():
@@ -906,6 +935,9 @@ def predict(
             )
 
             # final prediction head to map back to physical space
-            preds_tokens += [checkpoint(self.pred_heads[ii], tc_tokens, use_reentrant=False)]
+            if self.cf.pred_gradient_checkpoint_mode:
+                preds_tokens += [checkpoint(self.pred_heads[ii], tc_tokens, use_reentrant=False)]
+            else:
+                preds_tokens += [self.pred_heads[ii](tc_tokens)]
 
         return preds_tokens