File size: 1,061 Bytes
c1f5d12 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 | """Stage 3: full fine-tune of Qwen2.5-0.5B-Instruct into the Shaktimaan router (loss on the route 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/shaktimaan_routing").map(
lambda r: {"prompt": r["messages"][:2], "completion": r["messages"][2:]}).select_columns(["prompt", "completion"])
cfg = SFTConfig(
output_dir="ckpt", num_train_epochs=3, per_device_train_batch_size=8, gradient_accumulation_steps=4,
learning_rate=2e-5, lr_scheduler_type="cosine", warmup_steps=20, weight_decay=0.0, bf16=True,
max_length=3072, completion_only_loss=True, logging_steps=10, 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["test"])
trainer.train()
trainer.save_model("shaktimaan-0.5b")
trainer.processing_class.save_pretrained("shaktimaan-0.5b")
|