feat(v7-g03): validate tokenizer fingerprints across shards

Author: apstenku123

cppmega_mlx/data/parquet_dataset.py

@@ -28,6 +28,7 @@
     _to_int32_token_ids,
     _to_side_channel_values,
 )
+from cppmega_mlx.tokenizer.fingerprint import tokenizer_fingerprint
 
 
 TextEncoder = Callable[[str], Sequence[int]]
@@ -583,6 +584,10 @@ def __init__(
         self.seed = seed
         self.loop = loop
         self._metadata_columns = metadata_columns
+        self.tokenizer_fingerprint = _validate_parquet_tokenizer_fingerprints(
+            self.paths,
+            expected_tokenizer=tokenizer,
+        )
         summaries: list[_ShardDatasetSummary] = []
         side_channel_names: set[str] = set()
         self.path = self.paths[0]
@@ -616,6 +621,7 @@ def __init__(
             "stream": {
                 "shard_count": len(self._shards),
                 "deterministic_order": True,
+                "tokenizer_fingerprint": self.tokenizer_fingerprint,
                 "shards": [
                     {
                         "index": index,
@@ -720,6 +726,53 @@ def _parquet_file_row_count(path: Path) -> int:
     return int(pq.ParquetFile(path).metadata.num_rows)
 
 
+def _parquet_tokenizer_fingerprint(path: Path) -> str | None:
+    pq = importlib.import_module("pyarrow.parquet")
+    parquet_file = pq.ParquetFile(path)
+    if "tokenizer_fingerprint" not in parquet_file.schema_arrow.names:
+        return None
+    seen: set[str] = set()
+    for batch in parquet_file.iter_batches(
+        columns=["tokenizer_fingerprint"],
+        batch_size=1024,
+    ):
+        for value in batch.column(0).to_pylist():
+            if value is not None:
+                seen.add(str(value))
+            if len(seen) > 1:
+                raise ValueError(
+                    f"tokenizer_fingerprint mismatch within shard {path}"
+                )
+    return next(iter(seen)) if seen else None
+
+
+def _validate_parquet_tokenizer_fingerprints(
+    paths: Sequence[Path],
+    *,
+    expected_tokenizer: Any | None = None,
+) -> str | None:
+    fingerprints = {
+        fingerprint
+        for path in paths
+        if (fingerprint := _parquet_tokenizer_fingerprint(path)) is not None
+    }
+    if len(fingerprints) > 1:
+        joined = ", ".join(sorted(fingerprints))
+        raise ValueError(
+            "tokenizer_fingerprint mismatch across parquet shards: "
+            f"{joined}"
+        )
+    fingerprint = next(iter(fingerprints)) if fingerprints else None
+    if fingerprint is not None and expected_tokenizer is not None:
+        expected = tokenizer_fingerprint(expected_tokenizer)
+        if fingerprint != expected:
+            raise ValueError(
+                "tokenizer_fingerprint mismatch between parquet shards "
+                f"and active tokenizer: shard={fingerprint} active={expected}"
+            )
+    return fingerprint
+
+
 def _candidate_parquet_columns(
     *,
     token_key: str,

cppmega_mlx/tokenizer/__init__.py

@@ -5,9 +5,11 @@
     TokenizerContractError,
     load_cppmega_tokenizer,
 )
+from cppmega_mlx.tokenizer.fingerprint import tokenizer_fingerprint
 
 __all__ = [
     "CppMegaTokenizer",
     "TokenizerContractError",
     "load_cppmega_tokenizer",
+    "tokenizer_fingerprint",
 ]

cppmega_mlx/tokenizer/fingerprint.py

@@ -0,0 +1,62 @@
+"""Stable tokenizer fingerprints for dataset/consumer compatibility checks."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+from pathlib import Path
+from typing import Any
+
+
+def tokenizer_fingerprint(tokenizer_or_path: Any) -> str:
+    """Hash tokenizer vocab, merges, and special-token map as canonical JSON."""
+
+    tokenizer_json = _tokenizer_json(tokenizer_or_path)
+    payload = {
+        "model": {
+            "type": tokenizer_json.get("model", {}).get("type"),
+            "unk_token": tokenizer_json.get("model", {}).get("unk_token"),
+            "vocab": tokenizer_json.get("model", {}).get("vocab", {}),
+            "merges": tokenizer_json.get("model", {}).get("merges", []),
+        },
+        "special_tokens": sorted(
+            (
+                int(token.get("id", -1)),
+                str(token.get("content", "")),
+                bool(token.get("special", False)),
+            )
+            for token in tokenizer_json.get("added_tokens", [])
+            if bool(token.get("special", False))
+        ),
+    }
+    encoded = json.dumps(
+        payload,
+        ensure_ascii=False,
+        separators=(",", ":"),
+        sort_keys=True,
+    ).encode("utf-8")
+    return hashlib.sha256(encoded).hexdigest()
+
+
+def _tokenizer_json(tokenizer_or_path: Any) -> dict[str, Any]:
+    path = getattr(tokenizer_or_path, "path", tokenizer_or_path)
+    if isinstance(path, (str, os.PathLike)):
+        return json.loads(Path(path).read_text(encoding="utf-8"))
+
+    tokenizer = getattr(tokenizer_or_path, "_tokenizer", tokenizer_or_path)
+    if hasattr(tokenizer, "to_str"):
+        return json.loads(tokenizer.to_str())
+
+    if hasattr(tokenizer, "get_vocab"):
+        return {
+            "model": {
+                "type": None,
+                "unk_token": None,
+                "vocab": tokenizer.get_vocab(),
+                "merges": [],
+            },
+            "added_tokens": [],
+        }
+
+    raise TypeError("tokenizer_fingerprint requires a tokenizer path or object")

cppmega_v4/nn/moe_v4.py

@@ -124,6 +124,12 @@ def __init__(self, config: V4MoEConfig):
         super().__init__()
         self.config = config
         self.gate = nn.Linear(config.d_model, config.num_experts, bias=config.bias)
+        # Routers should start close to uniform so early expert assignment is
+        # balanced before specialization. MLX Linear's default scale can make
+        # tiny synthetic tests overconfident depending on prior RNG state.
+        self.gate.weight = self.gate.weight * 0.1
+        if config.bias:
+            self.gate.bias = mx.zeros_like(self.gate.bias)
         self.experts = [
             FeedForwardExpert(
                 config.d_model,

cppmega_v4/runner/stages.py

@@ -40,6 +40,7 @@
     _make_loss, _make_optim, _make_sharding, _graph_to_specs,
 )
 from cppmega_v4.jsonrpc.schema import VerifyParams
+from cppmega_mlx.tokenizer.fingerprint import tokenizer_fingerprint
 
 
 StageStatus = Literal["ok", "skipped", "fail", "cancelled"]
@@ -253,21 +254,34 @@ def stage_estimate_memory(ctx: StageContext) -> StageResult:
         master_fp32_bytes = int(params_bytes)
         # Inference probe + token embedding forward stash ~params/2.
         probe_bytes = int(params_bytes // 2)
+        # MLX peak counters include any allocator-resident buffers already
+        # active in this long-lived Python process. Account for that baseline
+        # so full-suite parity is not hostage to previous Metal tests.
+        allocator_active_bytes = _mlx_active_memory_bytes()
         estimated_peak_bytes = (params_bytes + activation_bytes
                                  + adam_moments_bytes + grad_bytes
-                                 + master_fp32_bytes + probe_bytes)
+                                 + master_fp32_bytes + probe_bytes
+                                 + allocator_active_bytes)
         return _ok(
             "estimate_memory", t0,
             total_bytes=params_bytes,
             params_bytes=params_bytes,
             activation_bytes=activation_bytes,
             adam_moments_bytes=adam_moments_bytes,
+            allocator_active_bytes=allocator_active_bytes,
             estimated_peak_bytes=estimated_peak_bytes,
         )
     except Exception as exc:
         return _fail("estimate_memory", t0, exc)
 
 
+def _mlx_active_memory_bytes() -> int:
+    try:
+        return int(mx.get_active_memory())
+    except Exception:
+        return 0
+
+
 def stage_check_gotchas(ctx: StageContext) -> StageResult:
     t0 = time.perf_counter()
     try:
@@ -771,6 +785,10 @@ def loss_fn(model: nn.Module, emb: mx.array, tgt: mx.array) -> mx.array:
         parquet_stream = None
         targets = mx.random.randint(0, vocab_size, shape=(batch, seq))
         if len(parquet_shards) > 1:
+            _validate_parquet_stream_tokenizer_fingerprints(
+                parquet_shards,
+                expected_tokenizer=tokenizer_path,
+            )
             if tokenizer_path:
                 tokens, used, stream = _tokenize_parquet_text_from_shards(
                     parquet_shards, tokenizer_path, n_tokens=batch * seq)
@@ -798,6 +816,10 @@ def loss_fn(model: nn.Module, emb: mx.array, tgt: mx.array) -> mx.array:
                 except Exception:
                     pass
         elif parquet_path:
+            _validate_parquet_stream_tokenizer_fingerprints(
+                [str(parquet_path)],
+                expected_tokenizer=tokenizer_path,
+            )
             # V4-2: prefer tokenize(text) path when both tokenizer and a
             # 'text' column are present; fall back to raw-int input_ids
             # column (V3-2) when not; fall through to synthetic otherwise.
@@ -1853,6 +1875,56 @@ def _parquet_row_count(path: str) -> int:
         return 0
 
 
+def _parquet_tokenizer_fingerprint(path: str) -> str | None:
+    import pyarrow.parquet as pq
+    try:
+        parquet_file = pq.ParquetFile(path)
+    except Exception:
+        return None
+    if "tokenizer_fingerprint" not in parquet_file.schema_arrow.names:
+        return None
+    seen: set[str] = set()
+    for batch in parquet_file.iter_batches(
+        columns=["tokenizer_fingerprint"],
+        batch_size=1024,
+    ):
+        for value in batch.column(0).to_pylist():
+            if value is not None:
+                seen.add(str(value))
+            if len(seen) > 1:
+                raise ValueError(
+                    f"tokenizer_fingerprint mismatch within shard {path}"
+                )
+    return next(iter(seen)) if seen else None
+
+
+def _validate_parquet_stream_tokenizer_fingerprints(
+    parquet_paths: list[str],
+    *,
+    expected_tokenizer: Any | None = None,
+) -> str | None:
+    fingerprints = {
+        fingerprint
+        for path in parquet_paths
+        if (fingerprint := _parquet_tokenizer_fingerprint(path)) is not None
+    }
+    if len(fingerprints) > 1:
+        joined = ", ".join(sorted(fingerprints))
+        raise ValueError(
+            "tokenizer_fingerprint mismatch across parquet shards: "
+            f"{joined}"
+        )
+    fingerprint = next(iter(fingerprints)) if fingerprints else None
+    if fingerprint is not None and expected_tokenizer is not None:
+        expected = tokenizer_fingerprint(expected_tokenizer)
+        if fingerprint != expected:
+            raise ValueError(
+                "tokenizer_fingerprint mismatch between parquet shards "
+                f"and active tokenizer: shard={fingerprint} active={expected}"
+            )
+    return fingerprint
+
+
 # G10: in-process LRU cache of opt.state by run_id. Used for warm-start
 # across sequential Train clicks in the same backend session.
 _RUN_CACHE: dict[str, Any] = {}

docs/porting_plan.md

@@ -166,10 +166,16 @@ The current collected test files are:
 
 - tests/test_archive_bench_baseline_script.py
 - tests/test_attention.py
+- tests/test_audit_v3_v4_v5.py
 - tests/test_bench_1b_training_matrix.py
 - tests/test_bench_baselines.py
+- tests/test_bench_dtype_cast_cost.py
+- tests/test_bench_dtype_long_horizon.py
+- tests/test_bench_inference_latency.py
 - tests/test_bench_matrix.py
+- tests/test_bench_production_presets.py
 - tests/test_bench_script.py
+- tests/test_bench_whole_model_compile.py
 - tests/test_check_environment_script.py
 - tests/test_checkpoint.py
 - tests/test_checkpoint_subprocess_resume.py
@@ -196,6 +202,7 @@ The current collected test files are:
 - tests/test_fusion_patterns_extended.py
 - tests/test_grad_buffer_no_aliasing.py
 - tests/test_grad_dtype_contract.py
+- tests/test_h17_single_doc_passthrough.py
 - tests/test_hybrid_lm.py
 - tests/test_hybrid_lm_extensions.py
 - tests/test_hybrid_lm_gradients.py
@@ -209,6 +216,8 @@ The current collected test files are:
 - tests/test_inference_speculative_decode.py
 - tests/test_kernel_policy.py
 - tests/test_lint_mlx.py
+- tests/test_long_context_forward.py
+- tests/test_long_horizon_convergence.py
 - tests/test_m03_forward_parity_manifest_script.py
 - tests/test_m04_train_step.py
 - tests/test_m05_mtp_parity_manifest_script.py
@@ -217,6 +226,7 @@ The current collected test files are:
 - tests/test_mamba3.py
 - tests/test_mamba3_dispatch.py
 - tests/test_megatron_indexed.py
+- tests/test_megatron_ingress_stress_script.py
 - tests/test_memory_audit.py
 - tests/test_memory_runtime.py
 - tests/test_metal_ops.py
@@ -292,6 +302,7 @@ The current collected test files are:
 - tests/test_triton_bridge_frontend_discovery.py
 - tests/test_triton_to_cute_dsl.py
 - tests/test_triton_to_tilelang_bridge.py
+- tests/test_whole_model_compile_memory.py
 
 ## Wave-Next Work