File size: 5,017 Bytes
c63bc31 | 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 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | """Fine-tuning entry point for CORTEX AI.
This script is written to run on GPU hardware. It is NOT executed in this
repository: the sandbox has 4 CPU cores, ~15 GiB of RAM and no accelerator, and
the checkpoint is 893 GB. Treat it as the documented path, not a verified run.
Usage::
python -m cortex_ai.training.train --config train_config.json
python -m cortex_ai.training.train --dataset data/cortex_sft.jsonl --epochs 1
Requires: torch, transformers, peft, datasets
"""
from __future__ import annotations
import argparse
import sys
from pathlib import Path
from .config import TrainConfig
from .dataset import dataset_stats, iter_encoded, load_jsonl
def build_arg_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description="Affinage de CORTEX AI")
parser.add_argument("--config", type=str, default="", help="fichier JSON de configuration")
parser.add_argument("--dataset", type=str, default="", help="dataset JSONL")
parser.add_argument("--output-dir", type=str, default="", help="dossier de sortie")
parser.add_argument("--epochs", type=int, default=0)
parser.add_argument("--dry-run", action="store_true", help="valider les donnees sans entrainer")
return parser
def main(argv: list[str] | None = None) -> int:
args = build_arg_parser().parse_args(argv)
config = TrainConfig.load(args.config) if args.config else TrainConfig()
if args.dataset:
config.dataset_path = args.dataset
if args.output_dir:
config.output_dir = args.output_dir
if args.epochs:
config.epochs = args.epochs
examples = load_jsonl(config.dataset_path)
stats = dataset_stats(examples)
print("=== Donnees ===")
for key, value in stats.items():
print(f" {key:20s} {value}")
# Encoding every example catches format problems before a GPU is booked.
encoded = list(iter_encoded(examples, config.thinking_mode))
print(f" {'exemples encodes':20s} {len(encoded)}")
print(f" {'prompt le plus long':20s} {max(len(e) for e in encoded)} caracteres")
if args.dry_run:
print("\n--dry-run : validation terminee, aucun entrainement lance.")
return 0
try:
import torch
from datasets import Dataset
from peft import LoraConfig, get_peft_model
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
DataCollatorForLanguageModeling,
Trainer,
TrainingArguments,
)
except ImportError as exc:
print(
f"\nDependance manquante : {exc.name}.\n"
"Installez torch, transformers, peft et datasets, puis relancez.\n"
"Sur cette machine (sans GPU, 893 Go de checkpoint), l'entrainement "
"n'est pas executable : utilisez --dry-run pour valider les donnees.",
file=sys.stderr,
)
return 2
if not torch.cuda.is_available():
print(
"\nAucun GPU detecte. L'affinage de ce checkpoint exige plusieurs "
"accelerateurs. Utilisez --dry-run pour valider les donnees.",
file=sys.stderr,
)
return 2
print("\n=== Chargement du modele ===")
tokenizer = AutoTokenizer.from_pretrained(config.model_id)
model = AutoModelForCausalLM.from_pretrained(
config.model_id, torch_dtype=torch.bfloat16, device_map="auto"
)
if config.method == "lora":
model = get_peft_model(
model,
LoraConfig(
r=config.lora.r,
lora_alpha=config.lora.alpha,
lora_dropout=config.lora.dropout,
target_modules=config.lora.target_modules,
task_type="CAUSAL_LM",
),
)
model.print_trainable_parameters()
dataset = Dataset.from_dict({"text": encoded})
def tokenize(batch):
return tokenizer(batch["text"], truncation=True, max_length=config.max_seq_length)
dataset = dataset.map(tokenize, batched=True, remove_columns=["text"])
training_args = TrainingArguments(
output_dir=config.output_dir,
num_train_epochs=config.epochs,
per_device_train_batch_size=config.batch_size,
gradient_accumulation_steps=config.gradient_accumulation_steps,
learning_rate=config.learning_rate,
warmup_ratio=config.warmup_ratio,
weight_decay=config.weight_decay,
logging_steps=config.logging_steps,
save_steps=config.save_steps,
bf16=True,
seed=config.seed,
report_to=[],
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=dataset,
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
trainer.train()
trainer.save_model(config.output_dir)
tokenizer.save_pretrained(config.output_dir)
print(f"\nModele enregistre dans {config.output_dir}")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|