fix underlying issue with `test_from_save_pretrained_dtype_inference` is that the `model.to(dtype)` cast at

[sayakpaul]

I think an underlying issue with test_from_save_pretrained_dtype_inference is that the model.to(dtype) cast at

diffusers/tests/models/testing_utils/common.py

Line 484 in 5d10b4d

model.to(dtype).save_pretrained(tmp_path)

is not dtype-aware (it will cast everything to dtype), but from_pretrained(..., torch_dtype=dtype) is dtype-aware (it respects _keep_in_fp32_modules, etc.). This causes the behavior of the two to diverge in several scenarios:

1. There are _keep_in_fp32_modules specified on the model (the more common case)
2. A non-persistent buffer like inv_freq is created with an explicit dtype (which is the case here)

which leads to a divergence between the outputs of model and model_loaded.

Originally posted by @dg845 in #13862 (comment)

Cc: @DN6

[akshan-main]

Had a go at this. I tried preparing the reference model the way from_pretrained(torch_dtype=dtype) does, instead of the plain model.to(dtype) , casting to dtype but keeping _keep_in_fp32_modules params/buffers in fp32 and restoring the non-persistent buffers like RoPE inv_freq:

# replacing `model.to(dtype).save_pretrained(tmp_path)`
non_persistent_buffers = {
    name: buffer.clone() for name, buffer in model.named_buffers() if name not in model.state_dict()
}
model.to(dtype)
for name, param in model.named_parameters():
    if fp32_modules and any(m in name.split(".") for m in fp32_modules):
        param.data = param.data.float()
for name, buffer in model.named_buffers():
    if name in non_persistent_buffers:
        buffer.data = non_persistent_buffers[name]
    elif fp32_modules and any(m in name.split(".") for m in fp32_modules):
        buffer.data = buffer.data.float()

model.save_pretrained(tmp_path)

With ideogram4's skip removed, here's what I got before vs after:

# before
E           AssertionError: Loaded model output differs for torch.bfloat16
E           Tensors not close! Mismatched elements: 113/128 (88.3%)
E             Max diff: 2.832031e-02 at index (0, 5, 15)
E             Actual:   1.630859e-01
E             Expected: 1.347656e-01
E             atol: 1.000000e-04, rtol: 0.000000e+00
======================== 2 failed, 2 warnings in 3.48s =========================

# after
tests/models/transformers/test_models_transformer_ideogram4.py::TestIdeogram4Transformer::test_from_save_pretrained_dtype_inference[fp16] PASSED [ 50%]
tests/models/transformers/test_models_transformer_ideogram4.py::TestIdeogram4Transformer::test_from_save_pretrained_dtype_inference[bf16] PASSED [100%]
======================== 2 passed, 2 warnings in 2.45s =========================

Running it across all the transformers + autoencoders dtype tests went from 6 failing to 4 (86 passing). It also seems to sort out chronoedit, which has scale_shift_table as a param under _keep_in_fp32_modules, so that one needs name-level matching rather than module-level. Nothing else regressed.

The 2 that still fail look like different things to me, not sure if they're in scope here:

autoencoder_tiny: reference and loaded weights came out identical but the outputs still differ, so it looks more like forward-pass numerics.
consistency_decoder_vae: its non-persistent buffers (means/stds) actually get cast to dtype by from_pretrained (loaded has them in bf16), unlike ideogram4's inv_freq which stays fp32, so restoring them isn't right for that one.

Happy to open a PR with this (and removing ideogram4's skip) if the approach looks reasonable.