Update arg_parser.py

arg_parser.py

@@ -66,6 +66,8 @@ def parse_commandline():
 
     # Multi GPU Support
     parser_training.add_argument("--local_rank", type=int, default=0, help="local rank if using torch.distributed.launch")
+    
+    parser_training.add_argument("--flash_attention", help="enables flash attention, can improve performance and reduce VRAM use")
 
     return vars(parser.parse_args())
 
@@ -102,4 +104,5 @@ def get_config() -> Finetune4bConfig:
         use_eos_token=args["use_eos_token"]!=0,
         groupsize=args["groupsize"],
         local_rank=args["local_rank"],
+        flash_attention=args["flash_attention"],
     )