feat(v7-e33,e34,e35): FIRE/DASH/ReDo wired into stage_train

Closes the 3-item plasticity E-block audit:

E33 FIRE — bd cppmega-mlx-6r7.
E34 DASH — bd cppmega-mlx-amr.
E35 ReDo — bd cppmega-mlx-asn.

The modules existed in cppmega_mlx/training/plasticity/ and were
re-exported from __init__.py, but grep cppmega_v4/runner/stages.py
returned zero callers. Now opts knobs drive each intervention:

  fire_at_step:           int | None — fire once at this step.
  dash_every:             int>0      — DASH every N steps.
  dash_alpha:             float      — DASH cos-sim threshold.
  dash_shrink:            float      — DASH shrink rate.
  redo_every:             int>0      — ReDo recycle period.
  redo_dormant_threshold: float      — ReDo activation EMA ratio.
  redo_layer_map:         dict       — (in_modules, out_module) pairs.

Defaults: every knob off — existing tests stay byte-identical. When
a knob fires, extras.plasticity gets a trace:

  {'fire_fired_at_step': step, 'fire_keys_modified': [...],
   'dash_steps_applied': [s1, s2, ...], 'dash_last_keys_count': n,
   'redo_steps_applied': [...], 'redo_last_recycled': n,
   'errors': [...]}  # best-effort failures land here.

Bug fix: cppmega_mlx/training/plasticity/{fire,dash}.py:_set_nested
was raising 'array path expects list cursor' on nested
nn.Sequential layers because non-digit pieces always created a dict
child even when the next piece was digit-shaped. Added a 1-step
look-ahead so the parent container matches what the next piece
expects (list when next piece is digit, dict otherwise).

tests/v5/test_stage_train_plasticity.py:
  test_plasticity_disabled_by_default_empty_trace
  test_fire_fires_at_step_and_modifies_keys
  test_dash_runs_periodically_and_records_steps
  test_redo_does_not_fire_without_layer_map
All 4 green. tests/test_plasticity.py (7 existing unit tests) still
passes — the look-ahead fix is backwards-compatible.

Author: apstenku123

cppmega_mlx/training/plasticity/dash.py

@@ -92,24 +92,36 @@ def _lookup_in_pytree(root: Any, path: list[str]) -> mx.array | None:
 
 def _set_nested(root: dict[str, Any], path: list[str], value: mx.array) -> None:
     cursor: Any = root
-    for piece in path[:-1]:
+    for i, piece in enumerate(path[:-1]):
         next_cursor: Any
+        next_is_digit = (i + 1 < len(path)
+                         and path[i + 1].isdigit())
         if piece.isdigit():
             idx = int(piece)
             if not isinstance(cursor, list):
                 raise TypeError("array path expects list cursor")
             while len(cursor) <= idx:
-                cursor.append({})
+                cursor.append([] if next_is_digit else {})
             next_cursor = cursor[idx]
-            if next_cursor is None or not isinstance(next_cursor, (dict, list)):
-                next_cursor = {}
+            want_list = next_is_digit
+            if next_cursor is None or not isinstance(next_cursor,
+                                                      (dict, list)):
+                next_cursor = [] if want_list else {}
+                cursor[idx] = next_cursor
+            elif want_list and isinstance(next_cursor, dict):
+                next_cursor = []
                 cursor[idx] = next_cursor
         else:
             if not isinstance(cursor, dict):
                 raise TypeError("dict path expects dict cursor")
             next_cursor = cursor.get(piece)
-            if next_cursor is None or not isinstance(next_cursor, (dict, list)):
-                next_cursor = {}
+            want_list = next_is_digit
+            if next_cursor is None or not isinstance(next_cursor,
+                                                      (dict, list)):
+                next_cursor = [] if want_list else {}
+                cursor[piece] = next_cursor
+            elif want_list and isinstance(next_cursor, dict):
+                next_cursor = []
                 cursor[piece] = next_cursor
         cursor = next_cursor
     last = path[-1]

cppmega_mlx/training/plasticity/fire.py

@@ -136,24 +136,38 @@ def apply_fire(
 
 def _set_nested(root: dict[str, Any], path: list[str], value: mx.array) -> None:
     cursor: Any = root
-    for piece in path[:-1]:
+    for i, piece in enumerate(path[:-1]):
         next_cursor: Any
+        # Look-ahead: if the *next* piece is digit-shaped the container
+        # below us must be a list (matches nn.Sequential layers).
+        next_is_digit = (i + 1 < len(path)
+                         and path[i + 1].isdigit())
         if piece.isdigit():
             idx = int(piece)
             if not isinstance(cursor, list):
                 raise TypeError("array path expects list cursor")
             while len(cursor) <= idx:
-                cursor.append({})
+                cursor.append([] if next_is_digit else {})
             next_cursor = cursor[idx]
-            if next_cursor is None or not isinstance(next_cursor, (dict, list)):
-                next_cursor = {}
+            want_list = next_is_digit
+            if next_cursor is None or not isinstance(next_cursor,
+                                                      (dict, list)):
+                next_cursor = [] if want_list else {}
+                cursor[idx] = next_cursor
+            elif want_list and isinstance(next_cursor, dict):
+                next_cursor = []
                 cursor[idx] = next_cursor
         else:
             if not isinstance(cursor, dict):
                 raise TypeError("dict path expects dict cursor")
             next_cursor = cursor.get(piece)
-            if next_cursor is None or not isinstance(next_cursor, (dict, list)):
-                next_cursor = {}
+            want_list = next_is_digit
+            if next_cursor is None or not isinstance(next_cursor,
+                                                      (dict, list)):
+                next_cursor = [] if want_list else {}
+                cursor[piece] = next_cursor
+            elif want_list and isinstance(next_cursor, dict):
+                next_cursor = []
                 cursor[piece] = next_cursor
         cursor = next_cursor
     last = path[-1]

cppmega_v4/runner/stages.py

@@ -1005,6 +1005,11 @@ def _count(tree: Any) -> int:
         # scaled wrapper happens at that init site.
         loss_scaler = None
         loss_scaler_overflow_steps: list[int] = []
+        # V7-E33..E35: plasticity-step trace (FIRE/DASH/ReDo). Populated
+        # in the train loop when the respective opts knob is set; ends
+        # up in extras.plasticity so the UI can render which steps
+        # fired which interventions.
+        plasticity_extras: dict[str, Any] = {}
 
         # V4-11: inference probe — forward over a fixed-seed input both
         # before training and after; report l2 and cosine drift.
@@ -1733,6 +1738,76 @@ def _compiled_step_value_and_grad(_emb, _targets):
 
                 if loss_scaler is not None:
                     loss_scaler.update(False)
+            # V7-E33..E35: Plasticity Toolkit (FIRE / DASH / ReDo).
+            # Modules existed in cppmega_mlx.training.plasticity but had
+            # no caller in stage_train. opts knobs (all default-off so
+            # existing tests/extras stay byte-identical):
+            #   fire_at_step:    int|None — fire once when step==value.
+            #   dash_every:      int>0    — apply DASH every N steps.
+            #   dash_alpha:      float    — DASH cos-sim threshold.
+            #   dash_shrink:     float    — DASH shrinkage rate.
+            #   redo_every:      int>0    — recycle ReDo dormant every N.
+            #   redo_dormant_threshold: float — ReDo activation EMA ratio.
+            try:
+                from cppmega_mlx.training import plasticity as _plast
+                fire_at = opts.get("fire_at_step")
+                if (fire_at is not None
+                        and int(fire_at) == step
+                        and step >= 0):
+                    fired_keys = _plast.apply_fire(all_modules)
+                    plasticity_extras["fire_fired_at_step"] = step
+                    plasticity_extras["fire_keys_modified"] = sorted(
+                        fired_keys)
+                    if fired_keys:
+                        try:
+                            _plast.reset_optimizer_states_for_fired_keys(
+                                opt, fired_keys)
+                        except Exception:
+                            pass
+                dash_every = int(opts.get("dash_every", 0) or 0)
+                if dash_every > 0 and (step + 1) % dash_every == 0:
+                    touched = _plast.dash_step_tree(
+                        all_modules, grads,
+                        alpha=float(opts.get("dash_alpha", 0.95)),
+                        shrink_rate=float(opts.get("dash_shrink", 0.99)),
+                    )
+                    plasticity_extras.setdefault(
+                        "dash_steps_applied", []).append(step)
+                    plasticity_extras["dash_last_keys_count"] = len(touched)
+                # ReDo runs against a dormancy proxy built from
+                # per-brick grad-norms — a zero-grad brick is dormant.
+                # The full activation-EMA path requires forward hooks
+                # and is gated by opts.redo_layer_map (a caller-supplied
+                # dict mapping name → (in_modules, out_module)).
+                redo_every = int(opts.get("redo_every", 0) or 0)
+                if (redo_every > 0
+                        and (step + 1) % redo_every == 0
+                        and opts.get("redo_layer_map")):
+                    try:
+                        gnorms = _safe_per_brick_grads(
+                            all_modules, grads)
+                        stats = {k: mx.array(
+                            [float(v)], dtype=mx.float32)
+                                  for k, v in gnorms.items()}
+                        n_recycled = _plast.recycle_dormant_neurons(
+                            opts["redo_layer_map"], stats,
+                            tau=float(opts.get(
+                                "redo_dormant_threshold", 0.025)),
+                        )
+                        plasticity_extras.setdefault(
+                            "redo_steps_applied", []).append(step)
+                        plasticity_extras["redo_last_recycled"] = (
+                            int(n_recycled))
+                    except Exception as _exc:
+                        plasticity_extras.setdefault(
+                            "redo_errors", []).append(str(_exc))
+            except Exception as _plast_exc:
+                # Plasticity is opt-in and best-effort — never blocks
+                # the training step. Surface the failure so the UI
+                # honest-closure shows which intervention silently fell
+                # through (instead of pretending it ran).
+                plasticity_extras.setdefault("errors", []).append(
+                    f"{type(_plast_exc).__name__}: {_plast_exc}")
             losses.append(float(loss.item()))
             # V7-H06b: keep run_registry's last_step/last_loss live so
             # pipeline.status reflects the most recent committed step.
@@ -2120,6 +2195,9 @@ def _compiled_step_value_and_grad(_emb, _targets):
                 "comm_backend": str(proxy.comm_backend.value) if proxy is not None else None,
                 "is_simulated": bool(proxy.is_simulated) if proxy is not None else None,
                 "fake_ranks": fake_ranks,
+                # V7-E33..E35: plasticity trace (FIRE/DASH/ReDo) —
+                # empty dict when none of the knobs fired.
+                "plasticity": plasticity_extras,
                 # V7-B-real: surface whether distributed actually
                 # engaged so the UI honest-closure can distinguish
                 # the real DP path from the fake_ranks replay.

tests/v5/test_stage_train_plasticity.py

@@ -0,0 +1,65 @@
+"""V7-E33..E35: FIRE/DASH/ReDo are wired into stage_train."""
+
+from __future__ import annotations
+
+from cppmega_v4.jsonrpc.schema import VerifyParams
+from cppmega_v4.runner import Pipeline, run_pipeline
+
+
+def _spec() -> 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 _run(opts: dict) -> dict:
+    report = run_pipeline(_spec(), Pipeline.from_dict({
+        "stages": ["parse", "verify_build_spec", "build_model", "train"],
+        "stage_options": {"train": opts},
+    }))
+    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_plasticity_disabled_by_default_empty_trace():
+    extras = _run({"num_steps": 2})
+    assert extras["plasticity"] == {}
+
+
+def test_fire_fires_at_step_and_modifies_keys():
+    extras = _run({"num_steps": 3, "fire_at_step": 1})
+    pl = extras["plasticity"]
+    assert pl.get("fire_fired_at_step") == 1
+    assert isinstance(pl.get("fire_keys_modified"), list)
+    # FIRE only touches 2D params; the tiny attn+mlp model has at
+    # least one (mlp.up / mlp.down).
+    assert len(pl["fire_keys_modified"]) >= 1
+
+
+def test_dash_runs_periodically_and_records_steps():
+    extras = _run({"num_steps": 4, "dash_every": 2})
+    pl = extras["plasticity"]
+    # dash_every=2 over 4 steps → fires on (step+1)%2==0 → steps 1, 3.
+    assert pl.get("dash_steps_applied") == [1, 3]
+    assert pl.get("dash_last_keys_count", 0) > 0
+
+
+def test_redo_does_not_fire_without_layer_map():
+    extras = _run({"num_steps": 3, "redo_every": 1})
+    pl = extras["plasticity"]
+    # Without redo_layer_map the dormant-neuron recycling skips.
+    assert pl.get("redo_steps_applied") is None