Spaces:
Running on Zero
Running on Zero
File size: 4,081 Bytes
0418f40 b9592cc 0418f40 | 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 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | """Preference alignment: DPO or ORPO on top of (or instead of) the SFT adapter.
python -m src.training.preference --method orpo --config configs/train_general.yaml
- ORPO: no reference model needed; can even replace the SFT stage (cheapest).
- DPO: run *after* SFT; TRL handles the frozen reference model implicitly when
training a PEFT adapter (the base weights are the reference).
Dataset must have prompt/chosen/rejected columns (conversational format works —
TRL extracts prompt/chosen/rejected from message lists).
Note: `pref_dataset` in the config defaults to a *general* preference set as a
placeholder. For real finance alignment, build pairs by sampling two model
answers per finance prompt and labeling the better one (human or strong-model
judge), then point pref_dataset at that repo.
"""
import argparse
import torch
import yaml
from datasets import load_dataset
from peft import LoraConfig, PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from trl import DPOConfig, DPOTrainer, ORPOConfig, ORPOTrainer
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--config", required=True)
ap.add_argument("--method", choices=["dpo", "orpo"], default="orpo")
ap.add_argument("--sft-adapter", default=None,
help="path/repo of the SFT adapter to continue from (default: config output_dir for dpo)")
args = ap.parse_args()
with open(args.config) as f:
cfg = yaml.safe_load(f)
bf16 = torch.cuda.is_bf16_supported()
dtype = torch.bfloat16 if bf16 else torch.float16
model = AutoModelForCausalLM.from_pretrained(
cfg["base_model"],
dtype=dtype,
device_map="auto",
quantization_config=BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_use_double_quant=True,
bnb_4bit_compute_dtype=dtype,
),
)
model.config.use_cache = False
# DPO continues from the SFT adapter; ORPO may start fresh from the base.
adapter = args.sft_adapter or (cfg["output_dir"] if args.method == "dpo" else None)
peft_config = None
if adapter:
model = PeftModel.from_pretrained(model, adapter, is_trainable=True)
print(f"[info] continuing from adapter: {adapter}")
else:
peft_config = LoraConfig(
r=cfg["lora_r"],
lora_alpha=cfg["lora_alpha"],
lora_dropout=cfg["lora_dropout"],
target_modules=cfg["target_modules"],
bias="none",
task_type="CAUSAL_LM",
)
tokenizer = AutoTokenizer.from_pretrained(cfg["base_model"])
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
pref_data = load_dataset(cfg["pref_dataset"], split="train")
out_dir = f"{cfg['output_dir']}-{args.method}"
common = dict(
output_dir=out_dir,
per_device_train_batch_size=2,
gradient_accumulation_steps=8,
learning_rate=float(cfg["pref_learning_rate"]),
max_steps=cfg["pref_max_steps"],
lr_scheduler_type="cosine",
warmup_ratio=0.05,
bf16=bf16,
fp16=not bf16,
gradient_checkpointing=True,
gradient_checkpointing_kwargs={"use_reentrant": False},
logging_steps=10,
save_steps=100,
report_to="none",
beta=cfg["pref_beta"],
)
if args.method == "dpo":
trainer = DPOTrainer(
model=model,
args=DPOConfig(**common),
train_dataset=pref_data,
processing_class=tokenizer,
peft_config=peft_config,
)
else:
trainer = ORPOTrainer(
model=model,
args=ORPOConfig(**common),
train_dataset=pref_data,
processing_class=tokenizer,
peft_config=peft_config,
)
trainer.train()
trainer.save_model(out_dir)
tokenizer.save_pretrained(out_dir)
print(f"[done] {args.method} adapter saved to {out_dir}")
if __name__ == "__main__":
main()
|