Fix action embedder MLP initialization after FSDP checkpoint load

README.md

@@ -22,6 +22,7 @@ The underlying architecture is a flow-based diffusion transformer that unifies T
 
 ## News
 
+- **[April 2026]** — Bugfix: properly initialize action embedder MLPs (`action_embedder_B_D` / `action_embedder_B_3D`) after loading the base checkpoint, so their weights are no longer left at zero under FSDP meta-device materialization.
 - **[March 2026]** — Released Cosmos-H-Surgical-Simulator for the [Open-H benchmark](#open-h-benchmark)
 
 For Cosmos-Predict2.5 updates, see the [upstream changelog](https://github.com/NVIDIA/Cosmos-Predict2.5).

cosmos_predict2/_src/imaginaire/trainer.py

@@ -40,6 +40,7 @@
 from cosmos_predict2._src.imaginaire.utils import callback, distributed, ema, log, misc
 from cosmos_predict2._src.imaginaire.utils.checkpointer import Checkpointer
 from cosmos_predict2._src.imaginaire.utils.misc import StragglerDetectorV2
+from cosmos_predict2._src.predict2.utils.dtensor_helper import broadcast_dtensor_model_states
 
 
 class ImaginaireTrainer:
@@ -170,6 +171,12 @@ def train(
         self.callbacks.on_optimizer_init_end()
         # Load the model checkpoint and get the starting iteration number.
         iteration = self.checkpointer.load(model, optimizer, scheduler, grad_scaler)
+
+        # Fix: re-initialize action embedder MLPs that were zeroed out by FSDP
+        # meta-device materialization when loading from a base checkpoint that
+        # lacks action_embedder keys.
+        self._reinitialize_action_embedders_if_needed(model, iteration)
+
         grad_accum_iter = 0
         log.critical(f"Distributed parallelism mode: {self.config.trainer.distributed_parallelism}")
         if self.config.trainer.distributed_parallelism == "ddp":
@@ -262,6 +269,87 @@ def train(
         distributed.barrier()
         self.callbacks.on_app_end()
 
+    @staticmethod
+    def _reinitialize_action_embedders_if_needed(model, iteration: int) -> None:
+        """Re-initialize action embedder MLPs if they are all-zero after checkpoint loading.
+
+        Handles three scenarios:
+        1. Fresh fine-tuning from base Cosmos checkpoint (iteration == 0,
+           weights all-zero) — reinitializes.
+        2. Resuming from a **broken** prior fine-tuning whose MLPs are still
+           all-zero (iteration > 0, weights all-zero) — reinitializes and
+           logs a recovery notice.
+        3. Resuming from a healthy fine-tuning with trained action embedders
+           (weights nonzero) — skips.
+
+        After reinit, copies the updated weights into ``model.net_ema`` using
+        the existing EMA updater and synchronizes both modules across ranks.
+        """
+        net = getattr(model, "net", None)
+        if net is None or not hasattr(net, "reinitialize_action_embedders"):
+            return
+
+        action_weight_params = [
+            (name, p)
+            for name, p in net.named_parameters()
+            if "action_embedder" in name
+            and name.endswith(".weight")
+            and (p.data.to_local() if hasattr(p.data, "to_local") else p.data).device.type != "meta"
+        ]
+        if not action_weight_params:
+            log.info("[ACTION-EMB-FIX] Skipped: no materialized action_embedder weight matrices found on model.net.")
+            return
+
+        all_zero = all(
+            (p.data.to_local() if hasattr(p.data, "to_local") else p.data).abs().max().item() == 0.0
+            for _, p in action_weight_params
+        )
+
+        if not all_zero:
+            log.info("[ACTION-EMB-FIX] Skipped: checkpoint already contains nonzero action embedder weights.")
+            return
+
+        if iteration == 0:
+            log.info(
+                "[ACTION-EMB-FIX] Fresh fine-tune from base checkpoint: "
+                "all action_embedder weights are zero at iteration 0, reinitializing (Kaiming uniform)."
+            )
+        else:
+            log.warning(
+                f"[ACTION-EMB-FIX] Recovery: all action_embedder weights are zero at iteration {iteration}; "
+                f"reinitializing (prior run was affected by the dead-MLP initialization bug)."
+            )
+
+        def _mirror_to_ema() -> None:
+            if not (hasattr(model, "net_ema") and model.net_ema is not None):
+                return
+            if hasattr(model, "net_ema_worker") and model.net_ema_worker is not None:
+                model.net_ema_worker.copy_to(src_model=model.net, tgt_model=model.net_ema)
+            else:
+                for tgt_param, src_param in zip(model.net_ema.parameters(), model.net.parameters(), strict=False):
+                    tgt_param.data.copy_(src_param.data)
+
+        fsdp_device_mesh = getattr(model, "fsdp_device_mesh", None)
+        if fsdp_device_mesh is not None:
+            # Under FSDP/DTensor, each rank reinitializes its own local shard,
+            # then broadcast_dtensor_model_states syncs replicated shards.
+            net.reinitialize_action_embedders()
+            _mirror_to_ema()
+            broadcast_dtensor_model_states(model.net, fsdp_device_mesh)
+            if hasattr(model, "net_ema") and model.net_ema is not None:
+                broadcast_dtensor_model_states(model.net_ema, fsdp_device_mesh)
+            log.info(
+                "[ACTION-EMB-FIX] Synchronized reinitialized weights (net + net_ema) across ranks via DTensor broadcast."
+            )
+        else:
+            if distributed.is_rank0():
+                net.reinitialize_action_embedders()
+                _mirror_to_ema()
+            distributed.sync_model_states(model.net, src=0)
+            if hasattr(model, "net_ema") and model.net_ema is not None:
+                distributed.sync_model_states(model.net_ema, src=0)
+            log.info("[ACTION-EMB-FIX] Synchronized reinitialized weights (net + net_ema) from rank 0.")
+
     def training_step(
         self,
         model_ddp: torch.nn.Module | distributed.DistributedDataParallel,

cosmos_predict2/_src/predict2/action/networks/action_conditioned_minimal_v1_lvg_dit.py

@@ -35,6 +35,29 @@ def __init__(self, in_features, hidden_features=None, out_features=None, act_lay
         self.fc2 = nn.Linear(hidden_features, out_features)
         self.drop = nn.Dropout(drop)
 
+    def init_weights(self) -> None:
+        """Initialize weights using Kaiming uniform (the ``nn.Linear`` default).
+
+        This must be called explicitly after FSDP meta-device materialization
+        because the default ``nn.Linear`` Kaiming init runs on the meta device
+        where it has no effect.
+
+        Handles FSDP DTensor-wrapped parameters by writing directly to the
+        local tensor data, bypassing the DTensor dispatch layer which can
+        silently intercept in-place operations like ``reset_parameters()``.
+        """
+        import torch.nn.init as init
+
+        for linear in (self.fc1, self.fc2):
+            w = linear.weight
+            b = linear.bias
+            local_w = w.to_local() if hasattr(w, "to_local") else w.data
+            local_b = b.to_local() if hasattr(b, "to_local") else b.data
+            init.kaiming_uniform_(local_w, a=5**0.5)
+            fan_in, _ = init._calculate_fan_in_and_fan_out(local_w)
+            bound = 1 / fan_in**0.5 if fan_in > 0 else 0
+            init.uniform_(local_b, -bound, bound)
+
     def forward(self, x):
         x = self.fc1(x)
         x = self.activation(x)
@@ -80,6 +103,31 @@ def __init__(self, *args, timestep_scale: float = 1.0, **kwargs):
             drop=0,
         )
 
+    def reinitialize_action_embedders(self) -> None:
+        """Re-initialize action embedder weights after FSDP checkpoint loading.
+
+        When fine-tuning from a pre-trained Cosmos base checkpoint that lacks
+        action_embedder keys, FSDP meta-device materialization leaves these
+        parameters as all-zeros (the normal ``nn.Linear`` Kaiming init runs on
+        the meta device where it has no effect).  This method reapplies
+        Kaiming uniform initialization directly on the local tensor data,
+        bypassing DTensor dispatch.
+        """
+        for name, mlp in [
+            ("action_embedder_B_D", self.action_embedder_B_D),
+            ("action_embedder_B_3D", self.action_embedder_B_3D),
+        ]:
+            mlp.init_weights()
+            fc1_w = mlp.fc1.weight
+            fc2_w = mlp.fc2.weight
+            local_fc1 = fc1_w.to_local() if hasattr(fc1_w, "to_local") else fc1_w.data
+            local_fc2 = fc2_w.to_local() if hasattr(fc2_w, "to_local") else fc2_w.data
+            log.info(
+                f"[ACTION-EMB-FIX] Reinitialized {name}: "
+                f"fc1.weight local_norm={local_fc1.norm().item():.6f}, "
+                f"fc2.weight local_norm={local_fc2.norm().item():.6f}"
+            )
+
     def forward(
         self,
         x_B_C_T_H_W: torch.Tensor,
@@ -232,6 +280,31 @@ def __init__(self, *args, timestep_scale: float = 1.0, **kwargs):
             drop=0,
         )
 
+    def reinitialize_action_embedders(self) -> None:
+        """Re-initialize action embedder weights after FSDP checkpoint loading.
+
+        When fine-tuning from a pre-trained Cosmos base checkpoint that lacks
+        action_embedder keys, FSDP meta-device materialization leaves these
+        parameters as all-zeros (the normal ``nn.Linear`` Kaiming init runs on
+        the meta device where it has no effect).  This method reapplies
+        Kaiming uniform initialization directly on the local tensor data,
+        bypassing DTensor dispatch.
+        """
+        for name, mlp in [
+            ("action_embedder_B_D", self.action_embedder_B_D),
+            ("action_embedder_B_3D", self.action_embedder_B_3D),
+        ]:
+            mlp.init_weights()
+            fc1_w = mlp.fc1.weight
+            fc2_w = mlp.fc2.weight
+            local_fc1 = fc1_w.to_local() if hasattr(fc1_w, "to_local") else fc1_w.data
+            local_fc2 = fc2_w.to_local() if hasattr(fc2_w, "to_local") else fc2_w.data
+            log.info(
+                f"[ACTION-EMB-FIX] Reinitialized {name}: "
+                f"fc1.weight local_norm={local_fc1.norm().item():.6f}, "
+                f"fc2.weight local_norm={local_fc2.norm().item():.6f}"
+            )
+
     def forward(
         self,
         x_B_C_T_H_W: torch.Tensor,