308 lines
9.6 KiB
Python
308 lines
9.6 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Fine-tuning QLoRA de gemma3:12b sur RunPod (A100 40/80GB).
|
|
|
|
Dataset V2 : 15.9K train, 53% raisonnement structuré (vs 95% lookups V1).
|
|
Sources : referentiels, pipeline, cocoa, ccam, cim10, reasoning,
|
|
negative, discrimination, fascicule_reasoning, guide_metho.
|
|
|
|
Usage sur RunPod :
|
|
1. Créer un pod A100 80GB (template PyTorch 2.4+ / CUDA 12.x)
|
|
2. Uploader les fichiers (train_runpod.py, setup.sh, data/)
|
|
3. bash setup.sh
|
|
4. python train_runpod.py [--epochs 3] [--export-gguf]
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
|
|
BASE = Path(__file__).resolve().parent
|
|
DATASETS = BASE / "data"
|
|
OUTPUT = BASE / "models"
|
|
OUTPUT.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
|
def check_prerequisites():
|
|
import torch
|
|
|
|
if not torch.cuda.is_available():
|
|
raise RuntimeError("CUDA non disponible.")
|
|
|
|
gpu_name = torch.cuda.get_device_name(0)
|
|
vram_total = torch.cuda.get_device_properties(0).total_memory / 1024**3
|
|
|
|
print(f"GPU: {gpu_name}")
|
|
print(f"VRAM: {vram_total:.1f} Go")
|
|
|
|
train_path = DATASETS / "pmsi_train.jsonl"
|
|
eval_path = DATASETS / "pmsi_eval.jsonl"
|
|
if not train_path.exists() or not eval_path.exists():
|
|
raise FileNotFoundError("Dataset non trouvé dans data/")
|
|
|
|
with open(train_path) as f:
|
|
n_train = sum(1 for _ in f)
|
|
with open(eval_path) as f:
|
|
n_eval = sum(1 for _ in f)
|
|
|
|
print(f"Dataset: {n_train} train + {n_eval} eval")
|
|
|
|
# Adapter le batch size à la VRAM
|
|
if vram_total >= 70:
|
|
suggested_batch = 8
|
|
elif vram_total >= 35:
|
|
suggested_batch = 4
|
|
else:
|
|
suggested_batch = 2
|
|
|
|
print(f"Batch size suggéré: {suggested_batch}")
|
|
return train_path, eval_path, suggested_batch
|
|
|
|
|
|
def load_model(model_name, max_seq_length, load_in_4bit=True):
|
|
from unsloth import FastLanguageModel
|
|
|
|
print(f"\nChargement de {model_name} (4-bit={load_in_4bit})...")
|
|
|
|
model, tokenizer = FastLanguageModel.from_pretrained(
|
|
model_name=model_name,
|
|
max_seq_length=max_seq_length,
|
|
dtype=None,
|
|
load_in_4bit=load_in_4bit,
|
|
)
|
|
|
|
print(f" Modèle chargé : {model.config._name_or_path}")
|
|
print(f" Paramètres : {model.num_parameters() / 1e9:.1f}B")
|
|
return model, tokenizer
|
|
|
|
|
|
def attach_lora(model, r=32, alpha=64, dropout=0.0):
|
|
from unsloth import FastLanguageModel
|
|
|
|
print(f"\nLoRA (r={r}, alpha={alpha})...")
|
|
|
|
model = FastLanguageModel.get_peft_model(
|
|
model,
|
|
r=r,
|
|
target_modules=[
|
|
"q_proj", "k_proj", "v_proj", "o_proj",
|
|
"gate_proj", "up_proj", "down_proj",
|
|
],
|
|
lora_alpha=alpha,
|
|
lora_dropout=dropout,
|
|
bias="none",
|
|
use_gradient_checkpointing="unsloth",
|
|
random_state=42,
|
|
)
|
|
|
|
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
|
total = sum(p.numel() for p in model.parameters())
|
|
print(f" Entraînables : {trainable / 1e6:.1f}M / {total / 1e9:.1f}B ({100 * trainable / total:.2f}%)")
|
|
return model
|
|
|
|
|
|
def load_dataset(train_path, eval_path):
|
|
from datasets import Dataset
|
|
|
|
def load_jsonl(path):
|
|
examples = []
|
|
with open(path) as f:
|
|
for line in f:
|
|
examples.append(json.loads(line.strip()))
|
|
return examples
|
|
|
|
train_ds = Dataset.from_list(load_jsonl(train_path))
|
|
eval_ds = Dataset.from_list(load_jsonl(eval_path))
|
|
|
|
print(f"\nDataset : {len(train_ds)} train + {len(eval_ds)} eval")
|
|
return train_ds, eval_ds
|
|
|
|
|
|
def format_chat(example, tokenizer):
|
|
text = tokenizer.apply_chat_template(
|
|
example["messages"],
|
|
tokenize=False,
|
|
add_generation_prompt=False,
|
|
)
|
|
return {"text": text}
|
|
|
|
|
|
def train(model, tokenizer, train_ds, eval_ds, args):
|
|
from trl import SFTTrainer, SFTConfig
|
|
|
|
print(f"\nConfig entraînement :")
|
|
print(f" Epochs : {args.epochs}")
|
|
print(f" LR : {args.lr}")
|
|
print(f" Batch : {args.batch} x grad_accum={args.grad_accum} = {args.batch * args.grad_accum}")
|
|
print(f" Max seq length : {args.max_seq_length}")
|
|
|
|
train_ds = train_ds.map(lambda x: format_chat(x, tokenizer), num_proc=4)
|
|
eval_ds = eval_ds.map(lambda x: format_chat(x, tokenizer), num_proc=4)
|
|
|
|
output_dir = OUTPUT / "pmsi-lora-checkpoints"
|
|
|
|
# Wandb optionnel
|
|
report = "none"
|
|
callbacks = []
|
|
try:
|
|
import wandb
|
|
wandb.init(project="pmsi-coder", name=f"v2-runpod-{args.epochs}ep-seq{args.max_seq_length}")
|
|
report = "wandb"
|
|
print(" Tracking : wandb")
|
|
except ImportError:
|
|
print(" Tracking : none (pip install wandb pour activer)")
|
|
|
|
training_args = SFTConfig(
|
|
output_dir=str(output_dir),
|
|
num_train_epochs=args.epochs,
|
|
per_device_train_batch_size=args.batch,
|
|
per_device_eval_batch_size=args.batch,
|
|
gradient_accumulation_steps=args.grad_accum,
|
|
learning_rate=args.lr,
|
|
weight_decay=0.01,
|
|
warmup_ratio=0.05,
|
|
lr_scheduler_type="cosine",
|
|
logging_steps=10,
|
|
eval_strategy="steps",
|
|
eval_steps=500,
|
|
save_strategy="steps",
|
|
save_steps=500,
|
|
save_total_limit=3,
|
|
fp16=False,
|
|
bf16=True,
|
|
max_seq_length=args.max_seq_length,
|
|
dataset_text_field="text",
|
|
seed=42,
|
|
report_to=report,
|
|
)
|
|
|
|
trainer = SFTTrainer(
|
|
model=model,
|
|
tokenizer=tokenizer,
|
|
train_dataset=train_ds,
|
|
eval_dataset=eval_ds,
|
|
args=training_args,
|
|
callbacks=callbacks,
|
|
)
|
|
|
|
total_steps = len(train_ds) * args.epochs // (args.batch * args.grad_accum)
|
|
print(f"\n Steps estimés : ~{total_steps}")
|
|
print(f" Démarrage...")
|
|
|
|
if args.resume:
|
|
trainer.train(resume_from_checkpoint=True)
|
|
else:
|
|
trainer.train()
|
|
|
|
final_dir = OUTPUT / "pmsi-lora-final"
|
|
model.save_pretrained(str(final_dir))
|
|
tokenizer.save_pretrained(str(final_dir))
|
|
print(f"\nLoRA sauvegardé : {final_dir}")
|
|
|
|
return trainer, final_dir
|
|
|
|
|
|
def export_merged_hf(model, tokenizer):
|
|
"""Sauvegarder le modèle mergé en 16-bit (HF format) pour conversion GGUF ultérieure."""
|
|
print(f"\nExport modèle mergé (16-bit HF)...")
|
|
|
|
merged_dir = OUTPUT / "pmsi-merged-hf"
|
|
merged_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
model.save_pretrained_merged(
|
|
str(merged_dir),
|
|
tokenizer,
|
|
save_method="merged_16bit",
|
|
)
|
|
|
|
size_gb = sum(f.stat().st_size for f in merged_dir.glob("*.safetensors")) / 1024**3
|
|
print(f" Modèle mergé : {merged_dir} ({size_gb:.1f} Go)")
|
|
print(f"\n Pour convertir en GGUF :")
|
|
print(f" python llama.cpp/convert_hf_to_gguf.py {merged_dir} --outfile pmsi-v2-q8.gguf --outtype q8_0")
|
|
print(f" llama-quantize pmsi-v2-q8.gguf pmsi-v2-q4km.gguf Q4_K_M")
|
|
return merged_dir
|
|
|
|
|
|
def export_gguf(model, tokenizer, final_dir, quantization="q4_k_m"):
|
|
"""Export GGUF via Unsloth (peut échouer — fallback sur export_merged_hf)."""
|
|
print(f"\nExport GGUF ({quantization})...")
|
|
|
|
gguf_dir = OUTPUT / "pmsi-gguf"
|
|
gguf_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
try:
|
|
model.save_pretrained_gguf(
|
|
str(gguf_dir),
|
|
tokenizer,
|
|
quantization_method=quantization,
|
|
)
|
|
|
|
gguf_files = list(gguf_dir.glob("*.gguf"))
|
|
if gguf_files:
|
|
gguf_path = gguf_files[0]
|
|
size_gb = gguf_path.stat().st_size / 1024**3
|
|
print(f" GGUF : {gguf_path.name} ({size_gb:.1f} Go)")
|
|
|
|
modelfile_path = gguf_dir / "Modelfile"
|
|
with open(modelfile_path, "w") as f:
|
|
f.write(f"FROM {gguf_path.name}\n\n")
|
|
f.write("PARAMETER temperature 0.3\n")
|
|
f.write("PARAMETER top_p 0.9\n")
|
|
f.write("PARAMETER num_ctx 8192\n")
|
|
f.write('PARAMETER stop "<end_of_turn>"\n')
|
|
f.write('PARAMETER stop "<eos>"\n')
|
|
|
|
print(f" Modelfile créé")
|
|
except Exception as e:
|
|
print(f" GGUF export échoué : {e}")
|
|
print(f" Fallback : export HF mergé...")
|
|
export_merged_hf(model, tokenizer)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Fine-tuning QLoRA RunPod")
|
|
|
|
parser.add_argument("--model", default="unsloth/gemma-3-12b-it-bnb-4bit")
|
|
parser.add_argument("--max-seq-length", type=int, default=2048)
|
|
|
|
parser.add_argument("--lora-r", type=int, default=32)
|
|
parser.add_argument("--lora-alpha", type=int, default=64)
|
|
parser.add_argument("--lora-dropout", type=float, default=0.0)
|
|
|
|
parser.add_argument("--epochs", type=int, default=3)
|
|
parser.add_argument("--lr", type=float, default=2e-4)
|
|
parser.add_argument("--batch", type=int, default=0, help="0=auto-detect")
|
|
parser.add_argument("--grad-accum", type=int, default=4)
|
|
|
|
parser.add_argument("--resume", action="store_true")
|
|
parser.add_argument("--export-gguf", action="store_true")
|
|
parser.add_argument("--gguf-quant", default="q4_k_m")
|
|
|
|
args = parser.parse_args()
|
|
|
|
train_path, eval_path, suggested_batch = check_prerequisites()
|
|
|
|
if args.batch == 0:
|
|
args.batch = suggested_batch
|
|
print(f"Batch auto-détecté : {args.batch}")
|
|
|
|
model, tokenizer = load_model(args.model, args.max_seq_length)
|
|
model = attach_lora(model, r=args.lora_r, alpha=args.lora_alpha, dropout=args.lora_dropout)
|
|
train_ds, eval_ds = load_dataset(train_path, eval_path)
|
|
trainer, final_dir = train(model, tokenizer, train_ds, eval_ds, args)
|
|
|
|
if args.export_gguf:
|
|
export_gguf(model, tokenizer, final_dir, args.gguf_quant)
|
|
|
|
print("\n" + "=" * 50)
|
|
print("Fine-tuning terminé !")
|
|
print(f" LoRA : {final_dir}")
|
|
if args.export_gguf:
|
|
print(f" GGUF : {OUTPUT / 'pmsi-gguf'}")
|
|
print("=" * 50)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|