better validation ux for attention mechanism selection

simpletuner/helpers/configuration/cmd_args.py

@@ -18,7 +18,11 @@
 
 from simpletuner.helpers.configuration.cli_utils import mapping_to_cli_args, normalize_lr_scheduler_value
 from simpletuner.helpers.logging import get_logger
-from simpletuner.helpers.training.attention_backend import AttentionBackendMode
+from simpletuner.helpers.training.attention_backend import (
+    AttentionBackendMode,
+    is_sageattention_available,
+    xformers_compute_capability_error,
+)
 from simpletuner.helpers.training.multi_process import should_log
 from simpletuner.helpers.training.optimizer_param import is_optimizer_deprecated, is_optimizer_grad_fp32
 from simpletuner.helpers.training.quantisation import MANUAL_QUANTIZATION_PRESETS, PIPELINE_QUANTIZATION_PRESETS
@@ -1282,6 +1286,18 @@ def _normalize_input_args(raw_args):
     if hasattr(args, "sageattention_usage"):
         args.sageattention_usage = AttentionBackendMode.from_raw(args.sageattention_usage)
 
+    attention_mech = getattr(args, "attention_mechanism", "diffusers")
+    if attention_mech == "xformers":
+        xformers_error = xformers_compute_capability_error()
+        if xformers_error:
+            raise ValueError(xformers_error)
+
+    if attention_mech.startswith("sage") and not is_sageattention_available():
+        raise ValueError(
+            f"SageAttention is not installed but --attention_mechanism={attention_mech} was requested. "
+            "Install it with: pip install sageattention"
+        )
+
     deprecated_options = {
         # how to deprecate options:
         # "flux_beta_schedule_alpha": "flow_beta_schedule_alpha",

simpletuner/helpers/training/attention_backend.py

@@ -93,6 +93,32 @@ class AttentionPhase(Enum):
     EVAL = "eval"
 
 
+def is_sageattention_available() -> bool:
+    """Check if sageattention package is importable."""
+    try:
+        import sageattention  # noqa: F401
+
+        return True
+    except ImportError:
+        return False
+
+
+def xformers_compute_capability_error() -> Optional[str]:
+    """Return an error message if xformers is unsupported on the current GPU, else None.
+
+    xformers does not support compute capability 9.0+ (Hopper architecture).
+    """
+    if not torch.cuda.is_available():
+        return None
+    major, _ = torch.cuda.get_device_capability()
+    if major >= 9:
+        return (
+            f"xformers is not supported on GPUs with compute capability 9.0+ "
+            f"(detected {major}.x). Use a different attention mechanism such as 'diffusers'."
+        )
+    return None
+
+
 SLAKey = Tuple[int, str, Optional[int], torch.dtype]
 
 

simpletuner/simpletuner_sdk/server/services/validation_service.py

@@ -14,6 +14,7 @@
 from typing import Any, Dict, List, Optional, Tuple
 
 from simpletuner.helpers.configuration.cli_utils import normalize_lr_scheduler_value
+from simpletuner.helpers.training.attention_backend import is_sageattention_available, xformers_compute_capability_error
 
 from ..services.field_registry_wrapper import lazy_field_registry
 
@@ -406,6 +407,20 @@ def _to_int(value: Any) -> Tuple[Optional[int], Optional[str]]:
         if warmup_error:
             result.add_error("lr_warmup_steps", "Warmup steps must be a whole number.")
 
+        # Attention mechanism availability checks
+        attention_mech = str(self._get_config_value(config, "attention_mechanism") or "diffusers")
+        if attention_mech == "xformers":
+            xformers_error = xformers_compute_capability_error()
+            if xformers_error:
+                result.add_error("attention_mechanism", xformers_error)
+
+        if attention_mech.startswith("sage") and not is_sageattention_available():
+            result.add_error(
+                "attention_mechanism",
+                f"SageAttention is not installed but '{attention_mech}' was selected. "
+                "Install it with: pip install sageattention",
+            )
+
     @staticmethod
     def _get_config_value(config: Dict[str, Any], field_name: str) -> Any:
         """Fetch a field value considering CLI-prefixed variants."""