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()})
# quick generate sanity
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"))