From ee85af17159617637293d011f6225c753fd98ce7 Mon Sep 17 00:00:00 2001 From: Volpeon Date: Tue, 16 May 2023 10:08:25 +0200 Subject: Patch xformers to cast dtypes --- train_lora.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) (limited to 'train_lora.py') diff --git a/train_lora.py b/train_lora.py index 70f0dc8..a58bef7 100644 --- a/train_lora.py +++ b/train_lora.py @@ -17,6 +17,7 @@ from accelerate.logging import get_logger from accelerate.utils import LoggerType, set_seed from peft import LoraConfig, get_peft_model # from diffusers.models.attention_processor import AttnProcessor +from diffusers.utils.import_utils import is_xformers_available import transformers import numpy as np @@ -33,7 +34,7 @@ from util.files import load_config, load_embeddings_from_dir # https://github.com/huggingface/peft/blob/main/examples/lora_dreambooth/train_dreambooth.py UNET_TARGET_MODULES_ORIG = ["to_q", "to_v", "query", "value"] -UNET_TARGET_MODULES = UNET_TARGET_MODULES_ORIG + ["to_out.0"] # ["to_k", key] +UNET_TARGET_MODULES = UNET_TARGET_MODULES_ORIG + ["to_out.0"] # ["to_k", "key"] TEXT_ENCODER_TARGET_MODULES_ORIG = ["q_proj", "v_proj"] TEXT_ENCODER_TARGET_MODULES = TEXT_ENCODER_TARGET_MODULES_ORIG + ["out_proj"] # ["k_proj"] TEXT_ENCODER_TARGET_MODULES_WITH_EMBEDDING = TEXT_ENCODER_TARGET_MODULES + ["token_embedding"] @@ -54,6 +55,16 @@ hidet.torch.dynamo_config.use_tensor_core(True) hidet.torch.dynamo_config.search_space(0) +if is_xformers_available(): + import xformers + import xformers.ops + + orig_xformers_memory_efficient_attention = xformers.ops.memory_efficient_attention + def xformers_memory_efficient_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, **kwargs): + return orig_xformers_memory_efficient_attention(query.to(key.dtype), key, value.to(key.dtype), **kwargs) + xformers.ops.memory_efficient_attention = xformers_memory_efficient_attention + + def parse_args(): parser = argparse.ArgumentParser( description="Simple example of a training script." -- cgit v1.2.3-54-g00ecf