diff --git a/config/jepa_multi_data_pipeline.yml b/config/jepa_multi_data_pipeline.yml
index f03e8a092..0ec19189f 100644
--- a/config/jepa_multi_data_pipeline.yml
+++ b/config/jepa_multi_data_pipeline.yml
@@ -24,51 +24,32 @@ stages:
   # Stage 1: JEPA pre-training with EMA teacher + deep SSL
   - name: pretrain
     command: train
-    base-config: ./config/config_jepa_multi_data_all_years.yml
-    config: ./config/config_jepa_multi_data_all_years.yml
-    # base-config: ./config/config_jepa_self_flow_multi_data.yml
-    # config: ./config/config_jepa_self_flow_multi_data.yml
+    base-config: config/jepa/space_jepa/space_jepa_and_reconstruction_pretraining_era5_o96_obs_o256_1979_2015.yml
+    config: config/jepa/space_jepa/space_jepa_and_reconstruction_pretraining_era5_o96_obs_o256_1979_2015.yml
     options: []
     chain: 2
     nodes: 2
 
-  # Stage 1: JEPA pre-training with EMA teacher + deep SSL
+  # Stage 2: JEPA pre-training with EMA teacher + deep SSL
   - name: pretrain-ft
-    command: train
-    config: ./config/config_jepa_multi_data_all_years_ft.yml
+    command: train-continue
+    config: config/jepa/space_jepa/space_jepa_and_reconstruction_pretraining_era5_o96_obs_o256_2016_2022.yml
     options: []
-    chain: 2
+    chain: 1
     nodes: 2
 
-  # # Stage 1.1: JEPA pre-training with EMA teacher + deep SSL
-  # - name: pretrain-cooldown
-  #   command: train-continue
-  #   from: pretrain
-  #   config: ./config/config_jepa_self_flow_cooldown.yml
-  #   options: []
-  #   chain: 1
-  #   nodes: 2
-
   # Stage 3: Forecasting finetuning (freezes encoder)
   - name: finetune-forecast
     command: train-continue
     from: pretrain-ft
-    config: ./config/config_jepa_multi_data_ft_forecast_all_years.yml
-    chain: 4
+    config: config/jepa/space_jepa/space_jepa_frozen_forecast_era5_o96_obs_o256.yml
+    chain: 2
     nodes: 2
 
-  # Stage 3: Forecasting finetuning (freezes encoder)
-  - name: finetune-forecast-ft
+  # Stage 3: Forecasting finetuning (unfreezes encoder)
+  - name: finetune-forecast-unfrozen
     command: train-continue
-    from: finetune-forecast
-    config: ./config/config_jepa_multi_data_ft_forecast.yml
-    chain: 3
+    from: pretrain-ft
+    config: config/jepa/space_jepa/space_jepa_unfrozen_forecast_era5_o96_obs_o256.yml
+    chain: 2
     nodes: 2
-
-  # # Stage 3: Forecasting finetuning (unfreezes encoder)
-  # - name: finetune-forecast-2
-  #   command: train-continue
-  #   from: finetune-forecast
-  #   config: ./config/config_jepa_multi_data_ft_forecast_unfreeze_encoder.yml
-  #   chain: 2
-  #   nodes: 2
diff --git a/config/streams/jepa_forecast_multi_data_all_years/avhrr.yml b/config/streams/jepa_forecast_multi_data_all_years/avhrr.yml
index 49597a39b..320d8ffae 100644
--- a/config/streams/jepa_forecast_multi_data_all_years/avhrr.yml
+++ b/config/streams/jepa_forecast_multi_data_all_years/avhrr.yml
@@ -8,7 +8,7 @@ METOP_ABC_AVHRR_IASI:
   forcing: True
   loss_weight: 1.0
   token_size: 512
-  tokenize_spacetime: True
+  tokenize_spacetime: False
   embed:
     net: transformer
     num_tokens: 1
diff --git a/config/streams/jepa_forecast_multi_data_all_years/synop.yml b/config/streams/jepa_forecast_multi_data_all_years/synop.yml
deleted file mode 100644
index a5dae3f9c..000000000
--- a/config/streams/jepa_forecast_multi_data_all_years/synop.yml
+++ /dev/null
@@ -1,29 +0,0 @@
-# obs_types
-# 0 : polar orbiting satellites
-# 1 : geostationay satellites
-# 2 : conventional observations
-
-SurfaceCombined :
-  type : obs
-  stream_id : 2
-  filenames : ['observations-ea-ofb-0001-1979-2022-combined-surface-v5-fixed-land-spatial80-min1km-lat60S70N-lsm09-min10.zarr', 'observations-ea-ofb-0001-2023-combined-surface-v5-fixed-land-heldout20-min1km-lat60S70N-lsm09-min10.zarr']
-  geoinfos: ['reportype', 'stalt', 'lsm']
-  forcing: True
-  token_size : 64
-  tokenize_spacetime : False
-  max_num_targets: -1
-  embed : 
-    net : transformer
-    num_tokens : 1
-    num_heads : 2
-    dim_embed : 256
-    num_blocks : 2
-  embed_target_coords : 
-    net : linear
-    dim_embed : 256
-  target_readout :
-    num_layers : 2
-    num_heads : 4
-  pred_head : 
-    ens_size : 1
-    num_layers : 1
diff --git a/config/streams/jepa_forecast_multi_data_od/avhrr.yml b/config/streams/jepa_forecast_multi_data_od/avhrr.yml
index 49597a39b..320d8ffae 100644
--- a/config/streams/jepa_forecast_multi_data_od/avhrr.yml
+++ b/config/streams/jepa_forecast_multi_data_od/avhrr.yml
@@ -8,7 +8,7 @@ METOP_ABC_AVHRR_IASI:
   forcing: True
   loss_weight: 1.0
   token_size: 512
-  tokenize_spacetime: True
+  tokenize_spacetime: False
   embed:
     net: transformer
     num_tokens: 1
diff --git a/config/streams/jepa_forecast_multi_data_od/synop.yml b/config/streams/jepa_forecast_multi_data_od/synop.yml
deleted file mode 100644
index a5dae3f9c..000000000
--- a/config/streams/jepa_forecast_multi_data_od/synop.yml
+++ /dev/null
@@ -1,29 +0,0 @@
-# obs_types
-# 0 : polar orbiting satellites
-# 1 : geostationay satellites
-# 2 : conventional observations
-
-SurfaceCombined :
-  type : obs
-  stream_id : 2
-  filenames : ['observations-ea-ofb-0001-1979-2022-combined-surface-v5-fixed-land-spatial80-min1km-lat60S70N-lsm09-min10.zarr', 'observations-ea-ofb-0001-2023-combined-surface-v5-fixed-land-heldout20-min1km-lat60S70N-lsm09-min10.zarr']
-  geoinfos: ['reportype', 'stalt', 'lsm']
-  forcing: True
-  token_size : 64
-  tokenize_spacetime : False
-  max_num_targets: -1
-  embed : 
-    net : transformer
-    num_tokens : 1
-    num_heads : 2
-    dim_embed : 256
-    num_blocks : 2
-  embed_target_coords : 
-    net : linear
-    dim_embed : 256
-  target_readout :
-    num_layers : 2
-    num_heads : 4
-  pred_head : 
-    ens_size : 1
-    num_layers : 1
diff --git a/packages/common/src/weathergen/common/config.py b/packages/common/src/weathergen/common/config.py
index e18de025a..a275c6882 100644
--- a/packages/common/src/weathergen/common/config.py
+++ b/packages/common/src/weathergen/common/config.py
@@ -422,6 +422,9 @@ def load_merge_configs(
     else:
         base_config = load_run_config(from_run_id, mini_epoch, None)
         from_run_id = get_run_id_from_config(base_config)
+        with open_dict(base_config):
+            # one-shot action key: must be set per stage, not inherited from the previous run
+            base_config.pop("reset_modules", None)
     with open_dict(base_config):
         base_config.from_run_id = from_run_id
     # use OmegaConf.unsafe_merge if too slow
diff --git a/src/weathergen/datasets/data_reader_obs.py b/src/weathergen/datasets/data_reader_obs.py
index 62b1dcfba..8aa5d196d 100644
--- a/src/weathergen/datasets/data_reader_obs.py
+++ b/src/weathergen/datasets/data_reader_obs.py
@@ -48,7 +48,7 @@ def __init__(
         # To read idx convert to a string, format e.g.: 197001010000
         base_date_str = dt_obj.strftime("%Y%m%d%H%M")
         self.hrly_index = self.z[f"idx_{base_date_str}_1"]
-        self.colnames = self.data.attrs["colnames"]
+        self.colnames = list(self.data.attrs["colnames"])
 
         data_colnames = [col for col in self.colnames if "obsvalue" in col]
         data_idx = [i for i, col in enumerate(self.colnames) if "obsvalue" in col]
@@ -74,7 +74,20 @@ def __init__(
         self.target_idx = np.array(self.target_idx, dtype=np.int64)
 
         # determine idx for coords and geoinfos
-        self.coords_idx = [self.colnames.index("lat"), self.colnames.index("lon")]
+        coords_channels = stream_info.get("coords_channels", ["lat", "lon"])
+        assert len(coords_channels) == 2, (
+            f"{stream_info['name']}: 'coords_channels' must be a list of exactly two "
+            f"names [lat, lon], got {coords_channels!r}."
+        )
+        lat_name, lon_name = coords_channels
+        for name in (lat_name, lon_name):
+            n = self.colnames.count(name)
+            assert n == 1, (
+                f"{stream_info['name']}: coordinate column not found in {self.filename}. "
+                f"Looked for '{lat_name}'/'{lon_name}'; available colnames: {self.colnames}. "
+                f"Set 'coords_channels' in the stream config to match data."
+            )
+        self.coords_idx = [self.colnames.index(lat_name), self.colnames.index(lon_name)]
 
         # geoinfo channels
         sname = stream_info["name"]
@@ -200,7 +213,7 @@ def _setup_sample_index(self) -> None:
                 self.indices_start = np.append(
                     self.indices_start,
                     np.ones(
-                        (diff_in_hours_end - self.hrly_index.shape[0] - 1) // step_hrs, dtype=int
+                        (diff_in_hours_end - (self.hrly_index.shape[0] - 1)) // step_hrs, dtype=int
                     )
                     * self.indices_start[-1],
                 )
@@ -209,7 +222,8 @@ def _setup_sample_index(self) -> None:
                     self.indices_end,
                     np.ones(
                         # add (len_hrs + 1) since above we also have diff_in_hours_start + len_hrs
-                        (diff_in_hours_end - self.hrly_index.shape[0] + (len_hrs + 1)) // step_hrs,
+                        (diff_in_hours_end - (self.hrly_index.shape[0] - 1) + (len_hrs + 1))
+                        // step_hrs,
                         dtype=int,
                     )
                     * self.indices_end[-1],
diff --git a/src/weathergen/datasets/masking.py b/src/weathergen/datasets/masking.py
index 4fc063bcb..2f94f496d 100644
--- a/src/weathergen/datasets/masking.py
+++ b/src/weathergen/datasets/masking.py
@@ -774,7 +774,7 @@ def apply_noise_to_data(
                     noise_level_per_cell = torch.zeros(
                         geoinfos.shape[0], dtype=geoinfos.dtype, device=geoinfos.device
                     )
-                    rd.geoinfos = torch.cat([geoinfos,noise_level_per_cell.unsqueeze(1)], dim=-1)
+                    rd.geoinfos = torch.cat([geoinfos, noise_level_per_cell.unsqueeze(1)], dim=-1)
                     input_data_geoinfo.append(rd)
                 input_data = input_data_geoinfo
             return input_data
@@ -824,7 +824,7 @@ def apply_noise_to_data(
                     noise_level_per_cell = torch.where(
                         point_noise_mask, noise_level_per_cell, noise_level_t
                     )
-                rd.geoinfos = torch.cat([rd.geoinfos,noise_level_per_cell.unsqueeze(1)], dim=-1)
+                rd.geoinfos = torch.cat([rd.geoinfos, noise_level_per_cell.unsqueeze(1)], dim=-1)
 
             noised.append(rd)
         return noised
diff --git a/src/weathergen/datasets/multi_stream_data_sampler.py b/src/weathergen/datasets/multi_stream_data_sampler.py
index d845d2e7d..97c7c4f1b 100644
--- a/src/weathergen/datasets/multi_stream_data_sampler.py
+++ b/src/weathergen/datasets/multi_stream_data_sampler.py
@@ -131,6 +131,7 @@ def __init__(self, cf: Config, mode_cfg: dict, stage: Stage):
             self.teacher_time_offset = 0
 
         self.batch_size = get_batch_size_from_config(mode_cfg)
+        self.num_workers = cf.data_loading.num_workers
         self.shuffle = mode_cfg.shuffle
 
         self.len_timedelta = mode_cfg.time_window_len
@@ -194,8 +195,13 @@ def check_samples(self, fsm: int):
 
         # streamlined calculation of length
         epoch_len = self.samples_per_mini_epoch
-        # adjust len to split loading across all workers and ensure it is multiple of batch_size
-        self.len = ((epoch_len // self.world_size) // self.batch_size) * self.batch_size
+        # adjust len to split loading across all workers and ensure it is multiple of batch_size;
+        # also account for num_workers so per-worker slice is a multiple of batch_size,
+        # preventing the range-loop in __iter__ from yielding extra batches via ceiling division
+        effective_workers = max(1, self.num_workers)
+        self.len = ((epoch_len // self.world_size) // (self.batch_size * effective_workers)) * (
+            self.batch_size * effective_workers
+        )
 
         n_duplicates = self.len * self.world_size - available_samples
         if not self.repeat_data:
@@ -732,7 +738,7 @@ def _get_batch(self, idx: int, num_forecast_steps: int):
                     input_data,
                     source_masks.metadata[sidx],
                     is_student=True,
-                    add_geoinfo_noise="noise_time" in stream_info.get("geoinfo_channels",[]),
+                    add_geoinfo_noise="noise_time" in stream_info.get("geoinfo_channels", []),
                 )
 
                 sdata = self._build_stream_data(
@@ -761,8 +767,10 @@ def _get_batch(self, idx: int, num_forecast_steps: int):
 
                 # Apply self-flow noise to teacher data (handled by masker)
                 input_data_target = self.masker.apply_noise_to_data(
-                    input_data_target_orig, target_masks.metadata[tidx], is_student=False,
-                    add_geoinfo_noise="noise_time" in stream_info.get("geoinfo_channels",[]),
+                    input_data_target_orig,
+                    target_masks.metadata[tidx],
+                    is_student=False,
+                    add_geoinfo_noise="noise_time" in stream_info.get("geoinfo_channels", []),
                 )
 
                 sdata = self._build_stream_data(
diff --git a/src/weathergen/datasets/stream_data.py b/src/weathergen/datasets/stream_data.py
index 0860d5ae8..606567458 100644
--- a/src/weathergen/datasets/stream_data.py
+++ b/src/weathergen/datasets/stream_data.py
@@ -462,7 +462,7 @@ def spoof(healpix_level: int, datetime, geoinfo_size, num_channels) -> IOReaderD
     coords = np.stack([lats.deg, lons.deg], axis=-1, dtype=np.float32)
     # spoof two tokens to avoid unnecessary computational load
     coords = coords[np.random.choice(coords.shape[0], size=2, replace=False)]
-    
+
     geoinfos = np.zeros((coords.shape[0], geoinfo_size), dtype=np.float32)
 
     data = np.zeros((coords.shape[0], num_channels), dtype=np.float32)
diff --git a/src/weathergen/datasets/tokenizer_utils.py b/src/weathergen/datasets/tokenizer_utils.py
index 1e2d9e591..02722475e 100644
--- a/src/weathergen/datasets/tokenizer_utils.py
+++ b/src/weathergen/datasets/tokenizer_utils.py
@@ -91,9 +91,9 @@ def encode_times_target(times, time_win) -> torch.tensor:
     dt = pd.to_datetime(times)
     dt_win = pd.to_datetime(time_win)
     # for target only provide local time
-    dt_delta = torch.tensor(
-        np.atleast_1d((dt - dt_win[0]).seconds), dtype=torch.float32
-    ).unsqueeze(1)
+    dt_delta = torch.tensor(np.atleast_1d((dt - dt_win[0]).seconds), dtype=torch.float32).unsqueeze(
+        1
+    )
     time_tensor = torch.cat(
         (
             dt_delta,
@@ -543,4 +543,11 @@ def get_target_coords_local(
     zi = 99
     a[..., (geoinfo_offset + zi) :] = target_coords[..., (geoinfo_offset + 2) :]
 
+    # Careful when merging develop-ssl into develop.
+    # This is not to be merged in to develop.
+    a[..., 98] = np.sin(coords[:, 0])
+    a[..., 97] = np.cos(coords[:, 0])
+    a[..., 96] = np.sin(coords[:, 1])
+    a[..., 95] = np.cos(coords[:, 1])
+
     return a
diff --git a/src/weathergen/model/ema.py b/src/weathergen/model/ema.py
index 63caeb6fd..cb17846d8 100644
--- a/src/weathergen/model/ema.py
+++ b/src/weathergen/model/ema.py
@@ -30,10 +30,10 @@ def __init__(
         self.rampup_ratio = rampup_ratio
         self.ema_model = empty_model
         self.is_model_sharded = is_model_sharded
-        self.batch_size = 1
-        # Build a name → param map once
+        # Build a name → param map once; used to precompute the update pairs.
         self.src_params = dict(self.original_model.named_parameters())
-
+        self._ema_update_params: list[torch.nn.Parameter] = []
+        self._src_update_params: list[torch.nn.Parameter] = []
         self.reset()
 
     @torch.no_grad()
@@ -56,13 +56,32 @@ def reset(self):
         if needs_strip:
             maybe_sharded_sd = {k.removeprefix("module."): v for k, v in maybe_sharded_sd.items()}
         mkeys, ukeys = self.ema_model.load_state_dict(maybe_sharded_sd, strict=False, assign=False)
+        self._rebuild_update_pairs()
         self.ema_model.eval()
 
-    def requires_grad_(self, flag: bool):
-        for p in self.ema_model.parameters():
-            p.requires_grad = flag
+    def _resolve_src_param(self, name: str):
+        p_src = self.src_params.get(name, None)
+        if p_src is None:
+            p_src = self.src_params.get("module." + name, None)
+        return p_src
+
+    def _rebuild_update_pairs(self):
+        """Precompute EMA/source parameter pairs once so update() stays branch-light."""
+        ema_params = []
+        src_params = []
+        for name, p_ema in self.ema_model.named_parameters():
+            p_src = self._resolve_src_param(name)
+            if p_src is None:
+                raise AssertionError(
+                    f"{name}: All parameters of the EMA model must be in the base model."
+                )
+            ema_params.append(p_ema)
+            src_params.append(p_src)
 
-    def get_current_beta(self, cur_step: int) -> float:
+        self._ema_update_params = ema_params
+        self._src_update_params = src_params
+
+    def get_current_beta(self, cur_step: int, batch_size: int) -> float:
         """
         Get current EMA beta value for monitoring.
 
@@ -78,30 +97,20 @@ def get_current_beta(self, cur_step: int) -> float:
         halflife_steps = self.halflife_steps
         if self.rampup_ratio is not None:
             halflife_steps = min(halflife_steps, cur_step * self.rampup_ratio)
-        beta = 0.5 ** (self.batch_size / max(halflife_steps, 1e-6))
+        beta = 0.5 ** (batch_size / max(halflife_steps, 1e-6))
         return beta
 
     @torch.no_grad()
-    def update(self, cur_step, batch_size):
+    def update(self, cur_step: int, batch_size: int):
         # ensure model remains sharded
         if self.is_model_sharded:
             self.ema_model.reshard()
-        # determine correct interpolation params
-        self.batch_size = batch_size
-        beta = self.get_current_beta(cur_step)
 
-        for name, p_ema in self.ema_model.named_parameters():
-            p_src = self.src_params.get(name, None)
-            # Due to DDP only being applied only to the student the names may missmatch
-            # Thus, we check for the alternate naming scheme
-            p_src = self.src_params.get("module." + name, None) if p_src is None else p_src
-            if "identity" in name.lower():
-                continue
-            if p_src is None:
-                # EMA-only param or intentionally excluded
-                assert False, f"{name}: All parameters of the EMA model must be in the base model."
+        # determine correct interpolation params
+        beta = self.get_current_beta(cur_step, batch_size)
 
-            p_ema.lerp_(p_src, 1.0 - beta)
+        torch._foreach_mul_(self._ema_update_params, beta)
+        torch._foreach_add_(self._ema_update_params, self._src_update_params, alpha=1.0 - beta)
 
     @torch.no_grad()
     def forward_eval(self, *args, **kwargs):
diff --git a/src/weathergen/model/engines.py b/src/weathergen/model/engines.py
index 3a5842b04..d8970ce52 100644
--- a/src/weathergen/model/engines.py
+++ b/src/weathergen/model/engines.py
@@ -32,6 +32,7 @@
 from weathergen.model.utils import ActivationFactory
 from weathergen.utils.utils import get_dtype
 
+
 class EmbeddingEngine(torch.nn.Module):
     name: "EmbeddingEngine"
 
@@ -104,7 +105,7 @@ def forward(self, batch, pe_embed):
             # skip empty stream
             if sdata.numel() == 0:
                 continue
-            
+
             # embedding from physical space to per patch latent representation
             x_embeds += [self.embeds[stream_name](sdata).flatten(0, 1)]
 
diff --git a/src/weathergen/model/model.py b/src/weathergen/model/model.py
index 7818612be..319ee9ef8 100644
--- a/src/weathergen/model/model.py
+++ b/src/weathergen/model/model.py
@@ -40,7 +40,7 @@
 from weathergen.model.layers import MLP, NamedLinear
 from weathergen.model.utils import get_num_parameters
 from weathergen.utils.distributed import is_root
-from weathergen.utils.utils import get_dtype, is_stream_forcing
+from weathergen.utils.utils import get_dtype, is_stream_reconstructed
 
 logger = logging.getLogger(__name__)
 
@@ -390,7 +390,8 @@ def create(self) -> "Model":
         )
 
         mode_cfg = cf.training_config
-        if cf.fe_num_blocks > 0:
+        fe_num_blocks = cf.get("fe_num_blocks", 0)
+        if fe_num_blocks > 0:
             self.forecast_engine = ForecastingEngine(cf, mode_cfg, self.num_healpix_cells)
         else:
             self.forecast_engine = IdentityEngine()
@@ -419,8 +420,9 @@ def create(self) -> "Model":
             for i_stream, si in enumerate(cf.streams):
                 stream_name = self.stream_names[i_stream]
 
-                # skip decoder if channels are empty
-                if is_stream_forcing(si):
+                # skip decoder for streams that are not physically reconstructed
+                # (forcing/input-only, or explicit reconstruct: false -> JEPA-only target)
+                if not is_stream_reconstructed(si):
                     continue
 
                 # skip for the moment to ensure target embedding and tte exist (ordering of
@@ -515,8 +517,9 @@ def create(self) -> "Model":
             for i_stream, si in enumerate(cf.streams):
                 stream_name = self.stream_names[i_stream]
 
-                # skip decoder if channels are empty
-                if is_stream_forcing(si):
+                # skip decoder for streams that are not physically reconstructed
+                # (forcing/input-only, or explicit reconstruct: false -> JEPA-only target)
+                if not is_stream_reconstructed(si):
                     continue
 
                 pred_spatial_shared = si.get("pred_spatial_shared")
@@ -779,10 +782,14 @@ def forward(self, model_params: ModelParams, batch: ModelBatch) -> ModelOutput:
 
             if self.forecast_engine:
                 tokens = self.forecast_engine(tokens, step, model_params.rope_coords)
-            # decoder predictions
-            output = self.predict_decoders(model_params, step, tokens, batch, output)
-            # latent predictions (raw and with SSL heads)
-            output = self.predict_latent(model_params, step, tokens, batch, output, intermediates)
+            if "masking" in self.cf.training_config.training_mode:
+                # decoder predictions
+                output = self.predict_decoders(model_params, step, tokens, batch, output)
+            if "student_teacher" in self.cf.training_config.training_mode:
+                # latent predictions (raw and with SSL heads)
+                output = self.predict_latent(
+                    model_params, step, tokens, batch, output, intermediates
+                )
 
         return output
 
@@ -890,6 +897,12 @@ def predict_decoders(
 
         # pair with tokens from assimilation engine to obtain target tokens
         for stream_name in self.stream_names:
+            # streams without a physical decoder (forcing, or reconstruct: false JEPA-only
+            # targets) have no embed_target_coords/target_token_engine. Skip them here even
+            # though they may still carry (unused) target coords on the student view.
+            if stream_name not in self.embed_target_coords:
+                continue
+
             # extract target coords for current stream and fstep and convert to one tensor
             t_coords = [
                 batch.samples[i_b].streams_data[stream_name].target_coords[step]
diff --git a/src/weathergen/model/model_interface.py b/src/weathergen/model/model_interface.py
index 6b42fca9a..121104666 100644
--- a/src/weathergen/model/model_interface.py
+++ b/src/weathergen/model/model_interface.py
@@ -29,7 +29,14 @@
 )
 from weathergen.model.layers import MLP
 from weathergen.model.model import Model, ModelParams
-from weathergen.model.utils import apply_fct_to_blocks, freeze_weights
+from weathergen.model.utils import (
+    apply_fct_to_blocks,
+    broadcast_matching_params,
+    check_reset_not_frozen,
+    freeze_weights,
+    log_trainable_summary,
+    reset_weights,
+)
 from weathergen.utils.distributed import is_root
 from weathergen.utils.utils import get_dtype
 
@@ -40,6 +47,20 @@
 type TrainingMode = str
 
 
+def _has_trainable_params(module: torch.nn.Module) -> bool:
+    """True if the module has at least one parameter with requires_grad=True.
+
+    FSDP2 raises "RuntimeError: _chunk_cat expects non-empty tensor" in the
+    backward reduce-scatter (foreach_reduce) when a fully_shard group contains
+    only frozen parameters, since there are no gradients to reduce. This happens
+    during fine-tuning (e.g. forecast fine-tuning freezes the encoder and
+    latent_heads). Skipping fully_shard for fully-frozen modules leaves their
+    parameters in the root FSDP group, which still has trainable parameters, so
+    they remain sharded without triggering the empty-gradient reduce.
+    """
+    return any(p.requires_grad for p in module.parameters())
+
+
 def init_model_and_shard(
     cf,
     dataset,
@@ -96,36 +117,36 @@ def init_model_and_shard(
         )
 
         for module in model.encoder.ae_local_engine.ae_local_blocks.modules():
-            if isinstance(module, modules_to_shard):
+            if isinstance(module, modules_to_shard) and _has_trainable_params(module):
                 fully_shard(module, **fsdp_kwargs)
 
         for module in model.encoder.ae_local_global_engine.ae_adapter.modules():
-            if isinstance(module, modules_to_shard):
+            if isinstance(module, modules_to_shard) and _has_trainable_params(module):
                 fully_shard(module, **fsdp_kwargs)
 
         for module in model.encoder.ae_global_engine.ae_global_blocks.modules():
-            if isinstance(module, modules_to_shard):
+            if isinstance(module, modules_to_shard) and _has_trainable_params(module):
                 fully_shard(module, **fsdp_kwargs)
 
         for module in model.forecast_engine.fe_blocks.modules():
-            if isinstance(module, modules_to_shard):
+            if isinstance(module, modules_to_shard) and _has_trainable_params(module):
                 # reshard_after_forward=False keeps FE parameters unsharded
                 # during the multi-step rollout loop.
                 # Needed for pushforward trick.
                 fully_shard(module, reshard_after_forward=False, **fsdp_kwargs)
 
         for module in model.latent_heads.modules():
-            if isinstance(module, modules_to_shard):
+            if isinstance(module, modules_to_shard) and _has_trainable_params(module):
                 fully_shard(module, **fsdp_kwargs)
 
         if model.deep_ssl_fusion is not None:
             for module in model.deep_ssl_fusion.modules():
-                if isinstance(module, modules_to_shard):
+                if isinstance(module, modules_to_shard) and _has_trainable_params(module):
                     fully_shard(module, **fsdp_kwargs)
 
         if model.deep_ssl_level_projections is not None:
             for module in model.deep_ssl_level_projections.modules():
-                if isinstance(module, modules_to_shard):
+                if isinstance(module, modules_to_shard) and _has_trainable_params(module):
                     fully_shard(module, **fsdp_kwargs)
 
         full_precision_fsdp_kwargs = {
@@ -140,7 +161,7 @@ def init_model_and_shard(
         }
 
         for module in model.target_token_engines.modules():
-            if isinstance(module, modules_to_shard):
+            if isinstance(module, modules_to_shard) and _has_trainable_params(module):
                 fully_shard(module, **full_precision_fsdp_kwargs)
 
     if with_ddp and with_fsdp:
@@ -158,22 +179,43 @@ def init_model_and_shard(
             torch.distributed.fsdp.register_fsdp_forward_method(embed, "forward_columns")
 
     # complete initalization and load model if inference/continuing a run
+    loaded_from_run_id = None
     if run_id_contd is not None:
         if is_root():
             logger.info(f"Continuing run with id={run_id_contd} at mini_epoch {mini_epoch_contd}.")
         model = load_model(cf, model, device, run_id_contd, mini_epoch_contd)
+        loaded_from_run_id = run_id_contd
     elif cf.get("load_chkpt", {}).get("run_id", None):
         run_id = cf.load_chkpt.run_id
         mini_epoch = cf.load_chkpt.get("mini_epoch", -1)
         if is_root():
             logger.info(f"Loading checkpoint from id={run_id} at mini_epoch {mini_epoch}.")
         model = load_model(cf, model, device, run_id, mini_epoch)
+        loaded_from_run_id = run_id
     else:
         if with_ddp and with_fsdp:
             model.to_empty(device="cuda")
             if with_fsdp:
                 model.reset_parameters()
 
+    # Reset specified modules when starting a new stage (e.g. pretrain -> finetune);
+    # skip when resuming the same run.
+    current_run_id = cf.general.run_id
+    if loaded_from_run_id is not None and loaded_from_run_id != current_run_id:
+        reset_modules = cf.get("reset_modules", "")
+        if reset_modules:
+            assert not with_fsdp, "reset_modules with FSDP-sharded parameters is not supported"
+            # a parameter that is both reset and frozen would stay random forever
+            check_reset_not_frozen(model, reset_modules)
+            if is_root():
+                logger.info(f"Resetting weights for modules matching: {reset_modules}")
+            apply_fct_to_blocks(model, reset_modules, reset_weights)
+            # each rank resets with its own RNG; sync to rank 0 like DDP does at wrap time
+            broadcast_matching_params(model, reset_modules, src=0)
+
+    if is_root():
+        log_trainable_summary(model)
+
     # model params
     model_params = ModelParams(cf).create(cf)
     model_params.reset_parameters(cf)
diff --git a/src/weathergen/model/utils.py b/src/weathergen/model/utils.py
index 7dd2060bb..b9fe5f3b2 100644
--- a/src/weathergen/model/utils.py
+++ b/src/weathergen/model/utils.py
@@ -29,6 +29,15 @@ def freeze_weights(block):
         p.requires_grad = False
 
 
+def reset_weights(block):
+    block_name = getattr(block, "name", type(block).__name__)
+    if hasattr(block, "reset_parameters"):
+        logger.info(f"Reset weights of block {block_name}")
+        block.reset_parameters()
+    else:
+        logger.info(f"Skip reset for block {block_name} (no reset_parameters)")
+
+
 def set_to_eval(block):
     if hasattr(block, "name"):
         logger.info(f"Set block {block.name} to eval mode")
@@ -51,6 +60,85 @@ def apply_fct_to_blocks(model, blocks, fct):
             fct(module)
 
 
+def broadcast_matching_params(model, blocks, src=0):
+    """
+    Broadcast parameters and buffers of blocks matching the regex from rank src to
+    all ranks. Needed after reset_parameters() under DDP: the reset draws from each
+    rank's own RNG, and DDP only syncs parameters at wrap time, so without a
+    broadcast the ranks train permanently diverged weights.
+    Args:
+        model : model instance with attribute named_modules
+        blocks : regex pattern to match block names
+        src : rank whose values are broadcast
+    """
+
+    if not (torch.distributed.is_available() and torch.distributed.is_initialized()):
+        return
+    seen = set()
+    tensors = []
+    for name, module in model.named_modules():
+        name = module.name if hasattr(module, "name") else name
+        if (re.fullmatch(blocks, name) is not None) and (name != ""):
+            for t in list(module.parameters()) + list(module.buffers()):
+                if id(t) not in seen:
+                    seen.add(id(t))
+                    tensors.append(t)
+    for t in tensors:
+        torch.distributed.broadcast(t.data, src=src)
+    logger.info(f"Broadcast {len(tensors)} reset tensors from rank {src}")
+
+
+def check_reset_not_frozen(model, reset_blocks):
+    """
+    Verify that no parameter about to be reset is frozen. A parameter that is reset
+    to random values but has requires_grad=False can never train, leaving random
+    dead weights in the model (almost never intended).
+    Args:
+        model : model instance with attribute named_modules
+        reset_blocks : regex pattern of block names that will be reset
+    Raises:
+        ValueError listing the frozen parameters that match the reset pattern.
+    """
+
+    frozen = []
+    for name, module in model.named_modules():
+        name = module.name if hasattr(module, "name") else name
+        if (re.fullmatch(reset_blocks, name) is not None) and (name != ""):
+            frozen += [f"{name}.{pn}" for pn, p in module.named_parameters() if not p.requires_grad]
+    if frozen:
+        frozen = sorted(set(frozen))
+        raise ValueError(
+            "reset_modules overlaps with frozen parameters; these would be reset to random "
+            "values but never trained. Remove them from freeze_modules or reset_modules: "
+            + ", ".join(frozen[:16])
+            + (" ..." if len(frozen) > 16 else "")
+        )
+
+
+def log_trainable_summary(model):
+    """
+    Log per-top-level-block parameter counts and trainable fractions.
+    """
+
+    # unwrap DDP for readable block names
+    block = model.module if isinstance(model, nn.parallel.DistributedDataParallel) else model
+    logger.info("Trainable parameter summary:")
+    for name, child in block.named_children():
+        n_total = sum(p.numel() for p in child.parameters())
+        if n_total == 0:
+            continue
+        n_train = sum(p.numel() for p in child.parameters() if p.requires_grad)
+        logger.info(
+            f"  {name}: {n_train:,} / {n_total:,} trainable ({100 * n_train / n_total:.1f}%)"
+        )
+    n_total = sum(p.numel() for p in block.parameters())
+    n_train = sum(p.numel() for p in block.parameters() if p.requires_grad)
+    if n_total > 0:
+        logger.info(
+            f"  total: {n_train:,} / {n_total:,} trainable ({100 * n_train / n_total:.1f}%)"
+        )
+
+
 class ActivationFactory:
     _registry = {
         "identity": nn.Identity,
diff --git a/src/weathergen/train/collapse_monitor.py b/src/weathergen/train/collapse_monitor.py
index fa5268ee4..b13453ba8 100644
--- a/src/weathergen/train/collapse_monitor.py
+++ b/src/weathergen/train/collapse_monitor.py
@@ -703,7 +703,7 @@ def extract_latent_tensor(
         for _calc_name, calculator in target_and_aux_calculators.items():
             if isinstance(calculator, EMATeacher):
                 step = batch_size * cf.general.istep
-                ema_beta = calculator.get_current_beta(step)
+                ema_beta = calculator.get_current_beta(step, batch_size)
                 break
 
         # Debug logging
diff --git a/src/weathergen/train/target_and_aux_ssl_teacher.py b/src/weathergen/train/target_and_aux_ssl_teacher.py
index 6afd89967..992ebaca1 100644
--- a/src/weathergen/train/target_and_aux_ssl_teacher.py
+++ b/src/weathergen/train/target_and_aux_ssl_teacher.py
@@ -91,8 +91,8 @@ def to_device(self, device) -> EncoderTeacher:
             module.to(device)
         return self
 
-    def get_current_beta(self, cur_step: int) -> float:
-        beta = self.ema_model.get_current_beta(cur_step)
+    def get_current_beta(self, cur_step: int, batch_size: int) -> float:
+        beta = self.ema_model.get_current_beta(cur_step, batch_size)
         return beta
 
 
@@ -118,9 +118,9 @@ def update_state_post_opt_step(self, istep, batch, model, **kwargs) -> None:
             self.ema_model.ema_model.reshard()
         self.ema_model.update(istep, self.batch_size)
 
-    def get_current_beta(self, cur_step: int) -> float:
+    def get_current_beta(self, cur_step: int, batch_size: int) -> float:
         """Return the current EMA interpolation beta for monitoring."""
-        return self.ema_model.get_current_beta(cur_step)
+        return self.ema_model.get_current_beta(cur_step, batch_size)
 
 
 class FrozenTeacher(EncoderTeacher):
diff --git a/src/weathergen/train/trainer.py b/src/weathergen/train/trainer.py
index 0b3c504db..5bda93797 100644
--- a/src/weathergen/train/trainer.py
+++ b/src/weathergen/train/trainer.py
@@ -256,7 +256,6 @@ def inference(self, cf, devices, run_id_contd, mini_epoch_contd):
 
         logger.info(f"Starting inference with id={self.cf.general.run_id}.")
 
-
         # inference validation set
         self.validate(0, self.test_cfg, self.batch_size_test_per_gpu)
         logger.info(f"Finished inference run with id: {cf.general.run_id}")
@@ -374,7 +373,7 @@ def run(self, cf, devices, run_id_contd=None, mini_epoch_contd=None):
         )
 
         # Restore optimizer momentum buffers when continuing from a checkpoint
-        if run_id_contd is not None:
+        if run_id_contd is not None and self.cf.general.istep != 0:
             self._load_optimizer_state(run_id_contd, mini_epoch_contd)
 
         if self.cf.general.istep > 0 and is_root():
diff --git a/src/weathergen/utils/utils.py b/src/weathergen/utils/utils.py
index 291ab1521..1bda83d9d 100644
--- a/src/weathergen/utils/utils.py
+++ b/src/weathergen/utils/utils.py
@@ -49,6 +49,23 @@ def is_stream_forcing(stream_cfg: dict, stage: Stage | None = None) -> bool:
     return is_forcing
 
 
+def is_stream_reconstructed(stream_cfg: dict, stage: Stage | None = None) -> bool:
+    """
+    Determine if a stream is physically reconstructed, i.e. has a decoder and contributes
+    to the physical (decoder) reconstruction loss.
+
+    A stream is NOT reconstructed if it is forcing (input-only) or if it explicitly opts
+    out via ``reconstruct: false``. The latter lets a stream still serve as a
+    student-teacher (JEPA) target while having no physical decoder, so JEPA can be trained
+    on all streams while only a subset is reconstructed in physical space. Note that, unlike
+    forcing streams, ``reconstruct: false`` streams keep a normal (non-empty) target mask,
+    so the teacher still encodes them.
+    """
+    if is_stream_forcing(stream_cfg, stage):
+        return False
+    return stream_cfg.get("reconstruct", True)
+
+
 def is_stream_diagnostic(stream_cfg: dict, stage: Stage | None = None) -> bool:
     """
     Determine if stream is diagnostic, i.e. does not contribute to model input