enable_group_offloading -> enable_group_offload for naming consistency

Author: a-r-r-o-w

docs/source/en/optimization/memory.md

@@ -162,7 +162,7 @@ In order to properly offload models after they're called, it is required to run
 
 Group offloading is the middle ground between sequential and model offloading. It works by offloading groups of internal layers (either `torch.nn.ModuleList` or `torch.nn.Sequential`), which uses less memory than model-level offloading. It is also faster than sequential-level offloading because the number of device synchronizations is reduced.
 
-To enable group offloading, call the [`~ModelMixin.enable_group_offloading`] method on the model if it is a Diffusers model implementation. For any other model implementation, use [`~hooks.group_offloading.apply_group_offloading`]:
+To enable group offloading, call the [`~ModelMixin.enable_group_offload`] method on the model if it is a Diffusers model implementation. For any other model implementation, use [`~hooks.group_offloading.apply_group_offloading`]:
 
 ```python
 import torch
@@ -175,8 +175,8 @@ onload_device = torch.device("cuda")
 offload_device = torch.device("cpu")
 pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", torch_dtype=torch.bfloat16)
 
-# We can utilize the enable_group_offloading method for Diffusers model implementations
-pipe.transformer.enable_group_offloading(onload_device=onload_device, offload_device=offload_device, offload_type="leaf_level", use_stream=True)
+# We can utilize the enable_group_offload method for Diffusers model implementations
+pipe.transformer.enable_group_offload(onload_device=onload_device, offload_device=offload_device, offload_type="leaf_level", use_stream=True)
 
 # For any other model implementations, the apply_group_offloading function can be used
 apply_group_offloading(pipe.text_encoder, onload_device=onload_device, offload_type="block_level", num_blocks_per_group=2)

src/diffusers/models/modeling_utils.py

@@ -446,7 +446,7 @@ def enable_layerwise_casting(
             self, storage_dtype, compute_dtype, skip_modules_pattern, skip_modules_classes, non_blocking
         )
 
-    def enable_group_offloading(
+    def enable_group_offload(
         self,
         onload_device: torch.device,
         offload_device: torch.device = torch.device("cpu"),
@@ -469,7 +469,7 @@ def enable_group_offloading(
             ...     "THUDM/CogVideoX-5b", subfolder="transformer", torch_dtype=torch.bfloat16
             ... )
 
-            >>> transformer.enable_group_offloading(
+            >>> transformer.enable_group_offload(
             ...     onload_device=torch.device("cuda"),
             ...     offload_device=torch.device("cpu"),
             ...     offload_type="leaf_level",

tests/hooks/test_group_offloading.py

@@ -109,25 +109,23 @@ def run_forward(model):
         self.model.to("cpu")
 
         model = self.get_model()
-        model.enable_group_offloading(torch_device, offload_type="block_level", num_blocks_per_group=3)
+        model.enable_group_offload(torch_device, offload_type="block_level", num_blocks_per_group=3)
         output_with_group_offloading1, mem1 = run_forward(model)
 
         model = self.get_model()
-        model.enable_group_offloading(torch_device, offload_type="block_level", num_blocks_per_group=1)
+        model.enable_group_offload(torch_device, offload_type="block_level", num_blocks_per_group=1)
         output_with_group_offloading2, mem2 = run_forward(model)
 
         model = self.get_model()
-        model.enable_group_offloading(
-            torch_device, offload_type="block_level", num_blocks_per_group=1, use_stream=True
-        )
+        model.enable_group_offload(torch_device, offload_type="block_level", num_blocks_per_group=1, use_stream=True)
         output_with_group_offloading3, mem3 = run_forward(model)
 
         model = self.get_model()
-        model.enable_group_offloading(torch_device, offload_type="leaf_level")
+        model.enable_group_offload(torch_device, offload_type="leaf_level")
         output_with_group_offloading4, mem4 = run_forward(model)
 
         model = self.get_model()
-        model.enable_group_offloading(torch_device, offload_type="leaf_level", use_stream=True)
+        model.enable_group_offload(torch_device, offload_type="leaf_level", use_stream=True)
         output_with_group_offloading5, mem5 = run_forward(model)
 
         # Precision assertions - offloading should not impact the output
@@ -144,12 +142,12 @@ def test_error_raised_if_streams_used_and_no_cuda_device(self):
         original_is_available = torch.cuda.is_available
         torch.cuda.is_available = lambda: False
         with self.assertRaises(ValueError):
-            self.model.enable_group_offloading(
+            self.model.enable_group_offload(
                 onload_device=torch.device("cuda"), offload_type="leaf_level", use_stream=True
             )
         torch.cuda.is_available = original_is_available
 
     def test_error_raised_if_supports_group_offloading_false(self):
         self.model._supports_group_offloading = False
         with self.assertRaisesRegex(ValueError, "does not support group offloading"):
-            self.model.enable_group_offloading(onload_device=torch.device("cuda"))
+            self.model.enable_group_offload(onload_device=torch.device("cuda"))

tests/models/test_modeling_common.py

@@ -1454,24 +1454,22 @@ def run_forward(model):
 
         torch.manual_seed(0)
         model = self.model_class(**init_dict)
-        model.enable_group_offloading(torch_device, offload_type="block_level", num_blocks_per_group=1)
+        model.enable_group_offload(torch_device, offload_type="block_level", num_blocks_per_group=1)
         output_with_group_offloading1 = run_forward(model)
 
         torch.manual_seed(0)
         model = self.model_class(**init_dict)
-        model.enable_group_offloading(
-            torch_device, offload_type="block_level", num_blocks_per_group=1, non_blocking=True
-        )
+        model.enable_group_offload(torch_device, offload_type="block_level", num_blocks_per_group=1, non_blocking=True)
         output_with_group_offloading2 = run_forward(model)
 
         torch.manual_seed(0)
         model = self.model_class(**init_dict)
-        model.enable_group_offloading(torch_device, offload_type="leaf_level")
+        model.enable_group_offload(torch_device, offload_type="leaf_level")
         output_with_group_offloading3 = run_forward(model)
 
         torch.manual_seed(0)
         model = self.model_class(**init_dict)
-        model.enable_group_offloading(torch_device, offload_type="leaf_level", use_stream=True)
+        model.enable_group_offload(torch_device, offload_type="leaf_level", use_stream=True)
         output_with_group_offloading4 = run_forward(model)
 
         self.assertTrue(torch.allclose(output_without_group_offloading, output_with_group_offloading1, atol=1e-5))

tests/pipelines/test_pipelines_common.py

@@ -2059,7 +2059,7 @@ def create_pipe():
             pipe.set_progress_bar_config(disable=None)
             return pipe
 
-        def enable_group_offloading_on_component(pipe, group_offloading_kwargs):
+        def enable_group_offload_on_component(pipe, group_offloading_kwargs):
             # We intentionally don't test VAE's here. This is because some tests enable tiling on the VAE. If
             # tiling is enabled and a forward pass is run, when cuda streams are used, the execution order of
             # the layers is not traced correctly. This causes errors. For apply group offloading to VAE, a
@@ -2077,9 +2077,9 @@ def enable_group_offloading_on_component(pipe, group_offloading_kwargs):
                 component = getattr(pipe, component_name)
                 if not getattr(component, "_supports_group_offloading", True):
                     continue
-                if hasattr(component, "enable_group_offloading"):
+                if hasattr(component, "enable_group_offload"):
                     # For diffusers ModelMixin implementations
-                    component.enable_group_offloading(torch.device(torch_device), **group_offloading_kwargs)
+                    component.enable_group_offload(torch.device(torch_device), **group_offloading_kwargs)
                 else:
                     # For other models not part of diffusers
                     apply_group_offloading(
@@ -2105,11 +2105,11 @@ def run_forward(pipe):
         output_without_group_offloading = run_forward(pipe)
 
         pipe = create_pipe()
-        enable_group_offloading_on_component(pipe, {"offload_type": "block_level", "num_blocks_per_group": 1})
+        enable_group_offload_on_component(pipe, {"offload_type": "block_level", "num_blocks_per_group": 1})
         output_with_group_offloading1 = run_forward(pipe)
 
         pipe = create_pipe()
-        enable_group_offloading_on_component(pipe, {"offload_type": "leaf_level"})
+        enable_group_offload_on_component(pipe, {"offload_type": "leaf_level"})
         output_with_group_offloading2 = run_forward(pipe)
 
         if torch.is_tensor(output_without_group_offloading):