fix(training): preserve packed row masks in stepper

Author: apstenku123

cppmega_mlx/training/compiled.py

@@ -61,7 +61,10 @@
 
 STABLE_BATCH_KEYS = (
     "tokens",
+    "target_tokens",
     "attention_mask",
+    "loss_mask",
+    "document_ids",
     "structure_ids",
     "dep_levels",
     "ast_depth_ids",
@@ -538,8 +541,8 @@ def normalize_compiled_batch(
     mx.compile keys off the Python input structure as well as array shapes
     and dtypes.  Keep every optional side channel present in the dict and use
     None for absent fields so callers do not alternate between different
-    dict key sets when switching between plain token batches and structured
-    batches.
+    dict key sets when switching among plain token batches, packed-row
+    training batches, and structured batches.
     """
 
     batch_dict = ensure_lm_batch(batch).as_dict()

tests/test_compiled_train.py

@@ -10,7 +10,7 @@
 from mlx.utils import tree_flatten
 import pytest
 
-from cppmega_mlx.data.batch import synthetic_token_batch
+from cppmega_mlx.data.batch import LMTokenBatch, synthetic_token_batch
 from cppmega_mlx.models.hybrid_lm import HybridTinyConfig, HybridTinyLM
 from cppmega_mlx.models.tiny_lm import TinyLM, TinyLMConfig
 from cppmega_mlx.training.compiled import (
@@ -355,12 +355,59 @@ def test_normalize_compiled_batch_uses_fixed_keys_for_optional_fields() -> None:
     plain_tokens = plain_dict["tokens"]
     assert plain_tokens is not None
     assert plain_tokens.shape == (2, 8)
+    assert plain_dict["target_tokens"] is None
     assert plain_dict["attention_mask"] is None
+    assert plain_dict["loss_mask"] is None
+    assert plain_dict["document_ids"] is None
     assert plain_dict["structure_ids"] is None
+    assert structured_dict["target_tokens"] is None
     assert structured_dict["attention_mask"] is not None
+    assert structured_dict["loss_mask"] is None
+    assert structured_dict["document_ids"] is None
     assert structured_dict["structure_ids"] is not None
 
 
+def test_normalize_compiled_batch_preserves_packed_row_training_fields() -> None:
+    tokens = mx.array([[1, 2, 3, 0]], dtype=mx.int32)
+    target_tokens = mx.array([[2, 3, 0, 0]], dtype=mx.int32)
+    loss_mask = mx.array([[1, 0, 0, 0]], dtype=mx.float32)
+    document_ids = mx.array([[0, 0, 1, 1]], dtype=mx.int32)
+    batch = LMTokenBatch(
+        tokens=tokens,
+        target_tokens=target_tokens,
+        loss_mask=loss_mask,
+        document_ids=document_ids,
+    )
+
+    normalized = normalize_compiled_batch(batch.as_dict())
+
+    assert tuple(normalized) == STABLE_BATCH_KEYS
+    assert normalized["tokens"] is tokens
+    assert normalized["target_tokens"] is target_tokens
+    assert normalized["loss_mask"] is loss_mask
+    assert normalized["document_ids"] is document_ids
+
+
+@pytest.mark.parametrize("compile_step", [False, True])
+def test_pretraining_step_honors_packed_row_loss_mask(compile_step: bool) -> None:
+    mx.random.seed(911)
+    model = TinyLM(_tiny_config())
+    optimizer = optim.AdamW(learning_rate=1e-3, weight_decay=0.0)
+    batch = LMTokenBatch(
+        tokens=mx.array([[1, 2, 3, 4]], dtype=mx.int32),
+        target_tokens=mx.array([[2, 3, 4, 0]], dtype=mx.int32),
+        loss_mask=mx.array([[1, 0, 0, 0]], dtype=mx.float32),
+    )
+
+    metrics = CompiledPretrainingStep(model, optimizer, compile=compile_step)(batch)
+
+    assert metrics.compiled is compile_step
+    assert metrics.ntokens == 1
+    assert metrics.trained_tokens == 1
+    assert math.isfinite(metrics.loss)
+    assert metrics.loss > 0
+
+
 def test_normalize_compiled_batch_does_not_mutate_mapping_inputs() -> None:
     batch = synthetic_token_batch(
         batch_size=2,

tests/test_train_hybrid_tiny_script.py

@@ -110,6 +110,27 @@ def _copy_real_parquet_head(
     pq.write_table(pa.Table.from_batches([batch]), sample_path)
 
 
+def _write_packed_row_parquet(path: Path) -> None:
+    pa = pytest.importorskip("pyarrow")
+    pq = pytest.importorskip("pyarrow.parquet")
+
+    table = pa.table(
+        {
+            "input_ids": pa.array([[1, 2, 3, 4]], type=pa.large_list(pa.int32())),
+            "target_ids": pa.array([[2, 3, 4, 0]], type=pa.large_list(pa.int32())),
+            "loss_mask": pa.array(
+                [[1.0, 0.0, 0.0, 0.0]],
+                type=pa.large_list(pa.float32()),
+            ),
+            "doc_ids": pa.array([[0, 0, 1, 1]], type=pa.large_list(pa.int32())),
+            "pack_id": pa.array([7], type=pa.int32()),
+            "valid_token_count": pa.array([4], type=pa.int32()),
+            "num_docs": pa.array([2], type=pa.int32()),
+        }
+    )
+    pq.write_table(table, path)
+
+
 def _tiny_route_args(symbol: str, *, steps: int = 1) -> list[str]:
     args = [
         "--json",
@@ -868,6 +889,75 @@ def test_real_gb10_parquet_cli_smoke_trains_token_only_after_retokenize(
     assert payload["step_metrics"][0]["ntokens"] == 63
 
 
+def test_packed_row_parquet_cli_smoke_honors_explicit_targets_and_loss_mask(
+    tmp_path: Path,
+) -> None:
+    parquet_path = tmp_path / "packed_rows.parquet"
+    _write_packed_row_parquet(parquet_path)
+
+    result = run_script(
+        str(parquet_path),
+        "--json",
+        "--data-format",
+        "parquet",
+        "--token-key",
+        "input_ids",
+        "--batch-size",
+        "1",
+        "--seq-len",
+        "4",
+        "--steps",
+        "1",
+        "--hidden-size",
+        "8",
+        "--num-attention-heads",
+        "1",
+        "--pattern",
+        "A",
+        "--depth",
+        "1",
+        "--vocab-size",
+        "16",
+        "--no-compile",
+    )
+
+    payload = _load_json_result(result)
+
+    assert payload["status"] == "ok"
+    assert payload["synthetic_npz"] is False
+    assert payload["dataset"]["path"] == str(parquet_path)
+    assert payload["dataset"]["metadata"]["source_format"] == "parquet"
+    assert payload["dataset"]["token_key"] == "input_ids"
+    assert payload["dataset"]["num_samples"] == 1
+    assert payload["dataset"]["num_batches"] == 1
+    assert payload["tokens_per_step"] == 1
+    assert payload["trained_tokens"] == 1
+    assert payload["step_metrics"][0]["ntokens"] == 1
+    assert payload["step_metrics"][0]["trained_tokens"] == 1
+    assert payload["final_loss"] > 0
+
+    receipt = payload["dataset"]["dataset_receipt"]["parquet_receipt"]
+    assert receipt["token_source"] == {
+        "mode": "token_column",
+        "column": "input_ids",
+        "type": "large_list<element: int32>",
+    }
+    assert receipt["training_column_sources"] == {
+        "target_tokens": {
+            "column": "target_ids",
+            "type": "large_list<element: int32>",
+        },
+        "loss_mask": {
+            "column": "loss_mask",
+            "type": "large_list<element: float>",
+        },
+        "document_ids": {
+            "column": "doc_ids",
+            "type": "large_list<element: int32>",
+        },
+    }
+
+
 def test_dry_run_json_accepts_explicit_no_compile() -> None:
     result = run_script(
         "--dry-run-json",