Diffuser test

[kiya00]

Before submitting

• Was this discussed/approved via a Github issue? (no need for typos and docs improvements)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure to update the docs?
• Did you write any new necessary tests?

What does this PR do?

Fixes #2075 .
Adds coverage tests for HF diffusers

Needs #2122: The bitsandbytes version needs to be updated to support HF Diffusers, as was done in PR #2122 .

cc @Borda

[kiya00]

bitsandbytes==0.44 fails with

RuntimeError: Failed to import diffusers.models.unets.unet_2d_condition because of the following error (look up to see its traceback):
No module named 'triton.ops'

[kiya00]

diffuser fails with bitsandbytes==0.45.0

diffuser works with bitsandbytes==0.45.1, but
FAILED thunder/tests/test_networks.py::test_quantization - ValueError: Blockwise quantization only supports 16/32-bit floats, but got torch.uint8

[t-vi]

@kiya00 the error is fixed in the bitsandbytes main branch, but we're waiting for the 0.47 release to bump. (@lianakoleva isolated this in #2238 )

[kiya00]

Hi @t-vi , thanks, I think it's fine to wait for 0.47 to merge this PR, I'll ask @IvanYashchuk on Monday

[IvanYashchuk]

bitsandbytes==0.44 fails with

RuntimeError: Failed to import diffusers.models.unets.unet_2d_condition because of the following error (look up to see its traceback):
No module named 'triton.ops'

What's the stack trace for this error?

I don't see triton.ops usage in https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/unets/unet_2d_condition.py

[kiya00]

root@4eb531844b44:/wayan/lightning-thunder# pytest thunder/tests/test_networks.py -k test_hf_diffusers[CompVis/stable-diffusion-v1-4]
================================================================== test session starts ===================================================================
platform linux -- Python 3.10.18, pytest-8.3.5, pluggy-1.6.0
benchmark: 5.1.0 (defaults: timer=torch.utils.benchmark.utils.timer.timer disable_gc=False min_rounds=5 min_time=0.000005 max_time=1.0 calibration_precision=10 warmup=True warmup_iterations=100000)
Test order randomisation NOT enabled. Enable with --random-order or --random-order-bucket=<bucket_type>
rootdir: /wayan/lightning-thunder
configfile: pyproject.toml
plugins: xdist-3.7.0, hypothesis-6.133.2, benchmark-5.1.0, timestamper-0.0.10, random-order-1.1.1, cov-6.1.1, timeout-2.4.0
timeout: 900.0s
timeout method: signal
timeout func_only: False
collected 52 items / 51 deselected / 1 selected

thunder/tests/test_networks.py F                                                                                                                   [100%]

======================================================================== FAILURES ========================================================================
____________________________________________________ test_hf_diffusers[CompVis/stable-diffusion-v1-4] ____________________________________________________

self = <module 'diffusers.models' from '/usr/local/lib/python3.10/dist-packages/diffusers/models/__init__.py'>, module_name = 'unets.unet_2d_condition'

    def _get_module(self, module_name: str):
        try:
>           return importlib.import_module("." + module_name, self.__name__)

/usr/local/lib/python3.10/dist-packages/diffusers/utils/import_utils.py:883:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
/usr/lib/python3.10/importlib/__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
/usr/local/lib/python3.10/dist-packages/diffusers/models/unets/__init__.py:5: in <module>
    from .unet_1d import UNet1DModel
/usr/local/lib/python3.10/dist-packages/diffusers/models/unets/unet_1d.py:24: in <module>
    from ..modeling_utils import ModelMixin
/usr/local/lib/python3.10/dist-packages/diffusers/models/modeling_utils.py:40: in <module>
    from ..quantizers import DiffusersAutoQuantizer, DiffusersQuantizer
/usr/local/lib/python3.10/dist-packages/diffusers/quantizers/__init__.py:19: in <module>
    from .auto import DiffusersAutoQuantizer
/usr/local/lib/python3.10/dist-packages/diffusers/quantizers/auto.py:22: in <module>
    from .bitsandbytes import BnB4BitDiffusersQuantizer, BnB8BitDiffusersQuantizer
/usr/local/lib/python3.10/dist-packages/diffusers/quantizers/bitsandbytes/__init__.py:2: in <module>
    from .utils import dequantize_and_replace, dequantize_bnb_weight, replace_with_bnb_linear
/usr/local/lib/python3.10/dist-packages/diffusers/quantizers/bitsandbytes/utils.py:32: in <module>
    import bitsandbytes as bnb
/usr/local/lib/python3.10/dist-packages/bitsandbytes/__init__.py:15: in <module>
    from .nn import modules
/usr/local/lib/python3.10/dist-packages/bitsandbytes/nn/__init__.py:21: in <module>
    from .triton_based_modules import (
/usr/local/lib/python3.10/dist-packages/bitsandbytes/nn/triton_based_modules.py:7: in <module>
    from bitsandbytes.triton.int8_matmul_mixed_dequantize import (
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

    import torch

    from bitsandbytes.triton.triton_utils import is_triton_available

    if not is_triton_available():

        def int8_matmul_mixed_dequantize(a, b, state_x, state_w, bias):
            return None
    else:
        import triton
        import triton.language as tl
>       from triton.ops.matmul_perf_model import early_config_prune, estimate_matmul_time
E       ModuleNotFoundError: No module named 'triton.ops'

/usr/local/lib/python3.10/dist-packages/bitsandbytes/triton/int8_matmul_mixed_dequantize.py:12: ModuleNotFoundError

The above exception was the direct cause of the following exception:

model_id = 'CompVis/stable-diffusion-v1-4'

    @thunder.tests.framework.requiresCUDA
    @pytest.mark.parametrize("model_id", hf_diffusers_unet2d_condition_model_ids)
    def test_hf_diffusers(model_id):
        from thunder.dynamo import thunderfx
>       from diffusers import UNet2DConditionModel

thunder/tests/test_networks.py:786:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
<frozen importlib._bootstrap>:1075: in _handle_fromlist
    ???
/usr/local/lib/python3.10/dist-packages/diffusers/utils/import_utils.py:874: in __getattr__
    value = getattr(module, name)
/usr/local/lib/python3.10/dist-packages/diffusers/utils/import_utils.py:873: in __getattr__
    module = self._get_module(self._class_to_module[name])
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = <module 'diffusers.models' from '/usr/local/lib/python3.10/dist-packages/diffusers/models/__init__.py'>, module_name = 'unets.unet_2d_condition'

    def _get_module(self, module_name: str):
        try:
            return importlib.import_module("." + module_name, self.__name__)
        except Exception as e:
>           raise RuntimeError(
                f"Failed to import {self.__name__}.{module_name} because of the following error (look up to see its"
                f" traceback):\n{e}"
            ) from e
E           RuntimeError: Failed to import diffusers.models.unets.unet_2d_condition because of the following error (look up to see its traceback):
E           No module named 'triton.ops'

/usr/local/lib/python3.10/dist-packages/diffusers/utils/import_utils.py:885: RuntimeError
================================================================ short test summary info =================================================================
FAILED thunder/tests/test_networks.py::test_hf_diffusers[CompVis/stable-diffusion-v1-4] - RuntimeError: Failed to import diffusers.models.unets.unet_2d_condition because of the following error (look up to see its traceback):
====================================================== 1 failed, 51 deselected, 1 warning in 0.36s =======================================================

[kiya00]

Hi @t-vi @IvanYashchuk ,it's ready to merge

[t-vi]

Thank you @kiya00

As much as we love broad coverage testing, we need to figure out how to reduce the test time.

148.69s call     thunder/tests/test_networks.py::test_hf_diffusers[ionet-official/bc8-alpha]
148.37s call     thunder/tests/test_networks.py::test_hf_diffusers[runwayml/stable-diffusion-v1-5]
148.29s call     thunder/tests/test_networks.py::test_hf_diffusers[stabilityai/stable-diffusion-xl-refiner-1.0]
147.92s call     thunder/tests/test_networks.py::test_hf_diffusers[runwayml/stable-diffusion-inpainting]
146.58s call     thunder/tests/test_networks.py::test_hf_diffusers[CompVis/stable-diffusion-v1-4]
145.77s call     thunder/tests/test_networks.py::test_hf_diffusers[stabilityai/stable-diffusion-xl-base-1.0]
144.68s call     thunder/tests/test_networks.py::test_hf_diffusers[stabilityai/sd-turbo]
144.28s call     thunder/tests/test_networks.py::test_hf_diffusers[diffusers/stable-diffusion-xl-1.0-inpainting-0.1]

[kiya00]

Hi @t-vi , Can I move it into test_coverage_trace.py? That way, it won't run on every commit, so the longer runtime would be acceptable. Would that work?

[Borda]

it seems there's no CUDA available to run it

That is expected, as you try to execute the CUDA test on a plain CPU machine/runner without any CUDA support, you would need to add it to Azure or LitCI jobs

[kiya00]

Hi @t-vi , @Borda , the diffuser test passed (the test_coverage_trace.py failed), could you review if it's fine to merge
https://dev.azure.com/Lightning-AI/lightning/_build/results?buildId=241223&view=logs&j=eb91186b-1c68-59cd-d743-62ff6863beba&t=9d433d3f-ed0b-5632-e22f-28f1e5b8f0f8

[Borda]

now failing with the missing argument in the used models:

FAILED thunder/tests/coverage_tests/test_coverage_trace.py::test_model_trace[facebook/bart-large] - NotImplementedError: casting a TensorProxy to a Python boolean is not suppo...
FAILED thunder/tests/coverage_tests/test_coverage_trace.py::test_model_trace[facebook/dinov2-small] - NotImplementedError: only modes 'nearest' and 'nearest-exact' are supported...
FAILED thunder/tests/coverage_tests/test_coverage_trace.py::test_model_trace[facebook/opt-1.3b] - NotImplementedError: casting a TensorProxy to a Python boolean is not suppo...
FAILED thunder/tests/coverage_tests/test_coverage_trace.py::test_model_trace[google/reformer-enwik8] - ValueError: If training, make sure that config.axial_pos_shape factors: (12...
FAILED thunder/tests/coverage_tests/test_coverage_trace.py::test_model_trace[google/vit-base-patch16-224] - TypeError: <function ViTModel.forward at 0x7ff650d7cd30>() got unexpected k...
FAILED thunder/tests/coverage_tests/test_coverage_trace.py::test_model_trace[t5-base] - ModuleNotFoundError: No module named 'fused_layer_norm_cuda'
FAILED thunder/tests/coverage_tests/test_coverage_trace.py::test_model_trace[xlnet-base-cased] - thunder.core.interpreter.InterpreterError: Encountered exception TypeError:...

[kiya00]

now failing with the missing argument in the used models:

FAILED thunder/tests/coverage_tests/test_coverage_trace.py::test_model_trace[facebook/bart-large] - NotImplementedError: casting a TensorProxy to a Python boolean is not suppo...
FAILED thunder/tests/coverage_tests/test_coverage_trace.py::test_model_trace[facebook/dinov2-small] - NotImplementedError: only modes 'nearest' and 'nearest-exact' are supported...
FAILED thunder/tests/coverage_tests/test_coverage_trace.py::test_model_trace[facebook/opt-1.3b] - NotImplementedError: casting a TensorProxy to a Python boolean is not suppo...
FAILED thunder/tests/coverage_tests/test_coverage_trace.py::test_model_trace[google/reformer-enwik8] - ValueError: If training, make sure that config.axial_pos_shape factors: (12...
FAILED thunder/tests/coverage_tests/test_coverage_trace.py::test_model_trace[google/vit-base-patch16-224] - TypeError: <function ViTModel.forward at 0x7ff650d7cd30>() got unexpected k...
FAILED thunder/tests/coverage_tests/test_coverage_trace.py::test_model_trace[t5-base] - ModuleNotFoundError: No module named 'fused_layer_norm_cuda'
FAILED thunder/tests/coverage_tests/test_coverage_trace.py::test_model_trace[xlnet-base-cased] - thunder.core.interpreter.InterpreterError: Encountered exception TypeError:...

Hi @Borda @t-vi , at first glance, each model appears to be failing for a different reason. Fixing them may take some time. Since this PR is only for the diffuser test, can we review and merge this one first, and then open a separate PR to fix test_coverage_trace.py (#2436) ?

[kiya00]

Thank you very much @Borda @t-vi for the help!