fix: add pull to solve the problem that fsdp2 traninng in shared will timeout

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,23 +509,78 @@ 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 _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")
         if os.path.exists(state_path):
             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):
-            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
             load_lora_weights(self.model, checkpoint_path)
-        
+
         if self.is_main_process:
             print(f"[Resume] Loaded from {checkpoint_path}, step={self.global_step}")