1from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
2from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
3from trl import SFTTrainer, SFTConfig
4from datasets import load_dataset
5
6# 1. Load 4-bit quantized base model
7bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype="bfloat16", bnb_4bit_quant_type="nf4")
8model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct", quantization_config=bnb, device_map="auto")
9tok = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
10
11# 2. Inject LoRA adapters (r=16)
12model = prepare_model_for_kbit_training(model)
13lora = LoraConfig(r=16, lora_alpha=32, target_modules=["q_proj","k_proj","v_proj","o_proj"], lora_dropout=0.05, bias="none")
14model = get_peft_model(model, lora)
15model.print_trainable_parameters() # ~0.4% of params
16
17# 3. Load + format dataset
18ds = load_dataset("json", data_files="train.jsonl", split="train")
19def format_chat(ex):
20 return tok.apply_chat_template(ex["messages"], tokenize=False)
21
22# 4. Train with TRL SFTTrainer
23trainer = SFTTrainer(
24 model=model,
25 train_dataset=ds.map(lambda e: {"text": format_chat(e)}),
26 args=SFTConfig(output_dir="out", num_train_epochs=3, per_device_train_batch_size=4,
27 gradient_accumulation_steps=4, learning_rate=2e-4, lr_scheduler_type="cosine",
28 bf16=True, logging_steps=10, save_strategy="epoch"),
29)
30trainer.train()
31
32# 5. Save adapter only (~50MB)
33trainer.model.save_pretrained("out/final-adapter")
