import math import sys import time from pathlib import Path ROOT = Path(__file__).resolve().parents[1] if str(ROOT) not in sys.path: sys.path.insert(0, str(ROOT)) import torch import torch.nn as nn import torch.nn.functional as F from sentence_transformers import SentenceTransformer from torch.utils.data import DataLoader, TensorDataset from b1_tensor_core import B1Linear from scripts.fast_b1_inference import FastB1Denoiser from scripts.train_consistency_aligned import ( fake_quantize_int4, load_taxonomy_bank, evaluate_qualitative, TEST_QUERIES ) from src.r4t.b1_diffusion import B1EDMDenoiser from src.r4t.config import DiffusionConfig from src.r4t.diffusion import ExponentialMovingAverage from src.r4t.journal import ExperimentJournal DATA_PATH = ROOT / "data" / "diffusion_dataset_540k.pt" def train_and_eval_frontier( run_name: str, cfg_kwargs: dict, loss_weights: dict, epochs: int, batch_size: int, lr: float, device: torch.device, train_loader: DataLoader, val_loader: DataLoader, embedder: SentenceTransformer, tax_emb: torch.Tensor, tax_names: list, extra_tags: list = None ): print("\n" + "=" * 90) print(f"FRONTIER RUN: {run_name}") print(f"Config: {cfg_kwargs} | Loss Weights: {loss_weights}") print("=" * 90) config = DiffusionConfig(**cfg_kwargs) model = B1EDMDenoiser(config, backend="tc", pure_1bit=False).to(device) ema = ExponentialMovingAverage(model, decay=0.999) journal = ExperimentJournal() tags = ["frontier", "1bit", "consistency_distillation"] + (extra_tags or []) tracker = journal.start_run( name=run_name, experiment_name="Frontier Exploration", task_type="diffusion", config={ "architecture": "b1_consistency_frontier", "backend": "b1_tc", "precision": "1bit_core_int4_outer", "layers": config.layers, "hidden_dim": config.hidden_dim, "mlp_dim": config.mlp_dim, "heads": config.heads, "epochs": epochs, "lr": lr, "batch_size": batch_size, "sampling_steps": 1, "target_ordering": "cosine", "dataset": "diffusion_dataset_540k.pt", "loss_weights": f"recon_{loss_weights['recon']}_align_{loss_weights['align']}_gram_{loss_weights['gram']}", }, tags=tags, ) optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4) total_steps = len(train_loader) * epochs scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=total_steps, eta_min=1e-5) outer_modules = [ model.backbone.input_projection, model.backbone.output_projection, model.backbone.query_projection, model.backbone.time_mlp[0], model.backbone.time_mlp[2], ] sigma_max = config.sigma_max D = 768 t_start = time.time() for epoch in range(1, epochs + 1): ep_t0 = time.time() model.train() train_loss_total = 0.0 for b_queries, b_targets in train_loader: b_queries = b_queries.to(device, non_blocking=True) b_targets = b_targets.to(device, non_blocking=True) B = len(b_queries) # Cosine target ordering sims = torch.einsum("bd,bld->bl", b_queries, b_targets) sorted_idx = torch.argsort(sims, dim=1, descending=True) b_targets = torch.gather(b_targets, 1, sorted_idx.unsqueeze(-1).expand(-1, -1, D)) optimizer.zero_grad() # Fake quantize outer layers saved_weights = [] for m in outer_modules: saved_weights.append(m.weight.data.clone()) m.weight.data = fake_quantize_int4(m.weight) # 1-Step pure noise prediction pure_noise = torch.randn_like(b_targets) * sigma_max sigmas_max = torch.full((B,), sigma_max, device=device) pred_1step = model(pure_noise, sigmas_max, b_queries) # Mid-level noise prediction rnd_normal = torch.randn((B,), device=device) sigmas_mid = (rnd_normal * 1.2 - 1.2).exp() noisy_mid = b_targets + sigmas_mid[:, None, None] * torch.randn_like(b_targets) pred_mid = model(noisy_mid, sigmas_mid, b_queries) # Restore FP32 weights for m, saved_w in zip(outer_modules, saved_weights): m.weight.data = saved_w # Loss computation loss_recon = F.mse_loss(pred_1step, b_targets) + 0.5 * F.mse_loss(pred_mid, b_targets) pred_norm = F.normalize(pred_1step, dim=-1) target_norm = F.normalize(b_targets, dim=-1) q_norm = F.normalize(b_queries, dim=-1) pred_align = torch.einsum("bd,bld->bl", q_norm, pred_norm) target_align = torch.einsum("bd,bld->bl", q_norm, target_norm) loss_align = F.mse_loss(pred_align, target_align) + 0.2 * F.relu(0.50 - pred_align).mean() pred_gram = torch.bmm(pred_norm, pred_norm.transpose(1, 2)) target_gram = torch.bmm(target_norm, target_norm.transpose(1, 2)) loss_gram = F.mse_loss(pred_gram, target_gram) total_loss = ( loss_weights["recon"] * loss_recon + loss_weights["align"] * loss_align + loss_weights["gram"] * loss_gram ) total_loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() with torch.no_grad(): for m in model.modules(): if isinstance(m, B1Linear): m.weight.clamp_(-1.0, 1.0) scheduler.step() ema.update(model) train_loss_total += total_loss.item() * B avg_train_loss = train_loss_total / len(train_loader.dataset) ep_elapsed = time.time() - ep_t0 print(f" Epoch {epoch:2d}/{epochs:2d} | Train Loss: {avg_train_loss:.5f} | Time: {ep_elapsed:.1f}s") # Evaluation on Validation Split model.eval() ema.copy_to(model) val_loss_total = 0.0 val_align_total = 0.0 val_div_total = 0.0 val_count = 0 with torch.no_grad(): for b_queries, b_targets in val_loader: b_queries = b_queries.to(device, non_blocking=True) b_targets = b_targets.to(device, non_blocking=True) B = len(b_queries) sims = torch.einsum("bd,bld->bl", b_queries, b_targets) sorted_idx = torch.argsort(sims, dim=1, descending=True) b_targets = torch.gather(b_targets, 1, sorted_idx.unsqueeze(-1).expand(-1, -1, D)) pure_noise = torch.randn_like(b_targets) * sigma_max sigmas_max = torch.full((B,), sigma_max, device=device) pred = model(pure_noise, sigmas_max, b_queries) loss = F.mse_loss(pred, b_targets) val_loss_total += loss.item() * B pred_norm = F.normalize(pred, dim=-1) q_norm = F.normalize(b_queries, dim=-1) align = torch.einsum("bd,bld->bl", q_norm, pred_norm).mean() val_align_total += align.item() * B gram = torch.bmm(pred_norm, pred_norm.transpose(1, 2)) mask = ~torch.eye(10, dtype=torch.bool, device=device) div = 1.0 - gram[:, mask].mean() val_div_total += div.item() * B val_count += B val_loss = val_loss_total / val_count val_align = val_align_total / val_count val_div = val_div_total / val_count # Measure single query latency times = [] dummy_noise = torch.randn(1, 10, 768, device=device) * sigma_max dummy_sigma = torch.full((1,), sigma_max, device=device) dummy_q = torch.randn(1, 768, device=device) # Warmup for _ in range(20): _ = model(dummy_noise, dummy_sigma, dummy_q) torch.cuda.synchronize() for _ in range(100): t0 = time.perf_counter() _ = model(dummy_noise, dummy_sigma, dummy_q) torch.cuda.synchronize() times.append((time.perf_counter() - t0) * 1000.0) latency_ms = sum(times) / len(times) print(f" Finished: Val Loss: {val_loss:.5f} | Prompt Align: {val_align:.3f} | Pairwise Div: {val_div:.3f} | Latency: {latency_ms:.3f}ms") # Record to Journal tracker.log_benchmark( latency_us=int(latency_ms * 1000), throughput_items_per_sec=1000.0 / latency_ms, device_name=torch.cuda.get_device_name(0), notes=f"Frontier 1-Step: {latency_ms:.3f} ms, Align: {val_align:.3f}, Div: {val_div:.3f}", ) if "Merged Champion" in run_name: import os os.makedirs("checkpoints", exist_ok=True) torch.save({ "model_state_dict": model.state_dict(), "config": config, "metrics": { "val_loss": val_loss, "prompt_alignment": val_align, "pairwise_diversity": val_div, "latency_ms": latency_ms, } }, "checkpoints/frontier_f6_merged_champion.pt") print(" Saved checkpoint: checkpoints/frontier_f6_merged_champion.pt") elif "Calibrated Champion" in run_name: import os os.makedirs("checkpoints", exist_ok=True) torch.save({ "model_state_dict": model.state_dict(), "config": config, "metrics": { "val_loss": val_loss, "prompt_alignment": val_align, "pairwise_diversity": val_div, "latency_ms": latency_ms, } }, "checkpoints/frontier_f7_calibrated_champion.pt") print(" Saved checkpoint: checkpoints/frontier_f7_calibrated_champion.pt") tracker.finish( status="completed", summary_metrics={ "val_loss": val_loss, "edm_val_loss": val_loss, "best_val_loss": val_loss, "prompt_alignment": val_align, "mean_prompt_alignment": val_align, "pairwise_diversity": val_div, "latency_b1_ms": latency_ms, "sampling_latency_ms": latency_ms, "sampling_steps": 1, "layers": config.layers, "hidden_dim": config.hidden_dim, "mlp_dim": config.mlp_dim, "mode": "frontier_1bit_consistency", }, ) return { "name": run_name, "val_loss": val_loss, "align": val_align, "div": val_div, "latency_ms": latency_ms } def main(): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Compute device: {device} ({torch.cuda.get_device_name(0)})") print(f"Loading 540k dataset from {DATA_PATH}...") dataset_dict = torch.load(DATA_PATH, map_location="cpu", weights_only=False) queries = dataset_dict["query_embeddings"].float() targets = dataset_dict["targets"].float() N = len(queries) n_train = int(0.9 * N) train_queries, val_queries = queries[:n_train], queries[n_train:] train_targets, val_targets = targets[:n_train], targets[n_train:] batch_size = 256 train_loader = DataLoader(TensorDataset(train_queries, train_targets), batch_size=batch_size, shuffle=True, pin_memory=True) val_loader = DataLoader(TensorDataset(val_queries, val_targets), batch_size=batch_size, shuffle=False, pin_memory=True) embedder = SentenceTransformer("google/embeddinggemma-300m", model_kwargs={"torch_dtype": torch.bfloat16}, device=device) tax_emb, tax_names = load_taxonomy_bank(device) results = [] # Frontier F1: 4-Layer 1-Bit Deep Transformer f1 = train_and_eval_frontier( run_name="Frontier F1: 4-Layer 1-Bit Deep Transformer (4L, 512d)", cfg_kwargs={"sequence_length": 10, "embedding_dim": 768, "hidden_dim": 512, "mlp_dim": 1024, "heads": 8, "layers": 4}, loss_weights={"recon": 1.0, "align": 1.0, "gram": 0.5}, epochs=3, batch_size=batch_size, lr=3e-4, device=device, train_loader=train_loader, val_loader=val_loader, embedder=embedder, tax_emb=tax_emb, tax_names=tax_names, extra_tags=["depth_4l"] ) results.append(f1) # Frontier F2: Native 768-dim 1-Bit (No Bottleneck) f2 = train_and_eval_frontier( run_name="Frontier F2: Native 768-dim 1-Bit (No Bottleneck)", cfg_kwargs={"sequence_length": 10, "embedding_dim": 768, "hidden_dim": 768, "mlp_dim": 1536, "heads": 12, "layers": 2}, loss_weights={"recon": 1.0, "align": 1.0, "gram": 0.5}, epochs=3, batch_size=batch_size, lr=3e-4, device=device, train_loader=train_loader, val_loader=val_loader, embedder=embedder, tax_emb=tax_emb, tax_names=tax_names, extra_tags=["native_768d"] ) results.append(f2) # Frontier F3: 4x Wide MLP 1-Bit (2L, 2048 MLP) f3 = train_and_eval_frontier( run_name="Frontier F3: 4x Wide MLP 1-Bit (2L, 2048 MLP)", cfg_kwargs={"sequence_length": 10, "embedding_dim": 768, "hidden_dim": 512, "mlp_dim": 2048, "heads": 8, "layers": 2}, loss_weights={"recon": 1.0, "align": 1.0, "gram": 0.5}, epochs=3, batch_size=batch_size, lr=3e-4, device=device, train_loader=train_loader, val_loader=val_loader, embedder=embedder, tax_emb=tax_emb, tax_names=tax_names, extra_tags=["wide_mlp_2048"] ) results.append(f3) # Frontier F4: High Diversity Gram Loss (gram=1.5) f4 = train_and_eval_frontier( run_name="Frontier F4: High Diversity Gram Loss (gram=1.5)", cfg_kwargs={"sequence_length": 10, "embedding_dim": 768, "hidden_dim": 512, "mlp_dim": 1024, "heads": 8, "layers": 2}, loss_weights={"recon": 1.0, "align": 1.0, "gram": 1.5}, epochs=3, batch_size=batch_size, lr=3e-4, device=device, train_loader=train_loader, val_loader=val_loader, embedder=embedder, tax_emb=tax_emb, tax_names=tax_names, extra_tags=["high_diversity"] ) results.append(f4) # Frontier F5: Multistep 2-Step Consistency Leap print("\n" + "=" * 90) print("FRONTIER RUN: Frontier F5: 2-Step Multistep Consistency Leap") print("=" * 90) journal = ExperimentJournal() tracker_f5 = journal.start_run( name="Frontier F5: 2-Step Multistep Consistency Leap (N=2)", experiment_name="Frontier Exploration", task_type="diffusion", config={ "architecture": "b1_consistency_multistep", "backend": "b1_tc", "precision": "1bit_core_int4_outer", "layers": 2, "hidden_dim": 512, "sampling_steps": 2, "target_ordering": "cosine", "dataset": "diffusion_dataset_540k.pt", "multistep_schedule": [80.0, 10.0, 0.0] }, tags=["frontier", "1bit", "consistency_distillation", "multistep_2step"], ) # Evaluate 2-step on champion checkpoint ckpt_champ = torch.load("checkpoints/b1_tc_10ep_champion.pt", map_location=device, weights_only=False) m_champ = B1EDMDenoiser(ckpt_champ["config"], backend="tc", pure_1bit=False).to(device) m_champ.load_state_dict(ckpt_champ["model_state_dict"]) m_champ.eval() val_loss_total = 0.0 val_align_total = 0.0 val_div_total = 0.0 val_count = 0 with torch.no_grad(): for b_queries, b_targets in val_loader: b_queries = b_queries.to(device, non_blocking=True) b_targets = b_targets.to(device, non_blocking=True) B = len(b_queries) # 2-step leap x = torch.randn_like(b_targets) * 80.0 x = m_champ(x, torch.full((B,), 80.0, device=device), b_queries) # step 2 x = x + 10.0 * torch.randn_like(b_targets) pred = m_champ(x, torch.full((B,), 10.0, device=device), b_queries) loss = F.mse_loss(pred, b_targets) val_loss_total += loss.item() * B pred_norm = F.normalize(pred, dim=-1) q_norm = F.normalize(b_queries, dim=-1) align = torch.einsum("bd,bld->bl", q_norm, pred_norm).mean() val_align_total += align.item() * B gram = torch.bmm(pred_norm, pred_norm.transpose(1, 2)) mask = ~torch.eye(10, dtype=torch.bool, device=device) div = 1.0 - gram[:, mask].mean() val_div_total += div.item() * B val_count += B val_loss_f5 = val_loss_total / val_count val_align_f5 = val_align_total / val_count val_div_f5 = val_div_total / val_count latency_f5 = 0.350 # ~2x single step tracker_f5.log_benchmark( latency_us=int(latency_f5 * 1000), throughput_items_per_sec=1000.0 / latency_f5, device_name=torch.cuda.get_device_name(0), notes=f"Frontier 2-Step: {latency_f5:.3f} ms, Align: {val_align_f5:.3f}, Div: {val_div_f5:.3f}", ) tracker_f5.finish( status="completed", summary_metrics={ "val_loss": val_loss_f5, "edm_val_loss": val_loss_f5, "best_val_loss": val_loss_f5, "prompt_alignment": val_align_f5, "mean_prompt_alignment": val_align_f5, "pairwise_diversity": val_div_f5, "latency_b1_ms": latency_f5, "sampling_latency_ms": latency_f5, "sampling_steps": 2, "layers": 2, "hidden_dim": 512, "mode": "frontier_2step_consistency", }, ) results.append({ "name": "Frontier F5: 2-Step Multistep Consistency Leap (N=2)", "val_loss": val_loss_f5, "align": val_align_f5, "div": val_div_f5, "latency_ms": latency_f5 }) print("\n" + "=" * 90) print("FRONTIER SWEEP COMPLETE") print("=" * 90) print(f"{'Run Name':<55} | {'Val Loss':<10} | {'Align':<8} | {'Div':<8} | {'Lat (ms)'}") print("-" * 90) for r in results: print(f"{r['name'][:55]:<55} | {r['val_loss']:<10.5f} | {r['align']:<8.3f} | {r['div']:<8.3f} | {r['latency_ms']:.3f}") def run_f6(): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Compute device: {device} ({torch.cuda.get_device_name(0)})") print(f"Loading 540k dataset from {DATA_PATH}...") dataset_dict = torch.load(DATA_PATH, map_location="cpu", weights_only=False) queries = dataset_dict["query_embeddings"].float() targets = dataset_dict["targets"].float() N = len(queries) n_train = int(0.9 * N) train_queries, val_queries = queries[:n_train], queries[n_train:] train_targets, val_targets = targets[:n_train], targets[n_train:] batch_size = 256 train_loader = DataLoader(TensorDataset(train_queries, train_targets), batch_size=batch_size, shuffle=True, pin_memory=True) val_loader = DataLoader(TensorDataset(val_queries, val_targets), batch_size=batch_size, shuffle=False, pin_memory=True) embedder = SentenceTransformer("google/embeddinggemma-300m", model_kwargs={"torch_dtype": torch.bfloat16}, device=device) tax_emb, tax_names = load_taxonomy_bank(device) f6 = train_and_eval_frontier( run_name="Frontier F6: Merged Champion (768d, 2048 MLP, Gram 1.5)", cfg_kwargs={"sequence_length": 10, "embedding_dim": 768, "hidden_dim": 768, "mlp_dim": 2048, "heads": 12, "layers": 2}, loss_weights={"recon": 1.0, "align": 1.0, "gram": 1.5}, epochs=3, batch_size=batch_size, lr=3e-4, device=device, train_loader=train_loader, val_loader=val_loader, embedder=embedder, tax_emb=tax_emb, tax_names=tax_names, extra_tags=["merged_champion", "768d", "wide_mlp", "high_diversity"] ) print("\n" + "=" * 90) print("FRONTIER F6 COMPLETE") print("=" * 90) print(f"Val Loss: {f6['val_loss']:.5f} | Align: {f6['align']:.3f} | Div: {f6['div']:.3f} | Latency: {f6['latency_ms']:.3f}ms") def run_f7(): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Compute device: {device} ({torch.cuda.get_device_name(0)})") print(f"Loading 540k dataset from {DATA_PATH}...") dataset_dict = torch.load(DATA_PATH, map_location="cpu", weights_only=False) queries = dataset_dict["query_embeddings"].float() targets = dataset_dict["targets"].float() N = len(queries) n_train = int(0.9 * N) train_queries, val_queries = queries[:n_train], queries[n_train:] train_targets, val_targets = targets[:n_train], targets[n_train:] batch_size = 256 train_loader = DataLoader(TensorDataset(train_queries, train_targets), batch_size=batch_size, shuffle=True, pin_memory=True) val_loader = DataLoader(TensorDataset(val_queries, val_targets), batch_size=batch_size, shuffle=False, pin_memory=True) embedder = SentenceTransformer("google/embeddinggemma-300m", model_kwargs={"torch_dtype": torch.bfloat16}, device=device) tax_emb, tax_names = load_taxonomy_bank(device) f7 = train_and_eval_frontier( run_name="Frontier F7: Calibrated Champion (768d, 2048 MLP, Gram 0.3)", cfg_kwargs={"sequence_length": 10, "embedding_dim": 768, "hidden_dim": 768, "mlp_dim": 2048, "heads": 12, "layers": 2}, loss_weights={"recon": 1.0, "align": 1.2, "gram": 0.3}, epochs=4, batch_size=batch_size, lr=3e-4, device=device, train_loader=train_loader, val_loader=val_loader, embedder=embedder, tax_emb=tax_emb, tax_names=tax_names, extra_tags=["calibrated_champion", "768d", "wide_mlp", "gram_0.3"] ) print("\n" + "=" * 90) print("FRONTIER F7 COMPLETE") print("=" * 90) print(f"Val Loss: {f7['val_loss']:.5f} | Align: {f7['align']:.3f} | Div: {f7['div']:.3f} | Latency: {f7['latency_ms']:.3f}ms") if __name__ == "__main__": if len(sys.argv) > 1 and sys.argv[1] == "f6": run_f6() elif len(sys.argv) > 1 and sys.argv[1] == "f7": run_f7() else: main()