Skip to content

Use device_map="auto" in single file tests to support large models on limited GPU memory #13816

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

Open
jiqing-feng wants to merge 2 commits into
base: main
Choose a base branch
from
+13 −7
Open

Use device_map="auto" in single file tests to support large models on limited GPU memory #13816

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9a5d226
fix flux tests OOM on 24G GPU
jiqing-feng May 27, 2026
2db10d5
revert wrong change
jiqing-feng May 27, 2026
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
10 changes: 5 additions & 5 deletions tests/models/testing_utils/single_file.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -107,8 +107,8 @@ def teardown_method(self):
backend_empty_cache(torch_device)

def test_single_file_model_config(self):
pretrained_kwargs = {"device": torch_device, **self.pretrained_model_kwargs}
single_file_kwargs = {"device": torch_device}
pretrained_kwargs = {"device_map": "auto", **self.pretrained_model_kwargs}
single_file_kwargs = {"device_map": "auto"}

if self.torch_dtype:
pretrained_kwargs["torch_dtype"] = self.torch_dtype
Expand All @@ -127,8 +127,8 @@ def test_single_file_model_config(self):
)

def test_single_file_model_parameters(self):
pretrained_kwargs = {"device_map": str(torch_device), **self.pretrained_model_kwargs}
single_file_kwargs = {"device": torch_device}
pretrained_kwargs = {"device_map": "auto", **self.pretrained_model_kwargs}
single_file_kwargs = {"device_map": "auto"}

if self.torch_dtype:
pretrained_kwargs["torch_dtype"] = self.torch_dtype
Expand Down Expand Up @@ -259,7 +259,7 @@ def test_checkpoint_variant_loading(self):
backend_empty_cache(torch_device)

def test_single_file_loading_with_device_map(self):
single_file_kwargs = {"device_map": torch_device}
single_file_kwargs = {"device_map": "auto"}

if self.torch_dtype:
single_file_kwargs["torch_dtype"] = self.torch_dtype
Expand Down
4 changes: 4 additions & 0 deletions tests/models/transformers/test_models_transformer_flux.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -344,6 +344,10 @@ def alternate_ckpt_paths(self):
def pretrained_model_name_or_path(self):
return "black-forest-labs/FLUX.1-dev"

@property
def torch_dtype(self):
return torch.bfloat16


class TestFluxTransformerBitsAndBytes(FluxTransformerTesterConfig, BitsAndBytesTesterMixin):
"""BitsAndBytes quantization tests for Flux Transformer."""
Expand Down
6 changes: 4 additions & 2 deletions tests/single_file/test_model_flux_transformer_single_file.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@

import gc

import torch

from diffusers import (
FluxTransformer2DModel,
)
Expand All @@ -38,9 +40,9 @@ class TestFluxTransformer2DModelSingleFile(SingleFileModelTesterMixin):
repo_id = "black-forest-labs/FLUX.1-dev"
subfolder = "transformer"

def test_device_map_cuda(self):
def test_device_map_auto(self):
backend_empty_cache(torch_device)
model = self.model_class.from_single_file(self.ckpt_path, device_map="cuda")
model = self.model_class.from_single_file(self.ckpt_path, device_map="auto", torch_dtype=torch.bfloat16)

del model
gc.collect()
Expand Down
Loading