V2: 12.97M params, HW optimizer, parallel BBPE, DPO, reasoning, inference, 10 bugs fixed
8594de8 verified | #!/usr/bin/env python3 | |
| """train_v2.py — Treino BiGRU_T_version V2 (15M params, HW optimizer, multimodal bug-hunt). | |
| Melhorias vs train_fast.py: | |
| - max_modules=10 (aumento de 4→10, memória permite), d_model=128 | |
| - ~14.45M params (objetivo 15M) | |
| - HamiltonianWasserstein optimizer ATIVADO (AdamW + W₂ + repulsão + LR cíclico) | |
| - TimeBudget + aggressive_cleanup (estilo Xavante) | |
| - Teste de multimodalidade (text/image/audio/video encoders) com poucos samples | |
| - Timeout aumentado (1800s total, 900s/epoch) | |
| - Verificação de OOM-Killer antes do treino | |
| - Smoke checks antes do treino principal | |
| Uso: | |
| export HF_TOKEN=hf_xxx # opcional, para datasets privados | |
| python3 scripts/train_v2.py | |
| """ | |
| from __future__ import annotations | |
| import os | |
| import sys | |
| import time | |
| import json | |
| import logging | |
| import gc | |
| import shutil | |
| from pathlib import Path | |
| # Adiciona src/ ao path | |
| sys.path.insert(0, str(Path(__file__).parent.parent / "src")) | |
| # Otimização Xeon | |
| os.environ.setdefault("OMP_NUM_THREADS", "2") | |
| os.environ.setdefault("MKL_NUM_THREADS", "2") | |
| os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") | |
| import torch | |
| torch.set_num_threads(2) | |
| try: | |
| from bigru_t.utils.xeon_runtime import optimize_xeon_environment | |
| optimize_xeon_environment() | |
| except Exception as e: | |
| logging.warning(f"Could not apply Xeon optimization: {e}") | |
| from bigru_t import ( | |
| UnifiedModel, UnifiedModelConfig, create_unified_model, | |
| BiGRU_T_Trainer, TrainerConfig, | |
| HamiltonianWassersteinOptimizer, | |
| CircularReasoningWasserstein, | |
| BiGRUTGenerator, | |
| aggressive_cleanup, get_rss_mb, TimeBudget, | |
| ) | |
| from bigru_t.data.streaming_datasets import stream_dataset | |
| logging.basicConfig( | |
| level=logging.INFO, | |
| format="%(asctime)s [%(levelname)s] %(message)s", | |
| datefmt="%H:%M:%S", | |
| handlers=[logging.StreamHandler()], | |
| ) | |
| logger = logging.getLogger(__name__) | |
| def check_oom_killer_risk() -> bool: | |
| """Verifica se há risco de OOM-Killer antes de iniciar treino. | |
| Checa: | |
| 1. dmesg por ocorrências recentes de OOM-Killer | |
| 2. Memória livre do sistema | |
| 3. Swap disponível | |
| Returns: | |
| True se seguro para prosseguir, False se risco alto. | |
| """ | |
| import subprocess | |
| # 1. dmesg OOM check | |
| try: | |
| result = subprocess.run( | |
| ["dmesg"], capture_output=True, text=True, timeout=5 | |
| ) | |
| if "Out of memory" in result.stdout or "Killed process" in result.stdout: | |
| logger.warning("⚠️ dmesg mostra OOM-Killer activity recente — risco alto") | |
| # Não aborta, apenas avisa (pode ser histórico antigo) | |
| except (FileNotFoundError, subprocess.TimeoutExpired): | |
| pass # dmesg pode não estar acessível | |
| # 2. Memória livre do sistema | |
| try: | |
| with open("/proc/meminfo", "r") as f: | |
| meminfo = dict(line.split(":")[1].strip().split()[0:2] for line in f if ":" in line) | |
| free_mb = int(meminfo.get("MemAvailable", "0")) / 1024 | |
| total_mb = int(meminfo.get("MemTotal", "0")) / 1024 | |
| logger.info(f" Sistema: {total_mb:.0f}MB total, {free_mb:.0f}MB disponível") | |
| if free_mb < 512: | |
| logger.error(f"❌ Memória disponível muito baixa ({free_mb:.0f}MB < 512MB)") | |
| return False | |
| except Exception: | |
| pass | |
| # 3. Disk free | |
| try: | |
| usage = shutil.disk_usage("/home/z/my-project") | |
| free_gb = usage.free / (1024**3) | |
| logger.info(f" Disk free: {free_gb:.2f}GB") | |
| if free_gb < 1.0: | |
| logger.error(f"❌ Disk livre muito baixo ({free_gb:.2f}GB < 1GB)") | |
| return False | |
| except Exception: | |
| pass | |
| return True | |
| def smoke_test_model(model, tokenizer) -> bool: | |
| """Smoke test rápido: forward + backward + inference.""" | |
| logger.info("Running smoke test (forward + backward + generate)...") | |
| try: | |
| import torch | |
| # Forward — COM grad (para permitir backward) | |
| dummy_ids = torch.randint(0, 100, (1, 16), dtype=torch.long) | |
| y_hat, delta = model(dummy_ids, temperature=1.0, use_hypothesis=False) | |
| assert y_hat.shape[0] == 1, f"batch mismatch: {y_hat.shape}" | |
| assert not torch.isnan(y_hat).any(), "NaN in y_hat" | |
| logger.info(f" Forward OK: y_hat {tuple(y_hat.shape)}, delta {tuple(delta.shape)}") | |
| # Backward | |
| target = torch.tensor([5], dtype=torch.long) | |
| loss = torch.nn.functional.cross_entropy(y_hat, target) | |
| loss.backward() | |
| grad_norm = sum(p.grad.norm().item() ** 2 for p in model.parameters() if p.grad is not None) ** 0.5 | |
| logger.info(f" Backward OK: loss={loss.item():.4f}, grad_norm={grad_norm:.4f}") | |
| model.zero_grad() | |
| # Generate (no grad) | |
| gen = BiGRUTGenerator(model, max_seq_len=64) | |
| with torch.no_grad(): | |
| out = gen.generate(dummy_ids, max_new_tokens=5, do_sample=False) | |
| logger.info(f" Generate OK: {dummy_ids.shape} -> {out.shape}") | |
| return True | |
| except Exception as e: | |
| import traceback | |
| logger.error(f" Smoke test FAILED: {e}") | |
| traceback.print_exc() | |
| return False | |
| def test_multimodal_encoders() -> dict: | |
| """Testa os encoders multimodais com dados sintéticos (bug detection). | |
| Não usa datasets reais (apenas poucos samples sintéticos para validar | |
| que os encoders não quebram). Retorna dict com status de cada encoder. | |
| BUGS ENCONTRADOS E CORRIGIDOS: | |
| 1. TextEncoder dependia de 3 módulos ausentes (attention_multimodal, | |
| embedding_reconfig, gru_hierarchy) — COPIADOS da fonte HF. | |
| 2. AudioEncoder: parâmetro `in_channels` é ignorado pelo Conv1d interno | |
| (sempre usa n_mels=80 como canal) — WORKAROUND: passar [B, 80, T]. | |
| 3. VideoEncoder.forward(frames, audio) espera frames=[B,T,C,H,W] (5D), | |
| NÃO [B,C,T,H,W] — corrigido no teste. | |
| 4. ModalRouter.forward(inputs: Dict[str, Tensor]) — espera dict nomeado. | |
| 5. FusionLayer.forward(modality_embs: list) — espera lista de embeddings. | |
| """ | |
| logger.info("Testing multimodal encoders (synthetic data, bug detection)...") | |
| results = {} | |
| try: | |
| from bigru_t.multimodal.text_encoder import TextEncoder | |
| enc = TextEncoder(vocab_size=16384, d_model=128, n_heads=4, n_gru_levels=2, max_seq_len=64) | |
| import torch | |
| x = torch.randint(0, 100, (2, 16), dtype=torch.long) | |
| out = enc(x) | |
| results["text_encoder"] = {"ok": True, "shape": str(tuple(out.shape))} | |
| logger.info(f" TextEncoder OK: {tuple(out.shape)}") | |
| except Exception as e: | |
| results["text_encoder"] = {"ok": False, "error": str(e)} | |
| logger.warning(f" TextEncoder FAIL: {e}") | |
| try: | |
| from bigru_t.multimodal.image_encoder import ImageEncoder | |
| enc = ImageEncoder(d_model=128, in_channels=3) | |
| import torch | |
| x = torch.randn(2, 3, 16, 16) | |
| out = enc(x) | |
| results["image_encoder"] = {"ok": True, "shape": str(tuple(out.shape))} | |
| logger.info(f" ImageEncoder OK: {tuple(out.shape)}") | |
| except Exception as e: | |
| results["image_encoder"] = {"ok": False, "error": str(e)} | |
| logger.warning(f" ImageEncoder FAIL: {e}") | |
| try: | |
| from bigru_t.multimodal.audio_encoder import AudioEncoder | |
| # BUG: in_channels é ignorado; Conv1d usa n_mels como canal | |
| enc = AudioEncoder(d_model=128, in_channels=80, n_mels=80) | |
| import torch | |
| # Shape correta: [B, n_mels=80, T_audio] | |
| x = torch.randn(2, 80, 100) | |
| out = enc(x) | |
| results["audio_encoder"] = {"ok": True, "shape": str(tuple(out.shape)), | |
| "note": "in_channels param unused; uses n_mels=80 as channel"} | |
| logger.info(f" AudioEncoder OK: {tuple(out.shape)}") | |
| except Exception as e: | |
| results["audio_encoder"] = {"ok": False, "error": str(e)} | |
| logger.warning(f" AudioEncoder FAIL: {e}") | |
| try: | |
| from bigru_t.multimodal.video_encoder import VideoEncoder | |
| enc = VideoEncoder(d_model=128, n_frames=4) | |
| import torch | |
| # BUG: frames deve ser [B, T, C, H, W] (5D, T antes de C) | |
| frames = torch.randn(2, 4, 3, 16, 16) # (B=2, T=4, C=3, H=16, W=16) | |
| audio = torch.randn(2, 80, 100) # (B, n_mels, T_audio) | |
| out = enc(frames, audio) | |
| results["video_encoder"] = {"ok": True, "shape": str(tuple(out.shape)), | |
| "note": "frames=[B,T,C,H,W], audio=[B,80,T]"} | |
| logger.info(f" VideoEncoder OK: {tuple(out.shape)}") | |
| except Exception as e: | |
| results["video_encoder"] = {"ok": False, "error": str(e)} | |
| logger.warning(f" VideoEncoder FAIL: {e}") | |
| try: | |
| from bigru_t.multimodal.modal_router import ModalRouter | |
| from bigru_t.multimodal.fusion_layer import FusionLayer | |
| router = ModalRouter(d_model=128) | |
| fusion = FusionLayer(d_model=128, n_modalities=4) | |
| import torch | |
| # ModalRouter.forward(inputs: Dict[str, Tensor]) — espera dict nomeado | |
| inputs = { | |
| "text": torch.randn(2, 128), | |
| "image": torch.randn(2, 128), | |
| "audio": torch.randn(2, 128), | |
| "video": torch.randn(2, 128), | |
| } | |
| routed = router(inputs) | |
| # FusionLayer.forward(modality_embs: list) — espera lista | |
| embs = [torch.randn(2, 128) for _ in range(4)] | |
| fused = fusion(embs) | |
| results["modal_router_fusion"] = {"ok": True, | |
| "router_shape": str(tuple(routed.shape)), | |
| "fusion_shape": str(tuple(fused.shape))} | |
| logger.info(f" ModalRouter OK: {tuple(routed.shape)} | Fusion OK: {tuple(fused.shape)}") | |
| except Exception as e: | |
| results["modal_router_fusion"] = {"ok": False, "error": str(e)} | |
| logger.warning(f" ModalRouter/Fusion FAIL: {e}") | |
| return results | |
| def main(): | |
| logger.info("=" * 70) | |
| logger.info("BiGRU_T_version V2 — 15M params + HW optimizer + multimodal bug-hunt") | |
| logger.info("=" * 70) | |
| # HF token (do ambiente; será apagado ao final) | |
| hf_token = os.environ.get("HF_TOKEN") or None | |
| if hf_token: | |
| logger.info("HF_TOKEN encontrado no ambiente") | |
| else: | |
| logger.info("HF_TOKEN não fornecido (datasets públicos apenas)") | |
| # 1. Verificação de OOM-Killer e recursos | |
| logger.info("\n[1/6] Verificando recursos do sistema...") | |
| if not check_oom_killer_risk(): | |
| logger.error("Recursos insuficientes. Abortando.") | |
| sys.exit(1) | |
| logger.info(f" RSS atual: {get_rss_mb():.1f}MB") | |
| # 2. Teste multimodal (bug detection) | |
| logger.info("\n[2/6] Testando multimodalidade (bug hunting)...") | |
| mm_results = test_multimodal_encoders() | |
| mm_ok = sum(1 for v in mm_results.values() if v.get("ok")) | |
| mm_total = len(mm_results) | |
| logger.info(f" Multimodal: {mm_ok}/{mm_total} encoders OK") | |
| # 3. Config do modelo (12.97M params — max_modules=8, d_model=128) | |
| logger.info("\n[3/6] Criando UnifiedModel (max_modules=8, d_model=128)...") | |
| model_config = UnifiedModelConfig( | |
| vocab_size=16384, | |
| d_model=128, | |
| max_seq_len=32, | |
| pad_token_id=1, | |
| max_modules=8, # base spec do usuário (8) | |
| bigru_hidden=32, | |
| d_transformer=64, | |
| nhead_tu=4, | |
| d_ff_tu=128, | |
| output_dim_u8cell=64, | |
| cache_len=16, | |
| d_cache=128, | |
| nhead_orq=8, | |
| d_ff_orq=256, | |
| trainT_dim=128, | |
| nhead_train=4, | |
| d_ff_train=256, | |
| num_layers_train=2, | |
| hypT_dim=128, | |
| nhead_hyp=4, | |
| d_ff_hyp=256, | |
| num_layers_hyp=2, | |
| num_bits=8, | |
| dropout=0.1, | |
| ) | |
| model, _ = create_unified_model(model_config) | |
| params = model.count_parameters() | |
| logger.info(f" params: {params['total']:,} ({params['total_M']:.2f}M)") | |
| logger.info(f" trainable: {params['trainable']:,} ({params['trainable_M']:.2f}M)") | |
| aggressive_cleanup(verbose=True) | |
| # 4. Tokenizer (usa o pré-treinado do source) | |
| logger.info("\n[4/6] Carregando tokenizer...") | |
| from tokenizers import Tokenizer | |
| tok_path = "/home/z/my-project/source/model_final/tokenizer/tokenizer.json" | |
| if not Path(tok_path).exists(): | |
| logger.error(f"Tokenizer não encontrado: {tok_path}") | |
| sys.exit(1) | |
| tokenizer = Tokenizer.from_file(tok_path) | |
| logger.info(f" Tokenizer loaded: vocab={tokenizer.get_vocab_size()}") | |
| # 5. Smoke test do modelo | |
| logger.info("\n[5/6] Smoke test do modelo...") | |
| if not smoke_test_model(model, tokenizer): | |
| logger.error("Smoke test falhou. Abortando.") | |
| sys.exit(1) | |
| aggressive_cleanup(verbose=True) | |
| # 6. Carrega datasets (streaming, poucos samples para bug detection) | |
| logger.info("\n[6/6] Carregando datasets (streaming, poucos samples)...") | |
| datasets = [ | |
| "CEIA-POSITIVO/ultrachat_br_clustred_balanced_v1", | |
| ] | |
| max_samples = 8 # poucos samples (bug detection, testes) | |
| all_samples = [] | |
| t0 = time.time() | |
| for ds_name in datasets: | |
| if time.time() - t0 > 180: | |
| logger.warning(f"Timeout carregando datasets após {len(all_samples)} amostras") | |
| break | |
| try: | |
| count_before = len(all_samples) | |
| for sample in stream_dataset(ds_name, max_samples=max_samples, hf_token=hf_token): | |
| all_samples.append(sample) | |
| if len(all_samples) >= max_samples * len(datasets): | |
| break | |
| if time.time() - t0 > 180: | |
| break | |
| logger.info(f" {ds_name}: +{len(all_samples)-count_before} samples") | |
| except Exception as e: | |
| logger.warning(f" {ds_name} falhou: {e}") | |
| if not all_samples: | |
| logger.error("Nenhuma amostra carregada de nenhum dataset") | |
| sys.exit(1) | |
| logger.info(f"Total: {len(all_samples)} samples em {time.time()-t0:.1f}s") | |
| # Split 90/10 | |
| import random | |
| random.seed(42) | |
| random.shuffle(all_samples) | |
| split = max(1, int(0.9 * len(all_samples))) | |
| train_samples = all_samples[:split] | |
| val_samples = all_samples[split:] or all_samples[:1] | |
| logger.info(f" train: {len(train_samples)} | val: {len(val_samples)}") | |
| aggressive_cleanup(verbose=True) | |
| # Trainer config — HW optimizer + time budget + cleanup agressivo | |
| trainer_config = TrainerConfig( | |
| epochs=2, | |
| datasets=",".join(datasets), | |
| max_samples_per_dataset=max_samples, | |
| max_seq_len=32, | |
| per_device_batch_size=1, | |
| grad_accum=2, | |
| lr=1e-3, | |
| optimizer_type="hamiltonian_wasserstein", # ATIVADO | |
| hw_lr_amp=0.3, | |
| hw_lr_freq=0.01, | |
| hw_sigma_w=1.0, | |
| hw_sigma_rep=0.1, | |
| hw_prune_every_n=0, # pruning off para bug detection | |
| use_hypothesis=True, | |
| stop_grad_hyp=True, | |
| meta_interval=5, | |
| max_total_time_s=1800.0, # 30 min total | |
| max_per_epoch_s=900.0, # 15 min/epoch | |
| cleanup_every_n_steps=5, | |
| log_every=1, | |
| save_temp_every=3, # salvar a cada 3 steps (garante checkpoint mesmo se OOM no fim) | |
| output_dir="/home/z/my-project/BiGRU_T_version/model_final", | |
| temp_dir="/home/z/my-project/BiGRU_T_version/_temp_checkpoints", | |
| keep_temp=False, | |
| ram_threshold_pct=90.0, | |
| disk_min_free_gb=1.0, | |
| loss_patience=30, | |
| ) | |
| trainer = BiGRU_T_Trainer( | |
| model=model, | |
| tokenizer=tokenizer, | |
| config=trainer_config, | |
| train_samples=train_samples, | |
| val_samples=val_samples, | |
| ) | |
| # Treino | |
| result = trainer.train() | |
| # Salva resultados extras (multimodal + smoke) | |
| extra_report = { | |
| "multimodal_test": mm_results, | |
| "smoke_test": {"passed": True}, | |
| "config": { | |
| "max_modules": model_config.max_modules, | |
| "d_model": model_config.d_model, | |
| "bigru_hidden": model_config.bigru_hidden, | |
| "d_transformer": model_config.d_transformer, | |
| "params_total": params["total"], | |
| "params_M": params["total_M"], | |
| "optimizer": "hamiltonian_wasserstein", | |
| }, | |
| } | |
| extra_path = Path("/home/z/my-project/BiGRU_T_version/multimodal_report.json") | |
| with open(extra_path, "w") as f: | |
| json.dump(extra_report, f, indent=2, default=str) | |
| logger.info(f"Relatório multimodal salvo: {extra_path}") | |
| print("\n=== Resultado Final ===") | |
| print(json.dumps(result, indent=2, default=str)) | |
| # Apaga HF token | |
| if "HF_TOKEN" in os.environ: | |
| del os.environ["HF_TOKEN"] | |
| logger.info("HF_TOKEN apagado do ambiente") | |
| sys.exit(0 if not result.get("killed") else 1) | |
| if __name__ == "__main__": | |
| main() | |