feat: remora PyTorch ZeroGPU layer-control research space

.gitignore

@@ -0,0 +1,6 @@
+__pycache__/
+*.py[cod]
+.venv/
+.env
+.env.*
+!.env.example

README.md

@@ -0,0 +1,59 @@
+---
+title: Remora Layer Lab
+emoji: 🐟
+colorFrom: blue
+colorTo: indigo
+sdk: gradio
+python_version: "3.12"
+sdk_version: "5.49.1"
+app_file: app.py
+models:
+  - LiquidAI/LFM2.5-1.2B-Instruct
+tags:
+  - remora
+  - layer-control
+  - lfm2
+  - zerogpu
+---
+
+# Remora Layer Lab
+
+A research-only Hugging Face ZeroGPU Space that runs `LiquidAI/LFM2.5-1.2B-Instruct` and applies explicit, reproducible controls to selected hidden-state layers.
+
+## What this is
+
+- A **PyTorch experimental adapter** for the existing Remora research method.
+- Captures per-layer hidden-state statistics for each generation run.
+- Applies a scalar gain to selected decoder-block outputs:
+
+  `controlled = baseline + strength × (baseline - token_mean)`
+
+  - `strength = 0`: observation only; output is unchanged.
+  - negative strength: damp the selected residual signal.
+  - positive strength: amplify it.
+- Produces a JSONL trace that can be compared with the Rust/llama.cpp Remora traces.
+
+## What this is not
+
+This is **not** the Rust production lane and does not replace local remora on Janus. The HF ZeroGPU Space is a bounded research surface: Gradio/PyTorch only, shared GPU allocation, daily quota, and no persistent runtime disk. The production route remains octopus → local `lfm-fast`.
+
+## Run protocol
+
+1. Use the default `strength = 0` to create a baseline trace.
+2. Repeat the same prompt with one selected layer and one non-zero strength.
+3. Compare output and trace statistics outside the Space. Save the downloaded JSONL with the prompt suite and run metadata.
+4. Keep a request below the declared GPU duration. Start with 32 output tokens.
+
+## Cost boundary
+
+For a PRO account, ZeroGPU has a 40-minute daily quota. Usage after the quota costs $1 per 10 minutes from pre-paid HF credits. The Space uses `large` (48 GB, one quota unit), not `xlarge`.
+
+## Local checks
+
+```bash
+python -m unittest discover -s tests -v
+```
+
+## Deploy
+
+Create an HF **Gradio ZeroGPU** Space called `gcorzo/remora-layer-lab`, select ZeroGPU hardware in its Settings, then push this directory to the Space repository. The HF token used for this must have Space write permission; no token is stored in this repository.

app.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+import json
+import tempfile
+import time
+from pathlib import Path
+
+import gradio as gr
+import spaces
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from remora_control import LayerController
+
+MODEL_ID = "LiquidAI/LFM2.5-1.2B-Instruct"
+MAX_OUTPUT_TOKENS = 128
+DEVICE = "cuda"
+
+# ZeroGPU's CUDA emulation makes module-level placement correct: a real GPU is
+# attached only when `run_experiment` enters the decorator.
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.bfloat16,
+    low_cpu_mem_usage=True,
+).to(DEVICE).eval()
+LAYER_COUNT = len(model.model.layers)
+
+
+def parse_layers(text: str) -> list[int]:
+    raw = [part.strip() for part in text.split(",") if part.strip()]
+    layers = sorted(set(int(part) for part in raw))
+    invalid = [value for value in layers if value < 0 or value >= LAYER_COUNT]
+    if invalid:
+        raise gr.Error(f"Layer values must be between 0 and {LAYER_COUNT - 1}; got {invalid}.")
+    return layers
+
+
+@spaces.GPU(duration=120, size="large")
+def run_experiment(prompt: str, selected_layers: str, gain: float, max_new_tokens: int, seed: int):
+    if not prompt.strip():
+        raise gr.Error("Enter a prompt.")
+    if max_new_tokens < 1 or max_new_tokens > MAX_OUTPUT_TOKENS:
+        raise gr.Error(f"Output tokens must be 1–{MAX_OUTPUT_TOKENS}.")
+
+    layers = parse_layers(selected_layers)
+    controller = LayerController(layers, gain)
+    generator = torch.Generator(device=DEVICE).manual_seed(int(seed))
+    started = time.perf_counter()
+    controller.attach(model)
+    try:
+        messages = [{"role": "user", "content": prompt}]
+        rendered = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        inputs = tokenizer(rendered, return_tensors="pt").to(DEVICE)
+        with torch.inference_mode():
+            output_ids = model.generate(
+                **inputs,
+                max_new_tokens=int(max_new_tokens),
+                do_sample=True,
+                temperature=0.1,
+                top_k=50,
+                repetition_penalty=1.05,
+                generator=generator,
+                use_cache=True,
+            )
+    finally:
+        controller.detach()
+
+    elapsed_s = time.perf_counter() - started
+    new_tokens = output_ids[0, inputs.input_ids.shape[1] :]
+    text = tokenizer.decode(new_tokens, skip_special_tokens=True)
+    trace = {
+        "schema": "remora-layer-lab/v1",
+        "model": MODEL_ID,
+        "layer_count": LAYER_COUNT,
+        "selected_layers": layers,
+        "gain": gain,
+        "seed": seed,
+        "prompt": prompt,
+        "generated_tokens": int(new_tokens.shape[0]),
+        "elapsed_s": round(elapsed_s, 4),
+        "tokens_per_second": round(float(new_tokens.shape[0]) / elapsed_s, 3) if elapsed_s else None,
+        "events": controller.json_events(),
+    }
+    artifact = Path(tempfile.mkstemp(prefix="remora-trace-", suffix=".json")[1])
+    artifact.write_text(json.dumps(trace, indent=2), encoding="utf-8")
+    status = (
+        f"{len(layers)} controlled layers; {len(trace['events'])} trace events; "
+        f"{trace['generated_tokens']} tokens in {elapsed_s:.2f}s "
+        f"({trace['tokens_per_second']} tok/s)."
+    )
+    return text, status, str(artifact)
+
+
+with gr.Blocks(title="Remora Layer Lab") as demo:
+    gr.Markdown(
+        "# Remora Layer Lab\n"
+        "PyTorch/ZeroGPU research adapter for controlled LFM2.5 decoder-layer experiments. "
+        "`gain=0` is an observation-only baseline. This is not the local Rust production lane."
+    )
+    with gr.Row():
+        prompt = gr.Textbox(label="Prompt", lines=6, value="Explain why water forecasting needs uncertainty.")
+        with gr.Column():
+            selected_layers = gr.Textbox(label=f"Layers (0–{LAYER_COUNT - 1}, comma-separated)", value="2,5")
+            gain = gr.Slider(-1.0, 1.0, value=0.0, step=0.05, label="Layer gain")
+            max_new_tokens = gr.Slider(1, MAX_OUTPUT_TOKENS, value=32, step=1, label="Max output tokens")
+            seed = gr.Number(value=42, precision=0, label="Random seed")
+            run = gr.Button("Run controlled experiment", variant="primary")
+    answer = gr.Textbox(label="Generated text", lines=8)
+    status = gr.Textbox(label="Run statistics")
+    trace = gr.File(label="Download trace JSON")
+    run.click(run_experiment, [prompt, selected_layers, gain, max_new_tokens, seed], [answer, status, trace])
+
+if __name__ == "__main__":
+    demo.launch()

docs/plans/2026-09-02-remora-layer-lab.md

@@ -0,0 +1,55 @@
+# Remora Layer Lab Implementation Plan
+
+> **For Hermes:** Use subagent-driven-development skill to implement this plan task-by-task.
+
+**Goal:** provide a bounded HF ZeroGPU laboratory for PyTorch layer-control experiments on public LFM2.5-1.2B weights, while retaining local Rust remora as the production path.
+
+**Architecture:** The Space runs `LiquidAI/LFM2.5-1.2B-Instruct` through Transformers and attaches forward hooks to selected `model.model.layers`. A deterministic gain control captures before/after RMS per token in a downloadable JSON trace. The Space uses Gradio plus `@spaces.GPU`, which is an explicitly bounded external Python research adapter and not fleet infrastructure.
+
+**Tech Stack:** Python 3.12, PyTorch 2.8+, Transformers 4.57.6, Gradio 5.49.1, Hugging Face ZeroGPU.
+
+---
+
+## Verified inputs
+
+- Model: `LiquidAI/LFM2.5-1.2B-Instruct`, public, Transformers-native, 16 layers (10 convolution + 6 attention).
+- ZeroGPU: Gradio only; declare GPU work with `@spaces.GPU`; `large` offers 48 GB; PRO quota is 40 minutes/day.
+- The account's active token cannot write a Space. Deployment requires a new token with Space-write permission.
+
+## Acceptance criteria
+
+- `gain=0` preserves exact hidden states while capturing a trace.
+- Selected layer indexes are validated against the 16-layer model.
+- Each run returns text, run speed and a downloadable JSON trace.
+- GPU work is bounded to 120 seconds/request on a 48 GB `large` allocation.
+- The project has unit tests, clear deployment instructions and no embedded credentials.
+
+## Tasks
+
+### Task 1: Core layer controller
+
+**Files:** `remora_control.py`, `tests/test_remora_control.py`
+
+Implement a hook controller over `model.model.layers`. Test zero-gain identity, non-zero gain effect, trace count and index validation.
+
+**Verification:** `python -m unittest discover -s tests -v`
+
+### Task 2: ZeroGPU application
+
+**Files:** `app.py`, `requirements.txt`, `README.md`
+
+Load the public 1.2B LFM model in bf16, create a Gradio prompt/control surface, decorate execution with `@spaces.GPU(duration=120, size="large")`, and return a trace artifact.
+
+**Verification:** `python -m py_compile app.py remora_control.py`
+
+### Task 3: Deploy and bounded benchmark
+
+**Files:** HF Space `gcorzo/remora-layer-lab` (external)
+
+Create a Gradio Space, choose ZeroGPU hardware in Settings, push this project and wait for build success. Run baseline and one non-zero gain with identical prompt/seed; download both traces and record runtime / generated tokens / quota use.
+
+**Verification:** Space build is `Running`; both trace files are downloadable; HF Billing shows resulting ZeroGPU compute usage.
+
+## Rollback
+
+Remove the Space or switch it to CPU Basic. No local router, model server or production octopus configuration changes are made by this project.

remora_control.py

@@ -0,0 +1,83 @@
+"""Layer-level Remora controls for PyTorch decoder stacks.
+
+This module is intentionally small and dependency-light. It controls the output
+of a selected decoder block and records scalar trace data; it does not claim
+that the output hook is identical to llama.cpp's internal MoE gate hook.
+"""
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass
+from typing import Any, Iterable
+
+import torch
+from torch import nn
+
+
+@dataclass(frozen=True)
+class TraceEvent:
+    layer: int
+    token_index: int
+    baseline_rms: float
+    controlled_rms: float
+    gain: float
+
+
+class LayerController:
+    """Attach controlled-output hooks to `model.model.layers`.
+
+    The control is deterministic for fixed model weights, prompt, sampling
+    parameters, and RNG seed. A gain of zero is a strict identity operation.
+    """
+
+    def __init__(self, layers: Iterable[int], gain: float) -> None:
+        self.layers = tuple(sorted(set(layers)))
+        self.gain = float(gain)
+        self.events: list[TraceEvent] = []
+        self._handles: list[Any] = []
+
+    def attach(self, model: nn.Module) -> None:
+        decoder_layers = getattr(getattr(model, "model", None), "layers", None)
+        if decoder_layers is None:
+            raise ValueError("model does not expose model.layers for layer control")
+        layer_count = len(decoder_layers)
+        invalid = [index for index in self.layers if index < 0 or index >= layer_count]
+        if invalid:
+            raise ValueError(f"layer indexes outside 0..{layer_count - 1}: {invalid}")
+        self.detach()
+        for index in self.layers:
+            self._handles.append(decoder_layers[index].register_forward_hook(self._hook(index)))
+
+    def detach(self) -> None:
+        for handle in self._handles:
+            handle.remove()
+        self._handles.clear()
+
+    def json_events(self) -> list[dict[str, float | int]]:
+        return [asdict(event) for event in self.events]
+
+    def _hook(self, layer: int):
+        def controlled_output(_module: nn.Module, _inputs: tuple[Any, ...], output: Any) -> Any:
+            hidden = output[0] if isinstance(output, tuple) else output
+            if not isinstance(hidden, torch.Tensor):
+                raise TypeError("decoder layer output must be a Tensor or tuple whose first value is a Tensor")
+
+            # [batch, sequence, hidden] -> one scalar event per generated/prompt token.
+            baseline_rms = hidden.float().pow(2).mean(dim=-1).sqrt()
+            token_mean = hidden.mean(dim=-1, keepdim=True)
+            controlled = hidden + self.gain * (hidden - token_mean)
+            controlled_rms = controlled.float().pow(2).mean(dim=-1).sqrt()
+            for token_index, (before, after) in enumerate(zip(baseline_rms[0], controlled_rms[0])):
+                self.events.append(
+                    TraceEvent(
+                        layer=layer,
+                        token_index=token_index,
+                        baseline_rms=float(before.detach().cpu()),
+                        controlled_rms=float(after.detach().cpu()),
+                        gain=self.gain,
+                    )
+                )
+            if isinstance(output, tuple):
+                return (controlled, *output[1:])
+            return controlled
+
+        return controlled_output

requirements.txt

@@ -0,0 +1,4 @@
+gradio==5.49.1
+huggingface-hub>=0.30,<1.0
+torch>=2.8.0
+transformers==4.57.6

tests/test_remora_control.py

@@ -0,0 +1,55 @@
+import unittest
+
+import torch
+from torch import nn
+
+from remora_control import LayerController
+
+
+class ToyModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.model = nn.Module()
+        self.model.layers = nn.ModuleList([nn.Identity(), nn.Identity(), nn.Identity()])
+
+    def forward(self, value):
+        for layer in self.model.layers:
+            value = layer(value)
+        return value
+
+
+class LayerControllerTests(unittest.TestCase):
+    def setUp(self):
+        self.model = ToyModel()
+        self.value = torch.tensor([[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]])
+
+    def test_zero_gain_is_identity_and_records_trace(self):
+        controller = LayerController([1], gain=0.0)
+        controller.attach(self.model)
+        try:
+            actual = self.model(self.value)
+        finally:
+            controller.detach()
+        self.assertTrue(torch.equal(actual, self.value))
+        self.assertEqual(2, len(controller.events))
+        self.assertEqual(1, controller.events[0].layer)
+        self.assertEqual(controller.events[0].baseline_rms, controller.events[0].controlled_rms)
+
+    def test_nonzero_gain_changes_selected_layer(self):
+        controller = LayerController([1], gain=0.5)
+        controller.attach(self.model)
+        try:
+            actual = self.model(self.value)
+        finally:
+            controller.detach()
+        self.assertFalse(torch.equal(actual, self.value))
+        self.assertEqual(2, len(controller.events))
+        self.assertGreater(controller.events[0].controlled_rms, controller.events[0].baseline_rms)
+
+    def test_invalid_layer_is_rejected(self):
+        with self.assertRaisesRegex(ValueError, "outside"):
+            LayerController([3], 0.0).attach(self.model)
+
+
+if __name__ == "__main__":
+    unittest.main()