feat(v8-r04): architectures.auto_fit + UI Auto-fit button

apstenku123 · apstenku123 · commit 5b877f4fecd8 · 2026-05-23T19:16:59.000+02:00
Closes cppmega-mlx-yz6n.4. New RPC architectures.auto_fit(preset,
host_info?) probes the platform (via platform.get_info when host_info
is null), picks total_hbm_bytes × headroom as the byte budget, runs
scale_down to land a fitting (H, L), then calls suggest_sharding with
a minimal default cross_entropy + adamw spec to get ranked sharding
proposals. Bundle returned: scaled + sharding + fits + reason.

UI: GalleryScaleDownSlider grows an "Auto-fit to my devbox" button
that fires the RPC and renders a result chip with the chosen topology
+ peak/budget reason string.

AC: llama3_8b on gb10_quarter fits the canonical 8B (51.67 GB / 137 GB)
end-to-end through scale_down + suggest_sharding.

Tests: 6/6 new pytest (fit/no-fit, unknown topology, platform-probe
fallback, dispatch e2e, headroom monotonicity), 1/1 new vitest case
(auto-fit button click → banner). Regression: 899 pytest + 459 vitest.
diff --git a/cppmega_v4/jsonrpc/auto_fit_method.py b/cppmega_v4/jsonrpc/auto_fit_method.py
@@ -0,0 +1,194 @@
+"""V8-R04: ``architectures.auto_fit`` — given a preset, pick the best
+``(hidden_size, num_layers)`` AND a sharding proposal for the detected
+(or supplied) devbox in one shot.
+
+Strategy:
+  1. Resolve the target topology — either ``host_info.topology`` from
+     the caller, or pick the first ``available_topologies`` entry from
+     ``platform.get_info``.
+  2. Compute ``target_bytes = total_hbm × headroom`` and run
+     :func:`architectures.scale_down.scale_down` to land a model that
+     fits the device.
+  3. Call :func:`suggest_sharding` against the scaled spec with a
+     minimal default ``loss=cross_entropy`` + ``optim=adamw`` so the
+     planner has enough to produce ranked proposals.
+  4. Return the bundle: scaled / sharding / fits / reason.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict
+
+from cppmega_v4.architectures.scale_down import scale_down as _scale_down
+from cppmega_v4.jsonrpc.cache import LRUCache
+from cppmega_v4.jsonrpc.methods import (
+    _cache_lookup, _cache_store, suggest_sharding,
+)
+from cppmega_v4.jsonrpc.schema import (
+    ScaleDownFromCanonical, ScaleDownResultModel,
+    SuggestShardingParams, SuggestShardingResult,
+)
+from cppmega_v4.parallelism import topology as _topo
+
+
+__all__ = [
+    "AutoFitParams",
+    "AutoFitHostInfo",
+    "AutoFitResult",
+    "auto_fit",
+    "TOPOLOGY_BUILDERS",
+]
+
+
+# Same table as memory_matrix_method but kept local to avoid a cycle.
+TOPOLOGY_BUILDERS: dict[str, Any] = {
+    "h100_8x":       lambda: _topo.h100_8x(),
+    "h200_8x":       lambda: _topo.h200_8x(),
+    "a100_8x":       lambda: _topo.a100_8x(),
+    "b100_8x":       lambda: _topo.b100_8x(),
+    "gb10_quarter":  lambda: _topo.gb10_quarter(),
+    "tpu_v6e_8":     lambda: _topo.tpu_v6e_8(),
+    "tpu_v5p_4":     lambda: _topo.tpu_v5p_4(),
+    "m3_ultra_solo": lambda: _topo.m3_ultra_solo(),
+}
+
+
+class AutoFitHostInfo(BaseModel):
+    """When supplied, overrides the server-side platform probe."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    topology: str
+    headroom: float = 0.9
+
+
+class AutoFitParams(BaseModel):
+    """Input — preset name + optional host overrides."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    preset: str
+    host_info: AutoFitHostInfo | None = None
+
+
+class AutoFitResult(BaseModel):
+    """Output — scale-down result + sharding proposals + fits verdict."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    scaled: ScaleDownResultModel
+    sharding: SuggestShardingResult
+    fits: bool
+    reason: str
+    topology: str
+    headroom: float
+
+
+def _resolve_topology(params: AutoFitParams) -> tuple[str, float]:
+    """Pick (topology_name, headroom) from host_info or platform probe."""
+    if params.host_info is not None:
+        if params.host_info.topology not in TOPOLOGY_BUILDERS:
+            raise ValueError(
+                f"unknown topology {params.host_info.topology!r}; "
+                f"choose from {sorted(TOPOLOGY_BUILDERS)}")
+        return params.host_info.topology, params.host_info.headroom
+    from cppmega_v4.runtime.platform_probe import probe_platform
+    info = probe_platform()
+    available = info.get("available_topologies") or []
+    for t in available:
+        if t in TOPOLOGY_BUILDERS:
+            return t, 0.9
+    return "m3_ultra_solo", 0.9   # final fallback
+
+
+def auto_fit(
+    params: AutoFitParams, *, cache: LRUCache | None = None,
+) -> AutoFitResult:
+    """Chain scale_down + suggest_sharding for the detected/given devbox."""
+    key, hit = _cache_lookup(cache, "architectures.auto_fit", params)
+    if hit is not None:
+        return hit
+
+    topo_name, headroom = _resolve_topology(params)
+    topo = TOPOLOGY_BUILDERS[topo_name]()
+    total_hbm = topo.total_hbm_bytes
+    target_bytes = int(total_hbm * headroom)
+    scaled = _scale_down(params.preset, target_bytes)
+
+    # Build minimal sharding query payload — use the scaled specs +
+    # cross_entropy + adamw defaults. We feed the scaled hidden_size
+    # so the suggester sees the right model size, and head_outputs to
+    # the last brick so verify_build_spec accepts the request.
+    graph_nodes = []
+    graph_edges = []
+    prev_name: str | None = None
+    for spec in scaled.specs:
+        if "parallel" in spec:
+            # Each branch becomes a leaf for the suggester.
+            for leaf in spec["parallel"]:
+                name = leaf.get("name") or leaf.get("kind")
+                graph_nodes.append({
+                    "id": name, "kind": leaf["kind"],
+                    "params": leaf.get("params", {}),
+                })
+                if prev_name is not None:
+                    graph_edges.append({"src": prev_name, "dst": name})
+                prev_name = name
+            continue
+        name = spec.get("name") or spec.get("kind")
+        graph_nodes.append({
+            "id": name, "kind": spec["kind"],
+            "params": spec.get("params", {}),
+        })
+        if prev_name is not None:
+            graph_edges.append({"src": prev_name, "dst": name})
+        prev_name = name
+
+    last = graph_nodes[-1]["id"] if graph_nodes else "out"
+    sharding_params = SuggestShardingParams.model_validate({
+        "graph": {"nodes": graph_nodes, "edges": graph_edges},
+        "dim_env": {"H": scaled.hidden_size, "B": 1, "S": 256},
+        "loss": {"kind": "cross_entropy", "head_outputs": [last]},
+        "optim": {"kind": "adamw", "groups": [
+            {"matcher": "all", "lr": 3e-4, "weight_decay": 0.01,
+             "betas": [0.9, 0.95]},
+        ], "gradient_clip_norm": 1.0, "mixed_precision": True},
+        "topology": {"factory": topo_name, "kwargs": {}},
+    })
+    sharding = suggest_sharding(sharding_params, cache=cache)
+
+    best_proposal = (
+        sharding.proposals[0] if sharding.proposals else None)
+    fits = bool(scaled.fits and (
+        best_proposal is None or best_proposal.fits))
+    axis_desc = (
+        ",".join(f"{a.axis_name}×{a.degree}"
+                 for a in best_proposal.sharding.axis_assignments)
+        if best_proposal else "<no proposal>")
+    reason = (
+        f"hidden={scaled.hidden_size}, layers={scaled.num_layers}, "
+        f"axis={axis_desc}, peak={scaled.estimated_bytes / 1e9:.2f} GB / "
+        f"{total_hbm / 1e9:.0f} GB")
+    out = AutoFitResult(
+        scaled=ScaleDownResultModel(
+            hidden_size=scaled.hidden_size,
+            num_layers=scaled.num_layers,
+            estimated_bytes=scaled.estimated_bytes,
+            target_bytes=scaled.target_bytes,
+            fits=scaled.fits,
+            scaled_down_from=ScaleDownFromCanonical(
+                hidden_size=scaled.scaled_down_from[0],
+                num_layers=scaled.scaled_down_from[1],
+            ),
+            specs=scaled.specs,
+        ),
+        sharding=sharding,
+        fits=fits,
+        reason=reason,
+        topology=topo_name,
+        headroom=headroom,
+    )
+    _cache_store(cache, key, out)
+    return out
diff --git a/cppmega_v4/jsonrpc/dispatcher.py b/cppmega_v4/jsonrpc/dispatcher.py
@@ -79,6 +79,10 @@
     MemoryMatrixParams,
     memory_matrix,
 )
+from cppmega_v4.jsonrpc.auto_fit_method import (
+    AutoFitParams,
+    auto_fit,
+)
 from cppmega_v4.jsonrpc.tokenizer_roundtrip_text_method import (
     TokenizerRoundtripTextParams,
     roundtrip_text,
@@ -137,6 +141,10 @@
         MemoryMatrixParams,
         lambda p, c: memory_matrix(p, cache=c),
     ),
+    "architectures.auto_fit": (
+        AutoFitParams,
+        lambda p, c: auto_fit(p, cache=c),
+    ),
     "probe.run": (
         ProbeRunParams,
         lambda p, c: probe_run(p, cache=c),
diff --git a/cppmega_v4/jsonrpc/schema.py b/cppmega_v4/jsonrpc/schema.py
@@ -809,6 +809,7 @@ class PipelineRunResult(BaseModel):
     "platform.get_info",
     "architectures.scale_down",
     "memory.matrix",
+    "architectures.auto_fit",
 })
 
 
diff --git a/tests/v4/test_auto_fit.py b/tests/v4/test_auto_fit.py
@@ -0,0 +1,89 @@
+"""V8-R04 unit tests: architectures.auto_fit RPC."""
+
+from __future__ import annotations
+
+import pytest
+
+from cppmega_v4.jsonrpc.auto_fit_method import (
+    AutoFitHostInfo, AutoFitParams, auto_fit,
+)
+from cppmega_v4.jsonrpc.dispatcher import dispatch
+from cppmega_v4.jsonrpc.schema import JsonRpcRequest
+
+
+def test_llama3_8b_on_gb10_fits_full_size():
+    """The canonical 8B preset fits on a single GB10 (137 GB unified)."""
+    r = auto_fit(AutoFitParams(
+        preset="llama3_8b",
+        host_info=AutoFitHostInfo(topology="gb10_quarter")))
+    assert r.topology == "gb10_quarter"
+    assert r.headroom == 0.9
+    assert r.fits is True
+    assert r.scaled.hidden_size == 4096   # canonical (no scale-down needed)
+    assert r.scaled.num_layers == 32
+    assert r.scaled.estimated_bytes <= int(137e9 * 0.9)
+    assert len(r.sharding.proposals) >= 1
+    assert "hidden=4096" in r.reason
+    assert "GB" in r.reason
+
+
+def test_huge_preset_on_tiny_topology_still_returns():
+    """When nothing fits, auto_fit still returns a result with fits=False
+    instead of raising."""
+    r = auto_fit(AutoFitParams(
+        preset="qwen3_235b_a22b",
+        host_info=AutoFitHostInfo(topology="gb10_quarter", headroom=0.5)))
+    # The full 235B doesn't fit even on a GB10 at H=8192 — but
+    # scale_down will pick the largest (H, L) ≤ target. If no cell on
+    # the search grid fits, fits=False.
+    assert r.topology == "gb10_quarter"
+    # Just check the contract is well-formed.
+    assert r.scaled.hidden_size >= 64
+    assert r.scaled.num_layers >= 1
+    assert isinstance(r.fits, bool)
+
+
+def test_unknown_topology_rejected():
+    with pytest.raises(ValueError, match="unknown topology"):
+        auto_fit(AutoFitParams(
+            preset="llama3_8b",
+            host_info=AutoFitHostInfo(topology="nope")))
+
+
+def test_host_info_omitted_uses_platform_probe():
+    """Without host_info, auto_fit picks a topology from platform.get_info."""
+    r = auto_fit(AutoFitParams(preset="llama3_8b"))
+    assert r.topology in {
+        "m3_ultra_solo", "gb10_quarter",
+        "h100_8x", "h200_8x", "a100_8x", "b100_8x",
+        "tpu_v6e_8", "tpu_v5p_4"}
+
+
+def test_dispatch_end_to_end():
+    req = JsonRpcRequest(
+        jsonrpc="2.0", id="t-r04",
+        method="architectures.auto_fit",
+        params={"preset": "qwen3_dense_4b",
+                "host_info": {"topology": "h100_8x"}},
+    )
+    resp = dispatch(req)
+    assert resp.error is None, resp.error
+    r = resp.result
+    assert r["topology"] == "h100_8x"
+    assert r["fits"] is True
+    assert r["scaled"]["hidden_size"] >= 64
+    assert isinstance(r["sharding"]["proposals"], list)
+    assert "GB" in r["reason"]
+
+
+def test_headroom_propagates_to_target_bytes():
+    """Lower headroom should produce a smaller target budget."""
+    r_high = auto_fit(AutoFitParams(
+        preset="llama3_8b",
+        host_info=AutoFitHostInfo(topology="gb10_quarter", headroom=0.9)))
+    r_low = auto_fit(AutoFitParams(
+        preset="llama3_8b",
+        host_info=AutoFitHostInfo(topology="gb10_quarter", headroom=0.3)))
+    # Lower headroom can never pick a *larger* model.
+    assert (r_low.scaled.hidden_size * r_low.scaled.num_layers
+            <= r_high.scaled.hidden_size * r_high.scaled.num_layers)
diff --git a/vbgui/src/components/GalleryScaleDownSlider.tsx b/vbgui/src/components/GalleryScaleDownSlider.tsx
@@ -52,6 +52,14 @@ function fmtBytes(n: number): string {
   return `${n} B`;
 }
 
+interface AutoFitResult {
+  scaled: ScaleDownPreview;
+  fits: boolean;
+  reason: string;
+  topology: string;
+  headroom: number;
+}
+
 export function GalleryScaleDownSlider({
   presets, rpc, onApply, initialBytes = ONE_GB, initialPreset,
 }: GalleryScaleDownSliderProps): JSX.Element {
@@ -61,6 +69,9 @@ export function GalleryScaleDownSlider({
   const [preview, setPreview] = useState<ScaleDownPreview | null>(null);
   const [err, setErr] = useState<string | null>(null);
   const [loading, setLoading] = useState(false);
+  const [autoFitResult, setAutoFitResult] =
+    useState<AutoFitResult | null>(null);
+  const [autoFitBusy, setAutoFitBusy] = useState(false);
   // Per-request token so stale responses are dropped.
   const reqIdRef = useRef(0);
 
@@ -90,6 +101,23 @@ export function GalleryScaleDownSlider({
     return () => clearTimeout(t);
   }, [preset, bytes, fetchPreview]);
 
+  const runAutoFit = useCallback(async () => {
+    if (!preset) return;
+    setAutoFitBusy(true);
+    setErr(null);
+    try {
+      const r = await rpc.call<AutoFitResult>(
+        "architectures.auto_fit", { preset });
+      setAutoFitResult(r);
+      // Mirror as the preview block so the user sees identical state.
+      setPreview(r.scaled);
+    } catch (e) {
+      setErr(e instanceof Error ? e.message : String(e));
+    } finally {
+      setAutoFitBusy(false);
+    }
+  }, [rpc, preset]);
+
   return (
     <div data-testid="gallery-scaledown" style={{ padding: 8,
          border: "1px solid #e5e7eb", borderRadius: 6, marginTop: 8,
@@ -120,6 +148,21 @@ export function GalleryScaleDownSlider({
       {loading && <span data-testid="gallery-scaledown-loading">…</span>}
       {err && <span data-testid="gallery-scaledown-error"
                     style={{ color: "#b91c1c" }}>{err}</span>}
+      <button
+        data-testid="gallery-auto-fit"
+        onClick={() => { void runAutoFit(); }}
+        disabled={autoFitBusy || !preset}
+        style={{ alignSelf: "flex-start", padding: "2px 8px" }}>
+        {autoFitBusy ? "auto-fitting…" : "Auto-fit to my devbox"}
+      </button>
+      {autoFitResult && (
+        <div data-testid="gallery-auto-fit-result"
+             style={{ fontSize: 11, color: "#374151",
+                      background: "#eff6ff", padding: 6,
+                      borderRadius: 4 }}>
+          <strong>{autoFitResult.topology}</strong> · {autoFitResult.reason}
+        </div>
+      )}
       {preview && (
         <div style={{ fontSize: 11, color: "#374151",
                       display: "flex", flexDirection: "column", gap: 2 }}>
diff --git a/vbgui/tests/GalleryScaleDownSlider.test.tsx b/vbgui/tests/GalleryScaleDownSlider.test.tsx
