tests: fix cuda model tests. (#14464)

* tests: fix cuda model tests.

* up

* up

* up

* add to docstrings.

* restore previous beehaviour.

---------

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

docker/diffusers-pytorch-cuda/Dockerfile

@@ -10,8 +10,11 @@ RUN apt-get -y update \
     && add-apt-repository ppa:deadsnakes/ppa && \
     apt-get update
 
+# cuda-cudart-dev provides the CUDA headers (cuda.h etc.) AOT-inductor compilation needs;
+# the runtime base image only ships the libraries.
 RUN apt install -y bash \
     build-essential \
+    cuda-cudart-dev-12-9 \
     git \
     git-lfs \
     curl \

src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py

@@ -24,7 +24,7 @@ from ...utils.accelerate_utils import apply_forward_hook
 from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward
 from ..attention_dispatch import dispatch_attention_fn
 from ..modeling_outputs import AutoencoderKLOutput
-from ..modeling_utils import ModelMixin
+from ..modeling_utils import ModelMixin, get_parameter_dtype
 from .vae import AutoencoderMixin, DecoderOutput, DiagonalGaussianDistribution
 
 
@@ -857,6 +857,9 @@ class AutoencoderKLMiniMaxH3(ModelMixin, ConfigMixin, AttentionMixin, Autoencode
             The latent distribution of the encoded videos. Note that MiniMax-H3 normalizes the sampled latents with
             `latents_mean` / `latents_std` afterwards.
         """
+        # Every module is pinned to float32 by `_keep_in_fp32_modules`, so a pipeline running in a lower `torch_dtype`
+        # hands over lower-precision pixels; align them with the weights, like the audio autoencoder does.
+        x = x.to(get_parameter_dtype(self.encoder))
         if self.use_slicing and x.shape[0] > 1:
             moments = torch.cat([self._encode(x_slice) for x_slice in x.split(1)])
         else:
@@ -881,6 +884,7 @@ class AutoencoderKLMiniMaxH3(ModelMixin, ConfigMixin, AttentionMixin, Autoencode
             [`~models.autoencoders.vae.DecoderOutput`] or `tuple`:
                 The decoded videos, shape `(batch_size, out_channels, num_frames, height, width)`.
         """
+        z = z.to(get_parameter_dtype(self.decoder))
         if self.use_slicing and z.shape[0] > 1:
             decoded = torch.cat([self._decode(z_slice) for z_slice in z.split(1)])
         else:

src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3_audio.py

@@ -524,6 +524,10 @@ class AutoencoderKLMiniMaxH3Audio(ModelMixin, ConfigMixin, AttentionMixin):
     """
 
     _supports_gradient_checkpointing = False
+    # `weight_norm` recomputes `weight` from `weight_g` / `weight_v` in a forward pre-hook, which runs before the
+    # leaf-level group offloading hook has onloaded them, so the convolution would see a CPU weight. Same reason the
+    # other weight-normalized audio autoencoders (`AutoencoderOobleck`, `Cosmos3AVAEAudioTokenizer`) opt out.
+    _supports_group_offloading = False
     # The released checkpoint is float32 and the DAC/BigVGAN stack (weight-normalized convolutions, Snake
     # activations) degrades audibly under bfloat16 (roughly 20 dB quieter decodes), so a pipeline-level
     # `torch_dtype=torch.bfloat16` must not downcast the weights.

src/diffusers/models/transformers/transformer_minimax_h3.py

@@ -46,12 +46,14 @@ class MiniMaxH3TransformerOutput(BaseOutput):
         sample (`torch.Tensor` of shape `(batch_size, num_video_tokens, in_channels * prod(patch_size))`):
             The video velocity prediction for the rows addressed by `video_indices`, in the same order. Conditioning
             rows are returned unmasked — masking them out before the scheduler step is the caller's job.
-        audio_sample (`torch.Tensor` of shape `(batch_size, num_audio_tokens, audio_in_channels)`):
-            The audio velocity prediction for the rows addressed by `audio_indices`, in the same order.
+        audio_sample (`torch.Tensor` of shape `(batch_size, num_audio_tokens, audio_in_channels)`, defaults to `None`):
+            The audio velocity prediction for the rows addressed by `audio_indices`, in the same order. `forward`
+            always populates it; it only defaults to `None` so that the output can be rebuilt from a plain dict of its
+            fields, which is how the accelerate offload hooks move a `BaseOutput` back to the input device.
     """
 
     sample: torch.Tensor
-    audio_sample: torch.Tensor
+    audio_sample: torch.Tensor | None = None
 
 
 def _apply_rotary_emb(hidden_states: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:

tests/models/autoencoders/test_models_autoencoder_kl_kvae_video.py

@@ -20,23 +20,19 @@ from diffusers import AutoencoderKLKVAEVideo
 from diffusers.utils.torch_utils import randn_tensor
 
 from ...testing_utils import enable_full_determinism, torch_device
-from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TrainingTesterMixin
+from ..testing_utils import (
+    BaseModelTesterConfig,
+    MemoryTesterMixin,
+    ModelTesterMixin,
+    TrainingTesterMixin,
+    run_nondeterministic,
+)
 from .testing_utils import NewAutoencoderTesterMixin
 
 
 enable_full_determinism()
 
 
-def _run_nondeterministic(fn):
-    # reflection_pad3d_backward_out_cuda has no deterministic CUDA implementation;
-    # temporarily relax the requirement for tests that do backward passes.
-    torch.use_deterministic_algorithms(False)
-    try:
-        fn()
-    finally:
-        torch.use_deterministic_algorithms(True)
-
-
 class AutoencoderKLKVAEVideoTesterConfig(BaseModelTesterConfig):
     @property
     def model_class(self):
@@ -91,14 +87,16 @@ class TestAutoencoderKLKVAEVideoTraining(AutoencoderKLKVAEVideoTesterConfig, Tra
         expected_set = {"KVAECachedEncoder3D", "KVAECachedDecoder3D"}
         super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
 
+    # reflection_pad3d_backward_out_cuda has no deterministic CUDA implementation, so every test with a backward
+    # pass runs with determinism relaxed.
     def test_training(self):
-        _run_nondeterministic(super().test_training)
+        run_nondeterministic(super().test_training)
 
     def test_training_with_ema(self):
-        _run_nondeterministic(super().test_training_with_ema)
+        run_nondeterministic(super().test_training_with_ema)
 
     def test_mixed_precision_training(self):
-        _run_nondeterministic(super().test_mixed_precision_training)
+        run_nondeterministic(super().test_mixed_precision_training)
 
     @pytest.mark.skip(
         "Gradient checkpointing recomputes the forward pass, but the model uses a stateful cache_dict "
@@ -112,8 +110,10 @@ class TestAutoencoderKLKVAEVideoTraining(AutoencoderKLKVAEVideoTesterConfig, Tra
 class TestAutoencoderKLKVAEVideoMemory(AutoencoderKLKVAEVideoTesterConfig, MemoryTesterMixin):
     """Memory optimization tests for AutoencoderKLKVAEVideo."""
 
+    # reflection_pad3d_backward_out_cuda has no deterministic CUDA implementation, so every test with a backward
+    # pass runs with determinism relaxed.
     def test_layerwise_casting_training(self):
-        _run_nondeterministic(super().test_layerwise_casting_training)
+        run_nondeterministic(super().test_layerwise_casting_training)
 
 
 class TestAutoencoderKLKVAEVideoSlicingTiling(AutoencoderKLKVAEVideoTesterConfig, NewAutoencoderTesterMixin):

tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py

@@ -25,8 +25,8 @@ from ..testing_utils import (
     BaseModelTesterConfig,
     MemoryTesterMixin,
     ModelTesterMixin,
-    TorchCompileTesterMixin,
     TrainingTesterMixin,
+    run_nondeterministic,
 )
 from .testing_utils import NewAutoencoderTesterMixin
 
@@ -134,6 +134,18 @@ class TestAutoencoderKLMiniMaxH3(AutoencoderKLMiniMaxH3TesterConfig, ModelTester
 class TestAutoencoderKLMiniMaxH3Memory(AutoencoderKLMiniMaxH3TesterConfig, MemoryTesterMixin):
     """Memory optimization tests for the MiniMax-H3 video autoencoder."""
 
+    @pytest.mark.skip(
+        "`_keep_in_fp32_modules` pins every module of this autoencoder, so layerwise casting has nothing to cast and "
+        "the memory footprint does not change."
+    )
+    def test_layerwise_casting_memory(self):
+        pass
+
+    # The encoder's reflect spatial padding has no deterministic CUDA backward
+    # (reflection_pad3d_backward_out_cuda), so every test with a backward pass runs with determinism relaxed.
+    def test_layerwise_casting_training(self):
+        run_nondeterministic(super().test_layerwise_casting_training)
+
 
 class TestAutoencoderKLMiniMaxH3Training(AutoencoderKLMiniMaxH3TesterConfig, TrainingTesterMixin):
     """Training tests for the MiniMax-H3 video autoencoder."""
@@ -143,14 +155,24 @@ class TestAutoencoderKLMiniMaxH3Training(AutoencoderKLMiniMaxH3TesterConfig, Tra
             expected_set={"MiniMaxH3VideoDownBlock3d", "MiniMaxH3VideoViTDecoder3d"}
         )
 
+    # The encoder's reflect spatial padding has no deterministic CUDA backward
+    # (reflection_pad3d_backward_out_cuda), so every test with a backward pass runs with determinism relaxed.
+    def test_training(self):
+        run_nondeterministic(super().test_training)
+
+    def test_training_with_ema(self):
+        run_nondeterministic(super().test_training_with_ema)
+
+    def test_mixed_precision_training(self):
+        run_nondeterministic(super().test_mixed_precision_training)
+
+    def test_gradient_checkpointing_equivalence(self):
+        run_nondeterministic(super().test_gradient_checkpointing_equivalence)
+
 
 class TestAutoencoderKLMiniMaxH3Attention(AutoencoderKLMiniMaxH3TesterConfig, AttentionTesterMixin):
     """Attention processor tests for the MiniMax-H3 video autoencoder."""
 
 
-class TestAutoencoderKLMiniMaxH3TorchCompile(AutoencoderKLMiniMaxH3TesterConfig, TorchCompileTesterMixin):
-    """Torch compile tests for the MiniMax-H3 video autoencoder."""
-
-
 class TestAutoencoderKLMiniMaxH3SlicingTiling(AutoencoderKLMiniMaxH3TesterConfig, NewAutoencoderTesterMixin):
     """Slicing and tiling tests for the MiniMax-H3 video autoencoder."""

tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3_audio.py

@@ -25,8 +25,8 @@ from ..testing_utils import (
     BaseModelTesterConfig,
     MemoryTesterMixin,
     ModelTesterMixin,
-    TorchCompileTesterMixin,
     TrainingTesterMixin,
+    run_nondeterministic,
 )
 
 
@@ -130,10 +130,33 @@ class TestAutoencoderKLMiniMaxH3Audio(AutoencoderKLMiniMaxH3AudioTesterConfig, M
 class TestAutoencoderKLMiniMaxH3AudioMemory(AutoencoderKLMiniMaxH3AudioTesterConfig, MemoryTesterMixin):
     """Memory optimization tests for the MiniMax-H3 audio autoencoder."""
 
+    @pytest.mark.skip(
+        "`_keep_in_fp32_modules` pins every module of this autoencoder, so layerwise casting has nothing to cast and "
+        "the memory footprint does not change."
+    )
+    def test_layerwise_casting_memory(self):
+        pass
+
+    # The causal-attention projection's `F.adaptive_avg_pool1d` has no deterministic CUDA backward
+    # (adaptive_avg_pool2d_backward_cuda), so every test with a backward pass runs with determinism relaxed.
+    def test_layerwise_casting_training(self):
+        run_nondeterministic(super().test_layerwise_casting_training)
+
 
 class TestAutoencoderKLMiniMaxH3AudioTraining(AutoencoderKLMiniMaxH3AudioTesterConfig, TrainingTesterMixin):
     """Training tests for the MiniMax-H3 audio autoencoder."""
 
+    # The causal-attention projection's `F.adaptive_avg_pool1d` has no deterministic CUDA backward
+    # (adaptive_avg_pool2d_backward_cuda), so every test with a backward pass runs with determinism relaxed.
+    def test_training(self):
+        run_nondeterministic(super().test_training)
+
+    def test_training_with_ema(self):
+        run_nondeterministic(super().test_training_with_ema)
+
+    def test_mixed_precision_training(self):
+        run_nondeterministic(super().test_mixed_precision_training)
+
 
 class TestAutoencoderKLMiniMaxH3AudioAttention(AutoencoderKLMiniMaxH3AudioTesterConfig, AttentionTesterMixin):
     """Attention processor tests for the MiniMax-H3 audio autoencoder."""
@@ -144,7 +167,3 @@ class TestAutoencoderKLMiniMaxH3AudioAttention(AutoencoderKLMiniMaxH3AudioTester
     )
     def test_attention_processor_count_mismatch_raises_error(self):
         pass
-
-
-class TestAutoencoderKLMiniMaxH3AudioTorchCompile(AutoencoderKLMiniMaxH3AudioTesterConfig, TorchCompileTesterMixin):
-    """Torch compile tests for the MiniMax-H3 audio autoencoder."""

tests/models/autoencoders/test_models_autoencoder_vidtok.py

@@ -19,23 +19,19 @@ from diffusers import AutoencoderVidTok
 from diffusers.utils.torch_utils import randn_tensor
 
 from ...testing_utils import enable_full_determinism, torch_device
-from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TrainingTesterMixin
+from ..testing_utils import (
+    BaseModelTesterConfig,
+    MemoryTesterMixin,
+    ModelTesterMixin,
+    TrainingTesterMixin,
+    run_nondeterministic,
+)
 from .testing_utils import NewAutoencoderTesterMixin
 
 
 enable_full_determinism()
 
 
-def _run_nondeterministic(fn):
-    # avg_pool3d_backward_cuda has no deterministic CUDA implementation;
-    # temporarily relax the requirement for tests that do backward passes.
-    torch.use_deterministic_algorithms(False)
-    try:
-        fn()
-    finally:
-        torch.use_deterministic_algorithms(True)
-
-
 class AutoencoderVidTokTesterConfig(BaseModelTesterConfig):
     @property
     def model_class(self):
@@ -90,24 +86,28 @@ class TestAutoencoderVidTokTraining(AutoencoderVidTokTesterConfig, TrainingTeste
         expected_set = {"VidTokEncoder3D", "VidTokDecoder3D"}
         super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
 
+    # avg_pool3d_backward_cuda has no deterministic CUDA implementation, so every test with a backward pass runs
+    # with determinism relaxed.
     def test_training(self):
-        _run_nondeterministic(super().test_training)
+        run_nondeterministic(super().test_training)
 
     def test_training_with_ema(self):
-        _run_nondeterministic(super().test_training_with_ema)
+        run_nondeterministic(super().test_training_with_ema)
 
     def test_mixed_precision_training(self):
-        _run_nondeterministic(super().test_mixed_precision_training)
+        run_nondeterministic(super().test_mixed_precision_training)
 
     def test_gradient_checkpointing_equivalence(self):
-        _run_nondeterministic(super().test_gradient_checkpointing_equivalence)
+        run_nondeterministic(super().test_gradient_checkpointing_equivalence)
 
 
 class TestAutoencoderVidTokMemory(AutoencoderVidTokTesterConfig, MemoryTesterMixin):
     """Memory optimization tests for AutoencoderVidTok."""
 
+    # avg_pool3d_backward_cuda has no deterministic CUDA implementation, so every test with a backward pass runs
+    # with determinism relaxed.
     def test_layerwise_casting_training(self):
-        _run_nondeterministic(super().test_layerwise_casting_training)
+        run_nondeterministic(super().test_layerwise_casting_training)
 
 
 class TestAutoencoderVidTokSlicingTiling(AutoencoderVidTokTesterConfig, NewAutoencoderTesterMixin):

tests/models/testing_utils/__init__.py

@@ -48,6 +48,7 @@ from .quantization import (
 )
 from .single_file import SingleFileTesterMixin
 from .training import TrainingTesterMixin
+from .utils import run_nondeterministic
 
 
 __all__ = [
@@ -94,6 +95,7 @@ __all__ = [
     "QuantoCompileTesterMixin",
     "QuantoConfigMixin",
     "QuantoTesterMixin",
+    "run_nondeterministic",
     "SDNQCompileTesterMixin",
     "SDNQConfigMixin",
     "SDNQTesterMixin",

tests/models/testing_utils/compile.py

@@ -82,7 +82,7 @@ class TorchCompileTesterMixin:
 
     @torch.no_grad()
     def test_torch_compile_repeated_blocks(self, recompile_limit=1):
-        if self.model_class._repeated_blocks is None:
+        if not self.model_class._repeated_blocks:
             pytest.skip("Skipping test as the model class doesn't have `_repeated_blocks` set.")
 
         init_dict = self.get_init_dict()

tests/models/testing_utils/utils.py

@@ -24,3 +24,21 @@ def _maybe_cast_to_bf16(backend, model, inputs_dict):
         for k, v in inputs_dict.items()
     }
     return model, inputs_dict
+
+
+def run_nondeterministic(fn):
+    """
+    Run `fn` with `enable_full_determinism`'s deterministic-algorithm requirement lifted, restoring the previous
+    setting afterwards.
+
+    Several models reach a backward kernel that has no deterministic CUDA implementation (reflection/replication
+    padding, average pooling), which makes every test doing a backward pass raise under
+    `torch.use_deterministic_algorithms(True)`. Wrap those tests instead of relaxing determinism for the whole module,
+    and name the offending op at the call site.
+    """
+    was_enabled = torch.are_deterministic_algorithms_enabled()
+    torch.use_deterministic_algorithms(False)
+    try:
+        fn()
+    finally:
+        torch.use_deterministic_algorithms(was_enabled)

tests/models/transformers/test_models_transformer_cosmos3.py

@@ -195,9 +195,11 @@ class TestCosmos3OmniTransformerModel(Cosmos3OmniTransformerTesterConfig, ModelT
         assert not torch.allclose(generation_before, generation_after)
 
     def test_cosmos3_edge_transformer_runs_action_workflow(self):
-        transformer = self.model_class(
-            **self.get_init_dict(), action_dim=3, action_gen=True, num_embodiment_domains=2
-        ).eval()
+        transformer = (
+            self.model_class(**self.get_init_dict(), action_dim=3, action_gen=True, num_embodiment_domains=2)
+            .to(torch_device)
+            .eval()
+        )
         inputs = self.get_dummy_inputs()
         inputs["position_ids"] = torch.zeros((3, 4), dtype=torch.long, device=torch_device)
         inputs["sequence_length"] = 4

tests/models/transformers/test_models_transformer_z_image.py

@@ -218,6 +218,11 @@ class TestZImageTransformerCompile(ZImageTransformerTesterConfig, TorchCompileTe
     def test_torch_compile_recompilation_and_graph_break(self):
         pass
 
+    def test_torch_compile_repeated_blocks(self):
+        # ZImageTransformerBlock is reused by noise_refiner (modulated), context_refiner (unmodulated), and the
+        # main layers (modulated, different sequence length), so the shared block forward compiles three times.
+        super().test_torch_compile_repeated_blocks(recompile_limit=3)
+
     @pytest.mark.skip("Fullgraph AoT is broken")
     def test_compile_works_with_aot(self, tmp_path):
         pass