mirror of
https://github.com/huggingface/diffusers.git
synced 2026-10-23 09:09:01 +03:00
[ci] fix compilation CI. (#15092)
* fix compilation CI. * fix attentionbackend name dispatch. * skip with a marker.
This commit is contained in:
@@ -418,7 +418,8 @@ def dispatch_attention_fn(
|
||||
# variable), or we use a custom backend based on whether user is using the `attention_backend` context manager
|
||||
backend_name, backend_fn = _AttentionBackendRegistry.get_active_backend()
|
||||
else:
|
||||
backend_name = AttentionBackendName(backend)
|
||||
# Skip the `Enum` call for members: older torch versions cannot always trace it in a full graph.
|
||||
backend_name = backend if isinstance(backend, AttentionBackendName) else AttentionBackendName(backend)
|
||||
backend_fn = _AttentionBackendRegistry._backends.get(backend_name)
|
||||
|
||||
kwargs = {
|
||||
|
||||
@@ -17,14 +17,11 @@ import pytest
|
||||
import torch
|
||||
|
||||
from diffusers import MMAudioVocoder
|
||||
from diffusers.utils import is_torch_version
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
|
||||
from ...testing_utils import enable_full_determinism, torch_device
|
||||
from ..testing_utils import (
|
||||
BaseModelTesterConfig,
|
||||
MemoryTesterMixin,
|
||||
ModelTesterMixin,
|
||||
)
|
||||
from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TorchCompileTesterMixin
|
||||
|
||||
|
||||
enable_full_determinism()
|
||||
@@ -104,3 +101,28 @@ class TestMMAudioVocoderMemory(MMAudioVocoderTesterConfig, MemoryTesterMixin):
|
||||
)
|
||||
def test_group_offloading_with_disk(self, tmp_path, record_stream, offload_type):
|
||||
super().test_group_offloading_with_disk(tmp_path, record_stream, offload_type)
|
||||
|
||||
|
||||
class TestMMAudioVocoderTorchCompile(MMAudioVocoderTesterConfig, TorchCompileTesterMixin):
|
||||
@pytest.mark.xfail(
|
||||
condition=is_torch_version("<", "2.12"),
|
||||
reason=(
|
||||
'With deterministic algorithms on, torch < 2.12 routes `F.pad(mode="replicate")` through `importlib`, '
|
||||
"which Dynamo cannot trace in a full graph."
|
||||
),
|
||||
raises=torch._dynamo.exc.Unsupported,
|
||||
strict=True,
|
||||
)
|
||||
def test_torch_compile_recompilation_and_graph_break(self):
|
||||
super().test_torch_compile_recompilation_and_graph_break()
|
||||
|
||||
@pytest.mark.xfail(
|
||||
reason=(
|
||||
"Block-level offloading hooks each `ups[i]` `ModuleList`, but `forward` calls `ups[i][0]` directly, so the "
|
||||
"hook never onloads the weights."
|
||||
),
|
||||
raises=RuntimeError,
|
||||
strict=True,
|
||||
)
|
||||
def test_compile_with_group_offloading(self):
|
||||
super().test_compile_with_group_offloading()
|
||||
|
||||
@@ -17,14 +17,11 @@ import pytest
|
||||
import torch
|
||||
|
||||
from diffusers import Kandinsky6SRLatentUpscalerBank
|
||||
from diffusers.utils import is_torch_version
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
|
||||
from ...testing_utils import enable_full_determinism, torch_device
|
||||
from ..testing_utils import (
|
||||
BaseModelTesterConfig,
|
||||
MemoryTesterMixin,
|
||||
ModelTesterMixin,
|
||||
)
|
||||
from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TorchCompileTesterMixin
|
||||
|
||||
|
||||
enable_full_determinism()
|
||||
@@ -116,3 +113,19 @@ class TestKandinsky6SRLatentUpscalerBankMemory(Kandinsky6SRLatentUpscalerBankTes
|
||||
@DISK_OFFLOAD_NUMERICS
|
||||
def test_disk_offload_with_safetensors(self, base_model_output, tmp_path):
|
||||
super().test_disk_offload_with_safetensors(base_model_output, tmp_path)
|
||||
|
||||
|
||||
class TestKandinsky6SRLatentUpscalerBankTorchCompile(
|
||||
Kandinsky6SRLatentUpscalerBankTesterConfig, TorchCompileTesterMixin
|
||||
):
|
||||
@pytest.mark.xfail(
|
||||
condition=is_torch_version("<", "2.12"),
|
||||
reason=(
|
||||
'With deterministic algorithms on, torch < 2.12 routes `F.pad(mode="replicate")` through `importlib`, '
|
||||
"which Dynamo cannot trace in a full graph."
|
||||
),
|
||||
raises=torch._dynamo.exc.Unsupported,
|
||||
strict=True,
|
||||
)
|
||||
def test_torch_compile_recompilation_and_graph_break(self):
|
||||
super().test_torch_compile_recompilation_and_graph_break()
|
||||
|
||||
@@ -2074,6 +2074,12 @@ class BFLTesterMixin(BFLConfigMixin, QuantizationTesterMixin):
|
||||
@is_torch_compile
|
||||
@require_accelerator
|
||||
@require_accelerate
|
||||
# Both schemes store E4M3 tensors (fp8r weights, nvfp4 block scales) and compute in bf16; Triton cannot compile either
|
||||
# below compute capability 8.9 (e.g. T4).
|
||||
@pytest.mark.skipif(
|
||||
torch_device == "cuda" and torch.cuda.get_device_capability() < (8, 9),
|
||||
reason="torch.compile of BFL models needs FP8 support in Triton (compute capability 8.9+)",
|
||||
)
|
||||
class BFLCompileTesterMixin(BFLConfigMixin, QuantizationCompileTesterMixin):
|
||||
"""
|
||||
Mixin class for testing torch.compile with BFL (`fp8r` / `nvfp4`) quantized models.
|
||||
|
||||
@@ -60,10 +60,11 @@ class Kandinsky6SRTransformerTesterConfig(BaseModelTesterConfig):
|
||||
"out_visual_dim": 8,
|
||||
"time_dim": 16,
|
||||
"patch_size": (1, 2, 2),
|
||||
"model_dim": 24,
|
||||
"model_dim": 32,
|
||||
"ff_dim": 32,
|
||||
"num_visual_blocks": 2,
|
||||
"axes_dims": (4, 4, 4),
|
||||
# head_dim = sum(axes_dims); Inductor lowers flex attention only for head dims of at least 16
|
||||
"axes_dims": (4, 6, 6),
|
||||
}
|
||||
|
||||
def _build_dummy_inputs(self, batch_size: int, num_frames: int, height: int, width: int) -> dict:
|
||||
@@ -106,6 +107,17 @@ class TestKandinsky6SRTransformerTorchCompile(Kandinsky6SRTransformerTesterConfi
|
||||
def get_dummy_inputs(self, height: int = 16, width: int = 16) -> dict:
|
||||
return self._build_dummy_inputs(batch_size=1, num_frames=2, height=height, width=width)
|
||||
|
||||
@pytest.mark.xfail(
|
||||
reason=(
|
||||
"The token grid is split into 8x8 blocks and `sliding_tile_mask` sizes its table by `max(...)` of the block "
|
||||
"counts, so the compiled graph specializes on the input shape and recompiles."
|
||||
),
|
||||
raises=torch._dynamo.exc.RecompileError,
|
||||
strict=True,
|
||||
)
|
||||
def test_compile_on_different_shapes(self):
|
||||
super().test_compile_on_different_shapes()
|
||||
|
||||
|
||||
@pytest.mark.skipif(torch_device == "cpu", reason="FlexAttention does not support backward on CPU.")
|
||||
class TestKandinsky6SRTransformerTraining(Kandinsky6SRTransformerTesterConfig, TrainingTesterMixin):
|
||||
|
||||
Reference in New Issue
Block a user