Skip to content

enable group_offloading and PipelineDeviceAndDtypeStabilityTests on XPU, all passed #11620

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
May 30, 2025
Merged

enable group_offloading and PipelineDeviceAndDtypeStabilityTests on XPU, all passed #11620

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
74a3164
enable group_offloading and PipelineDeviceAndDtypeStabilityTests on XPU,
yao-matrix May 28, 2025
2ed18b5
fix style
yao-matrix May 28, 2025
c0f96fe
fix
yao-matrix May 28, 2025
1d37294
Merge branch 'main' into ut-xpu
yao-matrix May 29, 2025
56de38c
Merge branch 'main' into ut-xpu
a-r-r-o-w May 30, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 26 additions & 19 deletions tests/hooks/test_group_offloading.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,13 @@
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.utils import get_logger
from diffusers.utils.import_utils import compare_versions
from diffusers.utils.testing_utils import require_torch_gpu, torch_device
from diffusers.utils.testing_utils import (
backend_empty_cache,
backend_max_memory_allocated,
backend_reset_peak_memory_stats,
require_torch_accelerator,
torch_device,
)


class DummyBlock(torch.nn.Module):
Expand Down Expand Up @@ -107,7 +113,7 @@ def __call__(self, x: torch.Tensor) -> torch.Tensor:
return x


@require_torch_gpu
@require_torch_accelerator
class GroupOffloadTests(unittest.TestCase):
in_features = 64
hidden_features = 256
Expand All @@ -125,8 +131,8 @@ def tearDown(self):
del self.model
del self.input
gc.collect()
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()
backend_empty_cache(torch_device)
backend_reset_peak_memory_stats(torch_device)

def get_model(self):
torch.manual_seed(0)
Expand All @@ -141,8 +147,8 @@ def test_offloading_forward_pass(self):
@torch.no_grad()
def run_forward(model):
gc.collect()
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()
backend_empty_cache(torch_device)
backend_reset_peak_memory_stats(torch_device)
self.assertTrue(
all(
module._diffusers_hook.get_hook("group_offloading") is not None
Expand All @@ -152,7 +158,7 @@ def run_forward(model):
)
model.eval()
output = model(self.input)[0].cpu()
max_memory_allocated = torch.cuda.max_memory_allocated()
max_memory_allocated = backend_max_memory_allocated(torch_device)
return output, max_memory_allocated

self.model.to(torch_device)
Expand Down Expand Up @@ -187,10 +193,10 @@ def run_forward(model):
self.assertTrue(torch.allclose(output_without_group_offloading, output_with_group_offloading5, atol=1e-5))

# Memory assertions - offloading should reduce memory usage
self.assertTrue(mem4 <= mem5 < mem2 < mem3 < mem1 < mem_baseline)
self.assertTrue(mem4 <= mem5 < mem2 <= mem3 < mem1 < mem_baseline)

def test_warning_logged_if_group_offloaded_module_moved_to_cuda(self):
if torch.device(torch_device).type != "cuda":
def test_warning_logged_if_group_offloaded_module_moved_to_accelerator(self):
if torch.device(torch_device).type not in ["cuda", "xpu"]:
return
self.model.enable_group_offload(torch_device, offload_type="block_level", num_blocks_per_group=3)
logger = get_logger("diffusers.models.modeling_utils")
Expand All @@ -199,8 +205,8 @@ def test_warning_logged_if_group_offloaded_module_moved_to_cuda(self):
self.model.to(torch_device)
self.assertIn(f"The module '{self.model.__class__.__name__}' is group offloaded", cm.output[0])

def test_warning_logged_if_group_offloaded_pipe_moved_to_cuda(self):
if torch.device(torch_device).type != "cuda":
def test_warning_logged_if_group_offloaded_pipe_moved_to_accelerator(self):
if torch.device(torch_device).type not in ["cuda", "xpu"]:
return
pipe = DummyPipeline(self.model)
self.model.enable_group_offload(torch_device, offload_type="block_level", num_blocks_per_group=3)
Expand All @@ -210,19 +216,20 @@ def test_warning_logged_if_group_offloaded_pipe_moved_to_cuda(self):
pipe.to(torch_device)
self.assertIn(f"The module '{self.model.__class__.__name__}' is group offloaded", cm.output[0])

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
def test_error_raised_if_streams_used_and_no_accelerator_device(self):
torch_accelerator_module = getattr(torch, torch_device, torch.cuda)
original_is_available = torch_accelerator_module.is_available
torch_accelerator_module.is_available = lambda: False
with self.assertRaises(ValueError):
self.model.enable_group_offload(
onload_device=torch.device("cuda"), offload_type="leaf_level", use_stream=True
onload_device=torch.device(torch_device), offload_type="leaf_level", use_stream=True
)
torch.cuda.is_available = original_is_available
torch_accelerator_module.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_offload(onload_device=torch.device("cuda"))
self.model.enable_group_offload(onload_device=torch.device(torch_device))

def test_error_raised_if_model_offloading_applied_on_group_offloaded_module(self):
pipe = DummyPipeline(self.model)
Expand All @@ -249,7 +256,7 @@ def test_error_raised_if_group_offloading_applied_on_sequential_offloaded_module
pipe.model.enable_group_offload(torch_device, offload_type="block_level", num_blocks_per_group=3)

def test_block_level_stream_with_invocation_order_different_from_initialization_order(self):
if torch.device(torch_device).type != "cuda":
if torch.device(torch_device).type not in ["cuda", "xpu"]:
return
model = DummyModelWithMultipleBlocks(
in_features=self.in_features,
Expand Down
10 changes: 5 additions & 5 deletions tests/pipelines/test_pipeline_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
UNet2DConditionModel,
)
from diffusers.pipelines.pipeline_loading_utils import is_safetensors_compatible, variant_compatible_siblings
from diffusers.utils.testing_utils import require_torch_gpu, torch_device
from diffusers.utils.testing_utils import require_torch_accelerator, torch_device


class IsSafetensorsCompatibleTests(unittest.TestCase):
Expand Down Expand Up @@ -850,9 +850,9 @@ def test_video_to_video(self):
self.assertTrue(stderr.getvalue() == "", "Progress bar should be disabled")


@require_torch_gpu
@require_torch_accelerator
class PipelineDeviceAndDtypeStabilityTests(unittest.TestCase):
expected_pipe_device = torch.device("cuda:0")
expected_pipe_device = torch.device(f"{torch_device}:0")
expected_pipe_dtype = torch.float64

def get_dummy_components_image_generation(self):
Expand Down Expand Up @@ -921,8 +921,8 @@ def test_deterministic_device(self):
pipe.to(device=torch_device, dtype=torch.float32)

pipe.unet.to(device="cpu")
pipe.vae.to(device="cuda")
pipe.text_encoder.to(device="cuda:0")
pipe.vae.to(device=torch_device)
pipe.text_encoder.to(device=f"{torch_device}:0")

pipe_device = pipe.device

Expand Down