Spaces:
Running
Running
| # type: ignore | |
| """ | |
| modules/gpu_trainer.py | |
| ================================================================================ | |
| TREINO QLoRA GPU / LoRA CPU — Space AKIRA | |
| ================================================================================ | |
| - GPU: base modelo 4-bit (bitsandbytes nf4) + LoRA (PEFT r=16) | |
| - CPU: base em FP32 + LoRA (PEFT r=16), apenas com RAM livre suficiente | |
| - Dataset: lista [{"text": ...}] (formato já produzido por | |
| ModelTrainer.prepare_dataset() em modules/treinamento_modelo.py) | |
| - Progresso persistido em <state_path> (JSON) → painel Gradio faz polling | |
| - Artefactos: adapter LoRA + tokenizer + training_report.json | |
| Env: | |
| AKIRA_TRAIN_MODEL base model (default AKIRA_LOCAL_MODEL) | |
| AKIRA_TRAIN_MAX_STEPS GPU default 100 | |
| AKIRA_TRAIN_BATCH_SIZE default 1 (L4 24GB: 1 + grad accum) | |
| AKIRA_TRAIN_GRAD_ACC default 4 | |
| AKIRA_TRAIN_LR default 2e-4 | |
| AKIRA_TRAIN_SEQ_LEN default 1024 | |
| AKIRA_TRAIN_OUT default /data/models/akira-tuned (ou ./models/...) | |
| AKIRA_CPU_TRAIN_MIN_RAM_GB default 48 | |
| AKIRA_CPU_TRAIN_MAX_STEPS default 10 | |
| AKIRA_CPU_TRAIN_SEQ_LEN default 256 | |
| Promotion only after finite, improved held-out validation loss | |
| ================================================================================ | |
| """ | |
| import os | |
| import json | |
| import time | |
| import threading | |
| import traceback | |
| import hashlib | |
| import math | |
| import random | |
| import sys | |
| import ctypes | |
| from typing import Optional, List, Dict, Any | |
| try: | |
| from loguru import logger # type: ignore | |
| except Exception: | |
| class _D: | |
| def info(self, *a, **k): pass | |
| def success(self, *a, **k): pass | |
| def warning(self, *a, **k): pass | |
| def error(self, *a, **k): pass | |
| def debug(self, *a, **k): pass | |
| logger = _D() # type: ignore | |
| DEFAULT_BASE_MODEL = os.getenv( | |
| "AKIRA_TRAIN_MODEL", | |
| os.getenv("AKIRA_LOCAL_MODEL", "Orenguteng/Llama-3.1-8B-Lexi-Uncensored-V2"), | |
| ) | |
| _train_lock = threading.Lock() # um treino de cada vez | |
| # ====================================================================== | |
| # Helpers de path / estado | |
| # ====================================================================== | |
| def _data_root() -> str: | |
| """Diretório persistente da Space (/data) quando existe.""" | |
| if os.path.isdir("/data") and os.access("/data", os.W_OK): | |
| return "/data" | |
| return os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "models") | |
| def default_output_dir() -> str: | |
| out = os.getenv("AKIRA_TRAIN_OUT", "").strip() | |
| if out: | |
| return out | |
| if _data_root() == "/data": | |
| return "/data/models/akira-tuned" | |
| return os.path.join(_data_root(), "akira-tuned") | |
| def fingerprint_dataset(dataset: List[Dict[str, Any]]) -> str: | |
| normalized = [ | |
| {"text": row.get("text", ""), "messages": row.get("messages", [])} | |
| for row in (dataset or []) | |
| if row and (row.get("text") or row.get("messages")) | |
| ] | |
| payload = json.dumps(normalized, ensure_ascii=False, separators=(",", ":")) | |
| return hashlib.sha256(payload.encode("utf-8")).hexdigest() | |
| def resolve_active_adapter(output_dir: Optional[str] = None, | |
| expected_base_model: Optional[str] = None): | |
| """Resolve only a validated adapter trained from the inference base model.""" | |
| root = os.path.abspath(output_dir or default_output_dir()) | |
| manifest_path = os.path.join(root, "active_adapter.json") | |
| if not os.path.isfile(manifest_path): | |
| return None, None | |
| try: | |
| with open(manifest_path, "r", encoding="utf-8") as f: | |
| manifest = json.load(f) | |
| relative_path = manifest.get("adapter_path") | |
| base_model = manifest.get("base_model") | |
| if not relative_path or not base_model: | |
| raise ValueError("manifesto sem adapter_path/base_model") | |
| adapter_path = os.path.abspath(os.path.join(root, relative_path)) | |
| if os.path.commonpath([root, adapter_path]) != root: | |
| raise ValueError("caminho do adaptador fora do diretório de treino") | |
| if expected_base_model and base_model != expected_base_model: | |
| raise ValueError( | |
| f"adaptador treinado em {base_model}, modelo de inferência é {expected_base_model}" | |
| ) | |
| if not os.path.isfile(os.path.join(adapter_path, "adapter_config.json")): | |
| raise ValueError("adapter_config.json não encontrado") | |
| with open(os.path.join(adapter_path, "adapter_config.json"), "r", encoding="utf-8") as f: | |
| adapter_config = json.load(f) | |
| adapter_base = adapter_config.get("base_model_name_or_path") | |
| if expected_base_model and adapter_base and adapter_base != expected_base_model: | |
| raise ValueError( | |
| f"config do adapter aponta para {adapter_base}, não {expected_base_model}" | |
| ) | |
| weights_exist = any( | |
| os.path.isfile(os.path.join(adapter_path, filename)) | |
| for filename in ("adapter_model.safetensors", "adapter_model.bin") | |
| ) | |
| if not weights_exist: | |
| raise ValueError("pesos do adaptador não encontrados") | |
| return adapter_path, None | |
| except Exception as e: | |
| return None, str(e) | |
| def training_state_path() -> str: | |
| if _data_root() == "/data": | |
| return "/data/training_state.json" | |
| return os.path.join(_data_root(), "training_state.json") | |
| def write_training_state(state: Dict[str, Any]) -> None: | |
| try: | |
| path = training_state_path() | |
| os.makedirs(os.path.dirname(path), exist_ok=True) | |
| tmp = path + ".tmp" | |
| with open(tmp, "w", encoding="utf-8") as f: | |
| json.dump(state, f, ensure_ascii=False, indent=2) | |
| os.replace(tmp, path) | |
| except Exception as e: | |
| logger.warning(f"[GPU-TRAIN] Falha ao gravar estado: {e}") | |
| def read_training_state() -> Dict[str, Any]: | |
| try: | |
| with open(training_state_path(), "r", encoding="utf-8") as f: | |
| return json.load(f) | |
| except Exception: | |
| return {"status": "idle"} | |
| def gpu_info() -> Dict[str, Any]: | |
| """Info da GPU para a aba Status do app.py.""" | |
| info: Dict[str, Any] = {"cuda_available": False, "device": None, "torch": None} | |
| try: | |
| import torch # type: ignore | |
| info["torch"] = torch.__version__ | |
| try: | |
| if torch.cuda.is_available(): | |
| # ZeroGPU: is_available() é emulado fora dos decorators → prova real | |
| torch.zeros(1, device="cuda") | |
| info["cuda_available"] = True | |
| except Exception as e: | |
| info["zerogpu_fora_do_decorator"] = str(e)[:180] | |
| if info["cuda_available"]: | |
| info["device"] = torch.cuda.get_device_name(0) | |
| props = torch.cuda.get_device_properties(0) | |
| info["vram_total_gb"] = round(props.total_mem / (1024 ** 3), 1) | |
| info["vram_allocated_gb"] = round(torch.cuda.memory_allocated(0) / (1024 ** 3), 2) | |
| info["vram_reserved_gb"] = round(torch.cuda.memory_reserved(0) / (1024 ** 3), 2) | |
| try: | |
| info["cuda_version"] = torch.version.cuda | |
| except Exception: | |
| pass | |
| except Exception as e: | |
| info["error"] = str(e) | |
| return info | |
| def is_gpu_training_available() -> bool: | |
| """True se CUDA real + transformers + peft + trl + datasets + bitsandbytes estiverem prontos. | |
| Em ZeroGPU, `torch.cuda.is_available()` devolve True FORA dos decorators | |
| @spaces.GPU mas qualquer operação CUDA levanta erro → aqui prova-se com | |
| uma operação mínima, para o painel dizer logo "precisa de L4 dedicada". | |
| """ | |
| try: | |
| import torch # type: ignore | |
| if not torch.cuda.is_available(): | |
| return False | |
| torch.zeros(1, device="cuda") | |
| except Exception: | |
| return False | |
| return _training_dependencies_available(require_bitsandbytes=True) | |
| def _training_dependencies_available(require_bitsandbytes: bool = False) -> bool: | |
| try: | |
| import torch # noqa: F401 # type: ignore | |
| import transformers # noqa: F401 # type: ignore | |
| import peft # noqa: F401 # type: ignore | |
| import trl # noqa: F401 # type: ignore | |
| import datasets # noqa: F401 # type: ignore | |
| if require_bitsandbytes: | |
| import bitsandbytes # noqa: F401 # type: ignore | |
| return True | |
| except Exception: | |
| return False | |
| def _available_memory_bytes() -> Optional[int]: | |
| """Available system RAM without adding a mandatory platform dependency.""" | |
| try: | |
| import psutil # type: ignore | |
| return int(psutil.virtual_memory().available) | |
| except ImportError: | |
| pass | |
| try: | |
| if sys.platform == "win32": | |
| class _MemoryStatus(ctypes.Structure): | |
| _fields_ = [ | |
| ("dwLength", ctypes.c_ulong), | |
| ("dwMemoryLoad", ctypes.c_ulong), | |
| ("ullTotalPhys", ctypes.c_ulonglong), | |
| ("ullAvailPhys", ctypes.c_ulonglong), | |
| ("ullTotalPageFile", ctypes.c_ulonglong), | |
| ("ullAvailPageFile", ctypes.c_ulonglong), | |
| ("ullTotalVirtual", ctypes.c_ulonglong), | |
| ("ullAvailVirtual", ctypes.c_ulonglong), | |
| ("ullAvailExtendedVirtual", ctypes.c_ulonglong), | |
| ] | |
| status = _MemoryStatus() | |
| status.dwLength = ctypes.sizeof(status) | |
| if ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(status)): | |
| return int(status.ullAvailPhys) | |
| if hasattr(os, "sysconf"): | |
| return int(os.sysconf("SC_AVPHYS_PAGES") * os.sysconf("SC_PAGE_SIZE")) | |
| except (AttributeError, OSError, ValueError): | |
| return None | |
| return None | |
| def _select_training_device(cuda_qlora_available: bool, | |
| available_ram_bytes: Optional[int], | |
| min_cpu_ram_bytes: int) -> Optional[str]: | |
| if cuda_qlora_available: | |
| return "cuda" | |
| if available_ram_bytes is not None and available_ram_bytes >= min_cpu_ram_bytes: | |
| return "cpu" | |
| return None | |
| def get_training_device() -> Optional[str]: | |
| """Prefer real CUDA QLoRA; otherwise allow bounded full-precision CPU LoRA.""" | |
| try: | |
| min_ram_gb = max(1, int(os.getenv("AKIRA_CPU_TRAIN_MIN_RAM_GB", "48"))) | |
| except ValueError: | |
| logger.warning("[TRAIN] AKIRA_CPU_TRAIN_MIN_RAM_GB inválido; usando 48 GB.") | |
| min_ram_gb = 48 | |
| device = _select_training_device( | |
| is_gpu_training_available(), | |
| _available_memory_bytes(), | |
| min_ram_gb * (1024 ** 3), | |
| ) | |
| if device == "cpu" and not _training_dependencies_available(): | |
| logger.warning("[TRAIN] LoRA CPU indisponível: faltam dependências de treino.") | |
| return None | |
| if device == "cpu": | |
| logger.info(f"[TRAIN] CUDA indisponível; treino LoRA CPU autorizado com >= {min_ram_gb} GB livres.") | |
| return device | |
| def is_training_available() -> bool: | |
| return get_training_device() is not None | |
| # ====================================================================== | |
| # Construção de argumentos compatível com várias versões de TRL/Transformers | |
| # ====================================================================== | |
| def _filter_kwargs(cls, kwargs: Dict[str, Any]) -> Dict[str, Any]: | |
| try: | |
| import dataclasses | |
| names = {f.name for f in dataclasses.fields(cls)} | |
| return {k: v for k, v in kwargs.items() if k in names} | |
| except Exception: | |
| return kwargs | |
| def _build_args(out_dir: str, base_model: str, start_ts: float, device: str = "cuda"): | |
| import os as _os | |
| is_cpu = device == "cpu" | |
| max_steps = int(_os.getenv( | |
| "AKIRA_CPU_TRAIN_MAX_STEPS" if is_cpu else "AKIRA_TRAIN_MAX_STEPS", | |
| "10" if is_cpu else "100", | |
| )) | |
| bs = int(_os.getenv("AKIRA_TRAIN_BATCH_SIZE", "1")) | |
| ga = int(_os.getenv("AKIRA_TRAIN_GRAD_ACC", "4")) | |
| lr = float(_os.getenv("AKIRA_TRAIN_LR", "1e-4" if is_cpu else "2e-4")) | |
| seq = int(_os.getenv( | |
| "AKIRA_CPU_TRAIN_SEQ_LEN" if is_cpu else "AKIRA_TRAIN_SEQ_LEN", | |
| "256" if is_cpu else "1024", | |
| )) | |
| common: Dict[str, Any] = dict( | |
| output_dir=out_dir, | |
| max_steps=max_steps, | |
| per_device_train_batch_size=bs, | |
| gradient_accumulation_steps=ga, | |
| learning_rate=lr, | |
| logging_steps=5, | |
| save_strategy="steps", | |
| save_steps=max(max_steps, 1), | |
| warmup_ratio=0.03, | |
| lr_scheduler_type="cosine", | |
| bf16=not is_cpu, | |
| fp16=False, | |
| use_cpu=is_cpu, | |
| no_cuda=is_cpu, | |
| gradient_checkpointing=True, | |
| report_to="none", | |
| remove_unused_columns=False, | |
| dataloader_num_workers=0, | |
| seed=42, | |
| ) | |
| # TRL SFTConfig (preferido — trata tokenização/formatting) | |
| try: | |
| from trl import SFTConfig # type: ignore | |
| kwargs = dict(common) | |
| fields = None | |
| try: | |
| import dataclasses | |
| fields = {f.name for f in dataclasses.fields(SFTConfig)} | |
| except Exception: | |
| fields = set() | |
| if "dataset_text_field" in fields: | |
| kwargs["dataset_text_field"] = "text" | |
| if "max_seq_length" in fields: | |
| kwargs["max_seq_length"] = seq | |
| elif "max_length" in fields: | |
| kwargs["max_length"] = seq | |
| if "packing" in fields: | |
| kwargs["packing"] = False | |
| return SFTConfig(**_filter_kwargs(SFTConfig, kwargs)) | |
| except Exception as e: | |
| logger.warning(f"[GPU-TRAIN] SFTConfig indisponível ({e}); a usar TrainingArguments") | |
| from transformers import TrainingArguments # type: ignore | |
| return TrainingArguments(**_filter_kwargs(TrainingArguments, common)) | |
| def _build_trainer(model, tokenizer, train_dataset, validation_dataset, args, peft_config=None): | |
| from trl import SFTTrainer # type: ignore | |
| try: | |
| return SFTTrainer(model=model, args=args, train_dataset=train_dataset, | |
| eval_dataset=validation_dataset, | |
| processing_class=tokenizer, peft_config=peft_config) | |
| except TypeError: | |
| # API antiga do TRL | |
| return SFTTrainer(model=model, args=args, train_dataset=train_dataset, | |
| eval_dataset=validation_dataset, | |
| tokenizer=tokenizer, peft_config=peft_config) | |
| # ====================================================================== | |
| # Callback → estado JSON para o painel Gradio | |
| # ====================================================================== | |
| def _make_state_callback(base_model: str, start_ts: float, device: str): | |
| from transformers import TrainerCallback # type: ignore | |
| class _StateCallback(TrainerCallback): | |
| def on_log(self, args, state, control, logs=None, **kwargs): | |
| if not logs: | |
| return | |
| write_training_state({ | |
| "status": "running", | |
| "base_model": base_model, | |
| "device": device, | |
| "training_mode": "QLoRA 4-bit" if device == "cuda" else "LoRA CPU full-precision", | |
| "step": int(state.global_step or 0), | |
| "max_steps": int(state.max_steps or 0), | |
| "epoch": round(float(state.epoch or 0.0), 2), | |
| "loss": logs.get("loss"), | |
| "learning_rate": logs.get("learning_rate"), | |
| "started_at": start_ts, | |
| "elapsed_s": round(time.time() - start_ts, 1), | |
| }) | |
| def on_train_end(self, args, state, control, **kwargs): | |
| write_training_state({ | |
| "status": "finalizing", | |
| "base_model": base_model, | |
| "device": device, | |
| "training_mode": "QLoRA 4-bit" if device == "cuda" else "LoRA CPU full-precision", | |
| "step": int(state.global_step or 0), | |
| "max_steps": int(state.max_steps or 0), | |
| "started_at": start_ts, | |
| "elapsed_s": round(time.time() - start_ts, 1), | |
| }) | |
| return _StateCallback() | |
| # ====================================================================== | |
| # Treino principal (síncrono — correr em thread) | |
| # ====================================================================== | |
| def train_lora(dataset: List[Dict[str, Any]], | |
| base_model: Optional[str] = None, | |
| out_dir: Optional[str] = None) -> Dict[str, Any]: | |
| """ | |
| Executa QLoRA sobre o dataset [{"text": ...}]. | |
| Devolve {"success": bool, ...} e mantém /data/training_state.json atualizado. | |
| """ | |
| base = base_model or DEFAULT_BASE_MODEL | |
| out = out_dir or default_output_dir() | |
| start_ts = time.time() | |
| if not _train_lock.acquire(blocking=False): | |
| return {"success": False, "error": "Treino já em execução."} | |
| try: | |
| # ---- Validações ---- | |
| rows = [r for r in (dataset or []) if r and (r.get("text") or r.get("messages"))] | |
| if len(rows) < 20: | |
| return {"success": False, "error": f"Dataset insuficiente para treino validado: {len(rows)} < 20 exemplos."} | |
| device = get_training_device() | |
| if device is None: | |
| required_gb = max(1, int(os.getenv("AKIRA_CPU_TRAIN_MIN_RAM_GB", "48"))) | |
| available_gb = _available_memory_bytes() | |
| available_text = ( | |
| f"{available_gb / (1024 ** 3):.1f} GB" | |
| if available_gb is not None else "indisponível" | |
| ) | |
| return { | |
| "success": False, | |
| "error": ( | |
| "Sem GPU QLoRA funcional e sem RAM suficiente para treino LoRA CPU " | |
| f"(necessários {required_gb} GB livres; disponíveis: {available_text})." | |
| ), | |
| } | |
| write_training_state({ | |
| "status": "starting", | |
| "base_model": base, | |
| "examples": len(rows), | |
| "started_at": start_ts, | |
| "out_dir": out, | |
| "device": device, | |
| "training_mode": "QLoRA 4-bit" if device == "cuda" else "LoRA CPU full-precision", | |
| }) | |
| import torch # type: ignore | |
| from transformers import AutoTokenizer, AutoModelForCausalLM # type: ignore | |
| from peft import LoraConfig # type: ignore | |
| from datasets import Dataset # type: ignore | |
| os.makedirs(out, exist_ok=True) | |
| is_cpu = device == "cpu" | |
| max_steps = int(os.getenv( | |
| "AKIRA_CPU_TRAIN_MAX_STEPS" if is_cpu else "AKIRA_TRAIN_MAX_STEPS", | |
| "10" if is_cpu else "100", | |
| )) | |
| mode = "LoRA CPU" if is_cpu else "QLoRA 4-bit GPU" | |
| logger.info(f"[TRAIN] A iniciar {mode}: {base} | {len(rows)} exemplos | max_steps={max_steps}") | |
| tokenizer = AutoTokenizer.from_pretrained(base, trust_remote_code=True) | |
| if tokenizer.pad_token is None: | |
| tokenizer.pad_token = tokenizer.eos_token | |
| if is_cpu: | |
| model = AutoModelForCausalLM.from_pretrained( | |
| base, | |
| torch_dtype=torch.float32, | |
| low_cpu_mem_usage=True, | |
| trust_remote_code=True, | |
| ) | |
| model.to("cpu") | |
| else: | |
| from transformers import BitsAndBytesConfig # type: ignore | |
| from peft import prepare_model_for_kbit_training # type: ignore | |
| quant = BitsAndBytesConfig( | |
| load_in_4bit=True, | |
| bnb_4bit_quant_type="nf4", | |
| bnb_4bit_compute_dtype=torch.bfloat16, | |
| bnb_4bit_use_double_quant=True, | |
| ) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| base, | |
| quantization_config=quant, | |
| device_map="auto", | |
| dtype=torch.bfloat16, | |
| trust_remote_code=True, | |
| ) | |
| model = prepare_model_for_kbit_training(model) | |
| model.config.use_cache = False | |
| lora = LoraConfig( | |
| r=16, | |
| lora_alpha=32, | |
| lora_dropout=0.05, | |
| target_modules="all-linear", | |
| bias="none", | |
| task_type="CAUSAL_LM", | |
| ) | |
| ds_texts: List[str] = [] | |
| for r in rows: | |
| txt = "" | |
| msgs = r.get("messages") | |
| # Preferir chat-template do modelo base (Qwen/Llama/...) via messages | |
| if msgs and getattr(tokenizer, "chat_template", None): | |
| try: | |
| txt = tokenizer.apply_chat_template( | |
| msgs, tokenize=False, add_generation_prompt=False | |
| ) | |
| except Exception: | |
| txt = "" | |
| if not txt: | |
| txt = r.get("text") or "" | |
| if txt and txt.strip(): | |
| ds_texts.append(txt) | |
| if len(ds_texts) < 20: | |
| return {"success": False, "error": f"Dataset após tokenização insuficiente: {len(ds_texts)} < 20 exemplos."} | |
| dataset_fingerprint = fingerprint_dataset(rows) | |
| validation_count = max(2, math.ceil(len(ds_texts) * 0.1)) | |
| validation_indices = list(range(len(ds_texts))) | |
| random.Random(42).shuffle(validation_indices) | |
| validation_indices = set(validation_indices[:validation_count]) | |
| validation_texts = [ | |
| text for index, text in enumerate(ds_texts) if index in validation_indices | |
| ] | |
| train_texts = [ | |
| text for index, text in enumerate(ds_texts) if index not in validation_indices | |
| ] | |
| if len(train_texts) < 18: | |
| return {"success": False, "error": "Dataset sem exemplos suficientes após separar a validação."} | |
| train_ds = Dataset.from_list([{"text": t} for t in train_texts]) | |
| validation_ds = Dataset.from_list([{"text": t} for t in validation_texts]) | |
| candidate_id = f"{int(start_ts)}-{dataset_fingerprint[:10]}" | |
| candidate_out = os.path.join(out, "versions", candidate_id) | |
| os.makedirs(candidate_out, exist_ok=False) | |
| args = _build_args(candidate_out, base, start_ts, device=device) | |
| trainer = _build_trainer( | |
| model, tokenizer, train_ds, validation_ds, args, peft_config=lora, | |
| ) | |
| trainer.add_callback(_make_state_callback(base, start_ts, device)) | |
| baseline_metrics = trainer.evaluate() | |
| baseline_loss = float(baseline_metrics.get("eval_loss", float("nan"))) | |
| if not math.isfinite(baseline_loss): | |
| raise ValueError("A validação inicial não produziu uma perda finita.") | |
| result = trainer.train() | |
| final_metrics = trainer.evaluate() | |
| validation_loss = float(final_metrics.get("eval_loss", float("nan"))) | |
| if not math.isfinite(validation_loss) or validation_loss >= baseline_loss: | |
| raise ValueError( | |
| f"Adaptador reprovado na validação: perda {validation_loss} " | |
| f"não melhorou face à referência {baseline_loss}." | |
| ) | |
| # ---- Guardar artefactos ---- | |
| trainer.save_model(candidate_out) | |
| try: | |
| tokenizer.save_pretrained(candidate_out) | |
| except Exception: | |
| pass | |
| if not os.path.isfile(os.path.join(candidate_out, "adapter_config.json")): | |
| raise ValueError("O treino terminou sem gerar adapter_config.json.") | |
| if not any(os.path.isfile(os.path.join(candidate_out, filename)) for filename in ( | |
| "adapter_model.safetensors", "adapter_model.bin" | |
| )): | |
| raise ValueError("O treino terminou sem gerar os pesos LoRA.") | |
| train_loss = None | |
| try: | |
| train_loss = float(result.training_loss) | |
| except Exception: | |
| pass | |
| report = { | |
| "success": True, | |
| "base_model": base, | |
| "out_dir": candidate_out, | |
| "device": device, | |
| "training_mode": mode, | |
| "examples": len(rows), | |
| "train_examples": len(train_texts), | |
| "validation_examples": len(validation_texts), | |
| "dataset_fingerprint": dataset_fingerprint, | |
| "global_step": int(getattr(result, "global_step", 0) or 0), | |
| "train_loss": train_loss, | |
| "baseline_validation_loss": baseline_loss, | |
| "validation_loss": validation_loss, | |
| "elapsed_s": round(time.time() - start_ts, 1), | |
| "finished_at": time.time(), | |
| } | |
| with open(os.path.join(candidate_out, "training_report.json"), "w", encoding="utf-8") as f: | |
| json.dump(report, f, ensure_ascii=False, indent=2) | |
| manifest = { | |
| "base_model": base, | |
| "adapter_path": os.path.relpath(candidate_out, out), | |
| "training_mode": mode, | |
| "dataset_fingerprint": dataset_fingerprint, | |
| "validation_loss": validation_loss, | |
| "activated_at": time.time(), | |
| } | |
| manifest_tmp = os.path.join(out, "active_adapter.json.tmp") | |
| with open(manifest_tmp, "w", encoding="utf-8") as f: | |
| json.dump(manifest, f, ensure_ascii=False, indent=2) | |
| os.replace(manifest_tmp, os.path.join(out, "active_adapter.json")) | |
| report_tmp = os.path.join(out, "training_report.json.tmp") | |
| with open(report_tmp, "w", encoding="utf-8") as f: | |
| json.dump(report, f, ensure_ascii=False, indent=2) | |
| os.replace(report_tmp, os.path.join(out, "training_report.json")) | |
| write_training_state({ | |
| "status": "done", | |
| "base_model": base, | |
| "out_dir": candidate_out, | |
| "device": device, | |
| "training_mode": mode, | |
| "step": report["global_step"], | |
| "max_steps": max_steps, | |
| "loss": train_loss, | |
| "validation_loss": validation_loss, | |
| "examples": len(rows), | |
| "elapsed_s": report["elapsed_s"], | |
| "started_at": start_ts, | |
| }) | |
| logger.success(f"[GPU-TRAIN] Concluído em {report['elapsed_s']}s → {out}") | |
| # ---- Libertar VRAM (modelo de treino não fica residente) ---- | |
| try: | |
| del trainer | |
| del model | |
| import gc | |
| gc.collect() | |
| torch.cuda.empty_cache() | |
| except Exception: | |
| pass | |
| return report | |
| except Exception as e: | |
| err = f"{e}" | |
| logger.error(f"[GPU-TRAIN] Falha: {err}\n{traceback.format_exc()}") | |
| write_training_state({ | |
| "status": "failed", | |
| "base_model": base, | |
| "device": locals().get("device"), | |
| "error": err[:500], | |
| "started_at": start_ts, | |
| "elapsed_s": round(time.time() - start_ts, 1), | |
| }) | |
| try: | |
| import torch # type: ignore | |
| import gc | |
| gc.collect() | |
| if torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |
| except Exception: | |
| pass | |
| return {"success": False, "error": err} | |
| finally: | |
| _train_lock.release() | |
| __all__ = [ | |
| "train_lora", | |
| "is_gpu_training_available", | |
| "get_training_device", | |
| "is_training_available", | |
| "_select_training_device", | |
| "gpu_info", | |
| "read_training_state", | |
| "write_training_state", | |
| "training_state_path", | |
| "default_output_dir", | |
| "resolve_active_adapter", | |
| "fingerprint_dataset", | |
| ] | |