set supports gradient checkpointing to true where necessary; add missing no split modules

Author: a-r-r-o-w

src/diffusers/models/modeling_utils.py

@@ -1196,7 +1196,7 @@ def _get_signature_keys(cls, obj):
     # Adapted from `transformers` modeling_utils.py
     def _get_no_split_modules(self, device_map: str):
         """
-        Get the modules of the model that should not be spit when using device_map. We iterate through the modules to
+        Get the modules of the model that should not be split when using device_map. We iterate through the modules to
         get the underlying `_no_split_modules`.
 
         Args:

src/diffusers/models/transformers/cogvideox_transformer_3d.py

@@ -210,6 +210,7 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
     """
 
     _supports_gradient_checkpointing = True
+    _no_split_modules = ["CogVideoXBlock", "CogVideoXPatchEmbed"]
 
     @register_to_config
     def __init__(

src/diffusers/models/transformers/latte_transformer_3d.py

@@ -27,6 +27,7 @@
 
 class LatteTransformer3DModel(ModelMixin, ConfigMixin):
     _supports_gradient_checkpointing = True
+    _no_split_modules = ["BasicTransformerBlock"]
 
     """
     A 3D Transformer model for video-like data, paper: https://arxiv.org/abs/2401.03048, offical code:

src/diffusers/models/transformers/transformer_allegro.py

@@ -221,6 +221,9 @@ class AllegroTransformer3DModel(ModelMixin, ConfigMixin):
             Scaling factor to apply in 3D positional embeddings across time dimension.
     """
 
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["AllegroTransformerBlock"]
+
     @register_to_config
     def __init__(
         self,

src/diffusers/models/transformers/transformer_cogview3plus.py

@@ -166,6 +166,7 @@ class CogView3PlusTransformer2DModel(ModelMixin, ConfigMixin):
     """
 
     _supports_gradient_checkpointing = True
+    _no_split_modules = ["CogView3PlusTransformerBlock", "CogView3PlusPatchEmbed"]
 
     @register_to_config
     def __init__(

src/diffusers/models/transformers/transformer_hunyuan_video.py

@@ -540,6 +540,12 @@ class HunyuanVideoTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
     """
 
     _supports_gradient_checkpointing = True
+    _no_split_modules = [
+        "HunyuanVideoTransformerBlock",
+        "HunyuanVideoSingleTransformerBlock",
+        "HunyuanVideoPatchEmbed",
+        "HunyuanVideoTokenRefiner",
+    ]
 
     @register_to_config
     def __init__(

src/diffusers/models/transformers/transformer_ltx.py

@@ -295,6 +295,7 @@ class LTXVideoTransformer3DModel(ModelMixin, ConfigMixin, FromOriginalModelMixin
     """
 
     _supports_gradient_checkpointing = True
+    _no_split_modules = ["LTXTransformerBlock"]
 
     @register_to_config
     def __init__(

src/diffusers/models/transformers/transformer_sd3.py

@@ -125,6 +125,7 @@ class SD3Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOrigi
     """
 
     _supports_gradient_checkpointing = True
+    _no_split_modules = ["JointTransformerBlock", "SD3SingleTransformerBlock"]
 
     @register_to_config
     def __init__(