| import os,sys,time,torch,json |
| os.environ.setdefault("HF_HOME","/workspace/.hf_home") |
| from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig |
| from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training |
| M="/workspace/models/gemma4-qat-unquantized"; rank=int(sys.argv[1]) |
| torch.cuda.init(); _=torch.zeros(1,device="cuda:0") |
| bnb=BitsAndBytesConfig(load_in_4bit=True,bnb_4bit_quant_type="nf4", |
| bnb_4bit_compute_dtype=torch.bfloat16,bnb_4bit_use_double_quant=True) |
| m=AutoModelForCausalLM.from_pretrained(M,quantization_config=bnb,dtype=torch.bfloat16, |
| device_map={"":0},attn_implementation="sdpa") |
| print(f"BASE_4BIT_WEIGHTS {torch.cuda.memory_allocated(0)/2**30:.2f} GiB",flush=True) |
| m.config.use_cache=False |
| m.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant":False}) |
| m=prepare_model_for_kbit_training(m,use_gradient_checkpointing=True) |
| mm=get_peft_model(m,LoraConfig(r=rank,lora_alpha=rank*2,lora_dropout=0.05,bias="none", |
| task_type="CAUSAL_LM",target_modules=["q_proj","k_proj","v_proj","o_proj","gate_proj","up_proj","down_proj"])) |
| tr=sum(p.numel() for p in mm.parameters() if p.requires_grad) |
| print(f"rank={rank} trainable={tr/1e6:.1f}M after-LoRA {torch.cuda.memory_allocated(0)/2**30:.2f} GiB",flush=True) |
| opt=torch.optim.AdamW([p for p in mm.parameters() if p.requires_grad],lr=1e-5) |
| res={} |
| for seq,mb in ((2048,1),(2048,2),(2048,4),(2048,8),(4096,4),(8192,2)): |
| try: |
| torch.cuda.empty_cache(); torch.cuda.reset_peak_memory_stats(0) |
| ids=torch.randint(0,250000,(mb,seq),device=0) |
| for i in range(3): |
| if i==1: torch.cuda.synchronize(); t=time.time() |
| out=mm(input_ids=ids,labels=ids); out.loss.backward() |
| opt.step(); opt.zero_grad(set_to_none=True) |
| torch.cuda.synchronize(); dt=(time.time()-t)/2; pk=torch.cuda.max_memory_allocated(0)/2**30 |
| res[f"s{seq}_mb{mb}"]={"tok_s":round(mb*seq/dt),"peak_gib":round(pk,1)} |
| print(f" seq={seq} mb={mb}: {mb*seq/dt:8.0f} tok/s/GPU peak {pk:5.1f} GiB {dt:.2f}s/step",flush=True) |
| del ids,out |
| except torch.cuda.OutOfMemoryError: |
| print(f" seq={seq} mb={mb}: OOM",flush=True); torch.cuda.empty_cache() |
| res["trainable_M"]=round(tr/1e6,1) |
| json.dump(res,open(f"/workspace/aditi/logs/bench_r{rank}.json","w"),indent=2) |
| print("BENCH_DONE",flush=True) |
|
|
