# type: ignore """ modules/gpu_trainer.py ================================================================================ TREINO LoRA/QLoRA REAL NA GPU — Space Akiragpu (NVIDIA L4 24GB) ================================================================================ - Base modelo 4-bit (bitsandbytes nf4) + LoRA (PEFT r=16) via TRL SFTTrainer - 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 Qwen/Qwen2.5-7B-Instruct) AKIRA_TRAIN_MAX_STEPS 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/...) ================================================================================ """ import os import json import time import threading import traceback 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", "Qwen/Qwen2.5-7B-Instruct") _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 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 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") import transformers # noqa: F401 # type: ignore import peft # noqa: F401 # type: ignore import trl # noqa: F401 # type: ignore import datasets # noqa: F401 # type: ignore return True except Exception: return False # ====================================================================== # 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): import os as _os max_steps = int(_os.getenv("AKIRA_TRAIN_MAX_STEPS", "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", "2e-4")) seq = int(_os.getenv("AKIRA_TRAIN_SEQ_LEN", "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=True, 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, dataset, args, peft_config=None): from trl import SFTTrainer # type: ignore try: return SFTTrainer(model=model, args=args, train_dataset=dataset, processing_class=tokenizer, peft_config=peft_config) except TypeError: # API antiga do TRL return SFTTrainer(model=model, args=args, train_dataset=dataset, tokenizer=tokenizer, peft_config=peft_config) # ====================================================================== # Callback → estado JSON para o painel Gradio # ====================================================================== def _make_state_callback(base_model: str, start_ts: float): 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, "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, "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) < 5: return {"success": False, "error": f"Dataset insuficiente: {len(rows)} < 5 exemplos."} if not is_gpu_training_available(): return {"success": False, "error": "GPU/tr_dependencies indisponíveis (torch.cuda/trl/peft)."} write_training_state({ "status": "starting", "base_model": base, "examples": len(rows), "started_at": start_ts, "out_dir": out, }) import torch # type: ignore from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig # type: ignore from peft import LoraConfig, prepare_model_for_kbit_training # type: ignore from datasets import Dataset # type: ignore os.makedirs(out, exist_ok=True) max_steps = int(os.getenv("AKIRA_TRAIN_MAX_STEPS", "100")) logger.info(f"[GPU-TRAIN] A iniciar QLoRA: {base} | {len(rows)} exemplos | max_steps={max_steps}") quant = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True, ) tokenizer = AutoTokenizer.from_pretrained(base, trust_remote_code=True) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token model = AutoModelForCausalLM.from_pretrained( base, quantization_config=quant, device_map="auto", dtype=torch.bfloat16, trust_remote_code=True, ) model.config.use_cache = False # QLoRA: necessário (trl 1.14 não chama isto sozinho) — norm em fp32 # + enable_input_require_grads p/ gradient checkpointing com LoRA model = prepare_model_for_kbit_training(model) 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) < 5: return {"success": False, "error": f"Dataset após tokenização: {len(ds_texts)} < 5 exemplos."} ds = Dataset.from_list([{"text": t} for t in ds_texts]) args = _build_args(out, base, start_ts) trainer = _build_trainer(model, tokenizer, ds, args, peft_config=lora) trainer.add_callback(_make_state_callback(base, start_ts)) result = trainer.train() # ---- Guardar artefactos ---- trainer.save_model(out) try: tokenizer.save_pretrained(out) except Exception: pass train_loss = None try: train_loss = float(result.training_loss) except Exception: pass report = { "success": True, "base_model": base, "out_dir": out, "examples": len(rows), "global_step": int(getattr(result, "global_step", 0) or 0), "train_loss": train_loss, "elapsed_s": round(time.time() - start_ts, 1), "finished_at": time.time(), } with open(os.path.join(out, "training_report.json"), "w", encoding="utf-8") as f: json.dump(report, f, ensure_ascii=False, indent=2) write_training_state({ "status": "done", "base_model": base, "out_dir": out, "step": report["global_step"], "max_steps": max_steps, "loss": train_loss, "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, "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", "gpu_info", "read_training_state", "write_training_state", "training_state_path", "default_output_dir", ]