Add FSDP option for Flux2

Author: leisuzz

examples/dreambooth/train_dreambooth_lora_flux2.py

@@ -47,6 +47,7 @@
 
 import numpy as np
 import torch
+import torch.distributed as dist
 import transformers
 from accelerate import Accelerator
 from accelerate.logging import get_logger
@@ -722,6 +723,7 @@ def parse_args(input_args=None):
     )
     parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank")
     parser.add_argument("--enable_npu_flash_attention", action="store_true", help="Enabla Flash Attention for NPU")
+    parser.add_argument("--fsdp_text_encoder", action="store_true", help="Use FSDP for text encoder")
 
     if input_args is not None:
         args = parser.parse_args(input_args)
@@ -1427,12 +1429,22 @@ def load_model_hook(models, input_dir):
     )
 
     def compute_text_embeddings(prompt, text_encoding_pipeline):
-        with torch.no_grad():
-            prompt_embeds, text_ids = text_encoding_pipeline.encode_prompt(
-                prompt=prompt,
-                max_sequence_length=args.max_sequence_length,
-                text_encoder_out_layers=args.text_encoder_out_layers,
-            )
+        if args.fsdp_text_encoder:
+            text_encoding_pipeline.text_encoder.eval()
+            with torch.no_grad():
+                prompt_embeds, text_ids = text_encoding_pipeline.encode_prompt(
+                    prompt=prompt,
+                    max_sequence_length=args.max_sequence_length,
+                    text_encoder_out_layers=args.text_encoder_out_layers,
+                    device=accelerator.device,
+                )
+        else:
+            with torch.no_grad():
+                prompt_embeds, text_ids = text_encoding_pipeline.encode_prompt(
+                    prompt=prompt,
+                    max_sequence_length=args.max_sequence_length,
+                    text_encoder_out_layers=args.text_encoder_out_layers,
+                )
         return prompt_embeds, text_ids
 
     def compute_remote_text_embeddings(prompts):
@@ -1507,6 +1519,38 @@ def _encode_single(prompt: str):
                     args.validation_prompt, text_encoding_pipeline
                 )
 
+    # Init FSDP for text encoder
+    if args.fsdp_text_encoder:
+        from functools import partial
+
+        from torch.distributed.fsdp import (
+            BackwardPrefetch,
+            CPUOffload,
+            ShardingStrategy,
+        )
+        from torch.distributed.fsdp import (
+            FullyShardedDataParallel as FSDP,
+        )
+        from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
+
+        original_text_encoder = text_encoding_pipeline.text_encoder
+        transformer_layer = type(original_text_encoder.model.language_model.layers[0])
+        auto_wrap_policy = partial(transformer_auto_wrap_policy, transformer_layer_cls={transformer_layer})
+
+        text_encoder_fsdp = FSDP(
+            original_text_encoder,
+            device_id=accelerator.device,
+            sharding_strategy=ShardingStrategy.FULL_SHARD,
+            cpu_offload=CPUOffload(offload_params=args.offload),
+            auto_wrap_policy=auto_wrap_policy,
+            backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
+            limit_all_gathers=True,
+            use_orig_params=True,
+        )
+
+        text_encoding_pipeline.text_encoder = text_encoder_fsdp
+        dist.barrier()
+
     # If custom instance prompts are NOT provided (i.e. the instance prompt is used for all images),
     # pack the statically computed variables appropriately here. This is so that we don't
     # have to pass them to the dataloader.
@@ -1536,6 +1580,8 @@ def _encode_single(prompt: str):
                 if train_dataset.custom_instance_prompts:
                     if args.remote_text_encoder:
                         prompt_embeds, text_ids = compute_remote_text_embeddings(batch["prompts"])
+                    elif args.fsdp_text_encoder:
+                        prompt_embeds, text_ids = compute_text_embeddings(batch["prompts"], text_encoding_pipeline)
                     else:
                         with offload_models(text_encoding_pipeline, device=accelerator.device, offload=args.offload):
                             prompt_embeds, text_ids = compute_text_embeddings(batch["prompts"], text_encoding_pipeline)
@@ -1836,15 +1882,42 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
 
     # Save the lora layers
     accelerator.wait_for_everyone()
+    is_fsdp = accelerator.state.fsdp_plugin is not None
+
+    if is_fsdp:
+        transformer = unwrap_model(transformer)
+        state_dict = accelerator.get_state_dict(transformer)
     if accelerator.is_main_process:
         modules_to_save = {}
-        transformer = unwrap_model(transformer)
-        if args.bnb_quantization_config_path is None:
-            if args.upcast_before_saving:
-                transformer.to(torch.float32)
-            else:
-                transformer = transformer.to(weight_dtype)
-        transformer_lora_layers = get_peft_model_state_dict(transformer)
+        if is_fsdp:
+            if args.bnb_quantization_config_path is None:
+                if args.upcast_before_saving:
+                    state_dict = {
+                        k: v.to(torch.float32) if isinstance(v, torch.Tensor) else v for k, v in state_dict.items()
+                    }
+                else:
+                    state_dict = {
+                        k: v.to(weight_dtype) if isinstance(v, torch.Tensor) else v for k, v in state_dict.items()
+                    }
+
+            transformer_lora_layers = get_peft_model_state_dict(
+                transformer,
+                state_dict=state_dict,
+            )
+            transformer_lora_layers = {
+                k: v.detach().cpu().contiguous() if isinstance(v, torch.Tensor) else v
+                for k, v in transformer_lora_layers.items()
+            }
+
+        else:
+            transformer = unwrap_model(transformer)
+            if args.bnb_quantization_config_path is None:
+                if args.upcast_before_saving:
+                    transformer.to(torch.float32)
+                else:
+                    transformer = transformer.to(weight_dtype)
+            transformer_lora_layers = get_peft_model_state_dict(transformer)
+
         modules_to_save["transformer"] = transformer
 
         Flux2Pipeline.save_lora_weights(

examples/dreambooth/train_dreambooth_lora_flux2_img2img.py

@@ -46,6 +46,7 @@
 
 import numpy as np
 import torch
+import torch.distributed as dist
 import transformers
 from accelerate import Accelerator
 from accelerate.logging import get_logger
@@ -691,6 +692,7 @@ def parse_args(input_args=None):
 
     parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank")
     parser.add_argument("--enable_npu_flash_attention", action="store_true", help="Enabla Flash Attention for NPU")
+    parser.add_argument("--fsdp_text_encoder", action="store_true", help="Use FSDP for text encoder")
 
     if input_args is not None:
         args = parser.parse_args(input_args)
@@ -1361,12 +1363,21 @@ def load_model_hook(models, input_dir):
     )
 
     def compute_text_embeddings(prompt, text_encoding_pipeline):
-        with torch.no_grad():
-            prompt_embeds, text_ids = text_encoding_pipeline.encode_prompt(
-                prompt=prompt, max_sequence_length=args.max_sequence_length
-            )
-            # prompt_embeds = prompt_embeds.to(accelerator.device)
-            # text_ids = text_ids.to(accelerator.device)
+        if args.fsdp_text_encoder:
+            text_encoding_pipeline.text_encoder.eval()
+            with torch.no_grad():
+                prompt_embeds, text_ids = text_encoding_pipeline.encode_prompt(
+                    prompt=prompt,
+                    max_sequence_length=args.max_sequence_length,
+                    device=accelerator.device,
+                )
+        else:
+            with torch.no_grad():
+                prompt_embeds, text_ids = text_encoding_pipeline.encode_prompt(
+                    prompt=prompt, max_sequence_length=args.max_sequence_length
+                )
+                # prompt_embeds = prompt_embeds.to(accelerator.device)
+                # text_ids = text_ids.to(accelerator.device)
         return prompt_embeds, text_ids
 
     def compute_remote_text_embeddings(prompts: str | list[str]):
@@ -1430,6 +1441,41 @@ def _encode_single(prompt: str):
                     args.validation_prompt, text_encoding_pipeline
                 )
 
+    # Init FSDP for text encoder
+    if args.fsdp_text_encoder:
+        from functools import partial
+
+        from torch.distributed.fsdp import (
+            BackwardPrefetch,
+            CPUOffload,
+            ShardingStrategy,
+        )
+        from torch.distributed.fsdp import (
+            FullyShardedDataParallel as FSDP,
+        )
+        from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
+
+        original_text_encoder = text_encoding_pipeline.text_encoder
+        transformer_layer = type(original_text_encoder.model.language_model.layers[0])
+        auto_wrap_policy = partial(
+            transformer_auto_wrap_policy,
+            transformer_layer_cls={transformer_layer}
+        )
+
+        text_encoder_fsdp = FSDP(
+            original_text_encoder,
+            device_id=accelerator.device,
+            sharding_strategy=ShardingStrategy.FULL_SHARD,
+            cpu_offload=CPUOffload(offload_params=args.offload),
+            auto_wrap_policy=auto_wrap_policy,
+            backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
+            limit_all_gathers=True,
+            use_orig_params=True,
+        )
+
+        text_encoding_pipeline.text_encoder = text_encoder_fsdp
+        dist.barrier()
+
     # If custom instance prompts are NOT provided (i.e. the instance prompt is used for all images),
     # pack the statically computed variables appropriately here. This is so that we don't
     # have to pass them to the dataloader.
@@ -1461,6 +1507,8 @@ def _encode_single(prompt: str):
                 if train_dataset.custom_instance_prompts:
                     if args.remote_text_encoder:
                         prompt_embeds, text_ids = compute_remote_text_embeddings(batch["prompts"])
+                    elif args.fsdp_text_encoder:
+                        prompt_embeds, text_ids = compute_text_embeddings(batch["prompts"], text_encoding_pipeline)
                     else:
                         with offload_models(text_encoding_pipeline, device=accelerator.device, offload=args.offload):
                             prompt_embeds, text_ids = compute_text_embeddings(batch["prompts"], text_encoding_pipeline)
@@ -1759,15 +1807,42 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
 
     # Save the lora layers
     accelerator.wait_for_everyone()
+    is_fsdp = accelerator.state.fsdp_plugin is not None
+
+    if is_fsdp:
+        transformer = unwrap_model(transformer)
+        state_dict = accelerator.get_state_dict(transformer)
     if accelerator.is_main_process:
         modules_to_save = {}
-        transformer = unwrap_model(transformer)
-        if args.bnb_quantization_config_path is None:
-            if args.upcast_before_saving:
-                transformer.to(torch.float32)
-            else:
-                transformer = transformer.to(weight_dtype)
-        transformer_lora_layers = get_peft_model_state_dict(transformer)
+        if is_fsdp:
+            if args.bnb_quantization_config_path is None:
+                if args.upcast_before_saving:
+                    state_dict = {
+                        k: v.to(torch.float32) if isinstance(v, torch.Tensor) else v for k, v in state_dict.items()
+                    }
+                else:
+                    state_dict = {
+                        k: v.to(weight_dtype) if isinstance(v, torch.Tensor) else v for k, v in state_dict.items()
+                    }
+
+            transformer_lora_layers = get_peft_model_state_dict(
+                transformer,
+                state_dict=state_dict,
+            )
+            transformer_lora_layers = {
+                k: v.detach().cpu().contiguous() if isinstance(v, torch.Tensor) else v
+                for k, v in transformer_lora_layers.items()
+            }
+
+        else:
+            transformer = unwrap_model(transformer)
+            if args.bnb_quantization_config_path is None:
+                if args.upcast_before_saving:
+                    transformer.to(torch.float32)
+                else:
+                    transformer = transformer.to(weight_dtype)
+            transformer_lora_layers = get_peft_model_state_dict(transformer)
+
         modules_to_save["transformer"] = transformer
 
         Flux2Pipeline.save_lora_weights(