Fix multi-head attention, checkpoint step persistence, and topology validation in simple_model.

pathwaysutils/extras/simple_model/config.py

@@ -14,7 +14,8 @@
 """Single universal configuration for the simple Transformer model."""
 
 import dataclasses
-import math
+
+from pathwaysutils import tpu_utils
 
 
 @dataclasses.dataclass(frozen=True)
@@ -54,5 +55,4 @@ class ModelConfig:
 
 def parse_topology(topology: str) -> int:
   """Returns the total number of chips in a topology string like '4x8' or '4x4x4'."""
-  dims = [int(x) for x in topology.strip().lower().split("x") if x.isdigit()]
-  return math.prod(dims) if dims else 1
+  return tpu_utils.parse_topology(topology)

pathwaysutils/extras/simple_model/model.py

@@ -60,13 +60,17 @@ def __init__(
     )
 
   def __call__(self, x: jax.Array) -> jax.Array:
+    b, t, _ = x.shape
     h = self.rmsnorm_1(x)
-    scores = (
-        jnp.einsum("btd,bsd->bts", self.attn_q(h), self.attn_k(h))
-        / jnp.sqrt(self.head_dim)
-    )
-    attn = jnp.einsum(
-        "bts,bsd->btd", jax.nn.softmax(scores, axis=-1), self.attn_v(h)
+    q = self.attn_q(h).reshape(b, t, self.num_heads, self.head_dim)
+    k = self.attn_k(h).reshape(b, t, self.num_heads, self.head_dim)
+    v = self.attn_v(h).reshape(b, t, self.num_heads, self.head_dim)
+    scores = jnp.einsum("bthd,bshd->bhts", q, k) / jnp.sqrt(self.head_dim)
+    mask = jnp.tril(jnp.ones((t, t), dtype=jnp.bool_))
+    scores = jnp.where(mask, scores, jnp.finfo(scores.dtype).min)
+    attn_weights = jax.nn.softmax(scores, axis=-1)
+    attn = jnp.einsum("bhts,bshd->bthd", attn_weights, v).reshape(
+        b, t, self.num_heads * self.head_dim
     )
     x = x + self.attn_o(attn)
     return x + self.mlp_down(geglu(self.mlp_gate_up(self.rmsnorm_2(x))))

pathwaysutils/extras/simple_model/simple_model_test.py

@@ -57,7 +57,10 @@ def tearDown(self):
   def test_config_dimensions_and_mesh(self):
     self.assertEqual(config_lib.parse_topology("4x8"), 32)
     self.assertEqual(config_lib.parse_topology("4x4x4"), 64)
-    self.assertEqual(config_lib.parse_topology("invalid"), 1)
+    with self.assertRaises(ValueError):
+      config_lib.parse_topology("invalid")
+    with self.assertRaises(ValueError):
+      config_lib.parse_topology("4xinvalid")
 
     self.assertIn("default", config_lib.CONFIGS)
     self.assertIn("test", config_lib.CONFIGS)
@@ -119,6 +122,14 @@ def test_model_forward_pass(self):
     scaled_logits = model(tokens)
     np.testing.assert_allclose(scaled_logits, logits * 2.0, rtol=1e-5)
 
+    # Verify causal attention masking: modifying a future token at t > 0 must
+    # not change logits at earlier positions.
+    tokens_modified = tokens.at[:, -1].set(1)
+    logits_modified = model(tokens_modified)
+    np.testing.assert_allclose(
+        logits_modified[:, :-1, :], scaled_logits[:, :-1, :], rtol=1e-5
+    )
+
   def test_synthetic_token_source(self):
     source = data_lib.SyntheticTokenSource(
         vocab_size=self.config.vocab_size,
@@ -231,6 +242,12 @@ def test_workload_defaults_and_snapshotter_init(self):
     self.assertEqual(workload.default_device, jax.devices()[0])
     self.assertEqual(workload.devices_per_slice, 0)
     self.assertIsNone(workload.snapshotter)
+    self.assertIsNone(workload.checkpoint_dir)
+
+    with self.assertRaisesRegex(
+        ValueError, "checkpoint_dir must be specified"
+    ):
+      workload_lib.Workload(config=self.config, checkpoint_interval=1)
 
     with self.assertRaises(RuntimeError):
       workload.run_step()
@@ -319,10 +336,12 @@ def test_snapshotter_integration(self):
       real_grain_st = workload.data_iterator.get_state()
       mock_snapshotter.load.return_value = {
           "params": fake_params,
+          "step": 2,
           "grain_state": real_grain_st,
       }
       workload.restore_from_snapshot()
       mock_snapshotter.load.assert_called_once()
+      self.assertEqual(workload.step, 2)
       self.assertEqual(workload.data_iterator.get_state(), real_grain_st)
     finally:
       workload.close()

pathwaysutils/extras/simple_model/workload.py

@@ -71,16 +71,18 @@ def __init__(
     self.total_steps = total_steps
     self.snapshot_interval = snapshot_interval
     self.checkpoint_interval = checkpoint_interval
-    self.checkpoint_dir = (
-        checkpoint_dir or f"/tmp/simple_model_ckpts_{int(time.time())}"
-    )
+    self.checkpoint_dir = checkpoint_dir
 
     self.snapshotter = snapshotter
     if self.snapshot_interval > 0 and self.snapshotter is None:
       self.snapshotter = ocp.v1.training.pathways.Snapshotter()
 
     self.checkpoint_manager: ocp.CheckpointManager | None = None
     if self.checkpoint_interval > 0:
+      if not self.checkpoint_dir:
+        raise ValueError(
+            "checkpoint_dir must be specified when checkpoint_interval > 0."
+        )
       self.checkpoint_manager = ocp.CheckpointManager(
           self.checkpoint_dir,
           options=ocp.CheckpointManagerOptions(enable_async_checkpointing=True),
@@ -157,7 +159,8 @@ def run_step(self) -> float:
       promoted = jax.device_put(self.params, self.artifacts.replicated_sharding)
       grain_st = self.data_iterator.get_state()
       self.snapshotter.save(
-          self.step, {"params": promoted, "grain_state": grain_st}
+          self.step,
+          {"params": promoted, "step": self.step, "grain_state": grain_st},
       )
       _logger.info("Saved in-memory snapshot at step %d.", self.step)
 
@@ -187,9 +190,12 @@ def restore_from_snapshot(self) -> None:
     self.snapshotter.join()
     payload = self.snapshotter.load({
         "params": self.artifacts.abstract_params,
+        "step": self.step,
         "grain_state": self.data_iterator.get_state(),
     })
     self.params = payload["params"]
+    if "step" in payload:
+      self.step = int(payload["step"])
     if "grain_state" in payload:
       self.data_iterator.set_state(payload["grain_state"])
     _logger.info("Restored state from in-memory snapshot.")
@@ -214,6 +220,9 @@ def close(self) -> None:
     """Cleans up resources and closes the data iterator and snapshotter."""
     if self.snapshotter is not None:
       self.snapshotter.join()
+    if self.checkpoint_manager is not None:
+      self.checkpoint_manager.close()
+      self.checkpoint_manager = None
     if self.data_iterator is not None:
       self.data_iterator.close()
       self.data_iterator = None