Diffuser test (#2141)

Co-authored-by: Thomas Viehmann <tv.code@beamnet.de>
Co-authored-by: Jirka Borovec <6035284+Borda@users.noreply.github.com>

Author: 3 people

.azure/gpu-coverage.yml

@@ -0,0 +1,77 @@
+trigger:
+  tags:
+    include: ["*"]
+  paths:
+    include:
+      - ".azure/gpu-coverage.yml"
+      - "requirements/coverage.txt"
+      - "thunder/tests/coverage/**"
+  branches:
+    include:
+      - "main"
+      - "release/*"
+      - "refs/tags/*"
+
+pr:
+  branches:
+    include: ["*"]
+
+jobs:
+  - job: coverage
+    strategy:
+      matrix:
+        "w/ torch 2.7.1":
+          docker-image: "ubuntu24.04-cuda12.6.3-cudnn-fe1.10.0-py3.10-pt_2.7.1-dev"
+    # how much time to give 'run always even if cancelled tasks' before stopping them
+    cancelTimeoutInMinutes: "2"
+    pool: "lit-rtx-3090"
+    variables:
+      DEVICES: $( python -c 'name = "$(Agent.Name)" ; gpus = name.split("_")[-1] if "_" in name else "0"; print(gpus)' )
+      TORCH_HOME: "/var/tmp/torch"
+      PIP_CACHE_DIR: "/var/tmp/pip"
+      PYTHONHASHSEED: "0"
+      NCCL_DEBUG: "INFO"
+      ALLOW_COVERAGE_TRACE: "1"
+    container:
+      image: "pytorchlightning/lightning-thunder:$(docker-image)"
+      options: "--gpus=all --shm-size=16g -v /var/tmp:/var/tmp"
+    workspace:
+      clean: all
+    steps:
+      - bash: |
+          echo $(DEVICES)
+          lspci | egrep 'VGA|3D'
+          dpkg-query -W -f='${Package} ${Version}\n' libnccl2 libnccl-dev
+          whereis nvidia
+          nvidia-smi
+          which python && which pip
+          python --version
+          pip --version
+          pip list
+          echo "##vso[task.setvariable variable=CUDA_VISIBLE_DEVICES]$(DEVICES)"
+        displayName: "Image info & NVIDIA"
+
+      - bash: |
+          set -ex
+          # drop pt from requirements so not to interfere with the existing one
+          bash scripts/remove-torch-lines.sh requirements/base.txt
+          cat requirements/base.txt
+
+          # double check on test requirements
+          pip install -U -r requirements/base.txt -r requirements/coverage.txt
+
+          # https://docs.codecov.com/docs/codecov-uploader
+          curl -Os https://uploader.codecov.io/latest/linux/codecov
+          chmod +x codecov
+
+          # install this package
+          python setup.py develop
+        displayName: "Install package & ..."
+
+      - bash: bash scripts/sanity-check.sh
+        displayName: "Sanity check / details"
+
+      - bash: |
+          PYTHONPATH=$(pwd)/thunder/tests pytest thunder/tests/coverage_tests
+        timeoutInMinutes: "45"
+        displayName: "Testing: coverage_tests"

.github/workflows/coverage_test.yml

@@ -0,0 +1,9 @@
+coverage ~=7.9.1
+pytest ==8.3.5
+pytest-cov ==6.2.1
+pytest-benchmark ==5.1.0
+transformers ==4.52.4
+lightning_sdk
+diffusers==0.34.0
+accelerate
+bitsandbytes==0.46.1

requirements/coverage.txt

@@ -20,6 +20,8 @@ jsonargparse # thunder/benchmarks/benchmark_litgpt.py
 bitsandbytes==0.46.1; 'arm' not in platform_machine and 'aarch' not in platform_machine
 bitsandbytes>=0.42,<0.43; 'arm' in platform_machine or 'aarch' in platform_machine
 transformers==4.52.4 # for test_networks.py
+diffusers==0.34.0 # for test_networks.py
+accelerate # for test_networks.py
 
 # Installs JAX on Linux and MacOS
 jaxlib; sys_platform == 'linux' or sys_platform == 'darwin'  # required for jax, see https://github.com/google/jax#installation

requirements/test.txt

@@ -0,0 +1,83 @@
+import torch
+import os
+import pytest
+
+
+if os.getenv("ALLOW_COVERAGE_TRACE") != "1":
+    pytest.skip("Skipping test_coverage_hf_diffusers.py in regular CI", allow_module_level=True)
+
+hf_diffusers_unet2d_condition_model_ids = [
+    "runwayml/stable-diffusion-v1-5",
+    "CompVis/stable-diffusion-v1-4",
+    "ionet-official/bc8-alpha",
+    "stabilityai/sd-turbo",
+    "runwayml/stable-diffusion-inpainting",
+    "stabilityai/stable-diffusion-xl-base-1.0",
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
+]
+
+from thunder.tests.framework import requiresCUDA
+
+
+@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
+
+    unet_config = UNet2DConditionModel.load_config(model_id, subfolder="unet", torch_dtype=torch.bfloat16)
+    unet = UNet2DConditionModel(unet_config)
+    in_channels = unet.config.in_channels
+    cross_attention_dim = unet.config.cross_attention_dim
+    addition_embed_type = unet.config.addition_embed_type
+
+    sample_size = 4
+    batch_size = 1
+    seq_length = 4
+
+    if "xl" in model_id:
+        time_ids_dim = 6
+        text_embeds_dim = 4
+        if "refiner" in model_id:
+            time_ids_dim = 2
+            text_embeds_dim = 4
+    else:
+        time_ids_dim = None
+        text_embeds_dim = None
+
+    input_shape = (batch_size, in_channels, sample_size, sample_size)
+    hidden_states_shape = (batch_size, seq_length, cross_attention_dim)
+
+    unet = unet.to("cuda", dtype=torch.bfloat16).requires_grad_(True)
+    compiled_model = thunderfx(unet)
+
+    def make_inputs(dtype=torch.bfloat16):
+        added_cond_kwargs = {}
+        with torch.device("cuda"):
+            input = torch.randn(input_shape, dtype=dtype)
+            hidden_states = torch.randn(hidden_states_shape, dtype=dtype)
+            timestep = torch.ones(batch_size, dtype=torch.long)
+            if addition_embed_type is not None:
+                assert text_embeds_dim is not None and time_ids_dim is not None
+                time_ids_shape = (batch_size, time_ids_dim)
+                text_embeds_shape = (batch_size, text_embeds_dim)
+                added_cond_kwargs["time_ids"] = torch.randn(time_ids_shape, device="cuda", dtype=dtype)
+                added_cond_kwargs["text_embeds"] = torch.randn(text_embeds_shape, device="cuda", dtype=dtype)
+        return (input, timestep, hidden_states), {"added_cond_kwargs": added_cond_kwargs}
+
+    compiled_args, compiled_kwargs = make_inputs(torch.bfloat16)
+    compiled_output = compiled_model(*compiled_args, **compiled_kwargs)
+
+    ref_output = unet(*compiled_args, **compiled_kwargs)
+
+    ref_output = ref_output.sample
+    compiled_output = compiled_output.sample
+
+    torch.testing.assert_close(compiled_output, ref_output, rtol=1e-2, atol=2e-1)
+
+    # TODO: Currently fails, needs investigation https://github.com/Lightning-AI/lightning-thunder/issues/2153
+    # loss_grad = torch.randn_like(compiled_output)
+    # grads_ref = torch.autograd.grad(ref_output, unet.parameters(), grad_outputs=loss_grad)
+    # grads_compiled = torch.autograd.grad(compiled_output, unet.parameters(), grad_outputs=loss_grad)
+    # torch.testing.assert_close(grads_ref, grads_compiled, rtol=1e-1, atol=1e-1)

thunder/tests/coverage_tests/__init__.py

@@ -17,7 +17,7 @@
     AutoModelForSeq2SeqLM,
     AutoModelForImageClassification,
 )
-from test_core import run_prologue
+from thunder.tests.test_core import run_prologue
 
 
 MODEL_LIST = [
@@ -59,6 +59,7 @@ def get_dummy_input(model_name, config):
         return {"input_ids": torch.randint(0, 1000, (1, 16), device="cpu")}
 
 
+@pytest.mark.skip(reason="https://github.com/Lightning-AI/lightning-thunder/issues/2436")
 @pytest.mark.parametrize("model_name", MODEL_LIST)
 def test_model_trace(model_name):
     print(f"\n=== Testing {model_name} ===")