File size: 7,098 Bytes
158177e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 | #!/usr/bin/env python3
"""Faz 2 egitim motoru: base safetensors LoRA + preprocessed veri."""
from __future__ import annotations
from pathlib import Path
from typing import Any
import numpy as np
import torch
from torch.utils.data import DataLoader, Subset
from meshai_train.base_weights import ensure_base_weights
from meshai_train.dataset import PreprocessedMeshDataset, collate_preprocessed
from meshai_train.faz2_models import (
FAZ2_VERSION,
build_faz2_from_weight_files,
faz2_loss,
load_faz2_checkpoint,
save_faz2_checkpoint,
)
def start_faz2_training(
*,
monitor: Any,
checkpoint_dir: Path,
output_dir: Path,
token: str | None,
epochs: int,
resume_from: Path | None,
validation_every: int,
val_ratio: float,
limit: int | None,
hf_repo: str,
data_root: Path | None,
log_fn: Any,
log_vram_fn: Any,
clear_gpu_fn: Any,
load_val_uids_fn: Any,
checkpoint_every: int = 50,
lora_rank: int = 8,
base_cache: Path | None = None,
) -> None:
if not token:
raise RuntimeError("Faz2 icin HF_TOKEN gerekli (base weight + data)")
device = "cuda" if torch.cuda.is_available() else "cpu"
cache = base_cache or Path("data/base_models")
log_fn(f"Pipeline surumu: {FAZ2_VERSION}")
log_fn(f"Veri: {hf_repo}" + (f" | limit={limit}" if limit else ""))
if torch.cuda.is_available():
log_fn(f"GPU: {torch.cuda.get_device_name(0)}")
log_vram_fn("startup")
weight_paths = ensure_base_weights(token=token, cache_dir=cache, log_fn=log_fn)
model = build_faz2_from_weight_files(weight_paths, lora_rank=lora_rank, log_fn=log_fn)
model = model.to(device)
# Faz1 hibrit ckpt varsa hybrid'e yukle
faz1 = checkpoint_dir / "latest_model.pt"
if resume_from and resume_from.exists():
pass
elif faz1.exists():
resume_from = faz1
dataset = PreprocessedMeshDataset(
token=token,
data_root=data_root,
hf_repo=hf_repo,
limit=limit,
)
val_uids = load_val_uids_fn()
val_idx = [i for i, o in enumerate(dataset.objects) if str(o.get("uid")) in val_uids]
if not val_idx:
val_count = max(1, int(len(dataset) * val_ratio))
val_idx = list(range(val_count))
train_idx = [i for i in range(len(dataset)) if i not in set(val_idx)]
if not train_idx:
train_idx = list(range(len(dataset)))
val_idx = train_idx[:1]
train_loader = DataLoader(
Subset(dataset, train_idx),
batch_size=1 if device == "cuda" else 1,
shuffle=True,
pin_memory=device == "cuda",
collate_fn=collate_preprocessed,
)
val_loader = DataLoader(
Subset(dataset, val_idx),
batch_size=1,
shuffle=False,
pin_memory=device == "cuda",
collate_fn=collate_preprocessed,
)
log_fn(f"Veri: {len(train_idx)} train + {len(val_idx)} val | LoRA rank={lora_rank}")
# Egitilebilir: LoRA + proj + hybrid (frozen W buffer)
trainable = [p for p in model.parameters() if p.requires_grad]
opt = torch.optim.AdamW(trainable, lr=5e-5, weight_decay=0.01)
latest = checkpoint_dir / "latest_faz2_model.pt"
global_step = 0
if resume_from and resume_from.exists():
global_step = load_faz2_checkpoint(resume_from, model, device)
log_fn(f"Resume: {resume_from} step={global_step}")
def _eval() -> float:
model.eval()
losses: list[float] = []
with torch.no_grad():
for batch in val_loader:
batch = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}
out = model(batch["geom_in"], batch["views"])
loss, _ = faz2_loss(out, batch)
losses.append(float(loss.item()))
model.train()
return float(np.mean(losses)) if losses else float("nan")
trainable_n = sum(p.numel() for p in trainable)
frozen_n = sum(p.numel() for p in model.parameters() if not p.requires_grad)
buffer_n = sum(b.numel() for b in model.buffers())
log_fn(f"Param: trainable={trainable_n:,} frozen={frozen_n:,} buffers(W)={buffer_n:,}")
for epoch in range(1, epochs + 1):
log_fn(f"--- Faz2 Epoch {epoch}/{epochs} ---")
model.train()
epoch_losses: list[float] = []
for batch in train_loader:
batch = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}
opt.zero_grad(set_to_none=True)
out = model(batch["geom_in"], batch["views"])
loss, parts = faz2_loss(out, batch)
if not torch.isfinite(loss):
monitor.note_nan_skip("faz2")
continue
loss.backward()
torch.nn.utils.clip_grad_norm_(trainable, 1.0)
opt.step()
global_step += 1
epoch_losses.append(float(loss.item()))
monitor.note_step(global_step, "faz2", parts["voxel"])
if global_step <= 3 or global_step % 20 == 0:
log_fn(
f"step={global_step} loss={float(loss.item()):.6f} "
f"voxel={parts['voxel']:.4f} trellis={parts['trellis_align']:.4f} "
f"hunyuan={parts['hunyuan_align']:.4f}"
)
if checkpoint_every > 0 and global_step % checkpoint_every == 0:
save_faz2_checkpoint(latest, epoch=epoch, global_step=global_step, model=model)
# Ayrica latest_model.pt olarak da yaz (orchestrator uyumu)
save_faz2_checkpoint(
checkpoint_dir / "latest_model.pt",
epoch=epoch,
global_step=global_step,
model=model,
)
log_fn(
f"CHECKPOINT_SAVED step={global_step} -> {latest} "
f"({latest.stat().st_size // (1024 * 1024)} MB)"
)
if global_step % validation_every == 0:
val = _eval()
monitor.note_validation(global_step, val, val)
log_vram_fn(f"step_{global_step}")
val = _eval()
mean = float(np.mean(epoch_losses)) if epoch_losses else float("nan")
monitor.note_epoch_end(epoch, epochs, mean, val, val, val)
save_faz2_checkpoint(latest, epoch=epoch, global_step=global_step, model=model)
save_faz2_checkpoint(
checkpoint_dir / "latest_model.pt",
epoch=epoch,
global_step=global_step,
model=model,
)
clear_gpu_fn()
log_fn(f"Epoch {epoch} kaydedildi -> {latest} ({latest.stat().st_size // (1024 * 1024)} MB)")
log_fn(f"CHECKPOINT_SAVED step={global_step} epoch={epoch} -> {latest}")
monitor.finish(ok=True)
log_fn("Faz2 TRELLIS/Hunyuan LoRA egitimi tamamlandi.")
log_fn(f"Cikti: {latest}")
|