From c916ed9e84543afb6810d6549c671e31915deb5d Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Thu, 13 Aug 2026 15:05:32 +0000 Subject: [PATCH 1/6] tests: fix cuda model tests. --- docker/diffusers-pytorch-cuda/Dockerfile | 3 ++ .../autoencoders/autoencoder_kl_minimax_h3.py | 11 ++++-- .../autoencoder_kl_minimax_h3_audio.py | 4 +++ .../models/transformers/transformer_ltx2.py | 6 ++-- .../transformers/transformer_minimax_h3.py | 6 ++-- .../test_models_autoencoder_kl_minimax_h3.py | 32 +++++++++++++++++ ..._models_autoencoder_kl_minimax_h3_audio.py | 35 +++++++++++++++++++ tests/models/testing_utils/compile.py | 2 +- .../test_models_transformer_cosmos3.py | 8 +++-- .../test_models_transformer_z_image.py | 5 +++ 10 files changed, 102 insertions(+), 10 deletions(-) diff --git a/docker/diffusers-pytorch-cuda/Dockerfile b/docker/diffusers-pytorch-cuda/Dockerfile index 28ae25717c54..345bef4648b8 100644 --- a/docker/diffusers-pytorch-cuda/Dockerfile +++ b/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 \ diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py b/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py index acc04f2c179c..a58bba8f56d8 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py @@ -24,7 +24,7 @@ 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,11 +857,15 @@ def encode(self, x: torch.Tensor, return_dict: bool = True) -> AutoencoderKLOutp The latent distribution of the encoded videos. Note that MiniMax-H3 normalizes the sampled latents with `latents_mean` / `latents_std` afterwards. """ + # The whole model is pinned to float32 (`_keep_in_fp32_modules`), so align the input with the encoder's + # parameter dtype and hand the caller back its own dtype, mirroring the audio autoencoder. + input_dtype = x.dtype + 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: moments = self._encode(x) - posterior = DiagonalGaussianDistribution(moments) + posterior = DiagonalGaussianDistribution(moments.to(input_dtype)) if not return_dict: return (posterior,) return AutoencoderKLOutput(latent_dist=posterior) @@ -881,10 +885,13 @@ def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | t [`~models.autoencoders.vae.DecoderOutput`] or `tuple`: The decoded videos, shape `(batch_size, out_channels, num_frames, height, width)`. """ + input_dtype = z.dtype + 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: decoded = self._decode(z) + decoded = decoded.to(input_dtype) if not return_dict: return (decoded,) return DecoderOutput(sample=decoded) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3_audio.py b/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3_audio.py index 957774f791bc..45d58457aad8 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3_audio.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3_audio.py @@ -524,6 +524,10 @@ class AutoencoderKLMiniMaxH3Audio(ModelMixin, ConfigMixin, AttentionMixin): """ _supports_gradient_checkpointing = False + # `torch.nn.utils.weight_norm` recomputes `weight` from `weight_g` / `weight_v` in a pre-forward hook, which + # runs before the group-offloading hook has onloaded the parameters, leaving the recomputed weight on the + # offload device (same reason the other weight-normed audio autoencoders disable it). + _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. diff --git a/src/diffusers/models/transformers/transformer_ltx2.py b/src/diffusers/models/transformers/transformer_ltx2.py index b5fdbcfb0b06..ca6947b3a535 100644 --- a/src/diffusers/models/transformers/transformer_ltx2.py +++ b/src/diffusers/models/transformers/transformer_ltx2.py @@ -97,8 +97,10 @@ class AudioVisualModelOutput(BaseOutput): The audio output of the audiovisual model. """ - sample: "torch.Tensor" # noqa: F821 - audio_sample: "torch.Tensor" # noqa: F821 + # The `None` defaults keep the class reconstructible from a plain `{field: value}` dict, which is how + # accelerate's `send_to_device` rebuilds model outputs when offloading hooks are attached. + sample: "torch.Tensor" = None # noqa: F821 + audio_sample: "torch.Tensor" = None # noqa: F821 class LTX2AdaLayerNormSingle(nn.Module): diff --git a/src/diffusers/models/transformers/transformer_minimax_h3.py b/src/diffusers/models/transformers/transformer_minimax_h3.py index 5b2be83bdbbd..2ac49454f3d1 100644 --- a/src/diffusers/models/transformers/transformer_minimax_h3.py +++ b/src/diffusers/models/transformers/transformer_minimax_h3.py @@ -50,8 +50,10 @@ class MiniMaxH3TransformerOutput(BaseOutput): The audio velocity prediction for the rows addressed by `audio_indices`, in the same order. """ - sample: torch.Tensor - audio_sample: torch.Tensor + # The `None` defaults keep the class reconstructible from a plain `{field: value}` dict, which is how + # accelerate's `send_to_device` rebuilds model outputs when offloading hooks are attached. + sample: torch.Tensor = None + audio_sample: torch.Tensor = None def _apply_rotary_emb(hidden_states: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py b/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py index cb6c4ebb12df..c01771a9283e 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py @@ -34,6 +34,16 @@ enable_full_determinism() +def _run_nondeterministic(fn): + # The encoder's reflect spatial padding has no deterministic CUDA backward + # (reflection_pad3d_backward_out_cuda); temporarily relax the requirement for tests that do backward passes. + torch.use_deterministic_algorithms(False) + try: + fn() + finally: + torch.use_deterministic_algorithms(True) + + # The temporal geometry is fixed by the released `clip_length` / `token_drop` pair: `17 * n + 5` pixel frames map to # `5 * n + 2` latent frames, so 22 frames is the shortest video this autoencoder can round-trip. NUM_FRAMES = 22 @@ -134,6 +144,16 @@ def test_encode_decode_temporal_geometry(self): 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 + + 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,6 +163,18 @@ def test_gradient_checkpointing_is_applied(self): expected_set={"MiniMaxH3VideoDownBlock3d", "MiniMaxH3VideoViTDecoder3d"} ) + 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.""" diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3_audio.py b/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3_audio.py index 676cae58d335..8c32dd9f10a0 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3_audio.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3_audio.py @@ -33,6 +33,16 @@ enable_full_determinism() +def _run_nondeterministic(fn): + # The causal-attention projection's `F.adaptive_avg_pool1d` has no deterministic CUDA backward + # (adaptive_avg_pool2d_backward_cuda); temporarily relax the requirement for tests that do backward passes. + torch.use_deterministic_algorithms(False) + try: + fn() + finally: + torch.use_deterministic_algorithms(True) + + # `prod(encoder_rates)` is the hop length, so a waveform of `HOP_LENGTH * NUM_LATENTS` samples encodes to # `NUM_LATENTS` latents. MiniMax-H3 carries stereo as two batch items of this mono autoencoder. HOP_LENGTH = 4 @@ -130,10 +140,29 @@ def test_encode_pads_to_the_hop_length(self): 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 + + 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.""" + 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.""" @@ -148,3 +177,9 @@ def test_attention_processor_count_mismatch_raises_error(self): class TestAutoencoderKLMiniMaxH3AudioTorchCompile(AutoencoderKLMiniMaxH3AudioTesterConfig, TorchCompileTesterMixin): """Torch compile tests for the MiniMax-H3 audio autoencoder.""" + + def test_torch_compile_recompilation_and_graph_break(self): + # Under `torch.use_deterministic_algorithms(True)`, `F.pad(mode="replicate")` (used by the anti-aliased + # resamplers) routes through an `importlib.import_module` call, which dynamo cannot trace with + # fullgraph=True on torch <= 2.10. + _run_nondeterministic(super().test_torch_compile_recompilation_and_graph_break) diff --git a/tests/models/testing_utils/compile.py b/tests/models/testing_utils/compile.py index 5d638ec0b6ff..4f41a2fc7beb 100644 --- a/tests/models/testing_utils/compile.py +++ b/tests/models/testing_utils/compile.py @@ -82,7 +82,7 @@ def test_torch_compile_recompilation_and_graph_break(self): @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() diff --git a/tests/models/transformers/test_models_transformer_cosmos3.py b/tests/models/transformers/test_models_transformer_cosmos3.py index bbe8d0755fce..bfe1d5b3be7f 100644 --- a/tests/models/transformers/test_models_transformer_cosmos3.py +++ b/tests/models/transformers/test_models_transformer_cosmos3.py @@ -195,9 +195,11 @@ def test_cosmos3_edge_generator_k_norm_does_not_change_causal_attention(self): 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 diff --git a/tests/models/transformers/test_models_transformer_z_image.py b/tests/models/transformers/test_models_transformer_z_image.py index 85aeb34c25c4..7b66bf5196aa 100644 --- a/tests/models/transformers/test_models_transformer_z_image.py +++ b/tests/models/transformers/test_models_transformer_z_image.py @@ -218,6 +218,11 @@ def get_dummy_inputs(self, height: int = 16, width: int = 16) -> dict[str, torch 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 From f20ff937e7ebcae74c119f78518c5a58286884f0 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Thu, 13 Aug 2026 15:27:04 +0000 Subject: [PATCH 2/6] up --- .../autoencoders/autoencoder_kl_minimax_h3.py | 9 ++---- .../autoencoder_kl_minimax_h3_audio.py | 6 ++-- .../models/transformers/transformer_ltx2.py | 8 ++--- .../transformers/transformer_minimax_h3.py | 8 ++--- .../test_models_autoencoder_kl_kvae_video.py | 30 ++++++++--------- .../test_models_autoencoder_kl_minimax_h3.py | 25 ++++++--------- ..._models_autoencoder_kl_minimax_h3_audio.py | 25 ++++++--------- .../test_models_autoencoder_vidtok.py | 32 +++++++++---------- tests/models/testing_utils/__init__.py | 2 ++ tests/models/testing_utils/utils.py | 16 ++++++++++ 10 files changed, 83 insertions(+), 78 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py b/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py index a58bba8f56d8..23ad2c725c00 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py @@ -857,15 +857,14 @@ def encode(self, x: torch.Tensor, return_dict: bool = True) -> AutoencoderKLOutp The latent distribution of the encoded videos. Note that MiniMax-H3 normalizes the sampled latents with `latents_mean` / `latents_std` afterwards. """ - # The whole model is pinned to float32 (`_keep_in_fp32_modules`), so align the input with the encoder's - # parameter dtype and hand the caller back its own dtype, mirroring the audio autoencoder. - input_dtype = x.dtype + # 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: moments = self._encode(x) - posterior = DiagonalGaussianDistribution(moments.to(input_dtype)) + posterior = DiagonalGaussianDistribution(moments) if not return_dict: return (posterior,) return AutoencoderKLOutput(latent_dist=posterior) @@ -885,13 +884,11 @@ def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | t [`~models.autoencoders.vae.DecoderOutput`] or `tuple`: The decoded videos, shape `(batch_size, out_channels, num_frames, height, width)`. """ - input_dtype = z.dtype 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: decoded = self._decode(z) - decoded = decoded.to(input_dtype) if not return_dict: return (decoded,) return DecoderOutput(sample=decoded) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3_audio.py b/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3_audio.py index 45d58457aad8..a70c947712ae 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3_audio.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3_audio.py @@ -524,9 +524,9 @@ class AutoencoderKLMiniMaxH3Audio(ModelMixin, ConfigMixin, AttentionMixin): """ _supports_gradient_checkpointing = False - # `torch.nn.utils.weight_norm` recomputes `weight` from `weight_g` / `weight_v` in a pre-forward hook, which - # runs before the group-offloading hook has onloaded the parameters, leaving the recomputed weight on the - # offload device (same reason the other weight-normed audio autoencoders disable it). + # `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 diff --git a/src/diffusers/models/transformers/transformer_ltx2.py b/src/diffusers/models/transformers/transformer_ltx2.py index ca6947b3a535..39850dc33d36 100644 --- a/src/diffusers/models/transformers/transformer_ltx2.py +++ b/src/diffusers/models/transformers/transformer_ltx2.py @@ -97,10 +97,10 @@ class AudioVisualModelOutput(BaseOutput): The audio output of the audiovisual model. """ - # The `None` defaults keep the class reconstructible from a plain `{field: value}` dict, which is how - # accelerate's `send_to_device` rebuilds model outputs when offloading hooks are attached. - sample: "torch.Tensor" = None # noqa: F821 - audio_sample: "torch.Tensor" = None # noqa: F821 + sample: "torch.Tensor" # noqa: F821 + # `forward` always populates `audio_sample`; the default is what lets the output be rebuilt from a plain dict of + # its fields, which is how the accelerate offload hooks move a `BaseOutput` back to the input device. + audio_sample: "torch.Tensor | None" = None # noqa: F821 class LTX2AdaLayerNormSingle(nn.Module): diff --git a/src/diffusers/models/transformers/transformer_minimax_h3.py b/src/diffusers/models/transformers/transformer_minimax_h3.py index 2ac49454f3d1..bbabe7c6c7fe 100644 --- a/src/diffusers/models/transformers/transformer_minimax_h3.py +++ b/src/diffusers/models/transformers/transformer_minimax_h3.py @@ -50,10 +50,10 @@ class MiniMaxH3TransformerOutput(BaseOutput): The audio velocity prediction for the rows addressed by `audio_indices`, in the same order. """ - # The `None` defaults keep the class reconstructible from a plain `{field: value}` dict, which is how - # accelerate's `send_to_device` rebuilds model outputs when offloading hooks are attached. - sample: torch.Tensor = None - audio_sample: torch.Tensor = None + sample: torch.Tensor + # `forward` always populates `audio_sample`; the default is what lets the output be rebuilt from a plain dict of + # its fields, which is how the accelerate offload hooks move a `BaseOutput` back to the input device. + audio_sample: torch.Tensor | None = None def _apply_rotary_emb(hidden_states: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_kvae_video.py b/tests/models/autoencoders/test_models_autoencoder_kl_kvae_video.py index a7097e128a13..6923def0d94d 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_kvae_video.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_kvae_video.py @@ -20,23 +20,19 @@ 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 @@ def test_gradient_checkpointing_is_applied(self): 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 @@ def test_gradient_checkpointing_equivalence(self): 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): diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py b/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py index c01771a9283e..92dd3a789dd3 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py @@ -27,6 +27,7 @@ ModelTesterMixin, TorchCompileTesterMixin, TrainingTesterMixin, + run_nondeterministic, ) from .testing_utils import NewAutoencoderTesterMixin @@ -34,16 +35,6 @@ enable_full_determinism() -def _run_nondeterministic(fn): - # The encoder's reflect spatial padding has no deterministic CUDA backward - # (reflection_pad3d_backward_out_cuda); temporarily relax the requirement for tests that do backward passes. - torch.use_deterministic_algorithms(False) - try: - fn() - finally: - torch.use_deterministic_algorithms(True) - - # The temporal geometry is fixed by the released `clip_length` / `token_drop` pair: `17 * n + 5` pixel frames map to # `5 * n + 2` latent frames, so 22 frames is the shortest video this autoencoder can round-trip. NUM_FRAMES = 22 @@ -151,8 +142,10 @@ class TestAutoencoderKLMiniMaxH3Memory(AutoencoderKLMiniMaxH3TesterConfig, Memor 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) + run_nondeterministic(super().test_layerwise_casting_training) class TestAutoencoderKLMiniMaxH3Training(AutoencoderKLMiniMaxH3TesterConfig, TrainingTesterMixin): @@ -163,17 +156,19 @@ def test_gradient_checkpointing_is_applied(self): 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) + 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 TestAutoencoderKLMiniMaxH3Attention(AutoencoderKLMiniMaxH3TesterConfig, AttentionTesterMixin): diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3_audio.py b/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3_audio.py index 8c32dd9f10a0..b1d5188d60ed 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3_audio.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3_audio.py @@ -27,22 +27,13 @@ ModelTesterMixin, TorchCompileTesterMixin, TrainingTesterMixin, + run_nondeterministic, ) enable_full_determinism() -def _run_nondeterministic(fn): - # The causal-attention projection's `F.adaptive_avg_pool1d` has no deterministic CUDA backward - # (adaptive_avg_pool2d_backward_cuda); temporarily relax the requirement for tests that do backward passes. - torch.use_deterministic_algorithms(False) - try: - fn() - finally: - torch.use_deterministic_algorithms(True) - - # `prod(encoder_rates)` is the hop length, so a waveform of `HOP_LENGTH * NUM_LATENTS` samples encodes to # `NUM_LATENTS` latents. MiniMax-H3 carries stereo as two batch items of this mono autoencoder. HOP_LENGTH = 4 @@ -147,21 +138,25 @@ class TestAutoencoderKLMiniMaxH3AudioMemory(AutoencoderKLMiniMaxH3AudioTesterCon 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) + 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) + 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) class TestAutoencoderKLMiniMaxH3AudioAttention(AutoencoderKLMiniMaxH3AudioTesterConfig, AttentionTesterMixin): @@ -182,4 +177,4 @@ def test_torch_compile_recompilation_and_graph_break(self): # Under `torch.use_deterministic_algorithms(True)`, `F.pad(mode="replicate")` (used by the anti-aliased # resamplers) routes through an `importlib.import_module` call, which dynamo cannot trace with # fullgraph=True on torch <= 2.10. - _run_nondeterministic(super().test_torch_compile_recompilation_and_graph_break) + run_nondeterministic(super().test_torch_compile_recompilation_and_graph_break) diff --git a/tests/models/autoencoders/test_models_autoencoder_vidtok.py b/tests/models/autoencoders/test_models_autoencoder_vidtok.py index 9810296a07d9..323d92b40572 100644 --- a/tests/models/autoencoders/test_models_autoencoder_vidtok.py +++ b/tests/models/autoencoders/test_models_autoencoder_vidtok.py @@ -19,23 +19,19 @@ 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 @@ def test_gradient_checkpointing_is_applied(self): 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): diff --git a/tests/models/testing_utils/__init__.py b/tests/models/testing_utils/__init__.py index 23aa871b80a2..67161159732a 100644 --- a/tests/models/testing_utils/__init__.py +++ b/tests/models/testing_utils/__init__.py @@ -48,6 +48,7 @@ ) from .single_file import SingleFileTesterMixin from .training import TrainingTesterMixin +from .utils import run_nondeterministic __all__ = [ @@ -94,6 +95,7 @@ "QuantoCompileTesterMixin", "QuantoConfigMixin", "QuantoTesterMixin", + "run_nondeterministic", "SDNQCompileTesterMixin", "SDNQConfigMixin", "SDNQTesterMixin", diff --git a/tests/models/testing_utils/utils.py b/tests/models/testing_utils/utils.py index 3beb59ed1a66..6c11552ae011 100644 --- a/tests/models/testing_utils/utils.py +++ b/tests/models/testing_utils/utils.py @@ -24,3 +24,19 @@ 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. + + 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. + """ + torch.use_deterministic_algorithms(False) + try: + fn() + finally: + torch.use_deterministic_algorithms(True) From 518c706a85441786f03e30f59c27c0c808460bf5 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Thu, 13 Aug 2026 15:35:09 +0000 Subject: [PATCH 3/6] up --- .../test_models_autoencoder_kl_minimax_h3.py | 5 ----- .../test_models_autoencoder_kl_minimax_h3_audio.py | 11 ----------- 2 files changed, 16 deletions(-) diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py b/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py index 92dd3a789dd3..a09fa57129ff 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py @@ -25,7 +25,6 @@ BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, - TorchCompileTesterMixin, TrainingTesterMixin, run_nondeterministic, ) @@ -175,9 +174,5 @@ class TestAutoencoderKLMiniMaxH3Attention(AutoencoderKLMiniMaxH3TesterConfig, At """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.""" diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3_audio.py b/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3_audio.py index b1d5188d60ed..f7b15be97770 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3_audio.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3_audio.py @@ -25,7 +25,6 @@ BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, - TorchCompileTesterMixin, TrainingTesterMixin, run_nondeterministic, ) @@ -168,13 +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.""" - - def test_torch_compile_recompilation_and_graph_break(self): - # Under `torch.use_deterministic_algorithms(True)`, `F.pad(mode="replicate")` (used by the anti-aliased - # resamplers) routes through an `importlib.import_module` call, which dynamo cannot trace with - # fullgraph=True on torch <= 2.10. - run_nondeterministic(super().test_torch_compile_recompilation_and_graph_break) From 6b9dfcb6aa998fae004766c04357663c87ba71f9 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Thu, 13 Aug 2026 15:49:45 +0000 Subject: [PATCH 4/6] up --- src/diffusers/models/transformers/transformer_ltx2.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_ltx2.py b/src/diffusers/models/transformers/transformer_ltx2.py index 39850dc33d36..b5fdbcfb0b06 100644 --- a/src/diffusers/models/transformers/transformer_ltx2.py +++ b/src/diffusers/models/transformers/transformer_ltx2.py @@ -98,9 +98,7 @@ class AudioVisualModelOutput(BaseOutput): """ sample: "torch.Tensor" # noqa: F821 - # `forward` always populates `audio_sample`; the default is what lets the output be rebuilt from a plain dict of - # its fields, which is how the accelerate offload hooks move a `BaseOutput` back to the input device. - audio_sample: "torch.Tensor | None" = None # noqa: F821 + audio_sample: "torch.Tensor" # noqa: F821 class LTX2AdaLayerNormSingle(nn.Module): From 3ac527a6b62b4ce5e01c971dc45bececcd47131c Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Sat, 15 Aug 2026 05:30:15 +0000 Subject: [PATCH 5/6] add to docstrings. --- .../models/transformers/transformer_minimax_h3.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_minimax_h3.py b/src/diffusers/models/transformers/transformer_minimax_h3.py index bbabe7c6c7fe..f49cdaca2eb6 100644 --- a/src/diffusers/models/transformers/transformer_minimax_h3.py +++ b/src/diffusers/models/transformers/transformer_minimax_h3.py @@ -46,13 +46,13 @@ 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 - # `forward` always populates `audio_sample`; the default is what lets the output be rebuilt from a plain dict of - # its fields, which is how the accelerate offload hooks move a `BaseOutput` back to the input device. audio_sample: torch.Tensor | None = None From 4c545670cbe4e24703fdbee9524f164908b4a735 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Sat, 15 Aug 2026 05:38:12 +0000 Subject: [PATCH 6/6] restore previous beehaviour. --- tests/models/testing_utils/utils.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/models/testing_utils/utils.py b/tests/models/testing_utils/utils.py index 6c11552ae011..07e4a38ddb21 100644 --- a/tests/models/testing_utils/utils.py +++ b/tests/models/testing_utils/utils.py @@ -28,15 +28,17 @@ def _maybe_cast_to_bf16(backend, model, inputs_dict): def run_nondeterministic(fn): """ - Run `fn` with `enable_full_determinism`'s deterministic-algorithm requirement lifted. + 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(True) + torch.use_deterministic_algorithms(was_enabled)