Update loader.py

This commit is contained in:
Daniel Han 2025-08-22 04:02:19 -07:00
commit a33ff972c1

View file

@ -437,10 +437,12 @@ class FastLanguageModel(FastLlamaModel):
if load_in_4bit:
# Fix up bitsandbytes config
config = model.config.to_dict()
torch_dtype = config.get("dtype") or config.get("torch_dtype")
quantization_config = \
{
# Sometimes torch_dtype is not a string!!
"bnb_4bit_compute_dtype" : model.config.to_dict()["torch_dtype"],
"bnb_4bit_compute_dtype" : torch_dtype,
"bnb_4bit_quant_type" : "nf4",
"bnb_4bit_use_double_quant" : True,
"llm_int8_enable_fp32_cpu_offload" : False,
@ -887,10 +889,12 @@ class FastModel(FastBaseModel):
if load_in_4bit:
# Fix up bitsandbytes config
config = model.config.to_dict()
torch_dtype = config.get("dtype") or config.get("torch_dtype")
quantization_config = \
{
# Sometimes torch_dtype is not a string!!
"bnb_4bit_compute_dtype" : model.config.to_dict()["torch_dtype"],
"bnb_4bit_compute_dtype" : torch_dtype,
"bnb_4bit_quant_type" : "nf4",
"bnb_4bit_use_double_quant" : True,
"llm_int8_enable_fp32_cpu_offload" : False,