tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

fix torch_dtype for model load

Browse Source
This commit is contained in:
Wing Lian
2023-05-14 08:44:22 -04:00
parent 0d28df0fd2
commit aef00b6c13
1 changed files with 6 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
src/axolotl/utils/models.py
View File

@@ -62,9 +62,12 @@ def load_model(
logging.info("patching with xformers attention") logging.info("patching with xformers attention")
hijack_llama_attention() hijack_llama_attention()
torch_dtype = ( if cfg.bf16:
torch.float16 if cfg.load_in_8bit or cfg.fp16 or cfg.bf16 else torch.float32 torch_dtype = torch.bfloat16
) elif cfg.load_in_8bit or cfg.fp16:
torch_dtype = torch.float16
else:
torch_dtype = torch.float32
try: try:
if cfg.load_4bit: if cfg.load_4bit:
from alpaca_lora_4bit.monkeypatch.peft_tuners_lora_monkey_patch import ( from alpaca_lora_4bit.monkeypatch.peft_tuners_lora_monkey_patch import (