423A35C7/LLaMA-Factory
1
0
Fork 0
mirror of https://github.com/hiyouga/LLaMA-Factory.git synced 2025-08-04 20:52:59 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix loading dtype

Browse Source 
Former-commit-id: c874e764b8334c18091233c3781009a39d67e794
This commit is contained in:
hiyouga 2023-10-14 20:15:24 +08:00
parent 1faba4b704
commit 089785c71b
1 changed files with 4 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
src/llmtuner/tuner/core/loader.py
View File

@ -88,11 +88,10 @@ def load_model_and_tokenizer(
tokenizer._pad = MethodType(PreTrainedTokenizerBase._pad, tokenizer)
# Set model dtype
if model_args.compute_dtype is not None:
if model_args.compute_dtype is not None: # for training
setattr(config, "torch_dtype", model_args.compute_dtype)
else: # priority: bf16 > fp16 > fp32
optim_dtype = infer_optim_dtype(model_dtype=getattr(config, "torch_dtype", None))
setattr(config, "torch_dtype", optim_dtype)
else: # for evaluation, priority: bf16 > fp16 > fp32
model_args.compute_dtype = infer_optim_dtype(model_dtype=getattr(config, "torch_dtype", None))
# Fix config (for Qwen)
if getattr(config, "model_type", None) == "qwen":
@ -185,7 +184,7 @@ def load_model_and_tokenizer(
model = AutoModelForCausalLM.from_pretrained(
model_to_load,
config=config,
torch_dtype=getattr(config, "torch_dtype"),
torch_dtype=model_args.compute_dtype,
low_cpu_mem_usage=(not is_deepspeed_zero3_enabled()),
**config_kwargs
)