feat(v4-9): hybrid optimizer bucket sizes in extras

Closes V4-9 / cppmega-mlx-9e7. Closes G8 from V4 audit (split was
wire-only proven).

Backend (stage_train):
  - When optimizer_kind === 'muon_adamw_hybrid', call
    split_param_groups on all_modules.parameters() to count parameter
    elements routed to Muon (2-D matmul) vs AdamW (1-D / 3-D+).
  - extras.muon_group_size + extras.adamw_group_size populated.
  - Both null for non-hybrid optimizers.

E2E 23_hybrid_optimizer_split.spec.ts: UI selects muon_adamw_hybrid,
Train, asserts both bucket sizes > 0 AND muon >= adamw (matmul
weights dominate norm scalars + biases by orders of magnitude).

1/1 green; stage_train pytest 27/27 still passing.

Author: apstenku123

cppmega_v4/runner/stages.py

@@ -457,6 +457,23 @@ def loss_fn(model: nn.Module, emb: mx.array, tgt: mx.array) -> mx.array:
                     pass
         all_modules = nn.Sequential(*modules, lm_head)
         opt, optimizer_kind = _build_optimizer(spec_optim, lr)
+        # V4-9: when hybrid, count params routed to each bucket so e2e can
+        # prove the split predicate actually saw 2D vs 1D/3D parameters.
+        muon_group_size: int | None = None
+        adamw_group_size: int | None = None
+        if optimizer_kind == "muon_adamw_hybrid":
+            try:
+                from cppmega_mlx.training.optimizers import split_param_groups
+                muon_t, adamw_t = split_param_groups(
+                    all_modules.parameters())
+                def _count(tree: Any) -> int:
+                    flat = dict(nn.utils.tree_flatten(tree))
+                    return sum(int(v.size) for v in flat.values()
+                               if hasattr(v, "size"))
+                muon_group_size = _count(muon_t)
+                adamw_group_size = _count(adamw_t)
+            except Exception:
+                pass
         loss_and_grad = nn.value_and_grad(all_modules, loss_fn)
 
         losses: list[float] = []
@@ -552,6 +569,8 @@ def loss_fn(model: nn.Module, emb: mx.array, tgt: mx.array) -> mx.array:
                     if getattr(ctx.spec, "loss", None) is not None
                     else "cross_entropy"
                 ),
+                "muon_group_size": muon_group_size,
+                "adamw_group_size": adamw_group_size,
                 "model_summary": _summarize_model(
                     ctx.spec, optimizer_kind, schedule_kind_label),
             },

vbgui/e2e/scenarios/23_hybrid_optimizer_split.spec.ts

@@ -0,0 +1,41 @@
+// V4-9: muon_adamw_hybrid optimizer routes 2D matmul weights to Muon
+// and 1D/3D+ tensors to AdamW. extras.muon_group_size + adamw_group_size
+// must both be > 0 to prove the split predicate actually fired.
+
+import { test, expect } from "@playwright/test";
+import { gotoApp, selectPreset, closeModal } from "../fixtures";
+import { readTrainExtras } from "../utils/train_extras";
+
+test("V4-9: hybrid optimizer split reports non-empty muon + adamw buckets",
+  async ({ page }) => {
+    test.setTimeout(60_000);
+    await gotoApp(page);
+    await selectPreset(page, "llama3_8b");
+
+    await page.getByTestId("sidebar-tab-optim").click();
+    await page.getByTestId("optim-kind").selectOption("muon_adamw_hybrid");
+    await page.getByTestId("optim-apply").click();
+
+    await page.getByTestId("run-pipeline-toggle").click();
+    await page.getByTestId("run-pipeline-train").click();
+    const modal = page.getByTestId("run-result-modal");
+    await modal.waitFor({ timeout: 60_000 });
+    const extras = await readTrainExtras(page);
+
+    expect(extras.optimizer_kind).toBe("muon_adamw_hybrid");
+
+    const muonSize = parseInt(
+      (await page.getByTestId("run-result-extras-train-muon_group_size")
+        .textContent()) ?? "0", 10);
+    const adamwSize = parseInt(
+      (await page.getByTestId("run-result-extras-train-adamw_group_size")
+        .textContent()) ?? "0", 10);
+
+    expect(muonSize).toBeGreaterThan(0);
+    expect(adamwSize).toBeGreaterThan(0);
+    // Sanity: Muon handles big matmul weights → its bucket should be
+    // at least as large as the AdamW bucket (norm scalars + biases).
+    expect(muonSize).toBeGreaterThanOrEqual(adamwSize);
+
+    await closeModal(page);
+  });