feat(v3-2): stage_train consumes real tokens from parquet when supplied

Closes V3-2 / cppmega-mlx-4pp. Fixes B2: stage_train hardcoded a
synthetic random-targets tensor, so the parquet shard + tokenizer the
user selected in the Data tab had zero effect on training.

Now: stage_options.train.parquet_path → _read_first_n_tokens(...) reads
the first batch*seq integers from input_ids / token_ids / tokens column
(reuses the column-detection convention from data_methods.py), clips
mod vocab_size, and uses them as the gradient targets. When the path
is absent, malformed, or lacks a token column the stage falls back to
the legacy random-targets path so E-4 matrix tests stay green.

extras gains:
  - data_source: "synthetic" | "parquet"
  - token_count: int (0 in synthetic mode)

These let Playwright assert the real-data code path actually executed
rather than silently degrading.

5 new pytest. Regression 85/85 across V3-1 / V3-2 / schedules / runner.

UI plumbing (passing parquet_path via stage_options.train from the
DataInspector selection) is a follow-up — backend exposure first.

Author: apstenku123

cppmega_v4/runner/stages.py

@@ -415,10 +415,32 @@ def loss_fn(model: nn.Module, emb: mx.array, tgt: mx.array) -> mx.array:
                 reduction="mean",
             )
 
-        # Build inputs once: synthetic Gaussian embeddings (no real
-        # tokenizer dependence — train matrix isolates the gradient path).
+        # V3-2: prefer real tokens from a parquet shard when opts
+        # supplies parquet_path. Falls back to synthetic random targets
+        # when absent (preserves E-4 matrix behaviour). data_source +
+        # token_count surface in extras so deep e2e can prove the real-
+        # data path actually executed instead of silently degrading.
         rng_key = mx.random.key(0)
-        targets = mx.random.randint(0, vocab_size, shape=(batch, seq))
+        data_source = "synthetic"
+        token_count = 0
+        parquet_path = opts.get("parquet_path")
+        if parquet_path:
+            try:
+                tokens = _read_first_n_tokens(parquet_path, n=batch * seq)
+                if len(tokens) >= batch * seq:
+                    tokens = [int(t) % vocab_size for t in tokens[:batch * seq]]
+                    targets = mx.array(tokens, dtype=mx.int32).reshape(
+                        batch, seq)
+                    data_source = "parquet"
+                    token_count = len(tokens)
+                else:
+                    targets = mx.random.randint(
+                        0, vocab_size, shape=(batch, seq))
+            except Exception:
+                targets = mx.random.randint(
+                    0, vocab_size, shape=(batch, seq))
+        else:
+            targets = mx.random.randint(0, vocab_size, shape=(batch, seq))
         all_modules = nn.Sequential(*modules, lm_head)
         opt, optimizer_kind = _build_optimizer(spec_optim, lr)
         loss_and_grad = nn.value_and_grad(all_modules, loss_fn)
@@ -508,6 +530,8 @@ def loss_fn(model: nn.Module, emb: mx.array, tgt: mx.array) -> mx.array:
                 "num_steps": n_steps,
                 "schedule_kind": schedule_kind_label,
                 "optimizer_kind": optimizer_kind,
+                "data_source": data_source,
+                "token_count": token_count,
                 "model_summary": _summarize_model(
                     ctx.spec, optimizer_kind, schedule_kind_label),
             },
@@ -521,6 +545,33 @@ def loss_fn(model: nn.Module, emb: mx.array, tgt: mx.array) -> mx.array:
 # ---------------------------------------------------------------------------
 
 
+def _read_first_n_tokens(parquet_path: str, n: int) -> list[int]:
+    """V3-2: read the first ``n`` token ids from a parquet shard so
+    stage_train can train against real corpus tokens instead of a
+    random target tensor. Reuses the column-detection convention from
+    cppmega_v4.jsonrpc.data_methods (input_ids / token_ids / tokens)."""
+    import pyarrow.parquet as pq
+    table = pq.read_table(parquet_path)
+    primary = None
+    for c in ("input_ids", "token_ids", "tokens"):
+        if c in table.column_names:
+            primary = c
+            break
+    if primary is None:
+        return []
+    column = table.column(primary)
+    out: list[int] = []
+    for chunk in column.chunks:
+        for cell in chunk.to_pylist():
+            if cell is None:
+                continue
+            for tok in cell:
+                out.append(int(tok))
+                if len(out) >= n:
+                    return out
+    return out
+
+
 def _build_optimizer(
     spec_optim: OptimSpec | None, base_lr: float,
 ) -> tuple[Any, str]:

tests/v4/test_stage_train_data.py

@@ -0,0 +1,112 @@
+"""V3-2: stage_train consumes real tokens from a parquet shard.
+
+Previously stage_train hardcoded a synthetic random-targets tensor and
+silently ignored any parquet/tokenizer the UI selected. Now opts may
+supply parquet_path; extras report data_source ('synthetic'|'parquet')
+and token_count for assertion.
+"""
+
+from __future__ import annotations
+
+import json
+import pathlib
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+import pytest
+
+from cppmega_v4.jsonrpc.schema import VerifyParams
+from cppmega_v4.runner import Pipeline, run_pipeline
+
+
+def _verify_params() -> VerifyParams:
+    return VerifyParams.model_validate({
+        "graph": {
+            "nodes": [
+                {"id": "attn", "kind": "attention", "params": {}},
+                {"id": "mlp", "kind": "mlp",
+                 "params": {"intermediate_size": 64, "activation": "swiglu"}},
+            ],
+            "edges": [{"src": "attn", "dst": "mlp"}],
+        },
+        "dim_env": {"B": 1, "S": 8, "H": 32, "nh": 2, "nkv": 1,
+                    "head_dim": 16},
+        "loss": {"kind": "cross_entropy", "head_outputs": ["mlp"]},
+        "optim": {"kind": "adamw",
+                  "groups": [{"matcher": "all", "lr": 1e-3,
+                              "weight_decay": 0.01,
+                              "betas": [0.9, 0.95]}]},
+    })
+
+
+def _write_parquet(tmp_path: pathlib.Path, column: str = "input_ids",
+                   n_rows: int = 4, n_tokens_per_row: int = 64) -> str:
+    rows = []
+    for r in range(n_rows):
+        rows.append(list(range(r * n_tokens_per_row,
+                               (r + 1) * n_tokens_per_row)))
+    table = pa.table({column: rows})
+    tmp_path.mkdir(parents=True, exist_ok=True)
+    path = tmp_path / "fixture.parquet"
+    pq.write_table(table, path)
+    return str(path)
+
+
+def _run(stage_options: dict) -> dict:
+    spec = _verify_params()
+    report = run_pipeline(spec, Pipeline.from_dict({
+        "stages": ["parse", "verify_build_spec", "build_model", "train"],
+        "stage_options": stage_options,
+    }))
+    train = next(s for s in report.stages if s.name == "train")
+    assert train.status == "ok", f"stage_train failed: {train.error}"
+    return train.extras
+
+
+def test_stage_train_synthetic_fallback_when_no_parquet():
+    """No parquet_path → falls back to synthetic random targets.
+    Backwards compat with E-4 train matrix."""
+    extras = _run(stage_options={"train": {"num_steps": 2}})
+    assert extras["data_source"] == "synthetic"
+    assert extras["token_count"] == 0
+
+
+def test_stage_train_reads_real_tokens_from_parquet(tmp_path):
+    """parquet_path → tokens loaded; extras.data_source='parquet'."""
+    path = _write_parquet(tmp_path)
+    extras = _run(stage_options={"train": {
+        "num_steps": 2, "parquet_path": path,
+    }})
+    assert extras["data_source"] == "parquet"
+    assert extras["token_count"] == 1 * 8  # batch=1, seq=8
+
+
+def test_stage_train_falls_back_when_parquet_missing_token_column(tmp_path):
+    """Parquet without input_ids/token_ids/tokens column → fallback."""
+    table = pa.table({"some_other_col": [[1, 2], [3, 4]]})
+    p = tmp_path / "no_tokens.parquet"
+    pq.write_table(table, p)
+    extras = _run(stage_options={"train": {
+        "num_steps": 2, "parquet_path": str(p),
+    }})
+    assert extras["data_source"] == "synthetic"
+    assert extras["token_count"] == 0
+
+
+def test_stage_train_falls_back_when_parquet_path_invalid():
+    """Bad parquet path → synthetic fallback (no exception leaks)."""
+    extras = _run(stage_options={"train": {
+        "num_steps": 2, "parquet_path": "/nonexistent/file.parquet",
+    }})
+    assert extras["data_source"] == "synthetic"
+    assert extras["token_count"] == 0
+
+
+def test_stage_train_uses_alt_column_names(tmp_path):
+    """token_ids and tokens are recognised in addition to input_ids."""
+    for col in ("token_ids", "tokens"):
+        p = _write_parquet(tmp_path / col, column=col)
+        extras = _run(stage_options={"train": {
+            "num_steps": 2, "parquet_path": p,
+        }})
+        assert extras["data_source"] == "parquet", f"col={col}"