wip: test that no new tokens were introduced (gibberish), take explicit QPos/KPos

Author: jyc

src/levanter/data/ul2r.py

@@ -637,17 +637,20 @@ def s_length() -> jnp.ndarray:
     return jax.lax.cond(task_kind == RX_TASK_KIND, rx_length, s_length)
 
 
-@functools.partial(jax.jit, static_argnames=("max_segments_per_example", "Pos"))
+@functools.partial(jax.jit, static_argnames=("max_segments_per_example", "QPos", "KPos"))
 def create_ul2r_example(
     key: PRNGKeyArray,
     task_params: jnp.ndarray,
     task_indices: jnp.ndarray,
     max_segments_per_example: int,
-    Pos: Axis,
+    QPos: Axis,
+    KPos: Axis,
     pad_token_id: int,
     tokens: hax.NamedArray,
     segment_ids: hax.NamedArray,
 ) -> LmExample:
+    jax.debug.print("create_ul2r_example start")
+
     # TODO Use NamedArrays more idiomatically
     # `unique_seg_ids = [3, 4, ..., -1, ...]`
     # Valid segment IDs come first, padded with -1.
@@ -710,7 +713,7 @@ def process_segment(key: PRNGKeyArray, id: int) -> tuple[jnp.ndarray, jnp.ndarra
         out_start = jnp.squeeze(out_starts[idx])
 
         segment = jnp.roll(tokens.array, -in_start)
-        inputs_len, denoising_tokens = to_ul2r_tokens(key, task_params[task_idx], segment, in_length, Pos.size)
+        inputs_len, denoising_tokens = to_ul2r_tokens(key, task_params[task_idx], segment, in_length, QPos.size)
 
         n_tokens = tokens.array.shape[0]
         input_mask = jnp.arange(n_tokens) < inputs_len
@@ -750,18 +753,17 @@ def loop(
     # TODO GreedyPrepackedDataset pads w/ zeros so can we end up with two
     # padding token IDs?
     loss_mask = ul2r_loss_mask(input_mask, out_seg_ids, denoising_tokens, pad_token_id)
-    loss_mask = hax.named(loss_mask, Pos)
+    loss_mask = hax.named(loss_mask, QPos)
 
-    KPos = Pos.alias("KPos")
     attn_mask = AttentionMask(
         is_causal=True,
         is_prefix=True,
-        input_mask=hax.named(input_mask, [Pos]),
-        segment_ids=(hax.named(out_seg_ids, [Pos]), hax.named(out_seg_ids, [KPos])),
+        input_mask=hax.named(input_mask, [QPos]),
+        segment_ids=(hax.named(out_seg_ids, [QPos]), hax.named(out_seg_ids, [KPos])),
     )
 
-    denoising_tokens = hax.named(denoising_tokens, Pos)
-    out_seg_ids = hax.named(out_seg_ids, Pos)
+    denoising_tokens = hax.named(denoising_tokens, QPos)
+    out_seg_ids = hax.named(out_seg_ids, QPos)
     return LmExample(tokens=denoising_tokens, loss_mask=loss_mask, attn_mask=attn_mask)
 
 
@@ -773,7 +775,8 @@ class Ul2rDataset(MappedAsyncDataset[tuple[TokenizedDict, TokenizedDict], LmExam
     def __init__(
         self,
         cache: TreeCache[TokenizedDict],
-        Pos: Axis,
+        QPos: Axis,
+        KPos: Axis,
         task_configs: typing.Dict[str, DenoisingConfig],
         task_probs: Dict[str, float],
         key: PRNGKeyArray,
@@ -826,14 +829,15 @@ def _compute_length(task_idx: jnp.ndarray, length: jnp.ndarray) -> int:
         # packed leaves and the second has the segment ids
         self.packed: GreedyPrepackedDataset[TokenizedDict] = GreedyPrepackedDataset(
             cache.store.tree,
-            Pos.size,
+            QPos.size,
             max_segments_per_example=max_segments_per_example,
             slice_strategy=slice_strategy,
             packing_lengths=out_lengths,
             # Reserve space for UL2R; denoising examples increase in length.
             pad_with_zeros=True,
         )
-        self.Pos = Pos
+        self.QPos = QPos
+        self.KPos = KPos
         self.pad_token_id = pad_token_id
 
         sharding = jax.sharding.SingleDeviceSharding(jax.local_devices(backend="cpu")[0])
@@ -842,14 +846,15 @@ def _compute_length(task_idx: jnp.ndarray, length: jnp.ndarray) -> int:
         @functools.partial(eqx.filter_jit, out_shardings=sharding)
         def _create_lm_example(e: tuple[TokenizedDict, TokenizedDict]) -> LmExample:
             example, seg_ids = e
-            tokens = hax.named(example["input_ids"], self.Pos)
-            segment_ids = hax.named(seg_ids["input_ids"], self.Pos)
+            tokens = hax.named(example["input_ids"], self.QPos)
+            segment_ids = hax.named(seg_ids["input_ids"], self.QPos)
             return create_ul2r_example(
                 key,
                 task_params,
                 task_indices,
                 max_segments_per_example,
-                self.Pos,
+                self.QPos,
+                self.KPos,
                 self.pad_token_id,
                 tokens,
                 segment_ids,

tests/test_ul2r.py

@@ -2,7 +2,7 @@
 from pathlib import Path
 import tempfile
 import typing
-from haliax import Axis
+from haliax import Axis, NamedArray
 from haliax.nn import hax
 import jax
 import jax.numpy as jnp
@@ -29,6 +29,9 @@
     XDenoisingConfig,
     SDenoisingConfig,
     Ul2rDataset,
+    R_TASK_TOKEN_ID,
+    X_TASK_TOKEN_ID,
+    S_TASK_TOKEN_ID,
 )
 from levanter.layers.attention import AttentionMask
 from levanter.models.lm_model import LmExample
@@ -398,7 +401,8 @@ def test_to_ul2r_s_tokens():
 
 
 def test_create_ul2r_example():
-    Pos = Axis("position", 128)
+    QPos = Axis("QPos", 128)
+    KPos = Axis("KPos", 128)
     pad_token_id = 0
     max_segments_per_example = 8
 
@@ -411,7 +415,7 @@ def test_create_ul2r_example():
             jnp.zeros(105, dtype=jnp.int32),  # padding
         ]
     )
-    tokens = hax.named(tokens, Pos)
+    tokens = hax.named(tokens, QPos)
 
     segment_ids = jnp.concatenate(
         [
@@ -421,7 +425,7 @@ def test_create_ul2r_example():
             jnp.full(105, -1),
         ]
     )
-    segment_ids = hax.named(segment_ids, Pos)
+    segment_ids = hax.named(segment_ids, QPos)
 
     task_configs = [
         RDenoisingConfig(mask_prob=0.15, mean_span_length=3.0),
@@ -438,15 +442,16 @@ def test_create_ul2r_example():
         task_params,
         task_indices,
         max_segments_per_example,
-        Pos,
+        QPos,
+        KPos,
         pad_token_id,
         tokens,
         segment_ids,
     )
 
     # Basic smoke checks
-    assert example.tokens.array.shape == (Pos.size,)
-    assert example.loss_mask.array.shape == (Pos.size,)
+    assert example.tokens.array.shape == (QPos.size,)
+    assert example.loss_mask.array.shape == (QPos.size,)
     assert example.attn_mask.is_causal
 
     # Should contain sentinel tokens after denoising
@@ -517,7 +522,8 @@ def test_ul2r_dataset_build(dummy_text_data, hf_tokenizer):
         cache = typing.cast(TreeCache[TokenizedDict], cache)
 
         # Test Ul2rDataset
-        Pos = hax.Axis("position", 128)
+        QPos = hax.Axis("QPos", 128)
+        KPos = hax.Axis("KPos", 128)
         task_configs = {
             "r": RDenoisingConfig(mask_prob=0.15, mean_span_length=3.0),
             "x": XDenoisingConfig(mask_prob=0.5, mean_span_length=3.0),
@@ -526,7 +532,8 @@ def test_ul2r_dataset_build(dummy_text_data, hf_tokenizer):
 
         dataset = Ul2rDataset(
             cache=cache,
-            Pos=Pos,
+            QPos=QPos,
+            KPos=KPos,
             task_configs=task_configs,
             task_probs={"r": 0.33, "x": 0.33, "s": 0.34},
             key=jax.random.PRNGKey(123),
@@ -541,8 +548,8 @@ def test_ul2r_dataset_build(dummy_text_data, hf_tokenizer):
 
         # Structure checks
         assert isinstance(ex, LmExample)
-        assert ex.tokens.axes == (Pos,)
-        assert ex.loss_mask.axes == (Pos,)
+        assert ex.tokens.axes == (QPos,)
+        assert ex.loss_mask.axes == (QPos,)
         assert isinstance(ex.attn_mask, AttentionMask)
         assert ex.attn_mask.is_causal
 
@@ -557,22 +564,40 @@ def test_ul2r_dataset_build(dummy_text_data, hf_tokenizer):
         assert not jnp.any(ex.loss_mask.array & ~non_padding)  # No loss on padding
         assert jnp.any(jnp.isin(ex.tokens.array, SENTINEL_TOKEN_IDS))  # Has sentinels from denoising
 
+        # Collect all original input tokens from the cache
+        original_tokens = set()
+        for item in cache_sync:
+            original_tokens.update(int(t) for t in item["input_ids"] if t != pad_id)
+
+        # Check that all output tokens (except pad, sentinels, and task tokens) were present in the input
+        # This helps verify we're not creating gibberish by overlapping spans
+        ul2r_special_tokens = set(SENTINEL_TOKEN_IDS.tolist()) | {
+            R_TASK_TOKEN_ID,
+            X_TASK_TOKEN_ID,
+            S_TASK_TOKEN_ID,
+            pad_id,
+        }
+        allowed_tokens = original_tokens | ul2r_special_tokens
+        output_tokens = set(int(t) for t in ex.tokens.array)
+        unexpected_tokens = output_tokens - allowed_tokens
+        assert len(unexpected_tokens) == 0, f"Found unexpected tokens not in input: {unexpected_tokens}"
+
         # Attention mask checks
-        if ex.attn_mask.prefix_mask is not None:
-            assert ex.attn_mask.prefix_mask.array.shape == (Pos.size, Pos.size)
-            # Materialize full attention mask (causal + prefix)
-            materialized = ex.attn_mask.materialize(Pos, Pos)
-            assert materialized is not None
-            # Diagonal should be True for all non-padding (tokens attend to themselves)
-            diag = jnp.diag(materialized.array)
-            assert jnp.all(diag[non_padding])
-            # Some off-diagonal should be True (bidirectional attention on input positions)
-            off_diag_sum = jnp.sum(materialized.array) - jnp.sum(diag)
-            assert off_diag_sum > 0, "Expected some bidirectional attention for input positions"
+        input_mask = typing.cast(NamedArray, ex.attn_mask.input_mask)
+        assert input_mask.array.shape == (QPos.size,)
+        # Materialize full attention mask (causal + prefix)
+        materialized = ex.attn_mask.materialize(QPos, KPos)
+        assert materialized is not None
+        # Diagonal should be True for all non-padding (tokens attend to themselves)
+        diag = jnp.diag(materialized.array)
+        assert jnp.all(diag[non_padding])
+        # Some off-diagonal should be True (bidirectional attention on input positions)
+        off_diag_sum = jnp.sum(materialized.array) - jnp.sum(diag)
+        assert off_diag_sum > 0, "Expected some bidirectional attention for input positions"
 
         # Check consistency across multiple examples
         for ex_i in [dataset_sync[i] for i in range(min(3, len(dataset_sync)))]:
-            assert ex_i.tokens.axes == (Pos,) and ex_i.loss_mask.axes == (Pos,)
+            assert ex_i.tokens.axes == (QPos,) and ex_i.loss_mask.axes == (QPos,)
             non_pad_i = jnp.sum(ex_i.tokens.array != pad_id)
             loss_i = jnp.sum(ex_i.loss_mask.array)
             assert 0 < loss_i < non_pad_i