File size: 3,116 Bytes
d18bbd3 | 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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | #!/usr/bin/env python3
"""KTO preference tuning on top of the Starburst 2.3 SFT model.
Desirable = winning entries + correct skips (837)
Undesirable = losing entries (384 real losers, held out of SFT)
KTO nudges the policy away from loss-like setups using per-sample binary labels —
no DPO-style pairs, no fabricated "rejected" data. Reference = the SFT model
(the merged 2.3) with the new LoRA disabled (PEFT auto-ref).
Free the GPU first: sudo systemctl stop starburst
Run: PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True unsloth_env/bin/python train_kto.py
"""
import os
import unsloth # noqa: F401 (import first to patch transformers/trl)
try:
from unsloth import PatchDPOTrainer
PatchDPOTrainer() # patches the DPO/KTO/ORPO trainers (no-op if already patched)
except Exception:
pass
from unsloth import FastModel
from datasets import load_dataset
from trl import KTOTrainer, KTOConfig
MODEL = os.environ.get("KTO_MODEL", "/home/sheen/trading-assistant/starburst_2.3_9b_merged")
MAXLEN = int(os.environ.get("KTO_MAXLEN", "2048"))
DATA = os.environ.get("KTO_DATA", "train_kto.jsonl")
OUT = os.environ.get("KTO_OUT", "qlora_kto_2_3")
EPOCHS = float(os.environ.get("KTO_EPOCHS", "1"))
LR = float(os.environ.get("KTO_LR", "1e-5"))
BETA = float(os.environ.get("KTO_BETA", "0.1"))
model, tokenizer = FastModel.from_pretrained(
MODEL, max_seq_length=MAXLEN, load_in_4bit=True, use_gradient_checkpointing="unsloth")
model = FastModel.get_peft_model(
model, r=16, lora_alpha=16, lora_dropout=0.0, bias="none",
finetune_vision_layers=False, finetune_language_layers=True,
finetune_attention_modules=True, finetune_mlp_modules=True,
use_gradient_checkpointing="unsloth", random_state=3407)
text_tok = getattr(tokenizer, "tokenizer", tokenizer)
# Use the SAME no-seed template the model is served with, so the completion the
# policy is scored on includes its own <think>...</think> (matches SFT + serving).
with open("/home/sheen/trading-assistant/starburst_chat_template.jinja") as f:
text_tok.chat_template = f.read()
ds = load_dataset("json", data_files=DATA, split="train")
trainer = KTOTrainer(
model=model,
ref_model=None, # PEFT: reference = base with adapter disabled
args=KTOConfig(
per_device_train_batch_size=2, # KTO needs actual batch > 1 (KL term)
gradient_accumulation_steps=4, # effective batch still 8
num_train_epochs=EPOCHS,
learning_rate=LR,
beta=BETA,
desirable_weight=1.0,
undesirable_weight=float(os.environ.get("KTO_UNDES_W", "1.8")), # 2.3: 837:384; per-run override
max_length=MAXLEN,
max_prompt_length=1024,
max_completion_length=1024,
warmup_ratio=0.05,
logging_steps=5,
optim="adamw_8bit",
weight_decay=0.0,
lr_scheduler_type="linear",
seed=3407,
output_dir=OUT,
report_to="none",
),
train_dataset=ds,
processing_class=text_tok,
)
trainer.train()
model.save_pretrained(OUT)
text_tok.save_pretrained(OUT)
print("KTO adapter saved to", OUT)
|