Migrate functionality into UI toggle

Author: zzlol63

modules/ui/TrainUI.py

@@ -29,6 +29,7 @@
 from modules.ui.TrainingTab import TrainingTab
 from modules.ui.VideoToolUI import VideoToolUI
 from modules.util import create
+from modules.util.attn.flash_attn_win import disable_flash_attn_win, enable_flash_attn_win
 from modules.util.callbacks.TrainCallbacks import TrainCallbacks
 from modules.util.commands.TrainCommands import TrainCommands
 from modules.util.config.TrainConfig import TrainConfig
@@ -133,6 +134,7 @@ def __init__(self):
         self.always_on_tensorboard_subprocess = None
         self.current_workspace_dir = self.train_config.workspace_dir
         self._check_start_always_on_tensorboard()
+        self._flash_attn_fallback_toggle()
 
         self.workspace_dir_trace_id = self.ui_state.add_var_trace("workspace_dir", self._on_workspace_dir_change_trace)
 
@@ -335,6 +337,10 @@ def create_general_tab(self, master):
                          tooltip="The device used to temporarily offload models while they are not used. Default:\"cpu\"")
         components.entry(frame, 16, 1, self.ui_state, "temp_device")
 
+        components.label(frame, 17, 0, "Use Flash-Attention Fallback",
+                         tooltip="Enables Flash-Attention fallback on Windows if native support is not available in PyTorch for a performance improvement during training/sampling.")
+        components.switch(frame, 17, 1, self.ui_state, "use_flash_attn_fallback", command=self._flash_attn_fallback_toggle)
+
         frame.pack(fill="both", expand=1)
         return frame
 
@@ -913,3 +919,9 @@ def _set_training_button_running(self):
 
     def _set_training_button_stopping(self):
         self._set_training_button_style("stopping")
+
+    def _flash_attn_fallback_toggle(self):
+        if self.train_config.use_flash_attn_fallback:
+            enable_flash_attn_win()
+        else:
+            disable_flash_attn_win()

modules/util/attn/flash_attn_win.py

@@ -146,7 +146,8 @@ def _flash_dynamic_scaled_dot_product_attention(query: torch.Tensor,
                                                         dropout_p: float = 0.0,
                                                         is_causal: bool = False,
                                                         scale: float | None = None,
-                                                        enable_gqa: bool = False):
+                                                        enable_gqa: bool = False,
+                                                        _fallback_sdpa = _scaled_dot_product_attention):
             if can_use_flash_attn(query, key, value, attn_mask, is_causal, enable_gqa):
                 # transpose(1,2) is equivalent to permute(0,2,1,3) for (B,H,L,D) -> (B,L,H,D)
                 q = query.transpose(1, 2)
@@ -161,7 +162,7 @@ def _flash_dynamic_scaled_dot_product_attention(query: torch.Tensor,
                 return out.transpose(1, 2)
 
             # Fallback
-            return _scaled_dot_product_attention(
+            return _fallback_sdpa(
                 query=query, key=key, value=value,
                 attn_mask=attn_mask, dropout_p=dropout_p,
                 is_causal=is_causal, scale=scale, enable_gqa=enable_gqa)

modules/util/config/TrainConfig.py

@@ -377,6 +377,7 @@ class TrainConfig(BaseConfig):
     loss_scaler: LossScaler
     learning_rate_scaler: LearningRateScaler
     clip_grad_norm: float
+    use_flash_attn_fallback: bool
 
     #layer filter
     layer_filter: str  # comma-separated
@@ -931,6 +932,7 @@ def default_values() -> 'TrainConfig':
         data.append(("loss_scaler", LossScaler.NONE, LossScaler, False))
         data.append(("learning_rate_scaler", LearningRateScaler.NONE, LearningRateScaler, False))
         data.append(("clip_grad_norm", 1.0, float, True))
+        data.append(("use_flash_attn_fallback", True, bool, False))
 
         # noise
         data.append(("offset_noise_weight", 0.0, float, False))

scripts/util/import_util.py

@@ -31,6 +31,3 @@ def script_imports(allow_zluda: bool = True):
             from modules.zluda import ZLUDA
 
             ZLUDA.initialize()
-
-    from modules.util.attn.flash_attn_win import enable_flash_attn_win
-    enable_flash_attn_win()