| import os, time, torch, json |
| os.environ.setdefault("HF_HOME","/workspace/.hf_home") |
| from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig |
| M="/workspace/models/gemma4-qat-unquantized" |
| torch.cuda.init() |
| _=torch.zeros(1,device="cuda:0") |
| torch.cuda.reset_peak_memory_stats(0) |
| bnb=BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", |
| bnb_4bit_compute_dtype=torch.bfloat16, |
| bnb_4bit_use_double_quant=True) |
| t=time.time() |
| tok=AutoTokenizer.from_pretrained(M) |
| m=AutoModelForCausalLM.from_pretrained(M, quantization_config=bnb, dtype=torch.bfloat16, |
| device_map={"":0}, attn_implementation="sdpa") |
| load=time.time()-t |
| w=torch.cuda.memory_allocated(0)/2**30 |
| print(f"LOAD {load:.0f}s WEIGHTS_RESIDENT {w:.2f} GiB peak {torch.cuda.max_memory_allocated(0)/2**30:.2f} GiB") |
| n=sum(p.numel() for p in m.parameters()); print(f"params(counted) {n/1e9:.2f}B") |
| print("dtypes:", {str(p.dtype) for p in m.parameters()}) |
| |
| m.eval() |
| msgs=[{"role":"user","content":"Reply with one JSON object: {\"answer\":\"A\"}"}] |
| ids=tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt").to(0) |
| with torch.no_grad(): |
| o=m.generate(ids, max_new_tokens=24, do_sample=False) |
| print("GEN:", repr(tok.decode(o[0][ids.shape[-1]:], skip_special_tokens=True))[:200]) |
| print(f"PEAK_AFTER_GEN {torch.cuda.max_memory_allocated(0)/2**30:.2f} GiB") |
| json.dump({"weights_gib":round(w,2),"load_s":round(load)}, open("/workspace/aditi/logs/bench4bit.json","w")) |
|
|
