fanout-diffusion / scripts /run_frontier_sweep.py

Add training pipelines, consistency distillation scripts, and interactive dashboard server

69dcf19 verified about 1 month ago

22 kB

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()

	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()