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..9df8991 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,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, 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..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 @@ -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,