mirror of
https://github.com/Druthulu/BFM-decomp
synced 2026-09-26 13:33:34 -04:00
Drew T
96 lines
4.8 KiB
Python
96 lines
4.8 KiB
Python
#!/usr/bin/env python3
|
|
"""train_lora.py — QLoRA fine-tune a code model into a BFM/gcc-2.7.2 MATCHING specialist.
|
|
|
|
Pilot recipe (docs/gen2-mips-matching-model.md). Trains on datasets/match_pairs/train_ft.jsonl
|
|
(produced by tools/format_finetune.py) and saves a LoRA adapter + a merged GGUF for LM Studio.
|
|
|
|
RUN THIS WHERE YOU HAVE A GPU + Unsloth (your 3080 Ti via WSL, or a cloud A100). NOT runnable in the
|
|
headless agent sandbox (no CUDA / no model download there). On a 12 GB card use the 7B base in 4-bit.
|
|
|
|
# one-time setup (WSL2 with CUDA, or cloud):
|
|
pip install "unsloth[cu121] @ git+https://github.com/unslothai/unsloth.git" trl peft datasets
|
|
|
|
python tools/train_lora.py # 7B pilot, ~3 epochs
|
|
python tools/train_lora.py --base unsloth/Qwen2.5-Coder-14B-Instruct-bnb-4bit --rank 32
|
|
|
|
Then load models/bfm-match-7b/*.gguf in LM Studio and eval LEAN:
|
|
LEAN=1 API_BASE=http://<host>:1234/v1 MODEL=bfm-match-7b \
|
|
.venv/bin/python tools/api_draft.py --targets .run/ab-exp/targets20.json --out .run/ab-exp/lora --iters 4
|
|
.venv/bin/python tools/ab_score.py --arms haiku local lora # beat stock-local's ~0
|
|
"""
|
|
import argparse, os
|
|
|
|
REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
ap.add_argument('--base', default='unsloth/Qwen2.5-Coder-7B-Instruct-bnb-4bit',
|
|
help='4-bit base (7B fits a 12GB card; 14B needs ~16GB)')
|
|
ap.add_argument('--data', default='datasets/match_pairs/train_ft.jsonl')
|
|
ap.add_argument('--out', default='models/bfm-match-7b')
|
|
ap.add_argument('--rank', type=int, default=16)
|
|
ap.add_argument('--epochs', type=float, default=3.0)
|
|
ap.add_argument('--max-steps', type=int, default=0, help='>0 overrides epochs (use for a smoke run)')
|
|
ap.add_argument('--lr', type=float, default=2e-4)
|
|
ap.add_argument('--maxlen', type=int, default=4096, help='token cap; lower to 2048 to avoid 12GB VRAM spill')
|
|
ap.add_argument('--batch', type=int, default=2, help='per-device batch; 1 halves activation memory')
|
|
ap.add_argument('--no-gguf', action='store_true', help='skip the merged-GGUF export step')
|
|
a = ap.parse_args()
|
|
|
|
# Imported here so --help works without a GPU/unsloth install.
|
|
from unsloth import FastLanguageModel
|
|
from unsloth.chat_templates import get_chat_template
|
|
from datasets import load_dataset
|
|
from trl import SFTTrainer, SFTConfig
|
|
|
|
model, tok = FastLanguageModel.from_pretrained(
|
|
model_name=a.base, max_seq_length=a.maxlen, load_in_4bit=True, dtype=None)
|
|
model = FastLanguageModel.get_peft_model(
|
|
model, r=a.rank, lora_alpha=a.rank * 2, lora_dropout=0.0, bias='none',
|
|
target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj',
|
|
'gate_proj', 'up_proj', 'down_proj'],
|
|
use_gradient_checkpointing='unsloth', random_state=3407)
|
|
tok = get_chat_template(tok, chat_template='qwen-2.5')
|
|
|
|
ds = load_dataset('json', data_files=os.path.join(REPO, a.data), split='train')
|
|
|
|
def fmt(ex):
|
|
return {'text': tok.apply_chat_template(ex['messages'], tokenize=False,
|
|
add_generation_prompt=False)}
|
|
ds = ds.map(fmt)
|
|
|
|
trainer = SFTTrainer(
|
|
model=model, tokenizer=tok, train_dataset=ds,
|
|
args=SFTConfig(
|
|
dataset_text_field='text', max_seq_length=a.maxlen,
|
|
per_device_train_batch_size=a.batch, gradient_accumulation_steps=max(1, 16 // a.batch),
|
|
warmup_ratio=0.03, num_train_epochs=a.epochs,
|
|
max_steps=(a.max_steps if a.max_steps > 0 else -1), learning_rate=a.lr,
|
|
logging_steps=1 if a.max_steps else 10, optim='adamw_8bit', weight_decay=0.01,
|
|
lr_scheduler_type='cosine', seed=3407,
|
|
output_dir=os.path.join(REPO, a.out, 'ckpt')))
|
|
trainer.train()
|
|
|
|
out = os.path.join(REPO, a.out)
|
|
model.save_pretrained(out); tok.save_pretrained(out)
|
|
print('LoRA adapter ->', a.out)
|
|
if not a.no_gguf:
|
|
model.save_pretrained_gguf(out, tok, quantization_method='q4_k_m')
|
|
# disk hygiene: GGUF conv leaves a ~15GB merged-16bit + ~15GB BF16 gguf; keep ONLY the q4
|
|
# (the ~30GB peak is what crashed C: last time — delete intermediates immediately after)
|
|
import glob as _g
|
|
for pat in (os.path.join(out, 'model-0000*.safetensors'),
|
|
os.path.join(out, 'model.safetensors.index.json'),
|
|
out + '_gguf/*BF16*.gguf', out + '_gguf/*bf16*.gguf', out + '_gguf/*F16*.gguf'):
|
|
for f in _g.glob(pat):
|
|
try:
|
|
os.remove(f); print('cleaned intermediate:', f)
|
|
except OSError:
|
|
pass
|
|
print('merged q4_k_m GGUF ->', a.out + '_gguf', '— intermediates cleaned')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|