AURORA v5 COMPLETO: atenção entre MoEs e variantes SOM (Teoremas 16.1–16.6), difusão multimodal (doc 17), quantização REAL sem fakes (torchao/QAT/FP8/BnB, doc 08 §§7–9), Agente Engenheiro com telemetria integral, vocab 16384, n_experts_fusao=4, treino do zero
96a4177 verified | # -*- coding: utf-8 -*- | |
| """Checkpoints de treino — v5 (doc 16 §8). | |
| Substitui `dados/estados_hf.py` (REMOVIDO a pedido — estados/fase-som-v4 | |
| apagados do Hub e do disco). Regra v5 motivada pelo bloqueio do Hub a | |
| uploads parciais em sequência curta: | |
| • durante o treino, checkpoints ficam SEMPRE LOCAIS (safetensors + SHA-256, | |
| protocolo idêntico ao doc 10 §9) e o diretório mantém só o último | |
| (poupa disco, item c do escopo); | |
| • NENHUM upload parcial durante o treino; | |
| • o snapshot FINAL entra no commit ÚNICO da publicação (scripts/ | |
| 06_publicar_hf.py chama `preparar_para_publicacao`, que copia o par | |
| (modelo.safetensors, meta.json) para o stage) — um único commit organizado. | |
| Teorema 16.9 (integridade do snapshot): seja σ o par gravado atomicamente | |
| com hash h(σ)=SHA-256. Na leitura, se h(σ')=h(σ) então σ'=σ (colisão | |
| desprezível, 2^−256), logo retomadas reproduzem EXATAMENTE o estado salvo. | |
| """ | |
| from __future__ import annotations | |
| import hashlib | |
| import json | |
| import os | |
| import shutil | |
| import torch | |
| from safetensors.torch import load as _st_load | |
| from safetensors.torch import save as _st_save | |
| def _sha256_bytes(b: bytes) -> str: | |
| return hashlib.sha256(b).hexdigest() | |
| class GestorCheckpoints: | |
| """Checkpoints locais atômicos + snapshot final para publicação única. | |
| Args: | |
| dir_local: diretório de checkpoints (mantém apenas o último). | |
| pasta_publicacao: pasta dentro do repo HF onde o snapshot final | |
| entra NO commit único (ex.: 'estados/fase-v5'). | |
| """ | |
| def __init__(self, dir_local: str = "/home/z/my-project/aurora/estados_local", | |
| pasta_publicacao: str = "estados/fase-v5"): | |
| self.dir_local = dir_local | |
| self.pasta_publicacao = pasta_publicacao | |
| self.eventos: list[str] = [] | |
| os.makedirs(dir_local, exist_ok=True) | |
| # ---------------- protocolo (doc 10 §9 preservado) ---------------- | |
| def _caminho_tag(self, tag: str) -> str: | |
| return os.path.join(self.dir_local, tag) | |
| def salvar(self, modelo, tag: str, meta: dict) -> str: | |
| """Serializa safetensors + meta com hash; mantém apenas a última tag.""" | |
| sd = {k: v.detach().contiguous() for k, v in modelo.state_dict().items()} | |
| # pesos empatados (tied): uma única cópia (load com strict=False) | |
| vistos: dict[int, str] = {} | |
| sd_unico: dict[str, torch.Tensor] = {} | |
| for k, v in sd.items(): | |
| ptr = v.data_ptr() | |
| if ptr in vistos: | |
| continue | |
| vistos[ptr] = k | |
| sd_unico[k] = v | |
| dados = bytes(_st_save(sd_unico, metadata={"formato": "aurora-v5", "tag": tag})) | |
| meta = dict(meta or {}) | |
| meta["hash_sha256"] = _sha256_bytes(dados) | |
| meta["tag"] = tag | |
| destino = self._caminho_tag(tag) | |
| os.makedirs(destino, exist_ok=True) | |
| with open(os.path.join(destino, "modelo.safetensors"), "wb") as f: | |
| f.write(dados) | |
| with open(os.path.join(destino, "meta.json"), "w", encoding="utf-8") as f: | |
| json.dump(meta, f, ensure_ascii=False, indent=2) | |
| # higiene: mantém SOMENTE a última tag (poupa disco, item c) | |
| removidos = 0 | |
| for nome in os.listdir(self.dir_local): | |
| if nome != tag: | |
| shutil.rmtree(os.path.join(self.dir_local, nome), | |
| ignore_errors=True) | |
| removidos += 1 | |
| if removidos: | |
| self.eventos.append(f"{removidos} checkpoint(s) antigo(s) apagado(s)") | |
| self.eventos.append(f"checkpoint {tag}: local ({len(dados)/1e6:.1f} MB)") | |
| return os.path.join(destino, "modelo.safetensors") | |
| def carregar(self, modelo, tag: str | None = None) -> bool: | |
| """Carrega tolerante a forma (mesma política da v3: tensores por | |
| chave com mesma forma; divergências registradas).""" | |
| tag = tag or self.ultima_tag() | |
| if tag is None: | |
| return False | |
| caminho = os.path.join(self._caminho_tag(tag), "modelo.safetensors") | |
| if not os.path.exists(caminho): | |
| self.eventos.append(f"checkpoint {tag}: ausente") | |
| return False | |
| with open(caminho, "rb") as f: | |
| dados = f.read() | |
| meta_path = os.path.join(self._caminho_tag(tag), "meta.json") | |
| if os.path.exists(meta_path): | |
| with open(meta_path, encoding="utf-8") as f: | |
| esperado = json.load(f).get("hash_sha256") | |
| if esperado and _sha256_bytes(dados) != esperado: | |
| self.eventos.append(f"checkpoint {tag}: HASH divergente — rejeitado") | |
| return False | |
| try: | |
| sd = _st_load(dados) | |
| modelo_sd = modelo.state_dict() | |
| próprios = {k: v for k, v in sd.items() | |
| if k in modelo_sd and modelo_sd[k].shape == v.shape} | |
| pulados = len(sd) - len(próprios) | |
| faltantes = [k for k in modelo_sd if k not in próprios and k not in sd] | |
| modelo.load_state_dict(próprios, strict=False) | |
| self.eventos.append( | |
| f"checkpoint {tag}: {len(próprios)} tensores, {pulados} pulados, " | |
| f"{len(faltantes)} novos") | |
| return True | |
| except Exception as e: | |
| self.eventos.append(f"carregar {tag}: {type(e).__name__}: {e}") | |
| return False | |
| def listar(self) -> list: | |
| """Tags disponíveis localmente, ordenadas lexicograficamente | |
| (tags canônicas `epoca-*`/`fase-*` são cronológicas).""" | |
| if not os.path.isdir(self.dir_local): | |
| return [] | |
| return sorted(d for d in os.listdir(self.dir_local) | |
| if os.path.isdir(os.path.join(self.dir_local, d)) | |
| and os.path.exists(os.path.join(self.dir_local, d, | |
| "meta.json"))) | |
| def ultima_tag(self) -> str | None: | |
| if not os.path.isdir(self.dir_local): | |
| return None | |
| tags = [d for d in os.listdir(self.dir_local) | |
| if os.path.isdir(os.path.join(self.dir_local, d)) | |
| and os.path.exists(os.path.join(self.dir_local, d, "meta.json"))] | |
| if not tags: | |
| return None | |
| # mais avançado = maior passo_global no meta (ordem cronológica) | |
| melhor, melhor_passo = None, -1 | |
| for tag in tags: | |
| try: | |
| with open(os.path.join(self.dir_local, tag, "meta.json"), | |
| encoding="utf-8") as f: | |
| passo = int(json.load(f).get("passo_global", -1)) | |
| except Exception: | |
| continue | |
| if passo > melhor_passo: | |
| melhor, melhor_passo = tag, passo | |
| return melhor | |
| def meta_de(self, tag: str | None = None) -> dict: | |
| tag = tag or self.ultima_tag() | |
| if tag is None: | |
| return {} | |
| caminho = os.path.join(self._caminho_tag(tag), "meta.json") | |
| if not os.path.exists(caminho): | |
| return {} | |
| with open(caminho, encoding="utf-8") as f: | |
| return json.load(f) | |
| # ---------------- publicação ÚNICA (doc 16 §8) ---------------- | |
| def preparar_para_publicacao(self, stage_raiz: str) -> str | None: | |
| """Copia o snapshot final para o stage do repo (commit único do | |
| script 06). Devolve o caminho do peso no stage.""" | |
| tag = self.ultima_tag() | |
| if tag is None: | |
| self.eventos.append("publicação: nenhum checkpoint local") | |
| return None | |
| origem = self._caminho_tag(tag) | |
| destino = os.path.join(stage_raiz, self.pasta_publicacao) | |
| os.makedirs(destino, exist_ok=True) | |
| shutil.copy2(os.path.join(origem, "modelo.safetensors"), | |
| os.path.join(destino, "modelo.safetensors")) | |
| meta = self.meta_de(tag) | |
| meta["publicado_em_commit_unico"] = True | |
| with open(os.path.join(destino, "meta.json"), "w", encoding="utf-8") as f: | |
| json.dump(meta, f, ensure_ascii=False, indent=2) | |
| self.eventos.append(f"snapshot {tag} → stage ({self.pasta_publicacao})") | |
| return os.path.join(destino, "modelo.safetensors") | |