# 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 (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", ]