feat(v3-1): stage_train honors OptimKind + extras.optimizer_kind + model_summary

Closes V3-1 / cppmega-mlx-wly. Fixes 3 confirmed UI→API→training bugs:

B1: stage_train hardcoded optim.AdamW. Lion/Lion8bit/Adam8bit/Muon/
    Muon_AdamW_Hybrid/SGD selected in UI OptimTab were silently ignored
    at train. Now _build_optimizer(spec_optim, lr) dispatches every
    OptimKind through cppmega_mlx.training.optimizers factories.

B1-bis: build_model raised "unsupported optimizer kind" for Lion/Lion8bit/
    Adam8bit. Pipeline orchestrator swallowed the BuildError and marked
    train "skipped" — UI showed superficial green. Now build_model wires
    these three through the same factories.

B1-tris: _make_optim dropped ns_steps and schedule when converting wire
    payload to OptimSpec. Muon-from-UI failed validation; scheduled groups
    silently lost their lr curve. Now both fields thread through with
    kind-aware defaults.

V3-3 (model_summary helper) lands alongside: _summarize_model surfaces
mlp_activation, attention pre/post_norm, mlp pre/post_norm, optimizer_kind,
schedule_kind, num_brick_kinds — enables UI-side assertions that a brick
context change reached training.

Also fixed: stage_train's forward chain silently zeroed gradients for
every q/k/v after attention because attention's zero-init-out (o_proj=
zeros for identity-at-init) killed downstream gradient flow. Real
transformer blocks have residuals; now stage_train adds residual when
shapes match so gradients propagate as in production.

19 new pytest in tests/v4/test_stage_train_optimizers.py. Regression
92/92 across pipeline + schedules + jsonrpc_methods + new file.

Author: apstenku123

cppmega_v4/buildspec/api.py

@@ -20,8 +20,8 @@
 from __future__ import annotations
 
 import time
-from collections.abc import Callable, Mapping, Sequence
-from dataclasses import dataclass, field
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
 from typing import Any
 
 import mlx.core as mx
@@ -35,7 +35,7 @@
 from cppmega_v4.buildspec.loss_spec import LossKind, LossSpec
 from cppmega_v4.buildspec.model_build_spec import ModelBuildSpec
 from cppmega_v4.buildspec.optim_spec import OptimKind, OptimSpec, ParamGroup
-from cppmega_v4.fusion.brick_graph import BrickGraph, BrickNode
+from cppmega_v4.fusion.brick_graph import BrickGraph
 from cppmega_v4.models.unified_superblock_v4 import BLOCK_BUILDERS
 
 
@@ -48,6 +48,14 @@
     # adapter-prefixed bricks are pure-shape rewires; their forward is
     # a passthrough (or handled separately via a future Stage F)
 })
+_AUTOGRAD_SAFE_KERNEL_KINDS: frozenset[str] = frozenset({"gdn", "kda"})
+
+
+def _autograd_safe_params(kind: str, params: Mapping[str, Any]) -> dict[str, Any]:
+    out = dict(params)
+    if kind in _AUTOGRAD_SAFE_KERNEL_KINDS and "kernel_path" not in out:
+        out["kernel_path"] = "path_a"
+    return out
 
 
 class BuiltSequentialModel(nn.Module):
@@ -67,9 +75,12 @@ def __init__(self, graph: BrickGraph, hidden_size: int):
         # them up as sub-modules (for parameter registration).
         self._node_attrs: dict[str, str] = {}
         for node in graph.nodes:
+            params = _autograd_safe_params(node.kind, node.params)
             module = node.module
+            if node.kind in _AUTOGRAD_SAFE_KERNEL_KINDS and "kernel_path" not in node.params:
+                module = None
             if module is None and node.kind in BLOCK_BUILDERS:
-                module = BLOCK_BUILDERS[node.kind](hidden_size, dict(node.params))
+                module = BLOCK_BUILDERS[node.kind](hidden_size, params)
             if module is None:
                 continue
             attr = f"brick_{_safe_name(node.name)}"
@@ -235,8 +246,33 @@ def _build_optimizer(spec: OptimSpec) -> object:
     nn.Module params is performed by callers (Stage F has the full
     matcher loop)."""
     if spec.kind is OptimKind.ADAMW:
+        from cppmega_mlx.training.optimizers import make_adamw
+        g = spec.groups[0]
+        return make_adamw(
+            learning_rate=g.lr,
+            betas=list(g.betas) if g.betas is not None else [0.9, 0.999],
+            weight_decay=g.weight_decay,
+        )
+    if spec.kind is OptimKind.LION:
+        from cppmega_mlx.training.optimizers import make_lion
+        g = spec.groups[0]
+        return make_lion(
+            learning_rate=g.lr,
+            betas=list(g.betas) if g.betas is not None else [0.9, 0.99],
+            weight_decay=g.weight_decay,
+        )
+    if spec.kind is OptimKind.LION_8BIT:
+        from cppmega_mlx.training.optimizers_quantized import make_lion8bit
+        g = spec.groups[0]
+        return make_lion8bit(
+            learning_rate=g.lr,
+            betas=list(g.betas) if g.betas is not None else [0.9, 0.99],
+            weight_decay=g.weight_decay,
+        )
+    if spec.kind is OptimKind.ADAM_8BIT:
+        from cppmega_mlx.training.optimizers_quantized import make_adam8bit
         g = spec.groups[0]
-        return optim.AdamW(
+        return make_adam8bit(
             learning_rate=g.lr,
             betas=list(g.betas) if g.betas is not None else [0.9, 0.999],
             weight_decay=g.weight_decay,

cppmega_v4/jsonrpc/methods.py

@@ -27,6 +27,7 @@
     OptimSpec,
     ParamGroup,
 )
+from cppmega_v4.buildspec.schedules import ScheduleSpec
 from cppmega_v4.fusion import from_block_specs
 from cppmega_v4.parallelism import (
     AxisAssignment,
@@ -43,18 +44,6 @@
     tpu_v6e_8,
     verify_distributed_plan,
 )
-
-
-_TOPOLOGY_FACTORIES = {
-    "h100_8x": h100_8x,
-    "h200_8x": h200_8x,
-    "a100_8x": a100_8x,
-    "b100_8x": b100_8x,
-    "gb10_quarter": gb10_quarter,
-    "tpu_v5p_4": tpu_v5p_4,
-    "tpu_v6e_8": tpu_v6e_8,
-    "m3_ultra_solo": m3_ultra_solo,
-}
 from cppmega_v4.probe import contract_probe as _contract_probe
 from cppmega_v4.probe import to_dict as _probe_to_dict
 from cppmega_v4.spec import (
@@ -94,6 +83,18 @@
 )
 
 
+_TOPOLOGY_FACTORIES = {
+    "h100_8x": h100_8x,
+    "h200_8x": h200_8x,
+    "a100_8x": a100_8x,
+    "b100_8x": b100_8x,
+    "gb10_quarter": gb10_quarter,
+    "tpu_v5p_4": tpu_v5p_4,
+    "tpu_v6e_8": tpu_v6e_8,
+    "m3_ultra_solo": m3_ultra_solo,
+}
+
+
 # ---------------------------------------------------------------------------
 # Coercion helpers: wire payload → backend dataclasses.
 # ---------------------------------------------------------------------------
@@ -123,16 +124,46 @@ def _make_loss(payload: LossSpecPayload) -> LossSpec:
 
 
 def _make_optim(payload: OptimSpecPayload) -> OptimSpec:
+    kind = OptimKind(payload.kind)
+
+    def _default_betas() -> tuple[float, float] | None:
+        if kind in (OptimKind.LION, OptimKind.LION_8BIT):
+            return (0.9, 0.99)
+        if kind is OptimKind.ADAM_8BIT:
+            return (0.9, 0.999)
+        if kind in (OptimKind.ADAMW, OptimKind.MUON_ADAMW_HYBRID):
+            return (0.9, 0.95)
+        return None
+
+    def _make_schedule(g) -> ScheduleSpec | None:
+        schedule = g.schedule
+        if schedule is None:
+            return None
+        return ScheduleSpec(
+            kind=schedule.kind,
+            warmup_steps=schedule.warmup_steps,
+            total_steps=schedule.total_steps,
+            min_lr_ratio=schedule.min_lr_ratio,
+            decay_steps=schedule.decay_steps,
+            power=schedule.power,
+        )
+
     groups = tuple(
         ParamGroup(
             matcher=g.matcher,
             lr=g.lr,
             weight_decay=g.weight_decay,
-            betas=g.betas if g.betas is not None else (0.9, 0.95),
+            betas=g.betas if g.betas is not None else _default_betas(),
+            ns_steps=(
+                g.ns_steps
+                if g.ns_steps is not None
+                else (5 if kind is OptimKind.MUON else None)
+            ),
+            schedule=_make_schedule(g),
         )
         for g in payload.groups
     )
-    return OptimSpec(kind=OptimKind(payload.kind), groups=groups)
+    return OptimSpec(kind=kind, groups=groups)
 
 
 def _make_topology(payload: TopologyPayload):