update

src/diffusers/models/_modeling_parallel.py

@@ -60,6 +60,16 @@ class ContextParallelConfig:
         rotate_method (`str`, *optional*, defaults to `"allgather"`):
             Method to use for rotating key/value states across devices in ring attention. Currently, only `"allgather"`
             is supported.
+        ulysses_anything (`bool`, *optional*, defaults to `False`):
+            Whether to enable "Ulysses Anything" mode, which supports arbitrary sequence lengths and head counts that
+            are not evenly divisible by `ulysses_degree`. When enabled, `ulysses_degree` must be greater than 1 and
+            `ring_degree` must be 1.
+        mesh (`torch.distributed.device_mesh.DeviceMesh`, *optional*):
+            A custom device mesh to use for context parallelism. If provided, this mesh will be used instead of
+            creating a new one. This is useful when combining context parallelism with other parallelism strategies
+            (e.g., FSDP, tensor parallelism) that share the same device mesh. The mesh must have both "ring" and
+            "ulysses" dimensions. Use size 1 for dimensions not being used (e.g., `mesh_shape=(2, 1, 4)` with
+            `mesh_dim_names=("ring", "ulysses", "fsdp")` for ring attention only with FSDP).
 
     """
 
@@ -68,6 +78,7 @@ class ContextParallelConfig:
     convert_to_fp32: bool = True
     # TODO: support alltoall
     rotate_method: Literal["allgather", "alltoall"] = "allgather"
+    mesh: torch.distributed.device_mesh.DeviceMesh | None = None
     # Whether to enable ulysses anything attention to support
     # any sequence lengths and any head numbers.
     ulysses_anything: bool = False
@@ -124,7 +135,7 @@ class ContextParallelConfig:
                 f"The product of `ring_degree` ({self.ring_degree}) and `ulysses_degree` ({self.ulysses_degree}) must not exceed the world size ({world_size})."
             )
 
-        self._flattened_mesh = self._mesh._flatten()
+        self._flattened_mesh = self._mesh["ring", "ulysses"]._flatten()
         self._ring_mesh = self._mesh["ring"]
         self._ulysses_mesh = self._mesh["ulysses"]
         self._ring_local_rank = self._ring_mesh.get_local_rank()

src/diffusers/models/attention_dispatch.py

@@ -62,8 +62,6 @@ _REQUIRED_FLEX_VERSION = "2.5.0"
 _REQUIRED_XLA_VERSION = "2.2"
 _REQUIRED_XFORMERS_VERSION = "0.0.29"
 
-logger = get_logger(__name__)  # pylint: disable=invalid-name
-
 _CAN_USE_FLASH_ATTN = is_flash_attn_available() and is_flash_attn_version(">=", _REQUIRED_FLASH_VERSION)
 _CAN_USE_FLASH_ATTN_3 = is_flash_attn_3_available()
 _CAN_USE_AITER_ATTN = is_aiter_available() and is_aiter_version(">=", _REQUIRED_AITER_VERSION)
@@ -75,18 +73,8 @@ _CAN_USE_XFORMERS_ATTN = is_xformers_available() and is_xformers_version(">=", _
 
 
 if _CAN_USE_FLASH_ATTN:
-    try:
-        from flash_attn import flash_attn_func, flash_attn_varlen_func
-        from flash_attn.flash_attn_interface import _wrapped_flash_attn_backward, _wrapped_flash_attn_forward
-    except (ImportError, OSError, RuntimeError) as e:
-        # Handle ABI mismatch or other import failures gracefully.
-        # This can happen when flash_attn was compiled against a different PyTorch version.
-        logger.warning(f"flash_attn is installed but failed to import: {e}. Falling back to native PyTorch attention.")
-        _CAN_USE_FLASH_ATTN = False
-        flash_attn_func = None
-        flash_attn_varlen_func = None
-        _wrapped_flash_attn_backward = None
-        _wrapped_flash_attn_forward = None
+    from flash_attn import flash_attn_func, flash_attn_varlen_func
+    from flash_attn.flash_attn_interface import _wrapped_flash_attn_backward, _wrapped_flash_attn_forward
 else:
     flash_attn_func = None
     flash_attn_varlen_func = None
@@ -95,47 +83,26 @@ else:
 
 
 if _CAN_USE_FLASH_ATTN_3:
-    try:
-        from flash_attn_interface import flash_attn_func as flash_attn_3_func
-        from flash_attn_interface import flash_attn_varlen_func as flash_attn_3_varlen_func
-    except (ImportError, OSError, RuntimeError) as e:
-        logger.warning(f"flash_attn_3 failed to import: {e}. Falling back to native attention.")
-        _CAN_USE_FLASH_ATTN_3 = False
-        flash_attn_3_func = None
-        flash_attn_3_varlen_func = None
+    from flash_attn_interface import flash_attn_func as flash_attn_3_func
+    from flash_attn_interface import flash_attn_varlen_func as flash_attn_3_varlen_func
 else:
     flash_attn_3_func = None
     flash_attn_3_varlen_func = None
 
 if _CAN_USE_AITER_ATTN:
-    try:
-        from aiter import flash_attn_func as aiter_flash_attn_func
-    except (ImportError, OSError, RuntimeError) as e:
-        logger.warning(f"aiter failed to import: {e}. Falling back to native attention.")
-        _CAN_USE_AITER_ATTN = False
-        aiter_flash_attn_func = None
+    from aiter import flash_attn_func as aiter_flash_attn_func
 else:
     aiter_flash_attn_func = None
 
 if _CAN_USE_SAGE_ATTN:
-    try:
-        from sageattention import (
-            sageattn,
-            sageattn_qk_int8_pv_fp8_cuda,
-            sageattn_qk_int8_pv_fp8_cuda_sm90,
-            sageattn_qk_int8_pv_fp16_cuda,
-            sageattn_qk_int8_pv_fp16_triton,
-            sageattn_varlen,
-        )
-    except (ImportError, OSError, RuntimeError) as e:
-        logger.warning(f"sageattention failed to import: {e}. Falling back to native attention.")
-        _CAN_USE_SAGE_ATTN = False
-        sageattn = None
-        sageattn_qk_int8_pv_fp8_cuda = None
-        sageattn_qk_int8_pv_fp8_cuda_sm90 = None
-        sageattn_qk_int8_pv_fp16_cuda = None
-        sageattn_qk_int8_pv_fp16_triton = None
-        sageattn_varlen = None
+    from sageattention import (
+        sageattn,
+        sageattn_qk_int8_pv_fp8_cuda,
+        sageattn_qk_int8_pv_fp8_cuda_sm90,
+        sageattn_qk_int8_pv_fp16_cuda,
+        sageattn_qk_int8_pv_fp16_triton,
+        sageattn_varlen,
+    )
 else:
     sageattn = None
     sageattn_qk_int8_pv_fp16_cuda = None
@@ -146,48 +113,26 @@ else:
 
 
 if _CAN_USE_FLEX_ATTN:
-    try:
-        # We cannot import the flex_attention function from the package directly because it is expected (from the
-        # pytorch documentation) that the user may compile it. If we import directly, we will not have access to the
-        # compiled function.
-        import torch.nn.attention.flex_attention as flex_attention
-    except (ImportError, OSError, RuntimeError) as e:
-        logger.warning(f"flex_attention failed to import: {e}. Falling back to native attention.")
-        _CAN_USE_FLEX_ATTN = False
-        flex_attention = None
-else:
-    flex_attention = None
+    # We cannot import the flex_attention function from the package directly because it is expected (from the
+    # pytorch documentation) that the user may compile it. If we import directly, we will not have access to the
+    # compiled function.
+    import torch.nn.attention.flex_attention as flex_attention
 
 
 if _CAN_USE_NPU_ATTN:
-    try:
-        from torch_npu import npu_fusion_attention
-    except (ImportError, OSError, RuntimeError) as e:
-        logger.warning(f"torch_npu failed to import: {e}. Falling back to native attention.")
-        _CAN_USE_NPU_ATTN = False
-        npu_fusion_attention = None
+    from torch_npu import npu_fusion_attention
 else:
     npu_fusion_attention = None
 
 
 if _CAN_USE_XLA_ATTN:
-    try:
-        from torch_xla.experimental.custom_kernel import flash_attention as xla_flash_attention
-    except (ImportError, OSError, RuntimeError) as e:
-        logger.warning(f"torch_xla failed to import: {e}. Falling back to native attention.")
-        _CAN_USE_XLA_ATTN = False
-        xla_flash_attention = None
+    from torch_xla.experimental.custom_kernel import flash_attention as xla_flash_attention
 else:
     xla_flash_attention = None
 
 
 if _CAN_USE_XFORMERS_ATTN:
-    try:
-        import xformers.ops as xops
-    except (ImportError, OSError, RuntimeError) as e:
-        logger.warning(f"xformers failed to import: {e}. Falling back to native attention.")
-        _CAN_USE_XFORMERS_ATTN = False
-        xops = None
+    import xformers.ops as xops
 else:
     xops = None
 
@@ -213,6 +158,8 @@ else:
     _register_fake = register_fake_no_op
 
 
+logger = get_logger(__name__)  # pylint: disable=invalid-name
+
 # TODO(aryan): Add support for the following:
 # - Sage Attention++
 # - block sparse, radial and other attention methods

src/diffusers/models/modeling_utils.py

@@ -1567,7 +1567,7 @@ class ModelMixin(torch.nn.Module, PushToHubMixin):
         mesh = None
         if config.context_parallel_config is not None:
             cp_config = config.context_parallel_config
-            mesh = torch.distributed.device_mesh.init_device_mesh(
+            mesh = cp_config.mesh or torch.distributed.device_mesh.init_device_mesh(
                 device_type=device_type,
                 mesh_shape=cp_config.mesh_shape,
                 mesh_dim_names=cp_config.mesh_dim_names,

src/diffusers/quantizers/quantization_config.py

@@ -36,7 +36,7 @@ from typing import Any, Callable
 
 from packaging import version
 
-from ..utils import deprecate, is_torch_available, is_torchao_available, is_torchao_version, logging
+from ..utils import is_torch_available, is_torchao_available, is_torchao_version, logging
 
 
 if is_torch_available():
@@ -844,8 +844,6 @@ class QuantoConfig(QuantizationConfigMixin):
         modules_to_not_convert: list[str] | None = None,
         **kwargs,
     ):
-        deprecation_message = "`QuantoConfig` is deprecated and will be removed in version 1.0.0."
-        deprecate("QuantoConfig", "1.0.0", deprecation_message)
         self.quant_method = QuantizationMethod.QUANTO
         self.weights_dtype = weights_dtype
         self.modules_to_not_convert = modules_to_not_convert

src/diffusers/quantizers/quanto/quanto_quantizer.py

@@ -3,7 +3,6 @@ from typing import TYPE_CHECKING, Any
 from diffusers.utils.import_utils import is_optimum_quanto_version
 
 from ...utils import (
-    deprecate,
     get_module_from_name,
     is_accelerate_available,
     is_accelerate_version,
@@ -43,9 +42,6 @@ class QuantoQuantizer(DiffusersQuantizer):
         super().__init__(quantization_config, **kwargs)
 
     def validate_environment(self, *args, **kwargs):
-        deprecation_message = "The Quanto quantizer is deprecated and will be removed in version 1.0.0."
-        deprecate("QuantoQuantizer", "1.0.0", deprecation_message)
-
         if not is_optimum_quanto_available():
             raise ImportError(
                 "Loading an optimum-quanto quantized model requires optimum-quanto library (`pip install optimum-quanto`)"

tests/models/testing_utils/parallelism.py

@@ -60,12 +60,7 @@ def _context_parallel_worker(rank, world_size, master_port, model_class, init_di
         model.eval()
 
         # Move inputs to device
-        inputs_on_device = {}
-        for key, value in inputs_dict.items():
-            if isinstance(value, torch.Tensor):
-                inputs_on_device[key] = value.to(device)
-            else:
-                inputs_on_device[key] = value
+        inputs_on_device = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()}
 
         # Enable context parallelism
         cp_config = ContextParallelConfig(**cp_dict)
@@ -89,6 +84,59 @@ def _context_parallel_worker(rank, world_size, master_port, model_class, init_di
             dist.destroy_process_group()
 
 
+def _custom_mesh_worker(
+    rank,
+    world_size,
+    master_port,
+    model_class,
+    init_dict,
+    cp_dict,
+    mesh_shape,
+    mesh_dim_names,
+    inputs_dict,
+    return_dict,
+):
+    """Worker function for context parallel testing with a user-provided custom DeviceMesh."""
+    try:
+        os.environ["MASTER_ADDR"] = "localhost"
+        os.environ["MASTER_PORT"] = str(master_port)
+        os.environ["RANK"] = str(rank)
+        os.environ["WORLD_SIZE"] = str(world_size)
+
+        dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)
+
+        torch.cuda.set_device(rank)
+        device = torch.device(f"cuda:{rank}")
+
+        model = model_class(**init_dict)
+        model.to(device)
+        model.eval()
+
+        inputs_on_device = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()}
+
+        # DeviceMesh must be created after init_process_group, inside each worker process.
+        mesh = torch.distributed.device_mesh.init_device_mesh(
+            "cuda", mesh_shape=mesh_shape, mesh_dim_names=mesh_dim_names
+        )
+        cp_config = ContextParallelConfig(**cp_dict, mesh=mesh)
+        model.enable_parallelism(config=cp_config)
+
+        with torch.no_grad():
+            output = model(**inputs_on_device, return_dict=False)[0]
+
+        if rank == 0:
+            return_dict["status"] = "success"
+            return_dict["output_shape"] = list(output.shape)
+
+    except Exception as e:
+        if rank == 0:
+            return_dict["status"] = "error"
+            return_dict["error"] = str(e)
+    finally:
+        if dist.is_initialized():
+            dist.destroy_process_group()
+
+
 @is_context_parallel
 @require_torch_multi_accelerator
 class ContextParallelTesterMixin:
@@ -126,3 +174,48 @@ class ContextParallelTesterMixin:
         assert return_dict.get("status") == "success", (
             f"Context parallel inference failed: {return_dict.get('error', 'Unknown error')}"
         )
+
+    @pytest.mark.parametrize(
+        "cp_type,mesh_shape,mesh_dim_names",
+        [
+            ("ring_degree", (2, 1, 1), ("ring", "ulysses", "fsdp")),
+            ("ulysses_degree", (1, 2, 1), ("ring", "ulysses", "fsdp")),
+        ],
+        ids=["ring-3d-fsdp", "ulysses-3d-fsdp"],
+    )
+    def test_context_parallel_custom_mesh(self, cp_type, mesh_shape, mesh_dim_names):
+        if not torch.distributed.is_available():
+            pytest.skip("torch.distributed is not available.")
+
+        if not hasattr(self.model_class, "_cp_plan") or self.model_class._cp_plan is None:
+            pytest.skip("Model does not have a _cp_plan defined for context parallel inference.")
+
+        world_size = 2
+        init_dict = self.get_init_dict()
+        inputs_dict = {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in self.get_dummy_inputs().items()}
+        cp_dict = {cp_type: world_size}
+
+        master_port = _find_free_port()
+        manager = mp.Manager()
+        return_dict = manager.dict()
+
+        mp.spawn(
+            _custom_mesh_worker,
+            args=(
+                world_size,
+                master_port,
+                self.model_class,
+                init_dict,
+                cp_dict,
+                mesh_shape,
+                mesh_dim_names,
+                inputs_dict,
+                return_dict,
+            ),
+            nprocs=world_size,
+            join=True,
+        )
+
+        assert return_dict.get("status") == "success", (
+            f"Custom mesh context parallel inference failed: {return_dict.get('error', 'Unknown error')}"
+        )