Adding support for blocks_to_swap option to gui

bmaltais · Sep 29, 2024 · a28026b · a28026b
1 parent 28dd25d
commit a28026b
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 15 deletions.
diff --git a/kohya_gui/class_flux1.py b/kohya_gui/class_flux1.py
@@ -202,22 +202,17 @@ def noise_offset_type_change(
                     )
 
                 with gr.Row(visible=True if finetuning else False):
-                    self.blockwise_fused_optimizers = gr.Checkbox(
-                        label="Blockwise Fused Optimizer",
-                        value=self.config.get(
-                            "flux1.blockwise_fused_optimizers", False
-                        ),
-                        info="Enable blockwise optimizers for fused backward pass and optimizer step. Any optimizer can be used.",
-                        interactive=True,
-                    )
-                    self.cpu_offload_checkpointing = gr.Checkbox(
-                        label="CPU Offload Checkpointing",
-                        value=self.config.get("flux1.cpu_offload_checkpointing", False),
-                        info="[Experimental] Enable offloading of tensors to CPU during checkpointing",
+                    self.blocks_to_swap = gr.Slider(
+                        label="Blocks to swap",
+                        value=self.config.get("flux1.blocks_to_swap", 0),
+                        info="The number of blocks to swap. The default is None (no swap). These options must be combined with --fused_backward_pass or --blockwise_fused_optimizers. The recommended maximum value is 36.",
+                        minimum=0,
+                        maximum=57,
+                        step=1,
                         interactive=True,
                     )
                     self.single_blocks_to_swap = gr.Slider(
-                        label="Single Blocks to swap",
+                        label="Single Blocks to swap (depercated)",
                         value=self.config.get("flux1.single_blocks_to_swap", 0),
                         info="[Experimental] Sets the number of 'single_blocks' (~320MB) to swap during the forward and backward passes.",
                         minimum=0,
@@ -226,14 +221,30 @@ def noise_offset_type_change(
                         interactive=True,
                     )
                     self.double_blocks_to_swap = gr.Slider(
-                        label="Double Blocks to swap",
+                        label="Double Blocks to swap (depercated)",
                         value=self.config.get("flux1.double_blocks_to_swap", 0),
                         info="[Experimental] Sets the number of 'double_blocks' (~640MB) to swap during the forward and backward passes.",
                         minimum=0,
                         maximum=38,
                         step=1,
                         interactive=True,
                     )
+
+                with gr.Row(visible=True if finetuning else False):
+                    self.blockwise_fused_optimizers = gr.Checkbox(
+                        label="Blockwise Fused Optimizer",
+                        value=self.config.get(
+                            "flux1.blockwise_fused_optimizers", False
+                        ),
+                        info="Enable blockwise optimizers for fused backward pass and optimizer step. Any optimizer can be used.",
+                        interactive=True,
+                    )
+                    self.cpu_offload_checkpointing = gr.Checkbox(
+                        label="CPU Offload Checkpointing",
+                        value=self.config.get("flux1.cpu_offload_checkpointing", False),
+                        info="[Experimental] Enable offloading of tensors to CPU during checkpointing",
+                        interactive=True,
+                    )
                     self.flux_fused_backward_pass = gr.Checkbox(
                         label="Fused Backward Pass",
                         value=self.config.get("flux1.fused_backward_pass", False),

diff --git a/kohya_gui/dreambooth_gui.py b/kohya_gui/dreambooth_gui.py
@@ -225,6 +225,7 @@ def save_configuration(
     blockwise_fused_optimizers,
     flux_fused_backward_pass,
     cpu_offload_checkpointing,
+    blocks_to_swap,
     single_blocks_to_swap,
     double_blocks_to_swap,
     mem_eff_save,
@@ -429,6 +430,7 @@ def open_configuration(
     blockwise_fused_optimizers,
     flux_fused_backward_pass,
     cpu_offload_checkpointing,
+    blocks_to_swap,
     single_blocks_to_swap,
     double_blocks_to_swap,
     mem_eff_save,
@@ -628,6 +630,7 @@ def train_model(
     blockwise_fused_optimizers,
     flux_fused_backward_pass,
     cpu_offload_checkpointing,
+    blocks_to_swap,
     single_blocks_to_swap,
     double_blocks_to_swap,
     mem_eff_save,
@@ -1043,6 +1046,7 @@ def train_model(
         "cpu_offload_checkpointing": (
             cpu_offload_checkpointing if flux1_checkbox else None
         ),
+        "blocks_to_swap": blocks_to_swap if flux1_checkbox else None,
         "single_blocks_to_swap": single_blocks_to_swap if flux1_checkbox else None,
         "double_blocks_to_swap": double_blocks_to_swap if flux1_checkbox else None,
         "mem_eff_save": mem_eff_save if flux1_checkbox else None,
@@ -1389,6 +1393,7 @@ def dreambooth_tab(
             flux1_training.blockwise_fused_optimizers,
             flux1_training.flux_fused_backward_pass,
             flux1_training.cpu_offload_checkpointing,
+            flux1_training.blocks_to_swap,
             flux1_training.single_blocks_to_swap,
             flux1_training.double_blocks_to_swap,
             flux1_training.mem_eff_save,

diff --git a/kohya_gui/finetune_gui.py b/kohya_gui/finetune_gui.py
@@ -233,6 +233,7 @@ def save_configuration(
     blockwise_fused_optimizers,
     flux_fused_backward_pass,
     cpu_offload_checkpointing,
+    blocks_to_swap,
     single_blocks_to_swap,
     double_blocks_to_swap,
     mem_eff_save,
@@ -443,11 +444,12 @@ def open_configuration(
     blockwise_fused_optimizers,
     flux_fused_backward_pass,
     cpu_offload_checkpointing,
+    blocks_to_swap,
     single_blocks_to_swap,
     double_blocks_to_swap,
     mem_eff_save,
-    training_preset,
     apply_t5_attn_mask,
+    training_preset,
 ):
     # Get list of function parameters and values
     parameters = list(locals().items())
@@ -659,6 +661,7 @@ def train_model(
     blockwise_fused_optimizers,
     flux_fused_backward_pass,
     cpu_offload_checkpointing,
+    blocks_to_swap,
     single_blocks_to_swap,
     double_blocks_to_swap,
     mem_eff_save,
@@ -1101,6 +1104,7 @@ def train_model(
         "cpu_offload_checkpointing": (
             cpu_offload_checkpointing if flux1_checkbox else None
         ),
+        "blocks_to_swap": blocks_to_swap if flux1_checkbox else None,
         "single_blocks_to_swap": single_blocks_to_swap if flux1_checkbox else None,
         "double_blocks_to_swap": double_blocks_to_swap if flux1_checkbox else None,
         "mem_eff_save": mem_eff_save if flux1_checkbox else None,
@@ -1521,6 +1525,7 @@ def list_presets(path):
             flux1_training.blockwise_fused_optimizers,
             flux1_training.flux_fused_backward_pass,
             flux1_training.cpu_offload_checkpointing,
+            flux1_training.blocks_to_swap,
             flux1_training.single_blocks_to_swap,
             flux1_training.double_blocks_to_swap,
             flux1_training.mem_eff_save,