Upload verify_port_parity.py with huggingface_hub

verify_port_parity.py

@@ -0,0 +1,256 @@
+"""Verify the PyTorch port matches the Flax model numerically (< 1e-3 max-abs-diff).
+
+Checks:
+  1. Encoder output for a fixed input_ids tensor
+  2. Decoder logits at step 0 (empty past_kv) for a fixed decoder_input_id
+     using the encoder output from step 1
+
+Tolerance: max(abs(flax_out - pt_out)) < 1e-3
+
+Flax is run in float32 to avoid bfloat16 precision noise masking real bugs.
+"""
+
+import sys
+from pathlib import Path
+import pickle
+
+import numpy as np
+import torch
+
+# Make the Cactus Flax package importable
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "external" / "needle"))
+# Make the PyTorch port importable
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+import jax
+import jax.numpy as jnp
+
+from needle.model.architecture import SimpleAttentionNetwork, TransformerConfig as FlaxTransformerConfig
+from needle_torch import NeedleModel, TransformerConfig
+
+ART = Path(__file__).resolve().parent / "artifacts"
+
+TOLERANCE = 1e-3
+
+
+# ---------------------------------------------------------------------------
+# Load helpers
+# ---------------------------------------------------------------------------
+
+def load_flax_checkpoint():
+    """Load the locally cached needle.pkl and return (flax_params, config_dict)."""
+    path = ART / "needle.pkl"
+    print(f"Loading Flax checkpoint from {path} ...", flush=True)
+    with open(path, "rb") as f:
+        ckpt = pickle.load(f)
+    return ckpt["params"], ckpt["config"]
+
+
+def cast_params_to_f32(params):
+    """Recursively cast all JAX arrays in a nested param tree to float32."""
+    if isinstance(params, dict):
+        return {k: cast_params_to_f32(v) for k, v in params.items()}
+    arr = np.asarray(params).astype(np.float32)
+    return jnp.array(arr)
+
+
+def load_pt_model(config_dict):
+    cfg = TransformerConfig(**{k: v for k, v in config_dict.items()
+                               if k in TransformerConfig.__dataclass_fields__})
+    m = NeedleModel(cfg)
+    m.eval()
+    state = torch.load(ART / "needle_torch.pt", map_location="cpu", weights_only=True)
+    m.load_state_dict(state, strict=True)
+    return m, cfg
+
+
+# ---------------------------------------------------------------------------
+# Bisection helper
+# ---------------------------------------------------------------------------
+
+def bisect_encoder(flax_model, flax_params_f32, pt_model, ids_np):
+    """Compare encoder layer-by-layer to find the first divergent layer."""
+    print("\n--- Encoder bisection ---", flush=True)
+    ids_jax = jnp.asarray(ids_np)
+
+    # Flax intermediates via capture_intermediates
+    _, state = flax_model.apply(
+        {'params': flax_params_f32},
+        ids_jax,
+        capture_intermediates=True,
+        method=flax_model.encode_text,
+    )
+
+    print("Flax intermediates structure (top level):")
+    def print_tree(d, prefix='', depth=0):
+        if depth > 5:
+            return
+        if isinstance(d, dict):
+            for k, v in d.items():
+                if isinstance(v, dict):
+                    print(f"{'  '*depth}{prefix}{k}/")
+                    print_tree(v, prefix='', depth=depth+1)
+                else:
+                    shape = getattr(v, 'shape', '?')
+                    print(f"{'  '*depth}{prefix}{k}: {shape}")
+    print_tree(state['intermediates'])
+
+    # PyTorch intermediates via hooks
+    pt_intermediates = {}
+    hooks = []
+    for i, layer in enumerate(pt_model.encoder.layers):
+        def make_hook(idx):
+            def hook(module, inp, output):
+                pt_intermediates[f'encoder_layer_{idx}'] = output.detach().cpu().numpy()
+            return hook
+        hooks.append(layer.register_forward_hook(make_hook(i)))
+
+    def final_norm_hook(module, inp, output):
+        pt_intermediates['encoder_final_norm'] = output.detach().cpu().numpy()
+    hooks.append(pt_model.encoder.final_norm.register_forward_hook(final_norm_hook))
+
+    with torch.no_grad():
+        _ = pt_model.encoder(torch.from_numpy(ids_np.astype(np.int64)))
+
+    for h in hooks:
+        h.remove()
+
+    print(f"PyTorch intermediates captured: {list(pt_intermediates.keys())}", flush=True)
+
+
+def bisect_decoder_step0(flax_model, flax_params_f32, pt_model, dec_id_np, flax_enc_out, pt_enc_out):
+    """Compare decoder step-0 layer by layer."""
+    print("\n--- Decoder step-0 bisection ---", flush=True)
+    dec_id_jax = jnp.asarray(dec_id_np)
+
+    _, state = flax_model.apply(
+        {'params': flax_params_f32},
+        dec_id_jax,
+        flax_enc_out,
+        capture_intermediates=True,
+        method=flax_model.decode,
+    )
+    print("Flax decoder intermediates (top-level):", list(state['intermediates'].keys()), flush=True)
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+def main():
+    flax_params, config_dict = load_flax_checkpoint()
+    print(f"Config: {config_dict}", flush=True)
+
+    # Cast Flax params to float32 to avoid bfloat16 precision differences
+    print("Casting Flax params to float32 ...", flush=True)
+    flax_params_f32 = cast_params_to_f32(flax_params)
+
+    # Build Flax model with float32 dtype
+    config_dict_f32 = dict(config_dict, dtype="float32")
+    flax_cfg = FlaxTransformerConfig(**config_dict_f32)
+    flax_model = SimpleAttentionNetwork(flax_cfg)
+
+    # Load PyTorch model
+    pt_model, pt_cfg = load_pt_model(config_dict)
+
+    # Fixed input token sequence
+    np.random.seed(0)
+    ids_np = np.array(
+        [[2, 100, 200, 300, 400, 500, 5, 600, 700, 800, 900, 1000, 1100, 1200, 1300, 1]],
+        dtype=np.int32,
+    )
+    ids_jax = jnp.asarray(ids_np)
+
+    # ── Check 1: Encoder ────────────────────────────────────────────────────
+    print("\n=== Check 1: Encoder ===", flush=True)
+
+    # Flax encode returns (encoder_out, mask)
+    flax_enc_out, flax_enc_mask = flax_model.apply(
+        {'params': flax_params_f32},
+        ids_jax,
+        method=flax_model.encode,
+    )
+    flax_enc_np = np.asarray(flax_enc_out).astype(np.float32)
+
+    with torch.no_grad():
+        pt_enc_out = pt_model.encoder(
+            torch.from_numpy(ids_np.astype(np.int64))
+        ).cpu().numpy()
+
+    print(f"Flax encoder output shape: {flax_enc_np.shape}, stats: "
+          f"min={flax_enc_np.min():.4f} max={flax_enc_np.max():.4f} "
+          f"mean={flax_enc_np.mean():.4f}", flush=True)
+    print(f"PT   encoder output shape: {pt_enc_out.shape}, stats: "
+          f"min={pt_enc_out.min():.4f} max={pt_enc_out.max():.4f} "
+          f"mean={pt_enc_out.mean():.4f}", flush=True)
+
+    enc_diff = float(np.max(np.abs(flax_enc_np - pt_enc_out)))
+    enc_mean_diff = float(np.mean(np.abs(flax_enc_np - pt_enc_out)))
+    print(f"\nencoder max-abs-diff:  {enc_diff:.6f}", flush=True)
+    print(f"encoder mean-abs-diff: {enc_mean_diff:.6f}", flush=True)
+
+    enc_ok = enc_diff < TOLERANCE
+    if not enc_ok:
+        print(f"encoder parity FAILED (diff={enc_diff:.6f} >= {TOLERANCE}) -- bisecting ...", flush=True)
+        bisect_encoder(flax_model, flax_params_f32, pt_model, ids_np)
+        sys.exit(1)
+    else:
+        print(f"encoder parity OK (diff={enc_diff:.6f} < {TOLERANCE})", flush=True)
+
+    # ── Check 2: Decoder step 0 ─────────────────────────────────────────────
+    print("\n=== Check 2: Decoder step 0 ===", flush=True)
+
+    dec_id_np = np.array([[1]], dtype=np.int32)
+    dec_id_jax = jnp.asarray(dec_id_np)
+
+    # Flax: decode(tgt, encoder_out) -> logits (B, T_dec, vocab_size)
+    flax_logits = flax_model.apply(
+        {'params': flax_params_f32},
+        dec_id_jax,
+        flax_enc_out,
+        method=flax_model.decode,
+    )
+    flax_logits_np = np.asarray(flax_logits).astype(np.float32)
+
+    with torch.no_grad():
+        past_kv = pt_model.decoder.initial_past_kv(batch=1)
+        pt_logits, _ = pt_model.decoder.step(
+            torch.from_numpy(dec_id_np.astype(np.int64)),
+            torch.from_numpy(pt_enc_out),
+            past_kv,
+        )
+        pt_logits_np = pt_logits.cpu().numpy()
+
+    print(f"Flax logits shape: {flax_logits_np.shape}, stats: "
+          f"min={flax_logits_np.min():.4f} max={flax_logits_np.max():.4f}", flush=True)
+    print(f"PT   logits shape: {pt_logits_np.shape}, stats: "
+          f"min={pt_logits_np.min():.4f} max={pt_logits_np.max():.4f}", flush=True)
+
+    logits_diff = float(np.max(np.abs(flax_logits_np - pt_logits_np)))
+    logits_mean_diff = float(np.mean(np.abs(flax_logits_np - pt_logits_np)))
+    print(f"\ndecoder step-0 logits max-abs-diff:  {logits_diff:.6f}", flush=True)
+    print(f"decoder step-0 logits mean-abs-diff: {logits_mean_diff:.6f}", flush=True)
+
+    dec_ok = logits_diff < TOLERANCE
+    if not dec_ok:
+        print(f"decoder parity FAILED (diff={logits_diff:.6f} >= {TOLERANCE}) -- bisecting ...", flush=True)
+        bisect_decoder_step0(flax_model, flax_params_f32, pt_model, dec_id_np, flax_enc_out, pt_enc_out)
+        sys.exit(1)
+    else:
+        print(f"decoder step-0 parity OK (diff={logits_diff:.6f} < {TOLERANCE})", flush=True)
+
+    # ── Summary ─────────────────────────────────────────────────────────────
+    print("\n" + "="*60, flush=True)
+    print("port parity OK (< 1e-3)", flush=True)
+    print(f"  encoder max-abs-diff:        {enc_diff:.6f}", flush=True)
+    print(f"  decoder step-0 max-abs-diff: {logits_diff:.6f}", flush=True)
+
+    flax_argmax = int(np.argmax(flax_logits_np[0, 0]))
+    pt_argmax = int(np.argmax(pt_logits_np[0, 0]))
+    print(f"  Flax argmax token: {flax_argmax}", flush=True)
+    print(f"  PT   argmax token: {pt_argmax}", flush=True)
+    print("="*60, flush=True)
+
+
+if __name__ == "__main__":
+    main()