File size: 1,881 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 | 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")
|