| """Tiny GPU smoke test: LoRA SFT 15 steps then generation sanity check.""" |
| import torch |
| from datasets import load_dataset |
| from peft import LoraConfig |
| from trl import SFTConfig, SFTTrainer |
| from transformers import AutoTokenizer |
| BASE='Qwen/Qwen2.5-1.5B-Instruct' |
| print('CUDA:',torch.cuda.is_available(), torch.cuda.get_device_name(0) if torch.cuda.is_available() else None) |
| ds=load_dataset('krystv/nomen-ai-sft',split='train[:100]') |
| peft=LoraConfig(r=16,lora_alpha=32,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']) |
| args=SFTConfig(output_dir='/tmp/smoke',max_steps=15,per_device_train_batch_size=8,gradient_accumulation_steps=1,learning_rate=2e-4,max_length=192,eos_token='<|im_end|>',fp16=True,gradient_checkpointing=True,logging_steps=5,logging_first_step=True,disable_tqdm=True,report_to='none') |
| trainer=SFTTrainer(model=BASE,args=args,train_dataset=ds,peft_config=peft); trainer.train(); print('TRAIN_OK', trainer.state.log_history[-1]) |
| m=trainer.model; tok=AutoTokenizer.from_pretrained(BASE) |
| msgs=[{'role':'system','content':'You are Nomen-AI, output only one brand name.'},{'role':'user','content':'[ROOT:japanese+nordic] [THEME:gaming] [SYL:3] [LEN:8] [CREATIVE:0.8] Generate one original brand name.'}] |
| text=tok.apply_chat_template(msgs,tokenize=False,add_generation_prompt=True); enc=tok(text,return_tensors='pt').to(m.device) |
| out=m.generate(**enc,max_new_tokens=12,do_sample=True,temperature=1.6,min_p=0.05,repetition_penalty=1.25,no_repeat_ngram_size=2,pad_token_id=tok.eos_token_id,eos_token_id=tok.eos_token_id) |
| print('GEN:',repr(tok.decode(out[0][enc['input_ids'].shape[1]:],skip_special_tokens=True))); print('VRAM_GB:',round(torch.cuda.max_memory_allocated()/1e9,2)); print('SMOKE_PASS') |
|
|
