Jedi / train.py
podstack's picture
Jedi v1.0
7a74082 verified
Raw History Blame Contribute Delete
951 Bytes
"""Stage 3: full fine-tune of Qwen2.5-0.5B-Instruct into Jedi (loss on the one-word verdict only)."""
import os
from datasets import load_from_disk
from trl import SFTConfig, SFTTrainer
BASE = os.environ.get("BASE", "Qwen/Qwen2.5-0.5B-Instruct")
ds = load_from_disk("data/jedi_judge").select_columns(["prompt", "completion"])
cfg = SFTConfig(
output_dir="ckpt", num_train_epochs=2, per_device_train_batch_size=8, gradient_accumulation_steps=4,
learning_rate=2e-5, lr_scheduler_type="cosine", warmup_steps=30, weight_decay=0.0, bf16=True,
max_length=2048, completion_only_loss=True, logging_steps=20, eval_strategy="epoch", save_strategy="epoch",
save_total_limit=1, report_to=[], gradient_checkpointing=True, dataloader_num_workers=4,
)
trainer = SFTTrainer(model=BASE, args=cfg, train_dataset=ds["train"], eval_dataset=ds["val"])
trainer.train()
trainer.save_model("jedi-0.5b")
trainer.processing_class.save_pretrained("jedi-0.5b")