Fix full-finetune save crash when checkpointing + layer offloading are enabled + diffusers output format

modules/util/checkpointing_util.py

@@ -1,3 +1,4 @@
+import copy
 import inspect
 from collections.abc import Callable
 from typing import Any
@@ -106,6 +107,17 @@ def __init__(self, orig_module: nn.Module, orig_forward, train_device: torch.dev
         self.conductor = conductor
         self.layer_index = layer_index
 
+    def __deepcopy__(self, memo):
+        # conductor holds torch.cuda.Stream/Event objects that cannot be deep-copied or pickled.
+        # deepcopy is only used at save time to build a dtype-converted CPU copy of the pipeline,
+        # where the conductor is never invoked, so share the existing instance instead of copying it.
+        cls = self.__class__
+        result = cls.__new__(cls)
+        memo[id(self)] = result
+        for key, value in self.__dict__.items():
+            result.__dict__[key] = value if key == "conductor" else copy.deepcopy(value, memo)
+        return result
+
     def __checkpointing_forward(self, dummy: torch.Tensor, call_id: int, *args):
         init_compile()  # workaround for https://github.com/pytorch/pytorch/issues/186537
         if self.layer_index == 0 and not torch.is_grad_enabled():