Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
8 changes: 4 additions & 4 deletions src/diffusers/models/modeling_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -208,10 +208,10 @@ def find_tensor_attributes(module: nn.Module) -> list[tuple[str, Tensor]]:

gen = parameter._named_members(get_members_fn=find_tensor_attributes)
last_tuple = None
for tuple in gen:
last_tuple = tuple
if tuple[1].is_floating_point():
return tuple[1].dtype
for t in gen:
last_tuple = t
if t[1].is_floating_point():
return t[1].dtype

if last_tuple is not None:
# fallback to the last dtype
Expand Down
28 changes: 28 additions & 0 deletions tests/models/test_modeling_common.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
from huggingface_hub.utils import HfHubHTTPError, is_jinja_available

from diffusers.models import FluxTransformer2DModel, SD3Transformer2DModel, UNet2DConditionModel
from diffusers.models.modeling_utils import get_parameter_dtype

from ..others.test_utils import TOKEN, USER, is_staging_test
from ..testing_utils import (
Expand Down Expand Up @@ -278,6 +279,33 @@ def get_dummy_inputs():

SD3Transformer2DModel._keep_in_fp32_modules = fp32_modules

@staticmethod
def _as_data_parallel_replica(module):
# `torch.nn.parallel.replicate` empties `_parameters` on every replica and re assigns the
# broadcast copies with `setattr`. Those copies are plain tensors rather than `nn.Parameter`,
# so they land in `replica.__dict__`, which is the state reproduced here.
replica = module._replicate_for_data_parallel()
for name, param in module._parameters.items():
setattr(replica, name, param.detach())
return replica

@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
def test_get_parameter_dtype_on_data_parallel_replica(self, dtype):
# A replica exposes no parameters and no buffers, so `get_parameter_dtype` falls through to
# the branch that scans `module.__dict__`. See https://github.com/huggingface/diffusers/issues/13789.
replica = self._as_data_parallel_replica(torch.nn.Linear(4, 4).to(dtype))

assert list(replica.named_parameters()) == []
assert list(replica.buffers()) == []
assert get_parameter_dtype(replica) == dtype

def test_get_parameter_dtype_falls_back_to_non_floating_point_tensor_attribute(self):
# Same branch, other exit: with no floating point tensor to report, the last one found wins.
module = torch.nn.Module()
module.token_ids = torch.zeros(4, dtype=torch.int64)

assert get_parameter_dtype(module) == torch.int64


class UNetTesterMixin:
@staticmethod
Expand Down
Loading