fix

examples/inference/basic/basic_dmd.py

@@ -21,6 +21,7 @@ def main():
         dit_cpu_offload=False,
         vae_cpu_offload=False,
         VSA_sparsity=0.8,
+        init_weights_from_safetensors="/mnt/weka/home/hao.zhang/wl/release/dmd_distill_1.3_4n_syn/checkpoint-900_weight_only/generator_inference_transformer"
     )
     load_end_time = time.perf_counter()
     load_time = load_end_time - load_start_time

fastvideo/training/checkpointing_utils.py

@@ -21,8 +21,12 @@ class ModelWrapper(torch.distributed.checkpoint.stateful.Stateful):
         state_dict = get_model_state_dict(
             self.model)  # type: ignore[no-any-return]
         # filter out non-trainable parameters
+        # Note: activation checkpointing adds ._checkpoint_wrapped_module. prefix
+        # to parameter names, but get_model_state_dict returns normalized keys.
+        # We need to normalize the parameter names for proper matching.
         param_requires_grad = set([
-            k for k, v in dict(self.model.named_parameters()).items()
+            k.replace("._checkpoint_wrapped_module.", ".")
+            for k, v in dict(self.model.named_parameters()).items()
             if v.requires_grad
         ])
         state_dict = {

fastvideo/training/training_utils.py

@@ -116,6 +116,44 @@ def get_sigmas(noise_scheduler,
     return sigma
 
 
+def _is_lora_training(transformer: torch.nn.Module) -> bool:
+    """Best-effort check for LoRA fine-tuning.
+
+    We treat a run as LoRA training when the model already contains LoRA
+    layers with trainable parameters. This avoids touching consolidated
+    full-model weights which are not updated during LoRA-only training.
+    """
+
+    state_keys = transformer.state_dict().keys()
+    return any("lora_" in name for name in state_keys)
+
+
+def _save_lora_adapter(cpu_state: dict[str, Any], adapter_path: str) -> None:
+    """Persist only the LoRA adapter weights for LoRA training runs.
+    
+    Args:
+        cpu_state: Pre-gathered state dict from gather_state_dict_on_cpu_rank0.
+                   Using the pre-gathered state dict avoids issues with invalid
+                   tensor storage pointers in FSDP-wrapped models.
+        adapter_path: Path to save the LoRA adapter safetensors file.
+    """
+
+    lora_state = {
+        name: tensor.detach().clone().cpu().contiguous()
+        for name, tensor in cpu_state.items() if "lora_" in name
+    }
+
+    if len(lora_state) == 0:
+        logger.warning(
+            "LoRA training detected but no LoRA parameters found to save.")
+        return
+
+    os.makedirs(os.path.dirname(adapter_path), exist_ok=True)
+    save_file(lora_state, adapter_path)
+    logger.info("Saved LoRA adapter with %d tensors to %s", len(lora_state),
+                adapter_path)
+
+
 def save_checkpoint(transformer,
                     rank,
                     output_dir,
@@ -161,35 +199,43 @@ def save_checkpoint(transformer,
 
     cpu_state = gather_state_dict_on_cpu_rank0(transformer, device=None)
     if rank == 0:
-        # Save model weights (consolidated)
+        if _is_lora_training(transformer):
-        transformer_save_dir = os.path.join(save_dir, "transformer")
+            adapter_path = os.path.join(save_dir, "lora_adapter.safetensors")
-        os.makedirs(transformer_save_dir, exist_ok=True)
+            _save_lora_adapter(cpu_state, adapter_path)
-        weight_path = os.path.join(transformer_save_dir,
+            logger.info(
-                                   "diffusion_pytorch_model.safetensors")
+                "LoRA training detected; saved adapter instead of consolidated weights."
-        logger.info("rank: %s, saving consolidated checkpoint to %s",
+            )
-                    rank,
+        else:
-                    weight_path,
+            # Save model weights (consolidated)
-                    local_main_process_only=False)
+            transformer_save_dir = os.path.join(save_dir, "transformer")
+            os.makedirs(transformer_save_dir, exist_ok=True)
+            weight_path = os.path.join(transformer_save_dir,
+                                       "diffusion_pytorch_model.safetensors")
+            logger.info("rank: %s, saving consolidated checkpoint to %s",
+                        rank,
+                        weight_path,
+                        local_main_process_only=False)
 
-        # Convert training format to diffusers format and save
+            # Convert training format to diffusers format and save
-        diffusers_state_dict = custom_to_hf_state_dict(
+            diffusers_state_dict = custom_to_hf_state_dict(
-            cpu_state, transformer.reverse_param_names_mapping)
+                cpu_state, transformer.reverse_param_names_mapping)
-        save_file(diffusers_state_dict, weight_path)
+            save_file(diffusers_state_dict, weight_path)
 
-        logger.info("rank: %s, consolidated checkpoint saved to %s",
+            logger.info("rank: %s, consolidated checkpoint saved to %s",
-                    rank,
+                        rank,
-                    weight_path,
+                        weight_path,
-                    local_main_process_only=False)
+                        local_main_process_only=False)
 
-        # Save model config
+            # Save model config
-        config_dict = transformer.hf_config
+            config_dict = transformer.hf_config
-        if "dtype" in config_dict:
+            if "dtype" in config_dict:
-            del config_dict["dtype"]  # TODO
+                del config_dict["dtype"]  # TODO
-        config_path = os.path.join(transformer_save_dir, "config.json")
+            config_path = os.path.join(transformer_save_dir, "config.json")
-        # save dict as json
+            # save dict as json
-        with open(config_path, "w") as f:
+            with open(config_path, "w") as f:
-            json.dump(config_dict, f, indent=4)
+                json.dump(config_dict, f, indent=4)
-        logger.info("--> checkpoint saved at step %s to %s", step, weight_path)
+            logger.info("--> checkpoint saved at step %s to %s", step,
+                        weight_path)
 
 
 def save_distillation_checkpoint(
@@ -421,37 +467,44 @@ def save_distillation_checkpoint(
                                                device=None)
 
     if rank == 0:
-        # Save generator model weights (consolidated) for inference
+        if _is_lora_training(generator_transformer):
-        os.makedirs(inference_save_dir, exist_ok=True)
+            adapter_path = os.path.join(save_dir, "lora_adapter.safetensors")
-        weight_path = os.path.join(inference_save_dir,
+            _save_lora_adapter(cpu_state, adapter_path)
-                                   "diffusion_pytorch_model.safetensors")
+            logger.info(
-        logger.info(
+                "LoRA training detected; saved adapter instead of consolidated generator weights."
-            "rank: %s, saving consolidated generator inference checkpoint to %s",
+            )
-            rank,
+        else:
-            weight_path,
+            # Save generator model weights (consolidated) for inference
-            local_main_process_only=False)
+            os.makedirs(inference_save_dir, exist_ok=True)
+            weight_path = os.path.join(inference_save_dir,
+                                       "diffusion_pytorch_model.safetensors")
+            logger.info(
+                "rank: %s, saving consolidated generator inference checkpoint to %s",
+                rank,
+                weight_path,
+                local_main_process_only=False)
 
-        # Convert training format to diffusers format and save
+            # Convert training format to diffusers format and save
-        diffusers_state_dict = custom_to_hf_state_dict(
+            diffusers_state_dict = custom_to_hf_state_dict(
-            cpu_state, generator_transformer.reverse_param_names_mapping)
+                cpu_state, generator_transformer.reverse_param_names_mapping)
-        save_file(diffusers_state_dict, weight_path)
+            save_file(diffusers_state_dict, weight_path)
 
-        logger.info(
+            logger.info(
-            "rank: %s, consolidated generator inference checkpoint saved to %s",
+                "rank: %s, consolidated generator inference checkpoint saved to %s",
-            rank,
+                rank,
-            weight_path,
+                weight_path,
-            local_main_process_only=False)
+                local_main_process_only=False)
 
-        # Save model config
+            # Save model config
-        config_dict = generator_transformer.hf_config
+            config_dict = generator_transformer.hf_config
-        if "dtype" in config_dict:
+            if "dtype" in config_dict:
-            del config_dict["dtype"]  # TODO
+                del config_dict["dtype"]  # TODO
-        config_path = os.path.join(inference_save_dir, "config.json")
+            config_path = os.path.join(inference_save_dir, "config.json")
-        # save dict as json
+            # save dict as json
-        with open(config_path, "w") as f:
+            with open(config_path, "w") as f:
-            json.dump(config_dict, f, indent=4)
+                json.dump(config_dict, f, indent=4)
-        logger.info("--> distillation checkpoint saved at step %s to %s", step,
+            logger.info("--> distillation checkpoint saved at step %s to %s",
-                    weight_path)
+                        step, weight_path)
 
         # Save generator_2 model weights (consolidated) for inference (MoE support)
         if generator_transformer_2 is not None:
@@ -599,8 +652,9 @@ def load_distillation_checkpoint(
                        checkpoint_path)
         return 0
 
-    # Extract step number from checkpoint path
+    # Extract step number from checkpoint path (normpath handles trailing slashes)
-    step = int(os.path.basename(checkpoint_path).split('-')[-1])
+    step = int(
+        os.path.basename(os.path.normpath(checkpoint_path)).split('-')[-1])
 
     if rank == 0:
         logger.info("Loading distillation checkpoint from step %s", step)