Allow Mochi Transformer to be split across multiple GPUs (huggingface#10300)

DN6 · Foundsheep · commit c4d6acd4104c · 2024-12-23T09:26:06.000+09:00
update
diff --git a/src/diffusers/models/transformers/transformer_mochi.py b/src/diffusers/models/transformers/transformer_mochi.py
@@ -335,6 +335,7 @@ class MochiTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOri
     """
 
     _supports_gradient_checkpointing = True
+    _no_split_modules = ["MochiTransformerBlock"]
 
     @register_to_config
     def __init__(
