fix the torch_dtype and quant_storage_dtype (#1614)

* fix the torch_dtype and quant_storage_dtype Co-Authored-By: Gabriel Altay <[email protected]> * quality --------- Co-authored-by: Gabriel Altay <[email protected]>
huggingface · Apr 4, 2024 · 8452d71 · 8452d71
1 parent 02b5aed
commit 8452d71
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/examples/sft/train.py b/examples/sft/train.py
@@ -40,7 +40,7 @@ class ModelArguments:
         metadata={"help": "Compute dtype for 4bit base models"},
     )
     bnb_4bit_quant_storage_dtype: Optional[str] = field(
-        default="float32",
+        default="uint8",
         metadata={"help": "Quantization storage dtype for 4bit base models"},
     )
     bnb_4bit_quant_type: Optional[str] = field(

diff --git a/examples/sft/utils.py b/examples/sft/utils.py
@@ -125,12 +125,15 @@ def create_and_prepare_model(args, data_args, training_args):
             load_in_4bit=args.use_4bit_quantization,
         )
     else:
+        torch_dtype = (
+            quant_storage_dtype if quant_storage_dtype and quant_storage_dtype.is_floating_point else torch.float32
+        )
         model = AutoModelForCausalLM.from_pretrained(
             args.model_name_or_path,
             quantization_config=bnb_config,
             trust_remote_code=True,
             attn_implementation="flash_attention_2" if args.use_flash_attn else "eager",
-            torch_dtype=quant_storage_dtype or torch.float32,
+            torch_dtype=torch_dtype,
         )
 
     peft_config = None