From a28026bca11ae0cb6b189ec2256e34c70fb88842 Mon Sep 17 00:00:00 2001 From: bmaltais Date: Sun, 29 Sep 2024 08:36:05 -0400 Subject: [PATCH] Adding support for blocks_to_swap option to gui --- kohya_gui/class_flux1.py | 39 ++++++++++++++++++++++++------------- kohya_gui/dreambooth_gui.py | 5 +++++ kohya_gui/finetune_gui.py | 7 ++++++- 3 files changed, 36 insertions(+), 15 deletions(-) diff --git a/kohya_gui/class_flux1.py b/kohya_gui/class_flux1.py index 2d74ae56..547e5193 100644 --- 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,7 +221,7 @@ 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, @@ -234,6 +229,22 @@ def noise_offset_type_change( 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 index 615db99b..f2291c06 100644 --- 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 index aa97a160..5dca363a 100644 --- 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,