add flash attention

requirements.txt

@@ -31,3 +31,4 @@ transformers
 accelerate
 diffusers
 bitsandbytes
+flash-attn

src/autotrain/cli/run_llm.py

@@ -325,6 +325,13 @@ def register_subcommand(parser: ArgumentParser):
                 "required": False,
                 "type": str,
             },
+            {
+                "arg": "--use_flash_attention_2",
+                "help": "Use flash attention 2",
+                "required": False,
+                "action": "store_true",
+                "alias": ["--use-flash-attention-2", "--use-fa2"],
+            },
         ]
         run_llm_parser = parser.add_parser("llm", description="✨ Run AutoTrain LLM")
         for arg in arg_list:
@@ -364,6 +371,7 @@ def __init__(self, args):
             "use_int8",
             "use_int4",
             "merge_adapter",
+            "use_flash_attention_2",
         ]
         for arg_name in store_true_arg_names:
             if getattr(self.args, arg_name) is None:
@@ -451,6 +459,7 @@ def run(self):
                 token=self.args.token,
                 merge_adapter=self.args.merge_adapter,
                 username=self.args.username,
+                use_flash_attention_2=self.args.use_flash_attention_2,
             )
 
             # space training

src/autotrain/trainers/clm/__main__.py

@@ -127,13 +127,15 @@ def train(config):
             torch_dtype=torch.float16,
             device_map={"": Accelerator().process_index} if torch.cuda.is_available() else None,
             trust_remote_code=True,
+            use_flash_attention_2=config.use_flash_attention_2,
         )
     else:
         model = AutoModelForCausalLM.from_pretrained(
             config.model,
             config=model_config,
             use_auth_token=config.token,
             trust_remote_code=True,
+            use_flash_attention_2=config.use_flash_attention_2,
         )
 
     model.resize_token_embeddings(len(tokenizer))

src/autotrain/trainers/clm/params.py

@@ -44,6 +44,7 @@ class LLMTrainingParams(BaseModel):
     target_modules: str = Field(None, title="Target modules")
     merge_adapter: bool = Field(False, title="Merge adapter")
     username: str = Field(None, title="Hugging Face Username")
+    use_flash_attention_2: bool = Field(False, title="Use flash attention 2")
 
     def save(self, output_dir):
         os.makedirs(output_dir, exist_ok=True)