Repository: huggingface/diffusers
File: tests/hooks/test_group_offloading.py (Python)

Code before:
        self.assertTrue(torch.allclose(out_ref, out, atol=1e-5), "Outputs do not match.")

        num_repeats = 4
        for i in range(num_repeats):
            out_ref = model_ref(x)
            out = model(x)

        self.assertTrue(torch.allclose(out_ref, out, atol=1e-5), "Outputs do not match after multiple invocations.")

        for (ref_name, ref_module), (name, module) in zip(model_ref.named_modules(), model.named_modules()):
            assert ref_name == name
            ref_outputs = (
                HookRegistry.check_if_exists_or_initialize(ref_module).get_hook("layer_output_tracker").outputs
            )
            outputs = HookRegistry.check_if_exists_or_initialize(module).get_hook("layer_output_tracker").outputs
            cumulated_absmax = 0.0
            for i in range(len(outputs)):
                diff = ref_outputs[0] - outputs[i]
                absdiff = diff.abs()
                absmax = absdiff.max().item()
                cumulated_absmax += absmax
            self.assertLess(
                cumulated_absmax, 1e-5, f"Output differences for {name} exceeded threshold: {cumulated_absmax:.5f}"
            )

    def test_vae_like_model_with_standalone_conv_layers(self):
        """Test that models with standalone Conv2d layers (like VAE) work with block-level offloading."""
        if torch.device(torch_device).type not in ["cuda", "xpu"]:
            return

        model = DummyVAELikeModel(in_features=64, hidden_features=128, out_features=64)

        model_ref = DummyVAELikeModel(in_features=64, hidden_features=128, out_features=64)
        model_ref.load_state_dict(model.state_dict(), strict=True)
        model_ref.to(torch_device)

        model.enable_group_offload(torch_device, offload_type="block_level", num_blocks_per_group=1, use_stream=True)

        x = torch.randn(2, 64).to(torch_device)

        with torch.no_grad():
            out_ref = model_ref(x)
            out = model(x)

        self.assertTrue(torch.allclose(out_ref, out, atol=1e-5), "Outputs do not match for VAE-like model.")

    def test_vae_like_model_without_streams(self):
        """Test VAE-like model with block-level offloading but without streams."""
        if torch.device(torch_device).type not in ["cuda", "xpu"]:
            return


Review comment (suggestion):
We can leverage an existing model implementations that were reported to be problematic with small configs and use them for testing here. 

Code after:
                )

    @parameterized.expand([("block_level",), ("leaf_level",)])
    def test_standalone_conv_layers_with_both_offload_types(self, offload_type: str):
        """Test that standalone Conv2d layers work correctly with both block-level and leaf-level offloading."""
        if torch.device(torch_device).type not in ["cuda", "xpu"]:
            return

        config = self.get_autoencoder_kl_config()
        model = AutoencoderKL(**config)

        model_ref = AutoencoderKL(**config)
        model_ref.load_state_dict(model.state_dict(), strict=True)
        model_ref.to(torch_device)

        model.enable_group_offload(torch_device, offload_type=offload_type, num_blocks_per_group=1, use_stream=True)

        x = torch.randn(2, 3, 32, 32).to(torch_device)

        with torch.no_grad():
            out_ref = model_ref(x).sample
            out = model(x).sample

        self.assertTrue(
            torch.allclose(out_ref, out, atol=1e-5),
            f"Outputs do not match for standalone Conv layers with {offload_type}.",
        )

    def test_multiple_invocations_with_vae_like_model(self):
        """Test that multiple forward passes work correctly with VAE-like model."""
        if torch.device(torch_device).type not in ["cuda", "xpu"]:
            return

        config = self.get_autoencoder_kl_config()
        model = AutoencoderKL(**config)

        model_ref = AutoencoderKL(**config)
        model_ref.load_state_dict(model.state_dict(), strict=True)
        model_ref.to(torch_device)

        model.enable_group_offload(torch_device, offload_type="block_level", num_blocks_per_group=1, use_stream=True)

        x = torch.randn(2, 3, 32, 32).to(torch_device)

        with torch.no_grad():
            for i in range(2):
                out_ref = model_ref(x).sample
                out = model(x).sample
                self.assertTrue(torch.allclose(out_ref, out, atol=1e-5), f"Outputs do not match at iteration {i}.")

    def test_nested_container_parameters_offloading(self):