File size: 2,527 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 | import os,sys,time,torch,json
os.environ.setdefault("HF_HOME","/workspace/.hf_home")
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model
M="/workspace/models/gemma4-qat-unquantized"; rank=int(sys.argv[1])
USE_LIGER = os.environ.get("LIGER","1")=="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")
if USE_LIGER:
try:
from liger_kernel.transformers import _apply_liger_kernel_to_instance
_apply_liger_kernel_to_instance(model=m); print("LIGER applied",flush=True)
except Exception as e: print("LIGER failed:",repr(e)[:200],flush=True)
m.config.use_cache=False
m.enable_input_require_grads() # doc 4: REQUIRED when checkpointing is on
if os.environ.get("CKPT","1")=="1":
m.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant":False})
else:
print("grad-ckpt OFF",flush=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"]))
mm.train()
print(f"rank={rank} liger={USE_LIGER} 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),(4096,2),(4096,4)):
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",flush=True)
del ids,out
except torch.cuda.OutOfMemoryError:
print(f" seq={seq} mb={mb}: OOM",flush=True); torch.cuda.empty_cache()
json.dump(res,open(f"/workspace/aditi/logs/bench4_r{rank}.json","w"),indent=2); print("DONE",flush=True)
|