diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py b/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py index acc04f2c179c..23ad2c725c00 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,6 +857,9 @@ 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. """ + # 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 @@ 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)`. """ + 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: 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..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,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. diff --git a/src/diffusers/models/transformers/transformer_minimax_h3.py b/src/diffusers/models/transformers/transformer_minimax_h3.py index 5b2be83bdbbd..bbabe7c6c7fe 100644 --- a/src/diffusers/models/transformers/transformer_minimax_h3.py +++ b/src/diffusers/models/transformers/transformer_minimax_h3.py @@ -51,7 +51,9 @@ class MiniMaxH3TransformerOutput(BaseOutput): """ sample: torch.Tensor - audio_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..e97dd8fd81ba 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 implementation, so every test below that runs a + # backward pass has to relax determinism. 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 " @@ -113,7 +111,7 @@ class TestAutoencoderKLKVAEVideoMemory(AutoencoderKLKVAEVideoTesterConfig, Memor """Memory optimization tests for AutoencoderKLKVAEVideo.""" 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 cb6c4ebb12df..497614b38141 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 @@ -134,10 +135,35 @@ 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 left to " + "cast and the memory footprint cannot go down." + ) + def test_layerwise_casting_memory(self): + pass + + # The encoder pads spatially with `mode="reflect"`, whose reflection_pad3d_backward_out_cuda has no deterministic + # implementation, so every test below that runs a backward pass has to relax determinism. + 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.""" + # See `TestAutoencoderKLMiniMaxH3Memory` for why these relax determinism. + 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) + def test_gradient_checkpointing_is_applied(self): super().test_gradient_checkpointing_is_applied( expected_set={"MiniMaxH3VideoDownBlock3d", "MiniMaxH3VideoViTDecoder3d"} 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..ab8956096664 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,8 +25,8 @@ BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, - TorchCompileTesterMixin, TrainingTesterMixin, + run_nondeterministic, ) @@ -130,10 +130,32 @@ 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 left to " + "cast and the memory footprint cannot go down." + ) + def test_layerwise_casting_memory(self): + pass + + # The latent projection pools with `F.adaptive_avg_pool1d`, whose adaptive_avg_pool2d_backward_cuda has no + # deterministic implementation, so every test below that runs a backward pass has to relax determinism. + 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.""" + # See `TestAutoencoderKLMiniMaxH3AudioMemory` for why these relax determinism. + 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 +166,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.""" diff --git a/tests/models/autoencoders/test_models_autoencoder_vidtok.py b/tests/models/autoencoders/test_models_autoencoder_vidtok.py index 9810296a07d9..61695b377f44 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,26 @@ 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 implementation, so every test below that runs a backward pass + # has to relax determinism. 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.""" 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)