File size: 2,755 Bytes
e7d8c8b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | import os, 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"
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)
tok=AutoTokenizer.from_pretrained(M)
m=AutoModelForCausalLM.from_pretrained(M,quantization_config=bnb,dtype=torch.bfloat16,
device_map={"":0},attn_implementation="sdpa")
base=torch.cuda.memory_allocated(0)/2**30
# sanity generate (transformers-5 returns BatchEncoding)
enc=tok.apply_chat_template([{"role":"user","content":"Reply with ONLY this JSON: {\"answer\":\"A\"}"}],
add_generation_prompt=True, return_dict=True, return_tensors="pt").to(0)
with torch.no_grad(): o=m.generate(**enc,max_new_tokens=24,do_sample=False)
print("GEN:",repr(tok.decode(o[0][enc['input_ids'].shape[-1]:],skip_special_tokens=True))[:160])
print(f"BASE_WEIGHTS {base:.2f} GiB")
m.config.use_cache=False
m=prepare_model_for_kbit_training(m, use_gradient_checkpointing=False)
TM=["q_proj","k_proj","v_proj","o_proj","gate_proj","up_proj","down_proj"]
res={}
for rank in (128, 256):
mm=get_peft_model(m, LoraConfig(r=rank,lora_alpha=rank*2,lora_dropout=0.05,bias="none",
task_type="CAUSAL_LM",target_modules=TM))
tr=sum(p.numel() for p in mm.parameters() if p.requires_grad)
for seq,mb in ((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)
opt=torch.optim.AdamW([p for p in mm.parameters() if p.requires_grad],lr=1e-5)
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
tps=mb*seq/dt
res[f"r{rank}_s{seq}_mb{mb}"]={"tok_s":round(tps),"peak_gib":round(pk,1),"s_step":round(dt,2)}
print(f"r={rank} seq={seq} mb={mb}: {tps:8.0f} tok/s peak {pk:5.1f} GiB {dt:.2f}s/step trainable {tr/1e6:.0f}M")
del opt,ids,out
except torch.cuda.OutOfMemoryError:
print(f"r={rank} seq={seq} mb={mb}: OOM"); torch.cuda.empty_cache()
mm=mm.unload() if hasattr(mm,'unload') else mm
m=getattr(mm,'base_model',m)
json.dump(res,open("/workspace/aditi/logs/bench_train.json","w"),indent=2)
|