| import os,sys,time,torch,json |
| os.environ.setdefault("HF_HOME","/workspace/.hf_home") |
| import torch.multiprocessing as mp |
| def worker(rank,q): |
| torch.cuda.set_device(rank) |
| from transformers import AutoModelForCausalLM, BitsAndBytesConfig |
| from peft import LoraConfig, get_peft_model |
| from liger_kernel.transformers import _apply_liger_kernel_to_instance |
| M="/workspace/models/gemma4-qat-unquantized" |
| 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={"":rank},attn_implementation="sdpa") |
| _apply_liger_kernel_to_instance(model=m); m.config.use_cache=False |
| m.enable_input_require_grads() |
| mm=get_peft_model(m,LoraConfig(r=128,lora_alpha=256,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() |
| opt=torch.optim.AdamW([p for p in mm.parameters() if p.requires_grad],lr=1e-5) |
| ids=torch.randint(0,250000,(1,2048),device=rank) |
| for i in range(6): |
| if i==2: 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)/4 |
| q.put((rank, round(2048/dt), round(torch.cuda.max_memory_allocated(rank)/2**30,1))) |
| if __name__=="__main__": |
| mp.set_start_method("spawn"); q=mp.Queue() |
| ps=[mp.Process(target=worker,args=(r,q)) for r in range(2)] |
| [p.start() for p in ps]; res=[q.get() for _ in range(2)]; [p.join() for p in ps] |
| tot=sum(r[1] for r in res) |
| for r in sorted(res): print(f" GPU{r[0]}: {r[1]} tok/s peak {r[2]} GiB") |
| print(f"AGGREGATE {tot} tok/s across 2 GPUs") |
|
|
