From 17e6aae6193d9ed5acb1996654d9014acd0c4cb1 Mon Sep 17 00:00:00 2001 From: Lukas Zbinden Date: Fri, 17 Apr 2026 13:23:28 +0200 Subject: [PATCH 1/2] Fix action embedder MLP initialization after FSDP checkpoint load --- README.md | 1 + cosmos_predict2/_src/imaginaire/trainer.py | 89 +++++++++++++++++++ .../action_conditioned_minimal_v1_lvg_dit.py | 69 ++++++++++++++ 3 files changed, 159 insertions(+) diff --git a/README.md b/README.md index 5a8c457..9b3b0a7 100644 --- a/README.md +++ b/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). diff --git a/cosmos_predict2/_src/imaginaire/trainer.py b/cosmos_predict2/_src/imaginaire/trainer.py index fac97df..89baf22 100644 --- a/cosmos_predict2/_src/imaginaire/trainer.py +++ b/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,88 @@ 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, diff --git a/cosmos_predict2/_src/predict2/action/networks/action_conditioned_minimal_v1_lvg_dit.py b/cosmos_predict2/_src/predict2/action/networks/action_conditioned_minimal_v1_lvg_dit.py index d4bbe13..924f881 100644 --- a/cosmos_predict2/_src/predict2/action/networks/action_conditioned_minimal_v1_lvg_dit.py +++ b/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,29 @@ 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 +278,29 @@ 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, From 5085dc57808578f8413105e306c9804106f8a914 Mon Sep 17 00:00:00 2001 From: Lukas Zbinden Date: Fri, 17 Apr 2026 15:13:06 +0200 Subject: [PATCH 2/2] fix ruff-format --- cosmos_predict2/_src/imaginaire/trainer.py | 15 +++++++-------- .../action_conditioned_minimal_v1_lvg_dit.py | 16 ++++++++++------ 2 files changed, 17 insertions(+), 14 deletions(-) diff --git a/cosmos_predict2/_src/imaginaire/trainer.py b/cosmos_predict2/_src/imaginaire/trainer.py index 89baf22..9df8991 100644 --- a/cosmos_predict2/_src/imaginaire/trainer.py +++ b/cosmos_predict2/_src/imaginaire/trainer.py @@ -292,13 +292,12 @@ def _reinitialize_action_embedders_if_needed(model, iteration: int) -> None: action_weight_params = [ (name, p) for name, p in net.named_parameters() - if "action_embedder" in name and name.endswith(".weight") + 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." - ) + log.info("[ACTION-EMB-FIX] Skipped: no materialized action_embedder weight matrices found on model.net.") return all_zero = all( @@ -307,9 +306,7 @@ def _reinitialize_action_embedders_if_needed(model, iteration: int) -> None: ) if not all_zero: - log.info( - "[ACTION-EMB-FIX] Skipped: checkpoint already contains nonzero action embedder weights." - ) + log.info("[ACTION-EMB-FIX] Skipped: checkpoint already contains nonzero action embedder weights.") return if iteration == 0: @@ -341,7 +338,9 @@ def _mirror_to_ema() -> None: 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.") + 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() diff --git a/cosmos_predict2/_src/predict2/action/networks/action_conditioned_minimal_v1_lvg_dit.py b/cosmos_predict2/_src/predict2/action/networks/action_conditioned_minimal_v1_lvg_dit.py index 924f881..cbfb627 100644 --- a/cosmos_predict2/_src/predict2/action/networks/action_conditioned_minimal_v1_lvg_dit.py +++ b/cosmos_predict2/_src/predict2/action/networks/action_conditioned_minimal_v1_lvg_dit.py @@ -53,9 +53,9 @@ def init_weights(self) -> None: 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) + 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 + bound = 1 / fan_in**0.5 if fan_in > 0 else 0 init.uniform_(local_b, -bound, bound) def forward(self, x): @@ -113,8 +113,10 @@ def reinitialize_action_embedders(self) -> None: 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)]: + 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 @@ -288,8 +290,10 @@ def reinitialize_action_embedders(self) -> None: 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)]: + 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