From 778035152d7ba63a2f980eea98142d5e6eec7fe3 Mon Sep 17 00:00:00 2001 From: yamashi Date: Fri, 7 Apr 2023 00:42:34 +0200 Subject: [PATCH] Update arg_parser.py --- arg_parser.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/arg_parser.py b/arg_parser.py index a985380..be18154 100644 --- a/arg_parser.py +++ b/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"], )