Eval v5: BitsAndBytesConfig for MoE 4-bit

eval_securecoder.py

@@ -6,6 +6,7 @@
 #     "torch",
 #     "transformers",
 #     "bitsandbytes",
+#     "accelerate",
 # ]
 # ///
 """Cheap evaluations for SecureCoder.
@@ -373,14 +374,24 @@ def main() -> int:
         log.error("HF_TOKEN not set"); return 1
     os.makedirs(args.out_dir, exist_ok=True)
     import torch
-    from transformers import AutoModelForCausalLM, AutoTokenizer
+    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
     random.seed(args.seed)
     log.info("loading merged model %s in 4-bit ...", args.merged_repo)
     tokenizer = AutoTokenizer.from_pretrained(args.merged_repo, token=token)
+    bnb = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.bfloat16,
+        bnb_4bit_use_double_quant=True,
+    )
     model = AutoModelForCausalLM.from_pretrained(
-        args.merged_repo, token=token, dtype=torch.bfloat16, load_in_4bit=True,
-        device_map="auto")
+        args.merged_repo,
+        token=token,
+        quantization_config=bnb,
+        dtype=torch.bfloat16,
+        device_map="auto",
+    )
     log.info("model loaded")
 
     sections = [