S-Dreamer/CodeCraftLab
0
1"""2Fine-tuning pipeline with structured logging and eval hooks.3 4Pipeline stages:5 1. Preflight validation — config, GPU, disk, token6 2. Dataset preparation — load, tokenize, split7 3. Model initialisation — base model + LoRA adapters8 4. Training — Trainer with custom callbacks9 5. Evaluation — post-training metric suite10 6. Checkpoint export — save + optional HF Hub push11 12Each stage emits structured log events. Eval hooks are composable and13run both during training (via TrainerCallback) and post-training.14"""15 16from __future__ import annotations17 18import json19import os20import shutil21import time22from dataclasses import dataclass, field23from pathlib import Path24from typing import Any25 26import structlog27import torch28from datasets import Dataset, DatasetDict, load_dataset29from peft import LoraConfig, TaskType, get_peft_model30from transformers import (31 AutoModelForCausalLM,32 AutoTokenizer,33 DataCollatorForLanguageModeling,34 PreTrainedModel,35 PreTrainedTokenizerBase,36 Trainer,37 TrainerCallback,38 TrainerControl,39 TrainerState,40 TrainingArguments,41)42 43from training.config import EvalMetric, EvalStrategy, TrainingJobConfig44from training.evaluators import (45 BleuEvaluator,46 ExecutionAccuracyEvaluator,47 ExactMatchEvaluator,48 PassAtKEvaluator,49)50 51log = structlog.get_logger(__name__)52 53 54# ---------------------------------------------------------------------------55# Eval result container56# ---------------------------------------------------------------------------57@dataclass58class EvalResults:59 job_name: str60 epoch: float61 step: int62 metrics: dict[str, float] = field(default_factory=dict)63 errors: list[str] = field(default_factory=list)64 duration_seconds: float = 0.065 66 def log(self, bound_log: structlog.BoundLogger) -> None:67 bound_log.info(68 "eval.completed",69 epoch=self.epoch,70 step=self.step,71 duration_seconds=round(self.duration_seconds, 2),72 **self.metrics,73 )74 for error in self.errors:75 bound_log.warning("eval.error", message=error)76 77 def to_dict(self) -> dict[str, Any]:78 return {79 "job_name": self.job_name,80 "epoch": self.epoch,81 "step": self.step,82 "metrics": self.metrics,83 "errors": self.errors,84 "duration_seconds": self.duration_seconds,85 }86 87 88# ---------------------------------------------------------------------------89# Eval hook registry90# ---------------------------------------------------------------------------91class EvalHookRunner:92 """93 Runs the configured evaluation metrics against a model + dataset.94 95 Evaluators are resolved from the job config at construction time.96 Each evaluator is independent; failures in one do not abort others.97 """98 99 def __init__(self, config: TrainingJobConfig, tokenizer: PreTrainedTokenizerBase) -> None:100 self._config = config101 self._tokenizer = tokenizer102 self._evaluators = self._build_evaluators()103 self._log = log.bind(job=config.job_name)104 105 def _build_evaluators(self) -> dict[EvalMetric, Any]:106 evals: dict[EvalMetric, Any] = {}107 eval_cfg = self._config.evaluation108 for metric in eval_cfg.metrics:109 match metric:110 case EvalMetric.PASS_AT_1:111 evals[metric] = PassAtKEvaluator(k=1, n=eval_cfg.num_samples_per_problem)112 case EvalMetric.PASS_AT_10:113 evals[metric] = PassAtKEvaluator(k=10, n=eval_cfg.num_samples_per_problem)114 case EvalMetric.BLEU:115 evals[metric] = BleuEvaluator()116 case EvalMetric.EXECUTION_ACCURACY:117 evals[metric] = ExecutionAccuracyEvaluator(118 timeout=self._config.evaluation.timeout_seconds119 )120 case EvalMetric.EXACT_MATCH:121 evals[metric] = ExactMatchEvaluator()122 return evals123 124 def run(125 self,126 model: PreTrainedModel,127 eval_dataset: Dataset,128 epoch: float,129 step: int,130 ) -> EvalResults:131 start = time.perf_counter()132 results = EvalResults(job_name=self._config.job_name, epoch=epoch, step=step)133 134 model.eval()135 with torch.no_grad():136 for metric, evaluator in self._evaluators.items():137 try:138 score = evaluator.evaluate(139 model=model,140 tokenizer=self._tokenizer,141 dataset=eval_dataset,142 )143 results.metrics[metric.value] = round(score, 4)144 self._log.info("eval.metric", metric=metric.value, score=score)145 except Exception as exc: # noqa: BLE001146 msg = f"{metric.value}: {exc}"147 results.errors.append(msg)148 self._log.warning("eval.metric_failed", metric=metric.value, error=str(exc))149 150 results.duration_seconds = time.perf_counter() - start151 results.log(self._log)152 return results153 154 155# ---------------------------------------------------------------------------156# Custom training callback157# ---------------------------------------------------------------------------158class CodeCraftLabCallback(TrainerCallback):159 """160 Injects structured logging and eval hooks into the HF Trainer loop.161 """162 163 def __init__(164 self,165 hook_runner: EvalHookRunner,166 eval_dataset: Dataset,167 results_path: Path,168 ) -> None:169 self._runner = hook_runner170 self._eval_dataset = eval_dataset171 self._results_path = results_path172 self._all_results: list[dict[str, Any]] = []173 self._log = log174 175 def on_epoch_end(176 self,177 args: TrainingArguments,178 state: TrainerState,179 control: TrainerControl,180 model: PreTrainedModel,181 **kwargs: Any,182 ) -> TrainerControl:183 self._log.info(184 "training.epoch_end",185 epoch=state.epoch,186 step=state.global_step,187 loss=state.log_history[-1].get("loss") if state.log_history else None,188 )189 results = self._runner.run(190 model=model,191 eval_dataset=self._eval_dataset,192 epoch=state.epoch or 0.0,193 step=state.global_step,194 )195 self._all_results.append(results.to_dict())196 self._persist_results()197 return control198 199 def on_log(200 self,201 args: TrainingArguments,202 state: TrainerState,203 control: TrainerControl,204 logs: dict[str, float],205 **kwargs: Any,206 ) -> TrainerControl:207 self._log.info("training.log", step=state.global_step, **logs)208 return control209 210 def on_train_end(211 self,212 args: TrainingArguments,213 state: TrainerState,214 control: TrainerControl,215 **kwargs: Any,216 ) -> TrainerControl:217 self._log.info(218 "training.completed",219 total_steps=state.global_step,220 total_flos=state.total_flos,221 )222 return control223 224 def _persist_results(self) -> None:225 self._results_path.write_text(226 json.dumps(self._all_results, indent=2), encoding="utf-8"227 )228 229 230# ---------------------------------------------------------------------------231# Pipeline232# ---------------------------------------------------------------------------233class FineTuningPipeline:234 """235 Orchestrates the full fine-tuning lifecycle.236 237 Usage:238 config = TrainingJobConfig.model_validate(raw_dict)239 pipeline = FineTuningPipeline(config)240 pipeline.run()241 """242 243 def __init__(self, config: TrainingJobConfig) -> None:244 self._config = config245 self._log = log.bind(job=config.job_name, model=config.base_model)246 self._output_dir = Path(config.checkpoint.output_dir) / config.job_name247 248 # ------------------------------------------------------------------249 # Public entry point250 # ------------------------------------------------------------------251 def run(self) -> Path:252 """Execute all pipeline stages. Returns the final checkpoint path."""253 self._log.info("pipeline.started")254 self._preflight()255 datasets = self._prepare_datasets()256 model, tokenizer = self._load_model()257 self._train(model, tokenizer, datasets)258 final_path = self._export(model, tokenizer)259 self._log.info("pipeline.finished", output=str(final_path))260 return final_path261 262 # ------------------------------------------------------------------263 # Stage 1: Preflight264 # ------------------------------------------------------------------265 def _preflight(self) -> None:266 self._log.info("pipeline.preflight")267 268 # Validate config (already done at submission, but be defensive)269 self._config.model_validate(self._config.model_dump())270 271 # GPU check272 if torch.cuda.is_available():273 device_name = torch.cuda.get_device_name(0)274 vram_gb = torch.cuda.get_device_properties(0).total_memory / 1e9275 self._log.info("preflight.gpu", device=device_name, vram_gb=round(vram_gb, 1))276 else:277 self._log.warning("preflight.no_gpu", message="Training on CPU — will be slow")278 279 # Disk space (rough check — 20 GB minimum)280 free_gb = shutil.disk_usage(self._output_dir.parent).free / 1e9281 if free_gb < 20:282 self._log.warning("preflight.disk_low", free_gb=round(free_gb, 1))283 284 # HF token if pushing285 if self._config.hub.push_to_hub and not os.environ.get("HF_TOKEN"):286 raise EnvironmentError("HF_TOKEN is required when hub.push_to_hub=true")287 288 self._output_dir.mkdir(parents=True, exist_ok=True)289 self._log.info("preflight.passed")290 291 # ------------------------------------------------------------------292 # Stage 2: Dataset preparation293 # ------------------------------------------------------------------294 def _prepare_datasets(self) -> DatasetDict:295 self._log.info("pipeline.dataset_prep")296 ds_cfg = self._config.dataset297 298 # Load — support both HF Hub paths and internal dataset IDs299 raw: Dataset300 if ds_cfg.dataset_id.startswith("ds_"):301 # Internal dataset — load from local store302 raw = Dataset.load_from_disk(f"./data/{ds_cfg.dataset_id}")303 else:304 raw = load_dataset(ds_cfg.dataset_id, split="train") # type: ignore[assignment]305 306 if ds_cfg.max_samples:307 raw = raw.select(range(min(ds_cfg.max_samples, len(raw))))308 309 if ds_cfg.shuffle:310 raw = raw.shuffle(seed=ds_cfg.shuffle_seed)311 312 n_train = int(len(raw) * ds_cfg.split_ratio)313 splits = DatasetDict(314 {315 "train": raw.select(range(n_train)),316 "eval": raw.select(range(n_train, len(raw))),317 }318 )319 self._log.info(320 "dataset.prepared",321 train_size=len(splits["train"]),322 eval_size=len(splits["eval"]),323 column=ds_cfg.text_column,324 )325 return splits326 327 # ------------------------------------------------------------------328 # Stage 3: Model initialisation329 # ------------------------------------------------------------------330 def _load_model(self) -> tuple[PreTrainedModel, PreTrainedTokenizerBase]:331 self._log.info("pipeline.model_load")332 hp = self._config.training333 334 dtype_map = {335 "fp32": torch.float32,336 "fp16": torch.float16,337 "bf16": torch.bfloat16,338 }339 torch_dtype = dtype_map.get(hp.precision.value, torch.bfloat16)340 341 tokenizer = AutoTokenizer.from_pretrained(self._config.base_model)342 if tokenizer.pad_token is None:343 tokenizer.pad_token = tokenizer.eos_token344 345 model = AutoModelForCausalLM.from_pretrained(346 self._config.base_model,347 torch_dtype=torch_dtype,348 device_map="auto" if torch.cuda.is_available() else "cpu",349 )350 351 if self._config.lora and self._config.lora.enabled:352 lora_cfg = self._config.lora353 peft_config = LoraConfig(354 task_type=TaskType.CAUSAL_LM,355 r=lora_cfg.r,356 lora_alpha=lora_cfg.alpha,357 lora_dropout=lora_cfg.dropout,358 target_modules=lora_cfg.target_modules,359 bias=lora_cfg.bias, # type: ignore[arg-type]360 )361 model = get_peft_model(model, peft_config)362 trainable, total = model.get_nb_trainable_parameters()363 self._log.info(364 "model.lora_applied",365 trainable_params=trainable,366 total_params=total,367 trainable_pct=round(100 * trainable / total, 2),368 )369 else:370 self._log.info("model.full_finetune")371 372 return model, tokenizer # type: ignore[return-value]373 374 # ------------------------------------------------------------------375 # Stage 4: Training376 # ------------------------------------------------------------------377 def _train(378 self,379 model: PreTrainedModel,380 tokenizer: PreTrainedTokenizerBase,381 datasets: DatasetDict,382 ) -> None:383 self._log.info("pipeline.training_start")384 hp = self._config.training385 ckpt = self._config.checkpoint386 eval_cfg = self._config.evaluation387 388 def tokenize(examples: dict[str, list[str]]) -> dict[str, Any]:389 return tokenizer(390 examples[self._config.dataset.text_column],391 truncation=True,392 max_length=hp.max_seq_length,393 padding=False,394 )395 396 tokenized = datasets.map(tokenize, batched=True, remove_columns=datasets["train"].column_names)397 398 training_args = TrainingArguments(399 output_dir=str(self._output_dir),400 num_train_epochs=hp.num_epochs,401 per_device_train_batch_size=hp.batch_size,402 per_device_eval_batch_size=hp.batch_size,403 gradient_accumulation_steps=hp.gradient_accumulation_steps,404 learning_rate=hp.learning_rate,405 weight_decay=hp.weight_decay,406 warmup_ratio=hp.warmup_ratio,407 max_grad_norm=hp.max_grad_norm,408 optim=hp.optimizer.value,409 lr_scheduler_type=hp.lr_scheduler,410 fp16=hp.precision.value == "fp16",411 bf16=hp.precision.value == "bf16",412 evaluation_strategy=eval_cfg.strategy.value,413 eval_steps=eval_cfg.eval_steps,414 save_strategy=ckpt.save_strategy.value,415 save_steps=ckpt.save_steps,416 save_total_limit=ckpt.save_total_limit,417 load_best_model_at_end=eval_cfg.load_best_model_at_end,418 metric_for_best_model=eval_cfg.metric_for_best_model.value,419 greater_is_better=eval_cfg.greater_is_better,420 seed=hp.seed,421 dataloader_num_workers=hp.dataloader_num_workers,422 report_to="none", # structlog handles all logging423 logging_steps=10,424 resume_from_checkpoint=ckpt.resume_from_checkpoint,425 push_to_hub=False, # push handled separately in export stage426 )427 428 hook_runner = EvalHookRunner(self._config, tokenizer)429 results_path = self._output_dir / "eval_results.json"430 callback = CodeCraftLabCallback(431 hook_runner=hook_runner,432 eval_dataset=datasets["eval"],433 results_path=results_path,434 )435 436 trainer = Trainer(437 model=model,438 args=training_args,439 train_dataset=tokenized["train"],440 eval_dataset=tokenized["eval"],441 data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),442 callbacks=[callback],443 )444 445 trainer.train(resume_from_checkpoint=ckpt.resume_from_checkpoint)446 447 # ------------------------------------------------------------------448 # Stage 5: Export + Hub push449 # ------------------------------------------------------------------450 def _export(self, model: PreTrainedModel, tokenizer: PreTrainedTokenizerBase) -> Path:451 self._log.info("pipeline.export")452 final_path = self._output_dir / "final"453 model.save_pretrained(str(final_path))454 tokenizer.save_pretrained(str(final_path))455 self._log.info("model.saved", path=str(final_path))456 457 hub_cfg = self._config.hub458 if hub_cfg.push_to_hub and hub_cfg.repo_id:459 self._log.info("hub.pushing", repo_id=hub_cfg.repo_id)460 model.push_to_hub(hub_cfg.repo_id, private=hub_cfg.private)461 tokenizer.push_to_hub(hub_cfg.repo_id, private=hub_cfg.private)462 self._log.info("hub.pushed", repo_id=hub_cfg.repo_id)463 464 return final_path465