fanout-diffusion / scripts /train_consistency_distillation_1step.py

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

f08972a verified about 2 months ago

15.3 kB

import argparse
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.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 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"
TAXONOMY_PATH = ROOT / "data" / "taxonomy_embeddings.pt"
BASE_CHECKPOINT = ROOT / "checkpoints" / "b1_tc_10ep_champion.pt"
OUTPUT_CHECKPOINT = ROOT / "checkpoints" / "champion_b1_consistency_1step.pt"

TEST_QUERIES = [
    "running shoes and athletic sneakers",
    "wireless noise cancelling headphones and audio gear",
    "organic gardening tools and indoor plant care",
    "python machine learning algorithms and deep neural networks",
    "espresso coffee machines and barista accessories",
]


def load_taxonomy_bank(device):
    if not TAXONOMY_PATH.exists():
        return None, None
    tax_data = torch.load(TAXONOMY_PATH, map_location="cpu", weights_only=False)
    emb = F.normalize(tax_data["embeddings"].float(), dim=-1).to(device)
    names = tax_data["names"]
    return emb, names


@torch.no_grad()
def evaluate_qualitative_1step(fast_model, embedder, tax_emb, tax_names, device, tracker):
    all_alignments = []
    all_diversities = []
    decoded_results = {}

    for q_idx, q_text in enumerate(TEST_QUERIES):
        z_q = embedder.encode([q_text], convert_to_tensor=True, normalize_embeddings=True, device=device).float()
        cached_mem = fast_model.precompute_cross_memory(z_q)

        # Single-step direct inference
        shape = (1, fast_model.config.sequence_length, fast_model.config.embedding_dim)
        noise = torch.randn(shape, device=device) * fast_model.config.sigma_max
        sigma = torch.full((1,), fast_model.config.sigma_max, device=device)

        fanout = fast_model.forward_with_cached_memory(noise, sigma, cached_mem)[0]
        fanout = F.normalize(fanout, dim=-1)

        # Prompt alignment
        align = (fanout @ z_q.T).squeeze(-1).mean().item()
        all_alignments.append(align)

        # Pairwise diversity
        sim_mat = fanout @ fanout.T
        mask = ~torch.eye(10, dtype=torch.bool, device=device)
        div = 1.0 - sim_mat[mask].mean().item()
        all_diversities.append(div)

        # Decode taxonomy
        terms = []
        if tax_emb is not None:
            top_matches = fanout @ tax_emb.T
            best_idx = top_matches.argmax(dim=-1).tolist()
            for s_idx, t_idx in enumerate(best_idx):
                sim_val = top_matches[s_idx, t_idx].item()
                terms.append(f"{tax_names[t_idx]} (sim: {sim_val:.3f})")
        else:
            terms = [f"Vector #{i+1}" for i in range(10)]

        decoded_results[q_text] = {
            "alignment": align,
            "diversity": div,
            "terms": terms,
        }

        if tracker:
            tracker.log_sample(
                step=200 + q_idx,
                sample_type="1step_consistency_fanout",
                input_payload={"query": q_text, "steps": 1, "mode": "consistency_distillation"},
                output_payload={"alignment": align, "diversity": div, "subqueries": terms},
            )

    mean_align = sum(all_alignments) / len(all_alignments)
    mean_div = sum(all_diversities) / len(all_diversities)
    return mean_align, mean_div, decoded_results


def benchmark_cuda_graph_1step(fast_model, device):
    D = fast_model.config.embedding_dim
    dummy_q = torch.randn(1, D, device=device)
    dummy_q = F.normalize(dummy_q, dim=-1)

    cached_mem = fast_model.precompute_cross_memory(dummy_q)
    shape = (1, fast_model.config.sequence_length, D)
    static_noise = torch.randn(shape, device=device) * fast_model.config.sigma_max
    static_sigma = torch.full((1,), fast_model.config.sigma_max, device=device)

    # Warmup graph stream
    g_stream = torch.cuda.Stream()
    g_stream.wait_stream(torch.cuda.current_stream())
    with torch.cuda.stream(g_stream):
        for _ in range(5):
            _ = fast_model.forward_with_cached_memory(static_noise, static_sigma, cached_mem)
    torch.cuda.current_stream().wait_stream(g_stream)

    # Capture graph
    g = torch.cuda.CUDAGraph()
    with torch.cuda.graph(g, stream=g_stream):
        _ = fast_model.forward_with_cached_memory(static_noise, static_sigma, cached_mem)

    # Benchmark Replay
    torch.cuda.synchronize()
    times = []
    for _ in range(100):
        t0 = time.perf_counter()
        g.replay()
        torch.cuda.synchronize()
        times.append((time.perf_counter() - t0) * 1000.0)

    avg_ms = sum(times) / len(times)
    p50_ms = sorted(times)[int(len(times) * 0.5)]
    p95_ms = sorted(times)[int(len(times) * 0.95)]
    qps = 1000.0 / avg_ms
    return avg_ms, p50_ms, p95_ms, qps


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, L, D = targets.shape

    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 = 128
    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)

    print(f"Loading Base Champion Checkpoint: {BASE_CHECKPOINT}...")
    ckpt = torch.load(BASE_CHECKPOINT, map_location=device, weights_only=False)
    config = ckpt["config"]
    model = B1EDMDenoiser(config, backend="tc", pure_1bit=False).to(device)

    if "ema_state_dict" in ckpt and "shadow" in ckpt["ema_state_dict"]:
        shadow = ckpt["ema_state_dict"]["shadow"]
        model.load_state_dict({k: shadow[k].to(device) for k in shadow})
    else:
        model.load_state_dict(ckpt["model_state_dict"])

    ema = ExponentialMovingAverage(model, decay=0.999)

    print("Loading SentenceTransformer and taxonomy bank...")
    embedder = SentenceTransformer("google/embeddinggemma-300m", model_kwargs={"torch_dtype": torch.bfloat16}, device=device)
    tax_emb, tax_names = load_taxonomy_bank(device)

    epochs = 4
    lr = 3e-4
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs * len(train_loader), eta_min=1e-5)

    journal = ExperimentJournal()
    tracker = journal.start_run(
        name="Champion B1: 1-Step Consistency Distillation (Tuned Diversity)",
        experiment_name="1-Bit Tensor Core Innovation",
        task_type="diffusion",
        config={
            "architecture": "consistency_distillation_1step",
            "backend": "b1_tc",
            "precision": "1bit_core_fp16_adapters",
            "layers": config.layers,
            "hidden_dim": config.hidden_dim,
            "epochs": epochs,
            "lr": lr,
            "batch_size": batch_size,
            "sampling_steps": 1,
            "target_ordering": "cosine",
        },
        tags=["1bit", "tensor_core", "consistency_distillation", "1step", "champion"],
    )

    print("\n" + "=" * 95)
    print("STARTING 1-STEP CONSISTENCY DISTILLATION TRAINING")
    print(f"Objective: Map pure Gaussian noise directly to semantic fanout in 1 single forward pass")
    print(f"Epochs: {epochs} | Batch Size: {batch_size} | Base Learning Rate: {lr}")
    print("=" * 95)

    sigma_max = config.sigma_max

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

            # 1. Direct 1-Step Prediction from pure noise (sigma_max)
            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)

            # 2. Intermediate noise level 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)

            # Losses:
            # Reconstruction Loss
            loss_recon = F.mse_loss(pred_1step, b_targets) + 0.5 * F.mse_loss(pred_mid, b_targets)

            # Centroid Prompt Alignment: Anchor the fanout cluster centroid to query
            pred_norm = F.normalize(pred_1step, dim=-1)
            centroid = F.normalize(pred_norm.mean(dim=1), dim=-1)
            align_cos = (centroid * b_queries).sum(dim=-1).mean()
            loss_align = 1.0 - align_cos

            # Pairwise Diversity Regularizer: Strongly penalize pairwise similarities above 0.30
            sim_mat = torch.bmm(pred_norm, pred_norm.transpose(1, 2))
            mask = ~torch.eye(L, dtype=torch.bool, device=device).unsqueeze(0).expand(B, -1, -1)
            pair_sims = sim_mat[mask]
            loss_div = F.relu(pair_sims - 0.30).mean()

            total_loss = loss_recon + 0.2 * loss_align + 1.2 * loss_div
            total_loss.backward()

            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()

            # STE weight clamp for B1Linear
            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

        train_loss = train_loss_total / len(train_queries)

        # Validation
        model.eval()
        val_loss_total = 0.0
        val_align_total = 0.0
        val_div_total = 0.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_1step = model(pure_noise, sigmas_max, b_queries)

                v_loss = F.mse_loss(pred_1step, b_targets)
                val_loss_total += v_loss.item() * B

                p_norm = F.normalize(pred_1step, dim=-1)
                align = (p_norm * b_queries[:, None, :]).sum(dim=-1).mean().item()
                val_align_total += align * B

                sim_mat = torch.bmm(p_norm, p_norm.transpose(1, 2))
                mask = ~torch.eye(L, dtype=torch.bool, device=device).unsqueeze(0).expand(B, -1, -1)
                div = (1.0 - sim_mat[mask].mean().item()) * B
                val_div_total += div

        val_loss = val_loss_total / len(val_queries)
        val_align = val_align_total / len(val_queries)
        val_div = val_div_total / len(val_queries)
        ep_time = time.time() - ep_t0

        print(f"Epoch [{epoch}/{epochs}] | Train Loss: {train_loss:.4f} | Val MSE: {val_loss:.4f} | 1-Step Align: {val_align:.3f} | Div: {val_div:.3f} | Time: {ep_time:.1f}s")

        tracker.log_metrics(
            step=epoch * len(train_loader),
            epoch=epoch,
            train_loss=train_loss,
            val_loss=val_loss,
            val_alignment=val_align,
            val_diversity=val_div,
            lr=scheduler.get_last_lr()[0],
        )

    # Save champion model
    print("\nSaving 1-Step Consistency Champion Checkpoint...")
    model.eval()
    model.freeze_for_inference()
    torch.save(
        {
            "epoch": epochs,
            "config": config,
            "model_state_dict": model.state_dict(),
            "val_loss": val_loss,
            "val_alignment": val_align,
            "val_diversity": val_div,
        },
        OUTPUT_CHECKPOINT,
    )
    print(f"Saved: {OUTPUT_CHECKPOINT} ({OUTPUT_CHECKPOINT.stat().st_size / (1024*1024):.2f} MB)")

    # Build FastB1Denoiser & CUDA Graph Benchmark
    print("\nBuilding FastB1Denoiser and Capturing CUDA Graph...")
    fast_model = FastB1Denoiser(model)
    avg_ms, p50_ms, p95_ms, qps = benchmark_cuda_graph_1step(fast_model, device)
    print(f"CUDA Graph 1-Step Latency: {avg_ms:.3f} ms ({p50_ms:.3f} ms P50, {p95_ms:.3f} ms P95) -> {qps:.1f} QPS (Batch=1)")

    # Qualitative Evaluation on Test Queries
    print("\nEvaluating Decoded Fanout Taxonomy Terms...")
    mean_align, mean_div, decoded = evaluate_qualitative_1step(fast_model, embedder, tax_emb, tax_names, device, tracker)

    for q, d in decoded.items():
        print(f"\nQuery: \"{q}\"")
        print(f"  Alignment: {d['alignment']:.3f} | Diversity: {d['diversity']:.3f}")
        for slot_idx, term in enumerate(d["terms"][:5]):
            print(f"   #{slot_idx+1}: {term}")

    tracker.log_benchmark(
        latency_us=int(avg_ms * 1000),
        throughput_items_per_sec=qps,
        device_name=torch.cuda.get_device_name(0),
        notes=f"1-Step Consistency Distillation: {avg_ms:.2f} ms (B=1), Align: {mean_align:.3f}",
    )

    tracker.finish(
        status="completed",
        summary_metrics={
            "sampling_latency_ms": avg_ms,
            "latency_ms": avg_ms,
            "prompt_alignment": mean_align,
            "pairwise_diversity": mean_div,
            "val_loss": val_loss,
            "best_val_loss": val_loss,
            "throughput_items_per_sec": qps,
            "sampling_steps": 1,
            "mode": "consistency_distillation_1step",
        },
    )

    print("\n" + "=" * 95)
    print(f"All done! 1-Step Latency: {avg_ms:.3f} ms | Alignment: {mean_align:.3f} | Diversity: {mean_div:.3f}")
    print("Run logged to journal.db and visible in web UI at http://127.0.0.1:8000")
    print("=" * 95)


if __name__ == "__main__":
    main()

	import argparse
	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.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 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"
	TAXONOMY_PATH = ROOT / "data" / "taxonomy_embeddings.pt"
	BASE_CHECKPOINT = ROOT / "checkpoints" / "b1_tc_10ep_champion.pt"
	OUTPUT_CHECKPOINT = ROOT / "checkpoints" / "champion_b1_consistency_1step.pt"

	TEST_QUERIES = [
	"running shoes and athletic sneakers",
	"wireless noise cancelling headphones and audio gear",
	"organic gardening tools and indoor plant care",
	"python machine learning algorithms and deep neural networks",
	"espresso coffee machines and barista accessories",
	]


	def load_taxonomy_bank(device):
	if not TAXONOMY_PATH.exists():
	return None, None
	tax_data = torch.load(TAXONOMY_PATH, map_location="cpu", weights_only=False)
	emb = F.normalize(tax_data["embeddings"].float(), dim=-1).to(device)
	names = tax_data["names"]
	return emb, names


	@torch.no_grad()
	def evaluate_qualitative_1step(fast_model, embedder, tax_emb, tax_names, device, tracker):
	all_alignments = []
	all_diversities = []
	decoded_results = {}

	for q_idx, q_text in enumerate(TEST_QUERIES):
	z_q = embedder.encode([q_text], convert_to_tensor=True, normalize_embeddings=True, device=device).float()
	cached_mem = fast_model.precompute_cross_memory(z_q)

	# Single-step direct inference
	shape = (1, fast_model.config.sequence_length, fast_model.config.embedding_dim)
	noise = torch.randn(shape, device=device) * fast_model.config.sigma_max
	sigma = torch.full((1,), fast_model.config.sigma_max, device=device)

	fanout = fast_model.forward_with_cached_memory(noise, sigma, cached_mem)[0]
	fanout = F.normalize(fanout, dim=-1)

	# Prompt alignment
	align = (fanout @ z_q.T).squeeze(-1).mean().item()
	all_alignments.append(align)

	# Pairwise diversity
	sim_mat = fanout @ fanout.T
	mask = ~torch.eye(10, dtype=torch.bool, device=device)
	div = 1.0 - sim_mat[mask].mean().item()
	all_diversities.append(div)

	# Decode taxonomy
	terms = []
	if tax_emb is not None:
	top_matches = fanout @ tax_emb.T
	best_idx = top_matches.argmax(dim=-1).tolist()
	for s_idx, t_idx in enumerate(best_idx):
	sim_val = top_matches[s_idx, t_idx].item()
	terms.append(f"{tax_names[t_idx]} (sim: {sim_val:.3f})")
	else:
	terms = [f"Vector #{i+1}" for i in range(10)]

	decoded_results[q_text] = {
	"alignment": align,
	"diversity": div,
	"terms": terms,
	}

	if tracker:
	tracker.log_sample(
	step=200 + q_idx,
	sample_type="1step_consistency_fanout",
	input_payload={"query": q_text, "steps": 1, "mode": "consistency_distillation"},
	output_payload={"alignment": align, "diversity": div, "subqueries": terms},
	)

	mean_align = sum(all_alignments) / len(all_alignments)
	mean_div = sum(all_diversities) / len(all_diversities)
	return mean_align, mean_div, decoded_results


	def benchmark_cuda_graph_1step(fast_model, device):
	D = fast_model.config.embedding_dim
	dummy_q = torch.randn(1, D, device=device)
	dummy_q = F.normalize(dummy_q, dim=-1)

	cached_mem = fast_model.precompute_cross_memory(dummy_q)
	shape = (1, fast_model.config.sequence_length, D)
	static_noise = torch.randn(shape, device=device) * fast_model.config.sigma_max
	static_sigma = torch.full((1,), fast_model.config.sigma_max, device=device)

	# Warmup graph stream
	g_stream = torch.cuda.Stream()
	g_stream.wait_stream(torch.cuda.current_stream())
	with torch.cuda.stream(g_stream):
	for _ in range(5):
	_ = fast_model.forward_with_cached_memory(static_noise, static_sigma, cached_mem)
	torch.cuda.current_stream().wait_stream(g_stream)

	# Capture graph
	g = torch.cuda.CUDAGraph()
	with torch.cuda.graph(g, stream=g_stream):
	_ = fast_model.forward_with_cached_memory(static_noise, static_sigma, cached_mem)

	# Benchmark Replay
	torch.cuda.synchronize()
	times = []
	for _ in range(100):
	t0 = time.perf_counter()
	g.replay()
	torch.cuda.synchronize()
	times.append((time.perf_counter() - t0) * 1000.0)

	avg_ms = sum(times) / len(times)
	p50_ms = sorted(times)[int(len(times) * 0.5)]
	p95_ms = sorted(times)[int(len(times) * 0.95)]
	qps = 1000.0 / avg_ms
	return avg_ms, p50_ms, p95_ms, qps


	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, L, D = targets.shape

	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 = 128
	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)

	print(f"Loading Base Champion Checkpoint: {BASE_CHECKPOINT}...")
	ckpt = torch.load(BASE_CHECKPOINT, map_location=device, weights_only=False)
	config = ckpt["config"]
	model = B1EDMDenoiser(config, backend="tc", pure_1bit=False).to(device)

	if "ema_state_dict" in ckpt and "shadow" in ckpt["ema_state_dict"]:
	shadow = ckpt["ema_state_dict"]["shadow"]
	model.load_state_dict({k: shadow[k].to(device) for k in shadow})
	else:
	model.load_state_dict(ckpt["model_state_dict"])

	ema = ExponentialMovingAverage(model, decay=0.999)

	print("Loading SentenceTransformer and taxonomy bank...")
	embedder = SentenceTransformer("google/embeddinggemma-300m", model_kwargs={"torch_dtype": torch.bfloat16}, device=device)
	tax_emb, tax_names = load_taxonomy_bank(device)

	epochs = 4
	lr = 3e-4
	optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
	scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs * len(train_loader), eta_min=1e-5)

	journal = ExperimentJournal()
	tracker = journal.start_run(
	name="Champion B1: 1-Step Consistency Distillation (Tuned Diversity)",
	experiment_name="1-Bit Tensor Core Innovation",
	task_type="diffusion",
	config={
	"architecture": "consistency_distillation_1step",
	"backend": "b1_tc",
	"precision": "1bit_core_fp16_adapters",
	"layers": config.layers,
	"hidden_dim": config.hidden_dim,
	"epochs": epochs,
	"lr": lr,
	"batch_size": batch_size,
	"sampling_steps": 1,
	"target_ordering": "cosine",
	},
	tags=["1bit", "tensor_core", "consistency_distillation", "1step", "champion"],
	)

	print("\n" + "=" * 95)
	print("STARTING 1-STEP CONSISTENCY DISTILLATION TRAINING")
	print(f"Objective: Map pure Gaussian noise directly to semantic fanout in 1 single forward pass")
	print(f"Epochs: {epochs} \| Batch Size: {batch_size} \| Base Learning Rate: {lr}")
	print("=" * 95)

	sigma_max = config.sigma_max

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

	# 1. Direct 1-Step Prediction from pure noise (sigma_max)
	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)

	# 2. Intermediate noise level 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)

	# Losses:
	# Reconstruction Loss
	loss_recon = F.mse_loss(pred_1step, b_targets) + 0.5 * F.mse_loss(pred_mid, b_targets)

	# Centroid Prompt Alignment: Anchor the fanout cluster centroid to query
	pred_norm = F.normalize(pred_1step, dim=-1)
	centroid = F.normalize(pred_norm.mean(dim=1), dim=-1)
	align_cos = (centroid * b_queries).sum(dim=-1).mean()
	loss_align = 1.0 - align_cos

	# Pairwise Diversity Regularizer: Strongly penalize pairwise similarities above 0.30
	sim_mat = torch.bmm(pred_norm, pred_norm.transpose(1, 2))
	mask = ~torch.eye(L, dtype=torch.bool, device=device).unsqueeze(0).expand(B, -1, -1)
	pair_sims = sim_mat[mask]
	loss_div = F.relu(pair_sims - 0.30).mean()

	total_loss = loss_recon + 0.2 * loss_align + 1.2 * loss_div
	total_loss.backward()

	torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
	optimizer.step()

	# STE weight clamp for B1Linear
	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

	train_loss = train_loss_total / len(train_queries)

	# Validation
	model.eval()
	val_loss_total = 0.0
	val_align_total = 0.0
	val_div_total = 0.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_1step = model(pure_noise, sigmas_max, b_queries)

	v_loss = F.mse_loss(pred_1step, b_targets)
	val_loss_total += v_loss.item() * B

	p_norm = F.normalize(pred_1step, dim=-1)
	align = (p_norm * b_queries[:, None, :]).sum(dim=-1).mean().item()
	val_align_total += align * B

	sim_mat = torch.bmm(p_norm, p_norm.transpose(1, 2))
	mask = ~torch.eye(L, dtype=torch.bool, device=device).unsqueeze(0).expand(B, -1, -1)
	div = (1.0 - sim_mat[mask].mean().item()) * B
	val_div_total += div

	val_loss = val_loss_total / len(val_queries)
	val_align = val_align_total / len(val_queries)
	val_div = val_div_total / len(val_queries)
	ep_time = time.time() - ep_t0

	print(f"Epoch [{epoch}/{epochs}] \| Train Loss: {train_loss:.4f} \| Val MSE: {val_loss:.4f} \| 1-Step Align: {val_align:.3f} \| Div: {val_div:.3f} \| Time: {ep_time:.1f}s")

	tracker.log_metrics(
	step=epoch * len(train_loader),
	epoch=epoch,
	train_loss=train_loss,
	val_loss=val_loss,
	val_alignment=val_align,
	val_diversity=val_div,
	lr=scheduler.get_last_lr()[0],
	)

	# Save champion model
	print("\nSaving 1-Step Consistency Champion Checkpoint...")
	model.eval()
	model.freeze_for_inference()
	torch.save(
	{
	"epoch": epochs,
	"config": config,
	"model_state_dict": model.state_dict(),
	"val_loss": val_loss,
	"val_alignment": val_align,
	"val_diversity": val_div,
	},
	OUTPUT_CHECKPOINT,
	)
	print(f"Saved: {OUTPUT_CHECKPOINT} ({OUTPUT_CHECKPOINT.stat().st_size / (1024*1024):.2f} MB)")

	# Build FastB1Denoiser & CUDA Graph Benchmark
	print("\nBuilding FastB1Denoiser and Capturing CUDA Graph...")
	fast_model = FastB1Denoiser(model)
	avg_ms, p50_ms, p95_ms, qps = benchmark_cuda_graph_1step(fast_model, device)
	print(f"CUDA Graph 1-Step Latency: {avg_ms:.3f} ms ({p50_ms:.3f} ms P50, {p95_ms:.3f} ms P95) -> {qps:.1f} QPS (Batch=1)")

	# Qualitative Evaluation on Test Queries
	print("\nEvaluating Decoded Fanout Taxonomy Terms...")
	mean_align, mean_div, decoded = evaluate_qualitative_1step(fast_model, embedder, tax_emb, tax_names, device, tracker)

	for q, d in decoded.items():
	print(f"\nQuery: \"{q}\"")
	print(f" Alignment: {d['alignment']:.3f} \| Diversity: {d['diversity']:.3f}")
	for slot_idx, term in enumerate(d["terms"][:5]):
	print(f" #{slot_idx+1}: {term}")

	tracker.log_benchmark(
	latency_us=int(avg_ms * 1000),
	throughput_items_per_sec=qps,
	device_name=torch.cuda.get_device_name(0),
	notes=f"1-Step Consistency Distillation: {avg_ms:.2f} ms (B=1), Align: {mean_align:.3f}",
	)

	tracker.finish(
	status="completed",
	summary_metrics={
	"sampling_latency_ms": avg_ms,
	"latency_ms": avg_ms,
	"prompt_alignment": mean_align,
	"pairwise_diversity": mean_div,
	"val_loss": val_loss,
	"best_val_loss": val_loss,
	"throughput_items_per_sec": qps,
	"sampling_steps": 1,
	"mode": "consistency_distillation_1step",
	},
	)

	print("\n" + "=" * 95)
	print(f"All done! 1-Step Latency: {avg_ms:.3f} ms \| Alignment: {mean_align:.3f} \| Diversity: {mean_div:.3f}")
	print("Run logged to journal.db and visible in web UI at http://127.0.0.1:8000")
	print("=" * 95)


	if __name__ == "__main__":
	main()