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