From 66807edb61636d7ee4dcf737baf3b782b9f4afa2 Mon Sep 17 00:00:00 2001 From: leild Date: Mon, 29 Jun 2026 10:27:24 +0800 Subject: [PATCH 1/2] modify accelerate_trainer --- .../trainer/accelerate/accelerate_trainer.py | 33 +++++++++++++++---- 1 file changed, 26 insertions(+), 7 deletions(-) diff --git a/mova/engine/trainer/accelerate/accelerate_trainer.py b/mova/engine/trainer/accelerate/accelerate_trainer.py index 1f25d95..7b8df7f 100644 --- a/mova/engine/trainer/accelerate/accelerate_trainer.py +++ b/mova/engine/trainer/accelerate/accelerate_trainer.py @@ -146,8 +146,8 @@ def _setup(self): if self.use_fsdp: from accelerate import FullyShardedDataParallelPlugin from torch.distributed.fsdp.fully_sharded_data_parallel import ( - FullOptimStateDictConfig, - FullStateDictConfig, + ShardedOptimStateDictConfig, + ShardedStateDictConfig, ) for attr_name in ['text_encoder', 'prompter', 'video_vae', 'audio_vae']: @@ -170,13 +170,11 @@ def _setup(self): fsdp_config['reshard_after_forward'] = True fsdp_plugin = FullyShardedDataParallelPlugin( - state_dict_config=FullStateDictConfig( + state_dict_config=ShardedStateDictConfig( offload_to_cpu=True, - # rank0_only=True, ), - optim_state_dict_config=FullOptimStateDictConfig( + optim_state_dict_config=ShardedOptimStateDictConfig( offload_to_cpu=True, - # rank0_only=True, ), **fsdp_config ) @@ -511,6 +509,17 @@ def _save_checkpoint(self, final: bool = False): os.makedirs(os.path.join(step_dir, "accelerator"), exist_ok=True) self.accelerator.wait_for_everyone() self.accelerator.save_state(os.path.join(step_dir, "accelerator")) + if self.use_fsdp: + rank = self.accelerator.process_index + accel_dir = os.path.join(step_dir, "accelerator") + os.makedirs(accel_dir, exist_ok=True) + torch.save(self.optimizer.state_dict(), + os.path.join(accel_dir, f"optimizer_{rank}.bin")) + torch.save(self.lr_scheduler.state_dict(), + os.path.join(accel_dir, f"scheduler_{rank}.bin")) + else: + self.accelerator.save_state(os.path.join(step_dir, "accelerator")) + def _resume_checkpoint(self, checkpoint_path: str): """Resume checkpoint""" @@ -522,7 +531,17 @@ def _resume_checkpoint(self, checkpoint_path: str): accelerator_path = os.path.join(checkpoint_path, "accelerator") if os.path.exists(accelerator_path): - self.accelerator.load_state(accelerator_path) + if self.use_fsdp: + rank = self.accelerator.process_index + optim_path = os.path.join(accelerator_path, f"optimizer_{rank}.bin") + sched_path = os.path.join(accelerator_path, f"scheduler_{rank}.bin") + if os.path.exists(optim_path): + self.optimizer.load_state_dict(torch.load(optim_path, map_location="cpu")) + if os.path.exists(sched_path): + self.lr_scheduler.load_state_dict(torch.load(sched_path, map_location="cpu")) + else: + self.accelerator.load_state(accelerator_path) + if self.use_lora: from mova.engine.trainer.accelerate.lora_utils import load_lora_weights From bfe966d74748c9b49b426d736dec2d73c10c47a2 Mon Sep 17 00:00:00 2001 From: gygdh-001 Date: Fri, 10 Jul 2026 11:26:06 +0800 Subject: [PATCH 2/2] fix: load model weights during FSDP checkpoint resume MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously _resume_checkpoint() only restored optimizer and scheduler states but never loaded model weights. This caused resumed training to use the original pretrained weights instead of the checkpointed trained weights, making resume effectively a from-scratch restart. Fix: add _load_model_weights_from_checkpoint() to load per-module diffusion_pytorch_model.bin files for all trainable modules (video_dit, video_dit_2, audio_dit, dual_tower_bridge) before restoring optimizer and scheduler states. Verified by: 50-parameter anchor comparison (max_diff=0), loss continuity (step-6 loss=0.3047 → step-7 loss=0.3047), and optimizer state key audit (3735 keys match). Signed-off-by: gygdh-001 --- .../trainer/accelerate/accelerate_trainer.py | 40 +++++++++++++++++-- 1 file changed, 37 insertions(+), 3 deletions(-) diff --git a/mova/engine/trainer/accelerate/accelerate_trainer.py b/mova/engine/trainer/accelerate/accelerate_trainer.py index 7b8df7f..f7f66e9 100644 --- a/mova/engine/trainer/accelerate/accelerate_trainer.py +++ b/mova/engine/trainer/accelerate/accelerate_trainer.py @@ -521,6 +521,37 @@ def _save_checkpoint(self, final: bool = False): self.accelerator.save_state(os.path.join(step_dir, "accelerator")) + def _load_model_weights_from_checkpoint(self, checkpoint_path): + """ + Load trainable module weights saved by _save_checkpoint. + Weights are saved as {module_name}/diffusion_pytorch_model.bin per module. + """ + unwrapped_model = self.accelerator.unwrap_model(self.model) + full_state_dict = {} + loaded_modules = [] + + for module_name in self.train_modules: + bin_path = os.path.join(checkpoint_path, module_name, "diffusion_pytorch_model.bin") + if os.path.isfile(bin_path): + module_state = torch.load(bin_path, map_location="cpu") + for key, value in module_state.items(): + full_state_dict[f"{module_name}.{key}"] = value + loaded_modules.append(module_name) + else: + raise FileNotFoundError( + f"[Resume] FATAL: missing weight file for module '{module_name}': {bin_path}" + ) + + missing_keys, unexpected_keys = unwrapped_model.load_state_dict( + full_state_dict, strict=False + ) + + if self.is_main_process: + print(f"[Resume] Loaded {len(full_state_dict)} model parameters from {len(loaded_modules)} modules: {loaded_modules}") + print(f"[Resume] {len(missing_keys)} missing keys (expected: frozen modules), first 5: {list(missing_keys)[:5]}") + if unexpected_keys: + print(f"[Resume] ⚠️ {len(unexpected_keys)} unexpected keys: {unexpected_keys[:5]}") + def _resume_checkpoint(self, checkpoint_path: str): """Resume checkpoint""" state_path = os.path.join(checkpoint_path, "trainer_state.pt") @@ -528,7 +559,11 @@ def _resume_checkpoint(self, checkpoint_path: str): state = torch.load(state_path, map_location="cpu") self.global_step = state["global_step"] self.epoch = state["epoch"] - + + # Load model weights (must happen before optimizer/scheduler restore) + if not self.use_lora: + self._load_model_weights_from_checkpoint(checkpoint_path) + accelerator_path = os.path.join(checkpoint_path, "accelerator") if os.path.exists(accelerator_path): if self.use_fsdp: @@ -542,11 +577,10 @@ def _resume_checkpoint(self, checkpoint_path: str): else: self.accelerator.load_state(accelerator_path) - if self.use_lora: from mova.engine.trainer.accelerate.lora_utils import load_lora_weights load_lora_weights(self.model, checkpoint_path) - + if self.is_main_process: print(f"[Resume] Loaded from {checkpoint_path}, step={self.global_step}")