Update Feather training runtime image

Dockerfile

@@ -1,6 +1,4 @@
-FROM pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel
-
-ARG HTM_CUDA_ARCH=sm_86
+FROM pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel
 
 ENV DEBIAN_FRONTEND=noninteractive \
     PIP_NO_CACHE_DIR=1 \
@@ -107,22 +105,12 @@ COPY overlay /workspace/feather
 COPY entrypoint.py /app/entrypoint.py
 WORKDIR /workspace/feather
 
-RUN python - <<'PY'
-from pathlib import Path
-for sh in Path('/workspace/feather/scripts').glob('*.sh'):
-    raw = sh.read_bytes()
-    norm = raw.replace(b'\r\n', b'\n')
-    if norm != raw:
-        sh.write_bytes(norm)
-PY
-
-RUN python -m py_compile hydra/training.py prepare.py train.py && \
-    bash -n scripts/run_domain_expanded_pretrain.sh
+RUN python -m py_compile hydra/training.py prepare.py train.py && \
+    bash -n scripts/run_domain_expanded_pretrain.sh
 
-RUN export LD_LIBRARY_PATH=/usr/local/cuda/lib64:${LD_LIBRARY_PATH} && \
-    export HTM_CUDA_ARCH=${HTM_CUDA_ARCH} && \
-    export CARGO_BUILD_JOBS=1 && \
-    maturin build --release -j 1 --features gpu --manifest-path htm_rust/Cargo.toml && \
-    pip install htm_rust/target/wheels/htm_rust-*.whl
+RUN export LD_LIBRARY_PATH=/usr/local/cuda/lib64:${LD_LIBRARY_PATH} && \
+    export HTM_CUDA_ARCH=sm_90 && \
+    maturin build --release --features gpu --manifest-path htm_rust/Cargo.toml && \
+    pip install htm_rust/target/wheels/htm_rust-*.whl
 
 CMD ["python", "/app/entrypoint.py"]

entrypoint.py

@@ -68,11 +68,7 @@ try:
 except ImportError:
     print('[boot] triton_cache_setup not found; skipping cache hydrate', flush=True)
 
-from huggingface_hub import HfApi  # noqa: E402  (import after cuda kick)
-if '/workspace/feather' not in sys.path:  # noqa: E402
-    sys.path.insert(0, '/workspace/feather')
-from scripts.benchmark_assets import hydrate_benchmark_assets  # noqa: E402
-from subsystems.sdr_retina import build_retina  # noqa: E402
+from huggingface_hub import HfApi  # noqa: E402  (import after cuda kick)
 
 REPO_ROOT = Path('/workspace/feather')
 CACHE_ROOT = Path.home() / '.cache' / 'autoresearch'
@@ -114,7 +110,7 @@ def _start_health_server() -> HTTPServer:
     return server
 
 
-def upload_artifact(api: HfApi, path: Path, dest: str) -> None:
+def upload_artifact(api: HfApi, path: Path, dest: str) -> None:
     if not path.exists():
         print(f'[upload] skip missing {path}', flush=True)
         return
@@ -124,20 +120,7 @@ def upload_artifact(api: HfApi, path: Path, dest: str) -> None:
         repo_id=OUTPUT_REPO,
         repo_type='model',
     )
-    print(f'[upload] uploaded {path} -> {OUTPUT_REPO}/{dest}', flush=True)
-
-
-def build_benchmark_mode_command() -> list[str]:
-    return [
-        'python',
-        str(REPO_ROOT / 'scripts' / 'benchmark_runner.py'),
-        '--benchmark', os.environ.get('HYDRA_BENCHMARK_NAME', 'GSM8K'),
-        '--generator-mode', 'hydra',
-        '--variant', os.environ.get('HYDRA_BENCHMARK_VARIANT', 'hydra_full'),
-        '--seed', os.environ.get('HYDRA_SEED', '42'),
-        '--out', str(REPO_ROOT / 'benchmark_result.json'),
-        '--ledger', str(REPO_ROOT / 'benchmark_ledger.json'),
-    ] + sys.argv[1:]
+    print(f'[upload] uploaded {path} -> {OUTPUT_REPO}/{dest}', flush=True)
 
 
 def _wait_for_cuda_ready(timeout_s: int = 120) -> None:
@@ -175,7 +158,7 @@ def _wait_for_cuda_ready(timeout_s: int = 120) -> None:
     print(f'[job] CUDA still not ready after {timeout_s}s — continuing anyway (training will likely fail)', flush=True)
 
 
-def run_job_mode() -> int:
+def run_job_mode() -> int:
     os.chdir(REPO_ROOT)
     os.environ.setdefault('HYDRA_TIME_BUDGET', '43200')
     os.environ.setdefault('HYDRA_TARGET_SHARDS', '2048')
@@ -220,46 +203,7 @@ def run_job_mode() -> int:
     else:
         print('[upload] HF_TOKEN not set; skipping artifact upload', flush=True)
 
-    return proc.returncode
-
-
-def run_benchmark_mode() -> int:
-    os.chdir(REPO_ROOT)
-    os.environ.setdefault('HYDRA_USE_NEMOTRON', '1')
-    if TOKEN:
-        try:
-            hydrate_benchmark_assets(
-                cache_dir=CACHE_ROOT,
-                output_repo=OUTPUT_REPO,
-                tokenizer_repo=os.environ.get('HYDRA_TOKENIZER_CACHE_REPO', OUTPUT_REPO),
-                token=TOKEN,
-            )
-        except Exception as e:
-            print(f'[benchmark] asset hydrate warning: {type(e).__name__}: {e}', flush=True)
-        try:
-            build_retina()
-        except Exception as e:
-            print(f'[benchmark] retina materialize warning: {type(e).__name__}: {e}', flush=True)
-    cmd = build_benchmark_mode_command()
-    print(f'[benchmark] command={cmd}', flush=True)
-    proc = subprocess.run(cmd, check=False)
-
-    if TOKEN:
-        api = HfApi(token=TOKEN)
-        try:
-            api.create_repo(repo_id=OUTPUT_REPO, repo_type='model', private=True, exist_ok=True)
-        except Exception as e:
-            print(f'[upload] create_repo warning: {type(e).__name__}: {e}', flush=True)
-        prefix = f'jobs/{JOB_ID}'
-        try:
-            upload_artifact(api, REPO_ROOT / 'benchmark_result.json', f'{prefix}/benchmark_result.json')
-            upload_artifact(api, REPO_ROOT / 'benchmark_ledger.json', f'{prefix}/benchmark_ledger.json')
-        except Exception as e:
-            print(f'[upload] upload warning: {type(e).__name__}: {e}', flush=True)
-    else:
-        print('[upload] HF_TOKEN not set; skipping benchmark artifact upload', flush=True)
-
-    return proc.returncode
+    return proc.returncode
 
 
 def run_space_mode() -> int:
@@ -273,12 +217,10 @@ def run_space_mode() -> int:
         server.server_close()
 
 
-def main() -> int:
-    if RUNTIME_MODE == 'job':
-        return run_job_mode()
-    if RUNTIME_MODE == 'benchmark':
-        return run_benchmark_mode()
-    return run_space_mode()
+def main() -> int:
+    if RUNTIME_MODE == 'job':
+        return run_job_mode()
+    return run_space_mode()
 
 
 if __name__ == '__main__':

mamba_ssm_init.py

@@ -24,8 +24,8 @@ mamba_inner_fn = None
 # stub is never actually invoked at runtime because the codebase does not use
 # torch.compile — but importing torch._inductor.* still requires the symbol to
 # exist at module load time.
-import triton as _triton  # noqa: E402
-if not hasattr(_triton, "set_allocator"):
+import triton as _triton  # noqa: E402
+if not hasattr(_triton, "set_allocator"):
     def _noop_set_allocator(_fn):  # pragma: no cover
         return None
     _triton.set_allocator = _noop_set_allocator
@@ -53,39 +53,7 @@ if not hasattr(_tcc, "triton_key"):
     def _triton_key_shim():
         import triton as _t
         return f"triton-{_t.__version__}-shim"
-    _tcc.triton_key = _triton_key_shim
-
-# Triton 3.5 wheels can occasionally load with an empty backend registry in
-# HF Jobs environments (driver.active -> "0 active drivers"), even though the
-# NVIDIA backend module is present and CudaDriver.is_active() is True.
-# Patch _create_driver to directly select CudaDriver when registry discovery
-# returns empty.
-import importlib as _importlib  # noqa: E402
-_triton_driver_mod = _importlib.import_module("triton.runtime.driver")
-if getattr(_triton_driver_mod, "backends", None) == {}:
-    from triton.backends.nvidia import driver as _nvidia_driver  # noqa: E402
-
-    def _create_driver_shim():
-        if hasattr(_nvidia_driver, "CudaDriver") and _nvidia_driver.CudaDriver.is_active():
-            return _nvidia_driver.CudaDriver()
-        raise RuntimeError(
-            "Triton backend registry is empty and NVIDIA CudaDriver is not active"
-        )
-
-    _triton_driver_mod._create_driver = _create_driver_shim
-    if hasattr(_triton_driver_mod, "driver") and hasattr(_triton_driver_mod.driver, "reset_active"):
-        _triton_driver_mod.driver.reset_active()
-
-_triton_compiler_mod = _importlib.import_module("triton.compiler.compiler")
-if getattr(_triton_compiler_mod, "backends", None) == {}:
-    from triton.backends import Backend as _Backend  # noqa: E402
-    from triton.backends.nvidia.compiler import CUDABackend as _CUDABackend  # noqa: E402
-    from triton.backends.nvidia.driver import CudaDriver as _CudaDriver  # noqa: E402
-
-    _triton_compiler_mod.backends["nvidia"] = _Backend(
-        compiler=_CUDABackend,
-        driver=_CudaDriver,
-    )
+    _tcc.triton_key = _triton_key_shim
 
 # Suppress torch.compile/_dynamo errors globally — we don't rely on torch.compile
 # for performance in this codebase (Muon + mamba3 CUDA kernels already fused),

overlay/configs/harness_config.py

@@ -1,13 +1,13 @@
-"""Harness configuration for HYDRA's self-evolving outer loop."""
-from typing import Literal
-
-from pydantic import BaseModel, Field
-
-type GateThresholds = dict[str, float]
-type GateConfig = dict[str, GateThresholds]
+"""Harness configuration for HYDRA's self-evolving outer loop."""
+from typing import Literal
 
+from pydantic import BaseModel, Field
 
-class HarnessConfig(BaseModel):
+type GateThresholds = dict[str, float]
+type GateConfig = dict[str, GateThresholds]
+
+
+class HarnessConfig(BaseModel):
     """Configuration for the HYDRA harness behavior."""
 
     # Inner loop
@@ -50,19 +50,19 @@ class HarnessConfig(BaseModel):
         default=5.0, description="Max % regression from best known val_bpb"
     )
 
-    # Keep/discard criteria
-    primary_metric: str = "val_bpb"
-    secondary_metrics: GateConfig = Field(
-        default_factory=lambda: {
-            "mhc_spectral_norm": {"max": 2.0},
-            "engram_hit_rate": {"min": 0.1},
-            "factual_english_score": {"min": 0.5},
-            "instruction_following_score": {"min": 0.5},
-            "distinct_2": {"min": 0.1},
-            "repetition_rate": {"max": 0.2},
-            "hestia_quant_error": {"max": 0.05},
-        }
-    )
+    # Keep/discard criteria
+    primary_metric: str = "val_bpb"
+    secondary_metrics: GateConfig = Field(
+        default_factory=lambda: {
+            "mhc_spectral_norm": {"max": 2.0},
+            "engram_hit_rate": {"min": 0.1},
+            "factual_english_score": {"min": 0.5},
+            "instruction_following_score": {"min": 0.5},
+            "distinct_2": {"min": 0.1},
+            "repetition_rate": {"max": 0.2},
+            "hestia_quant_error": {"max": 0.05},
+        }
+    )
 
     # Experiment execution
     experiment_timeout: int = Field(
@@ -80,29 +80,29 @@ class HarnessConfig(BaseModel):
     gate_mhc_spectral_norm: float | None = Field(
         default=None, description="Max mhc_spectral_norm for keep (None=disabled)"
     )
-    gate_engram_hit_rate: float | None = Field(
-        default=None, description="Min engram_hit_rate for keep (None=disabled)"
-    )
-    gate_tps_median: float | None = Field(
-        default=None,
-        description="Min steady-state tps_median for keep (None=disabled)",
-    )
-    gate_tps_p10: float | None = Field(
-        default=None,
-        description="Min steady-state tps_p10 for keep (None=disabled)",
-    )
-
-    def to_secondary_gates(self) -> GateConfig:
-        """Build active keep/discard gates from defaults plus gate_* overrides."""
-        gates = {metric: thresholds.copy() for metric, thresholds in self.secondary_metrics.items()}
-
-        if self.gate_mhc_spectral_norm is not None:
-            gates.setdefault("mhc_spectral_norm", {})["max"] = self.gate_mhc_spectral_norm
-        if self.gate_engram_hit_rate is not None:
-            gates.setdefault("engram_hit_rate", {})["min"] = self.gate_engram_hit_rate
-        if self.gate_tps_median is not None:
-            gates.setdefault("tps_median", {})["min"] = self.gate_tps_median
-        if self.gate_tps_p10 is not None:
-            gates.setdefault("tps_p10", {})["min"] = self.gate_tps_p10
-
-        return gates
+    gate_engram_hit_rate: float | None = Field(
+        default=None, description="Min engram_hit_rate for keep (None=disabled)"
+    )
+    gate_tps_median: float | None = Field(
+        default=None,
+        description="Min steady-state tps_median for keep (None=disabled)",
+    )
+    gate_tps_p10: float | None = Field(
+        default=None,
+        description="Min steady-state tps_p10 for keep (None=disabled)",
+    )
+
+    def to_secondary_gates(self) -> GateConfig:
+        """Build active keep/discard gates from defaults plus gate_* overrides."""
+        gates = {metric: thresholds.copy() for metric, thresholds in self.secondary_metrics.items()}
+
+        if self.gate_mhc_spectral_norm is not None:
+            gates.setdefault("mhc_spectral_norm", {})["max"] = self.gate_mhc_spectral_norm
+        if self.gate_engram_hit_rate is not None:
+            gates.setdefault("engram_hit_rate", {})["min"] = self.gate_engram_hit_rate
+        if self.gate_tps_median is not None:
+            gates.setdefault("tps_median", {})["min"] = self.gate_tps_median
+        if self.gate_tps_p10 is not None:
+            gates.setdefault("tps_p10", {})["min"] = self.gate_tps_p10
+
+        return gates

overlay/harness/eval_agent.py

@@ -1,15 +1,15 @@
-"""Eval agent: parse run.log and extract metrics from training runs."""
-import re
-import statistics
-from dataclasses import dataclass
-
-
-type GateThresholds = dict[str, float]
-type GateConfig = dict[str, GateThresholds]
+"""Eval agent: parse run.log and extract metrics from training runs."""
+import re
+import statistics
+from dataclasses import dataclass
+
+
+type GateThresholds = dict[str, float]
+type GateConfig = dict[str, GateThresholds]
 
 
 @dataclass
-class ExperimentResult:
+class ExperimentResult:
     """Parsed result from a single experiment run.
 
     All float fields default to 0.0; integer fields default to 0.
@@ -28,38 +28,38 @@ class ExperimentResult:
     peak_vram_mb: float = 0.0
     mfu_percent: float = 0.0
 
-    # Throughput
-    total_tokens_m: float = 0.0
-    num_steps: int = 0
-    tps_median: float = 0.0
-    tps_p10: float = 0.0
-    tps_min: float = 0.0
-    tps_max: float = 0.0
-    tps_samples: int = 0
+    # Throughput
+    total_tokens_m: float = 0.0
+    num_steps: int = 0
+    tps_median: float = 0.0
+    tps_p10: float = 0.0
+    tps_min: float = 0.0
+    tps_max: float = 0.0
+    tps_samples: int = 0
 
     # Model shape (echoed by train.py summary block)
     num_params_m: float = 0.0
     n_layer: int = 0
     d_model: int = 0
 
-    # Secondary health metrics
-    mhc_spectral_norm: float = 0.0
-    engram_hit_rate: float = 0.0
-    sr_bypass_rate: float = 0.0
-
-    # Evaluation breadth metrics
-    factual_english_score: float = 0.0
-    instruction_following_score: float = 0.0
-    distinct_1: float = 0.0
-    distinct_2: float = 0.0
-    repetition_rate: float = 0.0
-    repetition_bigram_rate: float = 0.0
-    calibration_ece: float = 0.0
-    calibration_brier: float = 0.0
-    calibration_accuracy: float = 0.0
-    calibration_tokens: int = 0
-    eval_seed: int = 0
-    eval_seed_group: str = ""
+    # Secondary health metrics
+    mhc_spectral_norm: float = 0.0
+    engram_hit_rate: float = 0.0
+    sr_bypass_rate: float = 0.0
+
+    # Evaluation breadth metrics
+    factual_english_score: float = 0.0
+    instruction_following_score: float = 0.0
+    distinct_1: float = 0.0
+    distinct_2: float = 0.0
+    repetition_rate: float = 0.0
+    repetition_bigram_rate: float = 0.0
+    calibration_ece: float = 0.0
+    calibration_brier: float = 0.0
+    calibration_accuracy: float = 0.0
+    calibration_tokens: int = 0
+    eval_seed: int = 0
+    eval_seed_group: str = ""
 
     # Status
     crashed: bool = False
@@ -80,48 +80,48 @@ _PATTERNS: dict[str, str] = {
     "n_layer": r"^n_layer:\s+(\d+)",
     "d_model": r"^d_model:\s+(\d+)",
     "mhc_spectral_norm": r"^mhc_spectral_norm:\s+([\d.]+)",
-    "engram_hit_rate": r"^engram_hit_rate:\s+([\d.]+)",
-    "sr_bypass_rate": r"^sr_bypass_rate:\s+([\d.]+)",
-    "factual_english_score": r"^factual_english_score:\s+([\d.]+)",
-    "instruction_following_score": r"^instruction_following_score:\s+([\d.]+)",
-    "distinct_1": r"^distinct_1:\s+([\d.]+)",
-    "distinct_2": r"^distinct_2:\s+([\d.]+)",
-    "repetition_rate": r"^repetition_rate:\s+([\d.]+)",
-    "repetition_bigram_rate": r"^repetition_bigram_rate:\s+([\d.]+)",
-    "calibration_ece": r"^calibration_ece:\s+([\d.]+)",
-    "calibration_brier": r"^calibration_brier:\s*([\d.]+)",
-    "calibration_accuracy": r"^calibration_accuracy:\s+([\d.]+)",
-    "calibration_tokens": r"^calibration_tokens:\s+(\d+)",
-    "eval_seed": r"^eval_seed:\s+(\d+)",
-    "eval_seed_group": r"^eval_seed_group:\s+(.+)",
-}
+    "engram_hit_rate": r"^engram_hit_rate:\s+([\d.]+)",
+    "sr_bypass_rate": r"^sr_bypass_rate:\s+([\d.]+)",
+    "factual_english_score": r"^factual_english_score:\s+([\d.]+)",
+    "instruction_following_score": r"^instruction_following_score:\s+([\d.]+)",
+    "distinct_1": r"^distinct_1:\s+([\d.]+)",
+    "distinct_2": r"^distinct_2:\s+([\d.]+)",
+    "repetition_rate": r"^repetition_rate:\s+([\d.]+)",
+    "repetition_bigram_rate": r"^repetition_bigram_rate:\s+([\d.]+)",
+    "calibration_ece": r"^calibration_ece:\s+([\d.]+)",
+    "calibration_brier": r"^calibration_brier:\s*([\d.]+)",
+    "calibration_accuracy": r"^calibration_accuracy:\s+([\d.]+)",
+    "calibration_tokens": r"^calibration_tokens:\s+(\d+)",
+    "eval_seed": r"^eval_seed:\s+(\d+)",
+    "eval_seed_group": r"^eval_seed_group:\s+(.+)",
+}
 
 # Attributes that should be parsed as int rather than float.
-_INT_ATTRS: frozenset[str] = frozenset(
-    {
-        "num_steps",
-        "n_layer",
-        "d_model",
-        "calibration_tokens",
-        "eval_seed",
-    }
-)
-_STR_ATTRS: frozenset[str] = frozenset({"eval_seed_group"})
-_STEP_TPS_PATTERN = re.compile(r"step=(\d+).*?\btps=(\d+)\b")
-_TPS_PATTERN = re.compile(r"\btps=(\d+)\b")
-
-
-def _percentile_linear(sorted_values: list[float], pct: float) -> float:
-    """Compute percentile via linear interpolation (0 <= pct <= 100)."""
-    if not sorted_values:
-        return 0.0
-    if len(sorted_values) == 1:
-        return sorted_values[0]
-    rank = (len(sorted_values) - 1) * (pct / 100.0)
-    lo = int(rank)
-    hi = min(lo + 1, len(sorted_values) - 1)
-    frac = rank - lo
-    return sorted_values[lo] * (1.0 - frac) + sorted_values[hi] * frac
+_INT_ATTRS: frozenset[str] = frozenset(
+    {
+        "num_steps",
+        "n_layer",
+        "d_model",
+        "calibration_tokens",
+        "eval_seed",
+    }
+)
+_STR_ATTRS: frozenset[str] = frozenset({"eval_seed_group"})
+_STEP_TPS_PATTERN = re.compile(r"step=(\d+).*?\btps=(\d+)\b")
+_TPS_PATTERN = re.compile(r"\btps=(\d+)\b")
+
+
+def _percentile_linear(sorted_values: list[float], pct: float) -> float:
+    """Compute percentile via linear interpolation (0 <= pct <= 100)."""
+    if not sorted_values:
+        return 0.0
+    if len(sorted_values) == 1:
+        return sorted_values[0]
+    rank = (len(sorted_values) - 1) * (pct / 100.0)
+    lo = int(rank)
+    hi = min(lo + 1, len(sorted_values) - 1)
+    frac = rank - lo
+    return sorted_values[lo] * (1.0 - frac) + sorted_values[hi] * frac
 
 
 def parse_run_log(log_path: str) -> ExperimentResult:
@@ -144,60 +144,60 @@ def parse_run_log(log_path: str) -> ExperimentResult:
         result.error_message = f"Log file not found: {log_path}"
         return result
 
-    # Detect crash signals in output. Keep this strict to avoid false positives
-    # from benign log lines that include "error" in a non-fatal context.
-    if (
-        "Traceback" in content
-        or "\nFAIL\n" in content
-        or "[TPS_GUARD] FAIL" in content
-        or "raise SystemExit(1)" in content
-    ):
-        result.crashed = True
-        lines = content.strip().splitlines()
-        result.error_message = "\n".join(lines[-20:])
-
-    for attr, pattern in _PATTERNS.items():
-        match = re.search(pattern, content, re.MULTILINE)
-        if match:
-            raw = match.group(1)
-            if attr in _INT_ATTRS:
-                setattr(result, attr, int(raw))
-            elif attr in _STR_ATTRS:
-                setattr(result, attr, raw.strip())
-            else:
-                setattr(result, attr, float(raw))
-
-    warmup_steps = 10
-    warmup_match = re.search(r"\[TPS_GUARD\] enabled .*?warmup_steps=(\d+)", content)
-    if warmup_match:
-        warmup_steps = int(warmup_match.group(1))
-
-    step_tps_samples: list[tuple[int, int]] = []
-    for m in _STEP_TPS_PATTERN.finditer(content):
-        step_tps_samples.append((int(m.group(1)), int(m.group(2))))
-
-    tps_values: list[float] = []
-    if step_tps_samples:
-        for step, tps in step_tps_samples:
-            if step >= warmup_steps:
-                tps_values.append(float(tps))
-        if not tps_values:
-            tps_values = [float(tps) for _, tps in step_tps_samples]
-    else:
-        tps_values = [float(m.group(1)) for m in _TPS_PATTERN.finditer(content)]
-
-    if tps_values:
-        sorted_tps = sorted(tps_values)
-        result.tps_samples = len(tps_values)
-        result.tps_median = float(statistics.median(tps_values))
-        result.tps_p10 = float(_percentile_linear(sorted_tps, 10.0))
-        result.tps_min = float(sorted_tps[0])
-        result.tps_max = float(sorted_tps[-1])
-
-    return result
-
-
-def check_secondary_alarms(result: ExperimentResult) -> list[str]:
+    # Detect crash signals in output. Keep this strict to avoid false positives
+    # from benign log lines that include "error" in a non-fatal context.
+    if (
+        "Traceback" in content
+        or "\nFAIL\n" in content
+        or "[TPS_GUARD] FAIL" in content
+        or "raise SystemExit(1)" in content
+    ):
+        result.crashed = True
+        lines = content.strip().splitlines()
+        result.error_message = "\n".join(lines[-20:])
+
+    for attr, pattern in _PATTERNS.items():
+        match = re.search(pattern, content, re.MULTILINE)
+        if match:
+            raw = match.group(1)
+            if attr in _INT_ATTRS:
+                setattr(result, attr, int(raw))
+            elif attr in _STR_ATTRS:
+                setattr(result, attr, raw.strip())
+            else:
+                setattr(result, attr, float(raw))
+
+    warmup_steps = 10
+    warmup_match = re.search(r"\[TPS_GUARD\] enabled .*?warmup_steps=(\d+)", content)
+    if warmup_match:
+        warmup_steps = int(warmup_match.group(1))
+
+    step_tps_samples: list[tuple[int, int]] = []
+    for m in _STEP_TPS_PATTERN.finditer(content):
+        step_tps_samples.append((int(m.group(1)), int(m.group(2))))
+
+    tps_values: list[float] = []
+    if step_tps_samples:
+        for step, tps in step_tps_samples:
+            if step >= warmup_steps:
+                tps_values.append(float(tps))
+        if not tps_values:
+            tps_values = [float(tps) for _, tps in step_tps_samples]
+    else:
+        tps_values = [float(m.group(1)) for m in _TPS_PATTERN.finditer(content)]
+
+    if tps_values:
+        sorted_tps = sorted(tps_values)
+        result.tps_samples = len(tps_values)
+        result.tps_median = float(statistics.median(tps_values))
+        result.tps_p10 = float(_percentile_linear(sorted_tps, 10.0))
+        result.tps_min = float(sorted_tps[0])
+        result.tps_max = float(sorted_tps[-1])
+
+    return result
+
+
+def check_secondary_alarms(result: ExperimentResult) -> list[str]:
     """Check secondary metrics against fixed alarm thresholds.
 
     Args:
@@ -216,44 +216,44 @@ def check_secondary_alarms(result: ExperimentResult) -> list[str]:
         alarms.append(
             f"engram_hit_rate={result.engram_hit_rate:.4f} < 0.1 (memory underused)"
         )
-    if 0 < result.mfu_percent < 10:
-        alarms.append(
-            f"mfu_percent={result.mfu_percent:.2f}% < 10% (GPU underutilized)"
-        )
-    if result.calibration_ece > 0.35:
-        alarms.append(
-            f"calibration_ece={result.calibration_ece:.4f} > 0.35 (poor calibration)"
-        )
-    if result.tps_median > 0 and result.tps_median < 50000:
-        alarms.append(
-            f"tps_median={result.tps_median:.0f} < 50000 (throughput below A10 objective)"
-        )
-
-    return alarms
-
-
-def _check_gate(
-    result: ExperimentResult,
-    gates: GateConfig,
-    metric: str,
-) -> tuple[bool, str] | None:
-    """Evaluate a single min/max gate against an ExperimentResult metric."""
-    gate = gates.get(metric, {})
-    value = getattr(result, metric)
-    max_value = gate.get("max")
-    if max_value is not None and value > max_value:
-        return False, f"{metric} {value:.4f} > gate {max_value}"
-    min_value = gate.get("min")
-    if min_value is not None and value < min_value:
-        return False, f"{metric} {value:.4f} < gate {min_value}"
-    return None
-
-
-def should_keep(
-    result: ExperimentResult,
-    best_bpb: float,
-    gates: GateConfig | None = None,
-) -> tuple[bool, str]:
+    if 0 < result.mfu_percent < 10:
+        alarms.append(
+            f"mfu_percent={result.mfu_percent:.2f}% < 10% (GPU underutilized)"
+        )
+    if result.calibration_ece > 0.35:
+        alarms.append(
+            f"calibration_ece={result.calibration_ece:.4f} > 0.35 (poor calibration)"
+        )
+    if result.tps_median > 0 and result.tps_median < 50000:
+        alarms.append(
+            f"tps_median={result.tps_median:.0f} < 50000 (throughput below A10 objective)"
+        )
+
+    return alarms
+
+
+def _check_gate(
+    result: ExperimentResult,
+    gates: GateConfig,
+    metric: str,
+) -> tuple[bool, str] | None:
+    """Evaluate a single min/max gate against an ExperimentResult metric."""
+    gate = gates.get(metric, {})
+    value = getattr(result, metric)
+    max_value = gate.get("max")
+    if max_value is not None and value > max_value:
+        return False, f"{metric} {value:.4f} > gate {max_value}"
+    min_value = gate.get("min")
+    if min_value is not None and value < min_value:
+        return False, f"{metric} {value:.4f} < gate {min_value}"
+    return None
+
+
+def should_keep(
+    result: ExperimentResult,
+    best_bpb: float,
+    gates: GateConfig | None = None,
+) -> tuple[bool, str]:
     """Decide whether to keep or discard an experiment.
 
     The primary criterion is strictly lower val_bpb than the current best.
@@ -277,24 +277,24 @@ def should_keep(
     if result.val_bpb >= best_bpb:
         return False, "discard"
 
-    # Secondary gate checks.
-    if gates:
-        gate_metrics = (
-            "mhc_spectral_norm",
-            "engram_hit_rate",
-            "factual_english_score",
-            "instruction_following_score",
-            "distinct_1",
-            "distinct_2",
-            "repetition_rate",
-            "repetition_bigram_rate",
-            "calibration_ece",
-            "tps_median",
-            "tps_p10",
-        )
-        for metric in gate_metrics:
-            gate_result = _check_gate(result, gates, metric)
-            if gate_result is not None:
-                return gate_result
-
-    return True, "keep"
+    # Secondary gate checks.
+    if gates:
+        gate_metrics = (
+            "mhc_spectral_norm",
+            "engram_hit_rate",
+            "factual_english_score",
+            "instruction_following_score",
+            "distinct_1",
+            "distinct_2",
+            "repetition_rate",
+            "repetition_bigram_rate",
+            "calibration_ece",
+            "tps_median",
+            "tps_p10",
+        )
+        for metric in gate_metrics:
+            gate_result = _check_gate(result, gates, metric)
+            if gate_result is not None:
+                return gate_result
+
+    return True, "keep"

overlay/harness/orchestrator.py

@@ -20,12 +20,12 @@ provides the infrastructure ("rails") that the autoresearch loop runs on.
 """
 import argparse
 import csv
-import os
-import subprocess
-import time
-
-from configs.harness_config import HarnessConfig
-from harness.eval_agent import ExperimentResult, check_secondary_alarms, parse_run_log, should_keep
+import os
+import subprocess
+import time
+
+from configs.harness_config import HarnessConfig
+from harness.eval_agent import ExperimentResult, check_secondary_alarms, parse_run_log, should_keep
 from harness.git_utils import REPO_DIR, commit_all, current_commit_short, reset_to
 from harness.health_monitor import check_health, reset_peak_stats
 from harness.meta_agent import run_meta_iteration
@@ -145,12 +145,12 @@ def run_experiment(timeout: int = 600) -> str:
 # ---------------------------------------------------------------------------
 
 
-def run_loop(
-    meta_interval: int = 20,
-    max_experiments: int | None = None,
-    experiment_timeout: int = 600,
-    secondary_gates: dict[str, dict[str, float]] | None = None,
-) -> None:
+def run_loop(
+    meta_interval: int = 20,
+    max_experiments: int | None = None,
+    experiment_timeout: int = 600,
+    secondary_gates: dict[str, dict[str, float]] | None = None,
+) -> None:
     """Run the HYDRA autoresearch loop.
 
     This function runs indefinitely (or until ``max_experiments`` is reached
@@ -163,10 +163,10 @@ def run_loop(
         secondary_gates: Optional gate thresholds forwarded to
             :func:`~harness.eval_agent.should_keep`.
     """
-    init_results_tsv()
-    if secondary_gates is None:
-        secondary_gates = HarnessConfig().to_secondary_gates()
-    best_bpb = _load_best_bpb()
+    init_results_tsv()
+    if secondary_gates is None:
+        secondary_gates = HarnessConfig().to_secondary_gates()
+    best_bpb = _load_best_bpb()
     experiment_num = count_experiments()
 
     print(

overlay/htm_rust/src/gpu/fused.rs

@@ -513,41 +513,41 @@ pub(super) fn launch_fused_batched_raw(
     assert_eq!(anom_per_region.len(), b);
     assert!(b >= 1, "need at least one region");
 
-    // Reset per-region step_scratch before each launch.
-    for &rp in region_ptrs.iter() {
-        let r = unsafe { &mut *rp };
-        let dev = r.sp_gpu.dev_ref().clone();
-        let fused = r
-            .fused_state
-            .as_mut()
-            .ok_or(DriverError(sys::CUresult::CUDA_ERROR_NOT_INITIALIZED))?;
-        dev.memset_zeros(&mut fused.step_scratch)?;
-        fused.iter_counter = fused.iter_counter.wrapping_add(1);
-    }
+    // Reset per-region step_scratch before each launch.
+    for &rp in region_ptrs.iter() {
+        let r = unsafe { &mut *rp };
+        let dev = r.sp_gpu.dev_ref().clone();
+        let fused = r
+            .fused_state
+            .as_mut()
+            .ok_or(DriverError(sys::CUresult::CUDA_ERROR_NOT_INITIALIZED))?;
+        dev.memset_zeros(&mut fused.step_scratch)?;
+        fused.iter_counter = fused.iter_counter.wrapping_add(1);
+    }
 
     // Shared config — all regions use identical sp/tm parameters.
-    let (grid_x, block_x, function_batched, cu_stream, cu_ctx) = {
-        let r0 = unsafe { &*region_ptrs[0] };
-        let fused = r0
-            .fused_state
-            .as_ref()
-            .ok_or(DriverError(sys::CUresult::CUDA_ERROR_NOT_INITIALIZED))?;
-        (
-            fused.grid_dim_x,
-            fused.block_dim_x,
-            fused.raw_kernel.function_batched,
-            *r0.sp_gpu.dev_ref().cu_stream(),
-            *r0.sp_gpu.dev_ref().cu_primary_ctx(),
-        )
+    let (grid_x, block_x, function_batched, cu_stream, cu_ctx) = {
+        let r0 = unsafe { &*region_ptrs[0] };
+        let fused = r0
+            .fused_state
+            .as_ref()
+            .ok_or(DriverError(sys::CUresult::CUDA_ERROR_NOT_INITIALIZED))?;
+        (
+            fused.grid_dim_x,
+            fused.block_dim_x,
+            fused.raw_kernel.function_batched,
+            *r0.sp_gpu.dev_ref().cu_stream(),
+            *r0.sp_gpu.dev_ref().cu_primary_ctx(),
+        )
     };
 
-    let cfg = {
-        let r = unsafe { &*region_ptrs[0] };
-        let fused = r
-            .fused_state
-            .as_ref()
-            .ok_or(DriverError(sys::CUresult::CUDA_ERROR_NOT_INITIALIZED))?;
-        FusedConfig {
+    let cfg = {
+        let r = unsafe { &*region_ptrs[0] };
+        let fused = r
+            .fused_state
+            .as_ref()
+            .ok_or(DriverError(sys::CUresult::CUDA_ERROR_NOT_INITIALIZED))?;
+        FusedConfig {
             input_bits: input_bits as u32,
             n_columns: r.sp_gpu.n_columns_accessor() as u32,
             synapses_per_col: r.sp_gpu.synapses_per_col_accessor() as u32,
@@ -572,42 +572,42 @@ pub(super) fn launch_fused_batched_raw(
             initial_perm_i16: r.tm_gpu.initial_perm_i16 as i32,
             t: t as u32,
             learn: if learn { 1 } else { 0 },
-            iter_seed: fused.iter_counter,
-            cooperative_grid_sync: 1,
-        }
-    };
+            iter_seed: fused.iter_counter,
+            cooperative_grid_sync: 1,
+        }
+    };
 
     // Build B FusedPtrs per-region.
-    let ptrs_vec: Vec<FusedPtrs> = (0..b)
-        .map(|i| {
-            let r = unsafe { &*region_ptrs[i] };
-            let fused = r
-                .fused_state
-                .as_ref()
-                .ok_or(DriverError(sys::CUresult::CUDA_ERROR_NOT_INITIALIZED))?;
-            Ok(FusedPtrs {
-                syn_bit: *r.sp_gpu.syn_bit_accessor().device_ptr(),
-                syn_perm: *r.sp_gpu.syn_perm_accessor().device_ptr(),
-                boost: *r.sp_gpu.boost_accessor().device_ptr(),
-                active_duty: *r.sp_gpu.active_duty_accessor().device_ptr(),
-                inhibition_threshold: *fused.inhibition_threshold.device_ptr(),
-                seg_cell_id: *r.tm_gpu.seg_cell_id_accessor().device_ptr(),
-                seg_syn_count: *r.tm_gpu.seg_syn_count_accessor().device_ptr(),
-                syn_presyn: *r.tm_gpu.syn_presyn_accessor().device_ptr(),
-                tm_syn_perm: *r.tm_gpu.syn_perm_accessor().device_ptr(),
-                cell_seg_count: *r.tm_gpu.cell_seg_count_accessor().device_ptr(),
-                cell_active_a: *fused.cell_active_bits_a.device_ptr(),
-                cell_active_b: *fused.cell_active_bits_b.device_ptr(),
-                cell_winner_a: *fused.cell_winner_bits_a.device_ptr(),
-                cell_winner_b: *fused.cell_winner_bits_b.device_ptr(),
-                inputs: inputs_per_region[i],
-                cols_out: cols_per_region[i],
-                anom_out: anom_per_region[i],
-                barrier_counters: 0u64,  // ABI-compat dummy; cluster barrier replaces DLB.
-                step_scratch: *fused.step_scratch.device_ptr(),
-            })
-        })
-        .collect::<Result<Vec<_>, DriverError>>()?;
+    let ptrs_vec: Vec<FusedPtrs> = (0..b)
+        .map(|i| {
+            let r = unsafe { &*region_ptrs[i] };
+            let fused = r
+                .fused_state
+                .as_ref()
+                .ok_or(DriverError(sys::CUresult::CUDA_ERROR_NOT_INITIALIZED))?;
+            Ok(FusedPtrs {
+                syn_bit: *r.sp_gpu.syn_bit_accessor().device_ptr(),
+                syn_perm: *r.sp_gpu.syn_perm_accessor().device_ptr(),
+                boost: *r.sp_gpu.boost_accessor().device_ptr(),
+                active_duty: *r.sp_gpu.active_duty_accessor().device_ptr(),
+                inhibition_threshold: *fused.inhibition_threshold.device_ptr(),
+                seg_cell_id: *r.tm_gpu.seg_cell_id_accessor().device_ptr(),
+                seg_syn_count: *r.tm_gpu.seg_syn_count_accessor().device_ptr(),
+                syn_presyn: *r.tm_gpu.syn_presyn_accessor().device_ptr(),
+                tm_syn_perm: *r.tm_gpu.syn_perm_accessor().device_ptr(),
+                cell_seg_count: *r.tm_gpu.cell_seg_count_accessor().device_ptr(),
+                cell_active_a: *fused.cell_active_bits_a.device_ptr(),
+                cell_active_b: *fused.cell_active_bits_b.device_ptr(),
+                cell_winner_a: *fused.cell_winner_bits_a.device_ptr(),
+                cell_winner_b: *fused.cell_winner_bits_b.device_ptr(),
+                inputs: inputs_per_region[i],
+                cols_out: cols_per_region[i],
+                anom_out: anom_per_region[i],
+                barrier_counters: 0u64,  // ABI-compat dummy; cluster barrier replaces DLB.
+                step_scratch: *fused.step_scratch.device_ptr(),
+            })
+        })
+        .collect::<Result<Vec<_>, DriverError>>()?;
 
     // Upload FusedPtrs array to device (B * sizeof(FusedPtrs) bytes).
     // FusedPtrs is repr(C) + DeviceRepr so htod_sync_copy handles it.
@@ -619,14 +619,14 @@ pub(super) fn launch_fused_batched_raw(
     // Grid = (grid_x, B, 1) with cluster_dim=(16,1,1): each region (Y slice)
     // occupies exactly one cluster of 16 blocks. All 8 clusters run concurrently
     // on the H200's 132 SMs (8 × 16 = 128 blocks ≤ 132 SMs).
-    let use_cluster = {
-        let r0 = unsafe { &*region_ptrs[0] };
-        let fused = r0
-            .fused_state
-            .as_ref()
-            .ok_or(DriverError(sys::CUresult::CUDA_ERROR_NOT_INITIALIZED))?;
-        fused.cluster_info.max_cluster_size > 0
-    };
+    let use_cluster = {
+        let r0 = unsafe { &*region_ptrs[0] };
+        let fused = r0
+            .fused_state
+            .as_ref()
+            .ok_or(DriverError(sys::CUresult::CUDA_ERROR_NOT_INITIALIZED))?;
+        fused.cluster_info.max_cluster_size > 0
+    };
 
     unsafe {
         result::ctx::set_current(cu_ctx)?;

overlay/hydra/eval.py

@@ -138,9 +138,6 @@ def _run_factual_english_gen(model, tokenizer, max_seq_len: int):
     num_samples = FACTUAL_SAMPLES
     batch = FACTUAL_BATCH
     gen_tokens = FACTUAL_GEN_TOKENS
-    # Optional fast incremental decode path for recurrence-capable backbones.
-    # If disabled, we preserve the original full-context re-forward behavior.
-    incremental_decode = os.environ.get("HYDRA_FACTUAL_GEN_INCREMENTAL", "1") == "1"
     temps = [0.7, 0.9, 1.1]
     hits = 0
 
@@ -157,18 +154,14 @@ def _run_factual_english_gen(model, tokenizer, max_seq_len: int):
                 temp = temps[batch_idx % len(temps)]
                 batch_idx += 1
                 ctx = torch.tensor([ids] * b, device="cuda", dtype=torch.long)
-                logits = model(ctx, targets=None)
                 for _ in range(gen_tokens):
+                    logits = model(ctx, targets=None)
                     next_logits = logits[:, -1, :] if logits.dim() == 3 else logits
                     probs = torch.softmax(next_logits.float() / temp, dim=-1)
                     next_id = torch.multinomial(probs, num_samples=1)
                     ctx = torch.cat([ctx, next_id], dim=1)
                     if ctx.size(1) >= max_seq_len:
                         break
-                    if incremental_decode:
-                        logits = model(ctx[:, -1:], targets=None)
-                    else:
-                        logits = model(ctx, targets=None)
                 # Transfer to CPU in one shot, no per-row sync
                 all_rows.extend(ctx.cpu().tolist())
                 samples_done += b

overlay/hydra/model.py

@@ -32,58 +32,58 @@ from __future__ import annotations
 
 import os
 
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-try:
-    from mamba_ssm import Mamba3
-except Exception:  # pragma: no cover - depends on optional runtime install
-    Mamba3 = None  # type: ignore[assignment]
-
-
-def _get_mamba3_cls():
-    global Mamba3
-    if Mamba3 is None:
-        try:
-            from mamba_ssm import Mamba3 as _Mamba3  # type: ignore
-            Mamba3 = _Mamba3  # type: ignore[assignment]
-        except Exception as exc:  # pragma: no cover - environment dependent
-            raise ImportError(
-                "mamba_ssm is required for Mamba-based HYDRA blocks. "
-                "Install mamba-ssm or use HYDRA_BASELINE_ARCH=transformer."
-            ) from exc
-    return Mamba3
-
-
-def _ensure_triton_cuda_backend_registered() -> None:
-    """Ensure Triton sees exactly one CUDA backend in HF Jobs A10 runtime.
-
-    In some Triton 3.5.1 environments, `triton.compiler.compiler.backends`
-    and `triton.runtime.driver.backends` are empty even though
-    `triton.backends.nvidia` is available and CUDA is active. When that
-    happens, Mamba3 layernorm path crashes at first forward with
-    "0 compatible backends for target (cuda)".
-    """
-    try:
-        import triton.compiler.compiler as cc
-        import triton.runtime.driver as rd
-        from triton.backends import Backend
-        from triton.backends.nvidia.compiler import CUDABackend
-        from triton.backends.nvidia.driver import CudaDriver
-
-        if hasattr(rd, "backends") and isinstance(rd.backends, dict) and not rd.backends:
-            rd.backends["nvidia"] = Backend(compiler=CUDABackend, driver=CudaDriver)
-
-        if hasattr(cc, "backends") and isinstance(cc.backends, dict) and not cc.backends:
-            cc.backends["nvidia"] = Backend(compiler=CUDABackend, driver=CudaDriver)
-    except Exception:
-        # Keep model construction resilient; runtime will raise explicit Triton
-        # errors later if backend setup is still invalid.
-        pass
-
-
-_ensure_triton_cuda_backend_registered()
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+try:
+    from mamba_ssm import Mamba3
+except Exception:  # pragma: no cover - depends on optional runtime install
+    Mamba3 = None  # type: ignore[assignment]
+
+
+def _get_mamba3_cls():
+    global Mamba3
+    if Mamba3 is None:
+        try:
+            from mamba_ssm import Mamba3 as _Mamba3  # type: ignore
+            Mamba3 = _Mamba3  # type: ignore[assignment]
+        except Exception as exc:  # pragma: no cover - environment dependent
+            raise ImportError(
+                "mamba_ssm is required for Mamba-based HYDRA blocks. "
+                "Install mamba-ssm or use HYDRA_BASELINE_ARCH=transformer."
+            ) from exc
+    return Mamba3
+
+
+def _ensure_triton_cuda_backend_registered() -> None:
+    """Ensure Triton sees exactly one CUDA backend in HF Jobs A10 runtime.
+
+    In some Triton 3.5.1 environments, `triton.compiler.compiler.backends`
+    and `triton.runtime.driver.backends` are empty even though
+    `triton.backends.nvidia` is available and CUDA is active. When that
+    happens, Mamba3 layernorm path crashes at first forward with
+    "0 compatible backends for target (cuda)".
+    """
+    try:
+        import triton.compiler.compiler as cc
+        import triton.runtime.driver as rd
+        from triton.backends import Backend
+        from triton.backends.nvidia.compiler import CUDABackend
+        from triton.backends.nvidia.driver import CudaDriver
+
+        if hasattr(rd, "backends") and isinstance(rd.backends, dict) and not rd.backends:
+            rd.backends["nvidia"] = Backend(compiler=CUDABackend, driver=CudaDriver)
+
+        if hasattr(cc, "backends") and isinstance(cc.backends, dict) and not cc.backends:
+            cc.backends["nvidia"] = Backend(compiler=CUDABackend, driver=CudaDriver)
+    except Exception:
+        # Keep model construction resilient; runtime will raise explicit Triton
+        # errors later if backend setup is still invalid.
+        pass
+
+
+_ensure_triton_cuda_backend_registered()
 
 from subsystems.hestia_mini import HestiaQAT
 from subsystems.htm import HTMLayer
@@ -98,30 +98,30 @@ from hydra.hyena_block import HyenaBlock
 from hydra.optimizer import MuonAdamW
 
 
-def norm(x: torch.Tensor) -> torch.Tensor:
-    """RMSNorm over the last dim — stateless, autocast-friendly."""
-    return F.rms_norm(x, (x.size(-1),))
-
-
-class TransformerBaselineBlock(nn.Module):
-    """Transformer-style delta block for matched baseline experiments.
-
-    This block returns a transformed delta tensor rather than owning the outer
-    residual connection, because ManifoldHyperConnection already handles stream
-    mixing and residual injection around the block function.
-    """
-
-    def __init__(self, d_model: int, n_heads: int, expand: int, dropout: float) -> None:
-        super().__init__()
-        self.self_attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True)
-        self.ff_in = nn.Linear(d_model, expand * d_model, bias=False)
-        self.ff_out = nn.Linear(expand * d_model, d_model, bias=False)
-        self.dropout = nn.Dropout(dropout)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        attn_out, _ = self.self_attn(x, x, x, need_weights=False)
-        ff = self.ff_out(F.gelu(self.ff_in(attn_out)))
-        return self.dropout(attn_out + ff)
+def norm(x: torch.Tensor) -> torch.Tensor:
+    """RMSNorm over the last dim — stateless, autocast-friendly."""
+    return F.rms_norm(x, (x.size(-1),))
+
+
+class TransformerBaselineBlock(nn.Module):
+    """Transformer-style delta block for matched baseline experiments.
+
+    This block returns a transformed delta tensor rather than owning the outer
+    residual connection, because ManifoldHyperConnection already handles stream
+    mixing and residual injection around the block function.
+    """
+
+    def __init__(self, d_model: int, n_heads: int, expand: int, dropout: float) -> None:
+        super().__init__()
+        self.self_attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True)
+        self.ff_in = nn.Linear(d_model, expand * d_model, bias=False)
+        self.ff_out = nn.Linear(expand * d_model, d_model, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        attn_out, _ = self.self_attn(x, x, x, need_weights=False)
+        ff = self.ff_out(F.gelu(self.ff_in(attn_out)))
+        return self.dropout(attn_out + ff)
 
 
 class PostSemClawModel(nn.Module):
@@ -136,12 +136,12 @@ class PostSemClawModel(nn.Module):
         model(x, y, reduction='mean')           -> scalar loss
     """
 
-    def __init__(self, config):
-        super().__init__()
-        _ensure_triton_cuda_backend_registered()
-        self.config = config
-        self._throughput_mode = os.environ.get("HYDRA_THROUGHPUT_MODE", "0") == "1"
-        self._baseline_arch = os.environ.get("HYDRA_BASELINE_ARCH", "mamba3").strip().lower()
+    def __init__(self, config):
+        super().__init__()
+        _ensure_triton_cuda_backend_registered()
+        self.config = config
+        self._throughput_mode = os.environ.get("HYDRA_THROUGHPUT_MODE", "0") == "1"
+        self._baseline_arch = os.environ.get("HYDRA_BASELINE_ARCH", "mamba3").strip().lower()
 
         # Token embedding
         self.wte = nn.Embedding(config.vocab_size, config.d_model)
@@ -163,31 +163,31 @@ class PostSemClawModel(nn.Module):
             print(f"[WARN] layers in both hyena_layers and gdn_layers; using Hyena: {sorted(_both)}", flush=True)
             _gdn_layer_set -= _hyena_layer_set
 
-        if _gdn_layer_set:
-            from hydra.gdn_block import GDNBlock  # requires `fla` package
-
-        def _build_block(i: int) -> nn.Module:
-            if self._baseline_arch == "transformer":
-                return TransformerBaselineBlock(
-                    d_model=config.d_model,
-                    n_heads=config.n_heads,
-                    expand=config.expand,
-                    dropout=float(os.environ.get("HYDRA_DROPOUT", "0.2")),
-                )
-            if i in _hyena_layer_set:
-                return HyenaBlock(
+        if _gdn_layer_set:
+            from hydra.gdn_block import GDNBlock  # requires `fla` package
+
+        def _build_block(i: int) -> nn.Module:
+            if self._baseline_arch == "transformer":
+                return TransformerBaselineBlock(
+                    d_model=config.d_model,
+                    n_heads=config.n_heads,
+                    expand=config.expand,
+                    dropout=float(os.environ.get("HYDRA_DROPOUT", "0.2")),
+                )
+            if i in _hyena_layer_set:
+                return HyenaBlock(
                     d_model=config.d_model,
                     seq_len=config.sequence_len,
                     order=int(os.environ.get("HYDRA_HYENA_ORDER", "2")),
                     filter_order=int(os.environ.get("HYDRA_HYENA_FILTER_DIM", "64")),
                 )
-            if i in _gdn_layer_set:
-                return GDNBlock(
-                    d_model=config.d_model,
-                    n_heads=config.n_heads,
-                )
-            mamba3_cls = _get_mamba3_cls()
-            return mamba3_cls(
+            if i in _gdn_layer_set:
+                return GDNBlock(
+                    d_model=config.d_model,
+                    n_heads=config.n_heads,
+                )
+            mamba3_cls = _get_mamba3_cls()
+            return mamba3_cls(
                 d_model=config.d_model,
                 d_state=config.d_state,
                 expand=config.expand,
@@ -201,43 +201,43 @@ class PostSemClawModel(nn.Module):
         self.blocks = nn.ModuleList([_build_block(i) for i in range(config.n_layer)])
 
         # Full-architecture SDR: offline semantic retina + STE (no-bypass).
-        if self._throughput_mode:
-            self.sdr_semantic = None
-            self.htm = None
-            self.htm_proj = None
-            self.htm_anom_proj = None
-            self.engram = None
-            self.engram_layer_idx = -1
-        else:
-            self.sdr_semantic = SemanticFoldingSDR(
-                vocab_size=config.vocab_size,
-                n_bits=config.sdr_n_bits,
-                target_active=config.sdr_target_active,
-                delta_rank=config.sdr_delta_rank,
-                som_warmup_steps=config.sdr_som_warmup,
-                som_update_interval=config.sdr_som_interval,
-            )
-
-            # HTM spatial pooler + temporal memory (Rust, Hebbian).
-            self.htm = HTMLayer(
-                input_bits=config.sdr_n_bits,
-                n_columns=config.htm_n_columns,
-                cells_per_column=config.htm_cells_per_column,
-                batch_size=1,
-                seed=42,
-                learn=True,
-                reset_each_forward=True,
-            )
-
-            self.htm_proj = nn.Linear(config.htm_n_columns, config.d_model, bias=False)
-            self.htm_anom_proj = nn.Linear(1, config.d_model, bias=False)
-
-            self.engram = GPUEngram(
-                d_model=config.d_model,
-                n_columns=config.engram_n_columns,
-                max_ngram=3,
-            )
-            self.engram_layer_idx = config.engram_layer_idx
+        if self._throughput_mode:
+            self.sdr_semantic = None
+            self.htm = None
+            self.htm_proj = None
+            self.htm_anom_proj = None
+            self.engram = None
+            self.engram_layer_idx = -1
+        else:
+            self.sdr_semantic = SemanticFoldingSDR(
+                vocab_size=config.vocab_size,
+                n_bits=config.sdr_n_bits,
+                target_active=config.sdr_target_active,
+                delta_rank=config.sdr_delta_rank,
+                som_warmup_steps=config.sdr_som_warmup,
+                som_update_interval=config.sdr_som_interval,
+            )
+
+            # HTM spatial pooler + temporal memory (Rust, Hebbian).
+            self.htm = HTMLayer(
+                input_bits=config.sdr_n_bits,
+                n_columns=config.htm_n_columns,
+                cells_per_column=config.htm_cells_per_column,
+                batch_size=1,
+                seed=42,
+                learn=True,
+                reset_each_forward=True,
+            )
+
+            self.htm_proj = nn.Linear(config.htm_n_columns, config.d_model, bias=False)
+            self.htm_anom_proj = nn.Linear(1, config.d_model, bias=False)
+
+            self.engram = GPUEngram(
+                d_model=config.d_model,
+                n_columns=config.engram_n_columns,
+                max_ngram=3,
+            )
+            self.engram_layer_idx = config.engram_layer_idx
 
         # Manifold-Constrained Hyper-Connections (one per Mamba-3 block).
         self.mhc = nn.ModuleList([
@@ -258,18 +258,18 @@ class PostSemClawModel(nn.Module):
         #   additional CE losses; no new params. Activated via HYDRA_MTP_K.
         self._mtp_k = max(1, int(os.environ.get("HYDRA_MTP_K", "1")))
 
-        # Learnability knob 3: gradient checkpointing on Mamba3 blocks.
-        self._grad_ckpt = os.environ.get("HYDRA_GRAD_CKPT", "0") == "1"
-
-        # Full-arch throughput knob: Engram remains in the architecture, but
-        # can run every N forwards and reuse its most recent residual delta on
-        # skipped forwards. This amortizes top-k Hopfield retrieval while still
-        # injecting Engram signal every microbatch. N=1 preserves exact legacy
-        # behavior.
-        self._engram_subsample = max(1, int(os.environ.get("HYDRA_ENGRAM_SUBSAMPLE", "1")))
-        self._engram_call_idx = 0
-        self._engram_delta_cache = None
-        self._engram_hit_rate_cache = None
+        # Learnability knob 3: gradient checkpointing on Mamba3 blocks.
+        self._grad_ckpt = os.environ.get("HYDRA_GRAD_CKPT", "0") == "1"
+
+        # Full-arch throughput knob: Engram remains in the architecture, but
+        # can run every N forwards and reuse its most recent residual delta on
+        # skipped forwards. This amortizes top-k Hopfield retrieval while still
+        # injecting Engram signal every microbatch. N=1 preserves exact legacy
+        # behavior.
+        self._engram_subsample = max(1, int(os.environ.get("HYDRA_ENGRAM_SUBSAMPLE", "1")))
+        self._engram_call_idx = 0
+        self._engram_delta_cache = None
+        self._engram_hit_rate_cache = None
 
         # Learnability knob 4: doc-separator BOS masking in packed sequences.
         self._doc_sep_mask = os.environ.get("HYDRA_DOC_SEP_MASK", "0") == "1"
@@ -373,8 +373,8 @@ class PostSemClawModel(nn.Module):
         # Required because to_empty() only moves params/buffers, and _retina_indices
         # is loaded from numpy (always CPU) by SemanticFoldingSDR.__init__.
         device = self.wte.weight.device
-        if self.sdr_semantic is not None and hasattr(self.sdr_semantic, '_retina_indices'):
-            self.sdr_semantic._retina_indices = self.sdr_semantic._retina_indices.to(device)
+        if self.sdr_semantic is not None and hasattr(self.sdr_semantic, '_retina_indices'):
+            self.sdr_semantic._retina_indices = self.sdr_semantic._retina_indices.to(device)
 
         # Embedding init: GPT-2 / LLaMA convention. std=1.0 was chosen for
         # vocab=8192; at larger vocabs, smaller std prevents logit blowup.
@@ -413,20 +413,20 @@ class PostSemClawModel(nn.Module):
                 ))
                 nn.init.normal_(block.out_proj.weight, mean=0.0, std=out_std)
 
-        if self.htm_proj is not None:
-            nn.init.normal_(self.htm_proj.weight, mean=0.0, std=s)
-        if self.htm_anom_proj is not None:
-            nn.init.normal_(self.htm_anom_proj.weight, mean=0.0, std=s)
+        if self.htm_proj is not None:
+            nn.init.normal_(self.htm_proj.weight, mean=0.0, std=s)
+        if self.htm_anom_proj is not None:
+            nn.init.normal_(self.htm_anom_proj.weight, mean=0.0, std=s)
 
         # Cast to bf16 to match Mamba3 dtype; Muon groups by shape so mixed
         # dtypes in the same shape group would break lerp_ dtype checks.
         self.wte.to(dtype=torch.bfloat16)
-        if self.htm_proj is not None:
-            self.htm_proj.to(dtype=torch.bfloat16)
-        if self.htm_anom_proj is not None:
-            self.htm_anom_proj.to(dtype=torch.bfloat16)
-        if self.engram is not None:
-            self.engram.to(dtype=torch.bfloat16)
+        if self.htm_proj is not None:
+            self.htm_proj.to(dtype=torch.bfloat16)
+        if self.htm_anom_proj is not None:
+            self.htm_anom_proj.to(dtype=torch.bfloat16)
+        if self.engram is not None:
+            self.engram.to(dtype=torch.bfloat16)
 
     def set_bos_token_id(self, bos_id: int) -> None:
         """Inform the model of the tokenizer's BOS id so doc-separator
@@ -472,10 +472,10 @@ class PostSemClawModel(nn.Module):
         wte = sum(p.numel() for p in self.wte.parameters())
         lm_head = sum(p.numel() for p in self.lm_head.parameters())
         blocks = sum(p.numel() for p in self.blocks.parameters())
-        sdr = sum(p.numel() for p in self.sdr_semantic.parameters()) if self.sdr_semantic is not None else 0
-        htm_proj = sum(p.numel() for p in self.htm_proj.parameters()) if self.htm_proj is not None else 0
-        htm_anom_proj = sum(p.numel() for p in self.htm_anom_proj.parameters()) if self.htm_anom_proj is not None else 0
-        engram = sum(p.numel() for p in self.engram.parameters()) if self.engram is not None else 0
+        sdr = sum(p.numel() for p in self.sdr_semantic.parameters()) if self.sdr_semantic is not None else 0
+        htm_proj = sum(p.numel() for p in self.htm_proj.parameters()) if self.htm_proj is not None else 0
+        htm_anom_proj = sum(p.numel() for p in self.htm_anom_proj.parameters()) if self.htm_anom_proj is not None else 0
+        engram = sum(p.numel() for p in self.engram.parameters()) if self.engram is not None else 0
         total = sum(p.numel() for p in self.parameters())
         return {
             'wte': wte, 'lm_head': lm_head, 'blocks': blocks,
@@ -544,19 +544,19 @@ class PostSemClawModel(nn.Module):
         # for p in self.sdr_semantic.parameters():
         #     if p.dim() == 2:
         #         matrix_params.append(p)
-        if self.htm_proj is not None:
-            for name, p in self.htm_proj.named_parameters():
-                if _muon_eligible(name, p):
-                    matrix_params.append(p)
-        if self.engram is not None:
-            for name, p in self.engram.named_parameters():
-                if _muon_eligible(name, p):
-                    matrix_params.append(p)
+        if self.htm_proj is not None:
+            for name, p in self.htm_proj.named_parameters():
+                if _muon_eligible(name, p):
+                    matrix_params.append(p)
+        if self.engram is not None:
+            for name, p in self.engram.named_parameters():
+                if _muon_eligible(name, p):
+                    matrix_params.append(p)
 
         # SDR params are intentionally not in any optimizer group — they
         # receive no gradient in the current forward, so any update would be
         # pure noise (weight_decay × lr on a zero-grad param).
-        sdr_param_ids = set(id(p) for p in self.sdr_semantic.parameters()) if self.sdr_semantic is not None else set()
+        sdr_param_ids = set(id(p) for p in self.sdr_semantic.parameters()) if self.sdr_semantic is not None else set()
         assigned = set(id(p) for p in embedding_params + lm_head_params + matrix_params)
         scalar_params = [
             p for p in self.parameters()
@@ -565,7 +565,7 @@ class PostSemClawModel(nn.Module):
 
         total_assigned = len(embedding_params) + len(lm_head_params) + len(matrix_params) + len(scalar_params)
         total_params = len(list(self.parameters()))
-        sdr_excluded = len(list(self.sdr_semantic.parameters())) if self.sdr_semantic is not None else 0
+        sdr_excluded = len(list(self.sdr_semantic.parameters())) if self.sdr_semantic is not None else 0
         assert total_assigned + sdr_excluded == total_params, (
             f"Parameter count mismatch: assigned {total_assigned} + sdr_excluded "
             f"{sdr_excluded} vs total {total_params}"
@@ -633,59 +633,59 @@ class PostSemClawModel(nn.Module):
         else:
             _t0 = None
 
-        dense_emb = self.wte(idx)  # (B, T, d_model) bf16
-        if self._throughput_mode:
-            self._last_sdr = None
-            sdr_active_bits = 0.0
-            htm_anomaly = dense_emb.new_tensor(0.0)
-            x = norm(dense_emb)
-            if _profile:
-                _t_htm_async = _ev()
-                _t_wte = _ev()
-                _t_htm_await = _ev()
-                _t_htm_proj = _ev()
-        else:
-            sdr_binary = self.sdr_semantic.binary_only(idx)
-            self._last_sdr = sdr_binary
-            _htm_sub = int(os.environ.get("HYDRA_HTM_SUBSAMPLE", "8"))
-            if not hasattr(self, '_htm_call_idx'):
-                self._htm_call_idx = 0
-
-            _run_htm = (self._htm_call_idx % _htm_sub == 0)
-            self._htm_call_idx += 1
-            if _run_htm:
-                htm_handle = self.htm.forward_async(sdr_binary)
-            else:
-                htm_handle = None
-
-            if _profile: _t_htm_async = _ev()
-            if _profile: _t_wte = _ev()
-
-            if _run_htm:
-                htm_out = self.htm.forward_await(htm_handle)
-                self._htm_cache = htm_out.detach()
-            elif hasattr(self, '_htm_cache') and self._htm_cache is not None and self._htm_cache.shape[0] == B and self._htm_cache.shape[1] == T:
-                htm_out = self._htm_cache
-            else:
-                htm_handle = self.htm.forward_async(sdr_binary)
-                htm_out = self.htm.forward_await(htm_handle)
-                self._htm_cache = htm_out.detach()
-
-            if _profile: _t_htm_await = _ev()
-            with torch.no_grad():
-                sdr_active_bits = float(self.sdr_semantic.target_active)
-                htm_anomaly = htm_out[..., -1].mean()
-            if self._htm_stop_grad:
-                htm_out = htm_out.detach()
-            htm_cols = htm_out[..., :-1].to(dense_emb.dtype)
-            htm_anom = htm_out[..., -1:].to(dense_emb.dtype)
-            htm_proj_out = self.htm_proj(htm_cols) + self.htm_anom_proj(htm_anom)
-            x = norm(dense_emb + htm_proj_out)
-            if _profile: _t_htm_proj = _ev()
-
-        # mHC-routed Mamba-3 stack with Engram injection at configured layer.
-        streams = self.mhc[0].init_streams(x)
-        _profile_layer_events = []
+        dense_emb = self.wte(idx)  # (B, T, d_model) bf16
+        if self._throughput_mode:
+            self._last_sdr = None
+            sdr_active_bits = 0.0
+            htm_anomaly = dense_emb.new_tensor(0.0)
+            x = norm(dense_emb)
+            if _profile:
+                _t_htm_async = _ev()
+                _t_wte = _ev()
+                _t_htm_await = _ev()
+                _t_htm_proj = _ev()
+        else:
+            sdr_binary = self.sdr_semantic.binary_only(idx)
+            self._last_sdr = sdr_binary
+            _htm_sub = int(os.environ.get("HYDRA_HTM_SUBSAMPLE", "8"))
+            if not hasattr(self, '_htm_call_idx'):
+                self._htm_call_idx = 0
+
+            _run_htm = (self._htm_call_idx % _htm_sub == 0)
+            self._htm_call_idx += 1
+            if _run_htm:
+                htm_handle = self.htm.forward_async(sdr_binary)
+            else:
+                htm_handle = None
+
+            if _profile: _t_htm_async = _ev()
+            if _profile: _t_wte = _ev()
+
+            if _run_htm:
+                htm_out = self.htm.forward_await(htm_handle)
+                self._htm_cache = htm_out.detach()
+            elif hasattr(self, '_htm_cache') and self._htm_cache is not None and self._htm_cache.shape[0] == B and self._htm_cache.shape[1] == T:
+                htm_out = self._htm_cache
+            else:
+                htm_handle = self.htm.forward_async(sdr_binary)
+                htm_out = self.htm.forward_await(htm_handle)
+                self._htm_cache = htm_out.detach()
+
+            if _profile: _t_htm_await = _ev()
+            with torch.no_grad():
+                sdr_active_bits = float(self.sdr_semantic.target_active)
+                htm_anomaly = htm_out[..., -1].mean()
+            if self._htm_stop_grad:
+                htm_out = htm_out.detach()
+            htm_cols = htm_out[..., :-1].to(dense_emb.dtype)
+            htm_anom = htm_out[..., -1:].to(dense_emb.dtype)
+            htm_proj_out = self.htm_proj(htm_cols) + self.htm_anom_proj(htm_anom)
+            x = norm(dense_emb + htm_proj_out)
+            if _profile: _t_htm_proj = _ev()
+
+        # mHC-routed Mamba-3 stack with Engram injection at configured layer.
+        streams = self.mhc[0].init_streams(x)
+        _profile_layer_events = []
 
         # Per-layer diagnostic panel. The pre-layer merged state h_pre lets us
         # measure residual contribution of each layer: delta_N = h_post - h_pre.
@@ -702,20 +702,20 @@ class PostSemClawModel(nn.Module):
                 h_pre = self.mhc[0].merge_streams(streams).detach().float()
             _run_svd = (self._diag_step % self._diag_svd_every) == 0
 
-        for i, (block, mhc_layer) in enumerate(zip(self.blocks, self.mhc)):
-            if _profile:
-                _t_layer_start = _ev()
-                _t_layer_mhc = None
-                _t_engram_start = None
-                _t_engram_end = None
-            else:
-                _t_layer_start = None
-                _t_layer_mhc = None
-                _t_engram_start = None
-                _t_engram_end = None
-
-            def _block_fn(h, _block=block):
-                return self.drop(_block(norm(h)))
+        for i, (block, mhc_layer) in enumerate(zip(self.blocks, self.mhc)):
+            if _profile:
+                _t_layer_start = _ev()
+                _t_layer_mhc = None
+                _t_engram_start = None
+                _t_engram_end = None
+            else:
+                _t_layer_start = None
+                _t_layer_mhc = None
+                _t_engram_start = None
+                _t_engram_end = None
+
+            def _block_fn(h, _block=block):
+                return self.drop(_block(norm(h)))
 
             # Learnability #3: gradient checkpointing. Wrap the block-fn so
             # the mhc layer's internal uses of it re-run the block in backward
@@ -726,31 +726,31 @@ class PostSemClawModel(nn.Module):
                 _raw_fn = _block_fn
                 def _block_fn(h, _raw=_raw_fn):  # noqa: E731
                     return _ckpt.checkpoint(_raw, h, use_reentrant=False)
-
-            streams = mhc_layer(streams, _block_fn)
-            if _profile:
-                _t_layer_mhc = _ev()
-
-            if self.engram is not None and i == self.engram_layer_idx:
-                if _profile: _t_engram_start = _ev()
-                x_mid = mhc_layer.merge_streams(streams)
-                _run_engram = (self._engram_call_idx % self._engram_subsample == 0)
-                self._engram_call_idx += 1
-                if _run_engram or self._engram_delta_cache is None or self._engram_delta_cache.shape != x_mid.shape:
-                    x_engram, hit_rate = self.engram(x_mid, idx)
-                    self._engram_delta_cache = (x_engram - x_mid).detach()
-                    self._engram_hit_rate_cache = hit_rate.detach() if torch.is_tensor(hit_rate) else hit_rate
-                    x_mid = x_engram
-                else:
-                    # Preserve gradient flow through the identity path while
-                    # reusing a detached Engram residual. The Engram module
-                    # still contributes to every forward; its expensive top-k
-                    # retrieval and parameter gradients run on the cadence.
-                    x_mid = x_mid + self._engram_delta_cache.to(dtype=x_mid.dtype, device=x_mid.device)
-                    hit_rate = self._engram_hit_rate_cache
-                streams = mhc_layer.init_streams(x_mid)
-                self._metrics['engram_hit_rate'] = hit_rate
-                if _profile: _t_engram_end = _ev()
+
+            streams = mhc_layer(streams, _block_fn)
+            if _profile:
+                _t_layer_mhc = _ev()
+
+            if self.engram is not None and i == self.engram_layer_idx:
+                if _profile: _t_engram_start = _ev()
+                x_mid = mhc_layer.merge_streams(streams)
+                _run_engram = (self._engram_call_idx % self._engram_subsample == 0)
+                self._engram_call_idx += 1
+                if _run_engram or self._engram_delta_cache is None or self._engram_delta_cache.shape != x_mid.shape:
+                    x_engram, hit_rate = self.engram(x_mid, idx)
+                    self._engram_delta_cache = (x_engram - x_mid).detach()
+                    self._engram_hit_rate_cache = hit_rate.detach() if torch.is_tensor(hit_rate) else hit_rate
+                    x_mid = x_engram
+                else:
+                    # Preserve gradient flow through the identity path while
+                    # reusing a detached Engram residual. The Engram module
+                    # still contributes to every forward; its expensive top-k
+                    # retrieval and parameter gradients run on the cadence.
+                    x_mid = x_mid + self._engram_delta_cache.to(dtype=x_mid.dtype, device=x_mid.device)
+                    hit_rate = self._engram_hit_rate_cache
+                streams = mhc_layer.init_streams(x_mid)
+                self._metrics['engram_hit_rate'] = hit_rate
+                if _profile: _t_engram_end = _ev()
 
             if _diag:
                 with torch.no_grad():
@@ -773,20 +773,20 @@ class PostSemClawModel(nn.Module):
                             self._metrics[f'layer_{i}_eff_rank'] = eff_rank
                         except Exception:
                             pass
-                    h_pre = h_post
-
-            if _profile:
-                _profile_layer_events.append(
-                    (i, _t_layer_start, _t_layer_mhc, _t_engram_start, _t_engram_end, _ev())
-                )
+                    h_pre = h_post
+
+            if _profile:
+                _profile_layer_events.append(
+                    (i, _t_layer_start, _t_layer_mhc, _t_engram_start, _t_engram_end, _ev())
+                )
 
         if _diag:
             self._diag_step += 1
 
         if _profile: _t_blocks = _ev()
 
-        self._metrics['sdr_active_bits'] = sdr_active_bits
-        self._metrics['htm_anomaly'] = htm_anomaly
+        self._metrics['sdr_active_bits'] = sdr_active_bits
+        self._metrics['htm_anomaly'] = htm_anomaly
 
         x = self.mhc[-1].merge_streams(streams)
         x = norm(x)
@@ -1000,8 +1000,8 @@ class PostSemClawModel(nn.Module):
                 _t_end = _ev()
                 torch.cuda.synchronize()
                 def _ms(a, b): return a.elapsed_time(b)
-                print(
-                    f"[PROFILE B={B} T={T}] "
+                print(
+                    f"[PROFILE B={B} T={T}] "
                     f"htm_launch={_ms(_t0, _t_htm_async):.2f} "
                     f"wte={_ms(_t_htm_async, _t_wte):.2f} "
                     f"htm_await={_ms(_t_wte, _t_htm_await):.2f} "
@@ -1010,23 +1010,23 @@ class PostSemClawModel(nn.Module):
                     f"merge={_ms(_t_blocks, _t_merge):.2f} "
                     f"lm_head_loss={_ms(_t_merge, _t_end):.2f} "
                     f"total={_ms(_t0, _t_end):.2f} ms",
-                    flush=True,
-                )
-                for _li, _start, _after_mhc, _engram_start, _engram_end, _end in _profile_layer_events:
-                    print(
-                        f"[PROFILE_LAYER B={B} T={T} layer={_li}] "
-                        f"mhc_block={_ms(_start, _after_mhc):.2f} "
-                        f"layer_total={_ms(_start, _end):.2f} ms",
-                        flush=True,
-                    )
-                    if _engram_start is not None and _engram_end is not None:
-                        print(
-                            f"[PROFILE_ENGRAM B={B} T={T} layer={_li}] "
-                            f"engram={_ms(_engram_start, _engram_end):.2f} "
-                            f"post_layer_total={_ms(_after_mhc, _end):.2f} ms",
-                            flush=True,
-                        )
-            return out
+                    flush=True,
+                )
+                for _li, _start, _after_mhc, _engram_start, _engram_end, _end in _profile_layer_events:
+                    print(
+                        f"[PROFILE_LAYER B={B} T={T} layer={_li}] "
+                        f"mhc_block={_ms(_start, _after_mhc):.2f} "
+                        f"layer_total={_ms(_start, _end):.2f} ms",
+                        flush=True,
+                    )
+                    if _engram_start is not None and _engram_end is not None:
+                        print(
+                            f"[PROFILE_ENGRAM B={B} T={T} layer={_li}] "
+                            f"engram={_ms(_engram_start, _engram_end):.2f} "
+                            f"post_layer_total={_ms(_after_mhc, _end):.2f} ms",
+                            flush=True,
+                        )
+            return out
 
         logits = self.lm_head(x).float()
         if _softcap_clamp:

overlay/hydra/training.py

@@ -4,20 +4,20 @@ Extracted from the monolithic train.py (W1 modularization). Semantics
 preserved. Public entrypoint: `main()`.
 """
 
-from __future__ import annotations
-
-import gc
-import hashlib
-import json
-import math
-import os
-import sys
-import threading
-import time
-from collections.abc import Mapping
-from dataclasses import asdict
-from pathlib import Path
-from typing import Any
+from __future__ import annotations
+
+import gc
+import hashlib
+import json
+import math
+import os
+import sys
+import threading
+import time
+from collections.abc import Mapping
+from dataclasses import asdict
+from pathlib import Path
+from typing import Any
 
 import torch
 
@@ -133,7 +133,7 @@ def _ckpt_snapshot_state_dicts(
     return msd, osd
 
 
-def save_ckpt(
+def save_ckpt(
     model: PostSemClawModel,
     optimizer: torch.optim.Optimizer,
     config: PostSemClawConfig,
@@ -214,233 +214,233 @@ def save_ckpt(
             target=_write, daemon=True, name=f"ckpt-save-{step}"
         )
         _CKPT_WORKER_THREAD.start()
-    except Exception as e:
-        print(f"[ckpt] SNAPSHOT FAILED {path}: {type(e).__name__}: {e}", flush=True)
-
-
-def _env_flag_enabled(env: Mapping[str, str], key: str) -> bool:
-    value = str(env.get(key, "0") or "0").strip().lower()
-    return value not in {"", "0", "false", "no", "off"}
-
-
-def _env_int(env: Mapping[str, str], key: str, default: int) -> int:
-    try:
-        return int(str(env.get(key, str(default)) or str(default)))
-    except ValueError:
-        return default
-
-
-def architecture_compliance_payload(env: Mapping[str, str]) -> dict[str, bool | int | str]:
-    throughput_mode = _env_flag_enabled(env, "HYDRA_THROUGHPUT_MODE")
-    fastpath = _env_flag_enabled(env, "HYDRA_FASTPATH")
-    force_htm_cpu = _env_flag_enabled(env, "HYDRA_FORCE_HTM_CPU")
-    inert_mamba = _env_flag_enabled(env, "HYDRA_INERT_MAMBA")
-    synthetic_retina = _env_flag_enabled(env, "HYDRA_ALLOW_SYNTHETIC_RETINA")
-    hyena_layers = str(env.get("HYDRA_HYENA_LAYERS", "") or "")
-    engram_subsample = _env_int(env, "HYDRA_ENGRAM_SUBSAMPLE", 1)
-    htm_subsample = _env_int(env, "HYDRA_HTM_SUBSAMPLE", 1)
-    full_arch_compliant = not any((
-        throughput_mode,
-        fastpath,
-        force_htm_cpu,
-        inert_mamba,
-        synthetic_retina,
-        bool(hyena_layers.strip()),
-    ))
-    return {
-        'full_arch_compliant': full_arch_compliant,
-        'throughput_mode': throughput_mode,
-        'fastpath': fastpath,
-        'force_htm_cpu': force_htm_cpu,
-        'inert_mamba': inert_mamba,
-        'synthetic_retina': synthetic_retina,
-        'hyena_layers': hyena_layers,
-        'engram_subsample': engram_subsample,
-        'htm_subsample': htm_subsample,
-    }
-
-
-def eval_attempt_batches(*, requested_batch: int, min_batch: int) -> list[int]:
-    requested = max(1, int(requested_batch))
-    minimum = max(1, int(min_batch))
-    batches: list[int] = []
-    current = requested
-    while current >= minimum:
-        if current not in batches:
-            batches.append(current)
-        if current == minimum:
-            break
-        next_batch = max(minimum, current // 2)
-        if next_batch == current:
-            break
-        current = next_batch
-    if minimum not in batches:
-        batches.append(minimum)
-    return batches
-
-
-def build_eval_plan(*, eval_tokens: int, requested_batch: int, max_seq_len: int, chunk_tokens: int, min_batch: int) -> dict[str, Any]:
-    effective_chunk_tokens = max(int(chunk_tokens), int(requested_batch) * int(max_seq_len))
-    chunk_count = max(1, math.ceil(int(eval_tokens) / effective_chunk_tokens))
-    return {
-        'eval_tokens': int(eval_tokens),
-        'eval_requested_batch': int(requested_batch),
-        'eval_chunk_tokens': int(effective_chunk_tokens),
-        'eval_chunk_count': int(chunk_count),
-        'eval_attempt_batches': eval_attempt_batches(requested_batch=requested_batch, min_batch=min_batch),
-        'eval_min_batch': int(max(1, min_batch)),
-    }
-
-
-def _fingerprint_descriptor(descriptor: Mapping[str, Any]) -> str:
-    payload = json.dumps(dict(descriptor), sort_keys=True, separators=(",", ":"))
-    return hashlib.sha1(payload.encode("utf-8")).hexdigest()[:12]
-
-
-def dataset_domain_payload(*, env: Mapping[str, str], prepare_module: Any, nemotron_module: Any | None) -> dict[str, Any]:
-    use_nemotron = _env_flag_enabled(env, "HYDRA_USE_NEMOTRON")
-    vocab_size = int(getattr(prepare_module, "VOCAB_SIZE", 0))
-
-    if use_nemotron and nemotron_module is not None:
-        use_full_blend = _env_flag_enabled(env, "HYDRA_USE_FULL_BLEND")
-        phase = str(env.get("HYDRA_NEMOTRON_PHASE", "phase1") or "phase1").strip().lower()
-        if use_full_blend:
-            train_weights = dict(getattr(nemotron_module, "FULL_BLEND_WEIGHTS", {}))
-            val_weights = dict(train_weights)
-        else:
-            train_weights = dict(
-                getattr(nemotron_module, "PHASE2_WEIGHTS", {}) if phase == "phase2" else getattr(nemotron_module, "PHASE1_WEIGHTS", {})
-            )
-            val_weights = {"Nemotron-Pretraining-Multiple-Choice": 1.0}
-        train_descriptor = {
-            "backend": "nemotron_stream",
-            "phase": "full_blend" if use_full_blend else phase,
-            "weights": train_weights,
-            "factual_inject_rate": _env_int(env, "HYDRA_FACTUAL_INJECT_RATE", 50),
-            "vocab_size": vocab_size,
-        }
-        val_descriptor = {
-            "backend": "nemotron_stream",
-            "phase": "full_blend" if use_full_blend else "val_multiple_choice",
-            "weights": val_weights,
-            "vocab_size": vocab_size,
-        }
-        data_backend = "nemotron_stream"
-    else:
-        all_files = list(getattr(prepare_module, "list_parquet_files", lambda: [])())
-        val_filename = str(getattr(prepare_module, "VAL_FILENAME", ""))
-        train_files = [str(path) for path in all_files if not str(path).endswith(val_filename)]
-        val_files = [str(path) for path in all_files if str(path).endswith(val_filename)]
-        train_descriptor = {
-            "backend": "climbmix_parquet",
-            "train_shard_count": len(train_files),
-            "train_shard_examples": sorted(Path(path).name for path in train_files[:3]),
-            "vocab_size": vocab_size,
-        }
-        val_descriptor = {
-            "backend": "climbmix_parquet",
-            "val_filename": val_filename,
-            "val_shard_count": len(val_files),
-            "vocab_size": vocab_size,
-        }
-        data_backend = "climbmix_parquet"
-
-    train_fingerprint = _fingerprint_descriptor(train_descriptor)
-    val_fingerprint = _fingerprint_descriptor(val_descriptor)
-    return {
-        "data_backend": data_backend,
-        "train_domain_descriptor": train_descriptor,
-        "val_domain_descriptor": val_descriptor,
-        "train_domain_fingerprint": train_fingerprint,
-        "val_domain_fingerprint": val_fingerprint,
-        "train_val_domain_match": train_fingerprint == val_fingerprint,
-    }
-
-
-def build_lineage_payload(
-    *,
-    env: Mapping[str, str],
-    seed: int,
-    resume_requested: bool,
-    resume_requested_path: str | None,
-    resume_loaded_path: str | None,
-    resume_step: int,
-    resume_epoch: int,
-) -> dict[str, Any]:
-    warmstart = _env_flag_enabled(env, "HYDRA_WARMSTART")
-    resume_applied = resume_loaded_path is not None and int(resume_step) > 0
-    if resume_applied and warmstart:
-        lineage_mode = "warmstart_resume"
-    elif resume_applied:
-        lineage_mode = "resume"
-    else:
-        lineage_mode = "fresh"
-    return {
-        "seed": int(seed),
-        "warmstart": warmstart,
-        "resume_requested": bool(resume_requested),
-        "resume_applied": resume_applied,
-        "resume_requested_path": resume_requested_path,
-        "resume_loaded_path": resume_loaded_path,
-        "resume_step": int(resume_step),
-        "resume_epoch": int(resume_epoch),
-        "lineage_mode": lineage_mode,
-    }
-
-
-def build_final_metrics_payload(
-    *,
-    secondary_metrics: dict[str, Any],
-    val_bpb: float | None,
-    val_ppl: float | None,
-    eval_status: str,
-    eval_error: str | None,
-    n_layer: int,
-    d_model: int,
-    num_params: int,
-    step: int,
-    total_tokens: int,
-    peak_vram_mb: float,
-    total_training_time: float,
-    sdr_target_active: int,
-    architecture_env: Mapping[str, str] | None = None,
-    eval_diagnostics: Mapping[str, Any] | None = None,
-    domain_fingerprints: Mapping[str, Any] | None = None,
-    lineage_payload: Mapping[str, Any] | None = None,
-) -> dict[str, Any]:
-    """Build final run metrics without conflating skipped eval and validation.
-
-    This helper deliberately preserves ``val_bpb=None`` when final eval did not
-    complete. HPO can then prune or explicitly label a fallback instead of
-    accidentally treating live training BPB as validation BPB.
-    """
-    payload = dict(secondary_metrics)
-    payload.update({
-        'eval_status': eval_status,
-        'eval_error': eval_error,
-        'objective_source': 'final_val' if val_bpb is not None else 'missing_final_val',
-        'val_bpb': float(val_bpb) if val_bpb is not None else None,
-        'val_ppl': float(val_ppl) if val_ppl is not None else None,
-        'n_layer': int(n_layer),
-        'd_model': int(d_model),
-        'num_params_M': float(num_params / 1e6),
-        'num_steps': int(step),
-        'total_tokens_M': float(total_tokens / 1e6),
-        'peak_vram_mb': float(peak_vram_mb),
-        'training_seconds': float(total_training_time),
-        'sdr_target_active': int(sdr_target_active),
-    })
-    payload.update(architecture_compliance_payload(architecture_env or dict(os.environ)))
-    if eval_diagnostics:
-        payload.update(dict(eval_diagnostics))
-    if domain_fingerprints:
-        payload.update(dict(domain_fingerprints))
-    if lineage_payload:
-        payload.update(dict(lineage_payload))
-    return payload
-
-
-def config_from_dict(cfg_dict: dict) -> PostSemClawConfig:
+    except Exception as e:
+        print(f"[ckpt] SNAPSHOT FAILED {path}: {type(e).__name__}: {e}", flush=True)
+
+
+def _env_flag_enabled(env: Mapping[str, str], key: str) -> bool:
+    value = str(env.get(key, "0") or "0").strip().lower()
+    return value not in {"", "0", "false", "no", "off"}
+
+
+def _env_int(env: Mapping[str, str], key: str, default: int) -> int:
+    try:
+        return int(str(env.get(key, str(default)) or str(default)))
+    except ValueError:
+        return default
+
+
+def architecture_compliance_payload(env: Mapping[str, str]) -> dict[str, bool | int | str]:
+    throughput_mode = _env_flag_enabled(env, "HYDRA_THROUGHPUT_MODE")
+    fastpath = _env_flag_enabled(env, "HYDRA_FASTPATH")
+    force_htm_cpu = _env_flag_enabled(env, "HYDRA_FORCE_HTM_CPU")
+    inert_mamba = _env_flag_enabled(env, "HYDRA_INERT_MAMBA")
+    synthetic_retina = _env_flag_enabled(env, "HYDRA_ALLOW_SYNTHETIC_RETINA")
+    hyena_layers = str(env.get("HYDRA_HYENA_LAYERS", "") or "")
+    engram_subsample = _env_int(env, "HYDRA_ENGRAM_SUBSAMPLE", 1)
+    htm_subsample = _env_int(env, "HYDRA_HTM_SUBSAMPLE", 1)
+    full_arch_compliant = not any((
+        throughput_mode,
+        fastpath,
+        force_htm_cpu,
+        inert_mamba,
+        synthetic_retina,
+        bool(hyena_layers.strip()),
+    ))
+    return {
+        'full_arch_compliant': full_arch_compliant,
+        'throughput_mode': throughput_mode,
+        'fastpath': fastpath,
+        'force_htm_cpu': force_htm_cpu,
+        'inert_mamba': inert_mamba,
+        'synthetic_retina': synthetic_retina,
+        'hyena_layers': hyena_layers,
+        'engram_subsample': engram_subsample,
+        'htm_subsample': htm_subsample,
+    }
+
+
+def eval_attempt_batches(*, requested_batch: int, min_batch: int) -> list[int]:
+    requested = max(1, int(requested_batch))
+    minimum = max(1, int(min_batch))
+    batches: list[int] = []
+    current = requested
+    while current >= minimum:
+        if current not in batches:
+            batches.append(current)
+        if current == minimum:
+            break
+        next_batch = max(minimum, current // 2)
+        if next_batch == current:
+            break
+        current = next_batch
+    if minimum not in batches:
+        batches.append(minimum)
+    return batches
+
+
+def build_eval_plan(*, eval_tokens: int, requested_batch: int, max_seq_len: int, chunk_tokens: int, min_batch: int) -> dict[str, Any]:
+    effective_chunk_tokens = max(int(chunk_tokens), int(requested_batch) * int(max_seq_len))
+    chunk_count = max(1, math.ceil(int(eval_tokens) / effective_chunk_tokens))
+    return {
+        'eval_tokens': int(eval_tokens),
+        'eval_requested_batch': int(requested_batch),
+        'eval_chunk_tokens': int(effective_chunk_tokens),
+        'eval_chunk_count': int(chunk_count),
+        'eval_attempt_batches': eval_attempt_batches(requested_batch=requested_batch, min_batch=min_batch),
+        'eval_min_batch': int(max(1, min_batch)),
+    }
+
+
+def _fingerprint_descriptor(descriptor: Mapping[str, Any]) -> str:
+    payload = json.dumps(dict(descriptor), sort_keys=True, separators=(",", ":"))
+    return hashlib.sha1(payload.encode("utf-8")).hexdigest()[:12]
+
+
+def dataset_domain_payload(*, env: Mapping[str, str], prepare_module: Any, nemotron_module: Any | None) -> dict[str, Any]:
+    use_nemotron = _env_flag_enabled(env, "HYDRA_USE_NEMOTRON")
+    vocab_size = int(getattr(prepare_module, "VOCAB_SIZE", 0))
+
+    if use_nemotron and nemotron_module is not None:
+        use_full_blend = _env_flag_enabled(env, "HYDRA_USE_FULL_BLEND")
+        phase = str(env.get("HYDRA_NEMOTRON_PHASE", "phase1") or "phase1").strip().lower()
+        if use_full_blend:
+            train_weights = dict(getattr(nemotron_module, "FULL_BLEND_WEIGHTS", {}))
+            val_weights = dict(train_weights)
+        else:
+            train_weights = dict(
+                getattr(nemotron_module, "PHASE2_WEIGHTS", {}) if phase == "phase2" else getattr(nemotron_module, "PHASE1_WEIGHTS", {})
+            )
+            val_weights = {"Nemotron-Pretraining-Multiple-Choice": 1.0}
+        train_descriptor = {
+            "backend": "nemotron_stream",
+            "phase": "full_blend" if use_full_blend else phase,
+            "weights": train_weights,
+            "factual_inject_rate": _env_int(env, "HYDRA_FACTUAL_INJECT_RATE", 50),
+            "vocab_size": vocab_size,
+        }
+        val_descriptor = {
+            "backend": "nemotron_stream",
+            "phase": "full_blend" if use_full_blend else "val_multiple_choice",
+            "weights": val_weights,
+            "vocab_size": vocab_size,
+        }
+        data_backend = "nemotron_stream"
+    else:
+        all_files = list(getattr(prepare_module, "list_parquet_files", lambda: [])())
+        val_filename = str(getattr(prepare_module, "VAL_FILENAME", ""))
+        train_files = [str(path) for path in all_files if not str(path).endswith(val_filename)]
+        val_files = [str(path) for path in all_files if str(path).endswith(val_filename)]
+        train_descriptor = {
+            "backend": "climbmix_parquet",
+            "train_shard_count": len(train_files),
+            "train_shard_examples": sorted(Path(path).name for path in train_files[:3]),
+            "vocab_size": vocab_size,
+        }
+        val_descriptor = {
+            "backend": "climbmix_parquet",
+            "val_filename": val_filename,
+            "val_shard_count": len(val_files),
+            "vocab_size": vocab_size,
+        }
+        data_backend = "climbmix_parquet"
+
+    train_fingerprint = _fingerprint_descriptor(train_descriptor)
+    val_fingerprint = _fingerprint_descriptor(val_descriptor)
+    return {
+        "data_backend": data_backend,
+        "train_domain_descriptor": train_descriptor,
+        "val_domain_descriptor": val_descriptor,
+        "train_domain_fingerprint": train_fingerprint,
+        "val_domain_fingerprint": val_fingerprint,
+        "train_val_domain_match": train_fingerprint == val_fingerprint,
+    }
+
+
+def build_lineage_payload(
+    *,
+    env: Mapping[str, str],
+    seed: int,
+    resume_requested: bool,
+    resume_requested_path: str | None,
+    resume_loaded_path: str | None,
+    resume_step: int,
+    resume_epoch: int,
+) -> dict[str, Any]:
+    warmstart = _env_flag_enabled(env, "HYDRA_WARMSTART")
+    resume_applied = resume_loaded_path is not None and int(resume_step) > 0
+    if resume_applied and warmstart:
+        lineage_mode = "warmstart_resume"
+    elif resume_applied:
+        lineage_mode = "resume"
+    else:
+        lineage_mode = "fresh"
+    return {
+        "seed": int(seed),
+        "warmstart": warmstart,
+        "resume_requested": bool(resume_requested),
+        "resume_applied": resume_applied,
+        "resume_requested_path": resume_requested_path,
+        "resume_loaded_path": resume_loaded_path,
+        "resume_step": int(resume_step),
+        "resume_epoch": int(resume_epoch),
+        "lineage_mode": lineage_mode,
+    }
+
+
+def build_final_metrics_payload(
+    *,
+    secondary_metrics: dict[str, Any],
+    val_bpb: float | None,
+    val_ppl: float | None,
+    eval_status: str,
+    eval_error: str | None,
+    n_layer: int,
+    d_model: int,
+    num_params: int,
+    step: int,
+    total_tokens: int,
+    peak_vram_mb: float,
+    total_training_time: float,
+    sdr_target_active: int,
+    architecture_env: Mapping[str, str] | None = None,
+    eval_diagnostics: Mapping[str, Any] | None = None,
+    domain_fingerprints: Mapping[str, Any] | None = None,
+    lineage_payload: Mapping[str, Any] | None = None,
+) -> dict[str, Any]:
+    """Build final run metrics without conflating skipped eval and validation.
+
+    This helper deliberately preserves ``val_bpb=None`` when final eval did not
+    complete. HPO can then prune or explicitly label a fallback instead of
+    accidentally treating live training BPB as validation BPB.
+    """
+    payload = dict(secondary_metrics)
+    payload.update({
+        'eval_status': eval_status,
+        'eval_error': eval_error,
+        'objective_source': 'final_val' if val_bpb is not None else 'missing_final_val',
+        'val_bpb': float(val_bpb) if val_bpb is not None else None,
+        'val_ppl': float(val_ppl) if val_ppl is not None else None,
+        'n_layer': int(n_layer),
+        'd_model': int(d_model),
+        'num_params_M': float(num_params / 1e6),
+        'num_steps': int(step),
+        'total_tokens_M': float(total_tokens / 1e6),
+        'peak_vram_mb': float(peak_vram_mb),
+        'training_seconds': float(total_training_time),
+        'sdr_target_active': int(sdr_target_active),
+    })
+    payload.update(architecture_compliance_payload(architecture_env or dict(os.environ)))
+    if eval_diagnostics:
+        payload.update(dict(eval_diagnostics))
+    if domain_fingerprints:
+        payload.update(dict(domain_fingerprints))
+    if lineage_payload:
+        payload.update(dict(lineage_payload))
+    return payload
+
+
+def config_from_dict(cfg_dict: dict) -> PostSemClawConfig:
     """Reconstruct a PostSemClawConfig from a checkpoint's asdict() payload.
 
     Newly-added fields (e.g. `hyena_layers`) are defaulted when absent in
@@ -500,14 +500,14 @@ def _try_load_ckpt(path: Path, model, optimizer, device):
     return step, total_training_time, smooth_train_loss, bpt_ema, epoch
 
 
-def maybe_resume_ckpt(
-    model: PostSemClawModel,
-    optimizer: torch.optim.Optimizer,
-    device: torch.device,
-) -> tuple[int, float, float, float, int, str | None]:
-    if not RESUME_CKPT or RESUME_CKPT.lower() == "none":
-        print("[ckpt] resume disabled; starting fresh", flush=True)
-        return 0, 0.0, 0.0, 0.0, 0, None
+def maybe_resume_ckpt(
+    model: PostSemClawModel,
+    optimizer: torch.optim.Optimizer,
+    device: torch.device,
+) -> tuple[int, float, float, float, int, str | None]:
+    if not RESUME_CKPT or RESUME_CKPT.lower() == "none":
+        print("[ckpt] resume disabled; starting fresh", flush=True)
+        return 0, 0.0, 0.0, 0.0, 0, None
 
     resume_path = Path(os.path.expanduser(RESUME_CKPT))
     # Try the primary path, then rotated backups. This is crucial because a
@@ -521,18 +521,18 @@ def maybe_resume_ckpt(
         if not cand.exists():
             continue
         try:
-            result = _try_load_ckpt(cand, model, optimizer, device)
-            if result is not None:
-                if cand != resume_path:
-                    print(f"[ckpt] fell back to rotation {cand.name}", flush=True)
-                step, total_training_time, smooth_train_loss, bpt_ema, epoch = result
-                return step, total_training_time, smooth_train_loss, bpt_ema, epoch, str(cand)
+            result = _try_load_ckpt(cand, model, optimizer, device)
+            if result is not None:
+                if cand != resume_path:
+                    print(f"[ckpt] fell back to rotation {cand.name}", flush=True)
+                step, total_training_time, smooth_train_loss, bpt_ema, epoch = result
+                return step, total_training_time, smooth_train_loss, bpt_ema, epoch, str(cand)
         except Exception as e:
             print(f"[ckpt] {cand.name} load failed: {type(e).__name__}: {e}", flush=True)
             continue
 
-    print(f"[ckpt] no usable checkpoint in {resume_path} + rotations; starting fresh", flush=True)
-    return 0, 0.0, 0.0, 0.0, 0, None
+    print(f"[ckpt] no usable checkpoint in {resume_path} + rotations; starting fresh", flush=True)
+    return 0, 0.0, 0.0, 0.0, 0, None
 
 
 # ---------------------------------------------------------------------------
@@ -561,14 +561,14 @@ def main() -> None:
 
     # Streaming path skips prepare.py (which normally trains the tokenizer
     # and builds the retina), so we must materialize both before model init.
-    if os.environ.get("HYDRA_USE_NEMOTRON", "0") == "1":
-        _p_nemo.ensure_tokenizer()
-        if os.environ.get("HYDRA_THROUGHPUT_MODE", "0") != "1":
-            # Retina: HF Hub cache hit for this (vocab, n_bits, target_active) combo
-            # returns in seconds; otherwise build_retina streams Nemotron docs to
-            # compute cooccurrence + train SOM, then uploads back to the cache.
-            import subsystems.sdr_retina as _sdr_retina
-            _sdr_retina.build_retina()
+    if os.environ.get("HYDRA_USE_NEMOTRON", "0") == "1":
+        _p_nemo.ensure_tokenizer()
+        if os.environ.get("HYDRA_THROUGHPUT_MODE", "0") != "1":
+            # Retina: HF Hub cache hit for this (vocab, n_bits, target_active) combo
+            # returns in seconds; otherwise build_retina streams Nemotron docs to
+            # compute cooccurrence + train SOM, then uploads back to the cache.
+            import subsystems.sdr_retina as _sdr_retina
+            _sdr_retina.build_retina()
     tokenizer = Tokenizer.from_directory()
     vocab_size = tokenizer.get_vocab_size()
     print(f"Vocab size: {vocab_size:,}")
@@ -614,18 +614,18 @@ def main() -> None:
         weight_decay=WEIGHT_DECAY,
     )
 
-    step, total_training_time, smooth_train_loss, bpt_ema, resume_epoch, resume_loaded_path = maybe_resume_ckpt(
-        model, optimizer, device,
-    )
-    lineage_payload = build_lineage_payload(
-        env=dict(os.environ),
-        seed=SEED,
-        resume_requested=bool(RESUME_CKPT and RESUME_CKPT.lower() != "none"),
-        resume_requested_path=RESUME_CKPT if RESUME_CKPT and RESUME_CKPT.lower() != "none" else None,
-        resume_loaded_path=resume_loaded_path,
-        resume_step=step,
-        resume_epoch=resume_epoch,
-    )
+    step, total_training_time, smooth_train_loss, bpt_ema, resume_epoch, resume_loaded_path = maybe_resume_ckpt(
+        model, optimizer, device,
+    )
+    lineage_payload = build_lineage_payload(
+        env=dict(os.environ),
+        seed=SEED,
+        resume_requested=bool(RESUME_CKPT and RESUME_CKPT.lower() != "none"),
+        resume_requested_path=RESUME_CKPT if RESUME_CKPT and RESUME_CKPT.lower() != "none" else None,
+        resume_loaded_path=resume_loaded_path,
+        resume_step=step,
+        resume_epoch=resume_epoch,
+    )
 
     # Learnability #4: inform the model of the BOS token id so it can mask
     # doc-separator positions in packed sequences. Always set (the mask only
@@ -1020,22 +1020,22 @@ def main() -> None:
     # does not benefit from overlap with backward). HYDRA_EVAL_TOKENS controls
     # how many val tokens to sweep (default 2 M, short enough for autoresearch
     # 5-min budgets).
-    val_bpb: float | None = None
-    val_ppl: float | None = None
-    eval_status = "not_started"
-    eval_error: str | None = None
-    _eval_B = int(os.environ.get("HYDRA_EVAL_BATCH", str(max(1, DEVICE_BATCH_SIZE // 2))))
-    _eval_tokens = int(os.environ.get("HYDRA_EVAL_TOKENS", str(2 * 524288)))
-    _eval_chunk_tokens = int(os.environ.get("HYDRA_EVAL_CHUNK_TOKENS", str(_eval_tokens)))
-    _eval_min_batch = int(os.environ.get("HYDRA_EVAL_MIN_BATCH", "1"))
-    eval_diagnostics = build_eval_plan(
-        eval_tokens=_eval_tokens,
-        requested_batch=_eval_B,
-        max_seq_len=MAX_SEQ_LEN,
-        chunk_tokens=_eval_chunk_tokens,
-        min_batch=_eval_min_batch,
-    )
-    try:
+    val_bpb: float | None = None
+    val_ppl: float | None = None
+    eval_status = "not_started"
+    eval_error: str | None = None
+    _eval_B = int(os.environ.get("HYDRA_EVAL_BATCH", str(max(1, DEVICE_BATCH_SIZE // 2))))
+    _eval_tokens = int(os.environ.get("HYDRA_EVAL_TOKENS", str(2 * 524288)))
+    _eval_chunk_tokens = int(os.environ.get("HYDRA_EVAL_CHUNK_TOKENS", str(_eval_tokens)))
+    _eval_min_batch = int(os.environ.get("HYDRA_EVAL_MIN_BATCH", "1"))
+    eval_diagnostics = build_eval_plan(
+        eval_tokens=_eval_tokens,
+        requested_batch=_eval_B,
+        max_seq_len=MAX_SEQ_LEN,
+        chunk_tokens=_eval_chunk_tokens,
+        min_batch=_eval_min_batch,
+    )
+    try:
         # Aggressive VRAM reclaim for 6GB cards. Peak training VRAM = 5.1GB
         # which leaves < 1GB for the eval forward — the driver can't satisfy
         # the allocation. Free EVERY tensor we don't strictly need:
@@ -1057,70 +1057,70 @@ def main() -> None:
             model._last_sdr = None
         import gc as _gc
         _gc.collect()
-        torch.cuda.empty_cache()
-        torch.cuda.synchronize()
-        try:
-            _free_mb = torch.cuda.mem_get_info()[0] / 1024 / 1024
-            eval_diagnostics["eval_free_vram_before_mb"] = float(_free_mb)
-            print(f"[VAL] free_vram_mb={_free_mb:.0f} (cleared optimizer state)", flush=True)
-        except Exception:
-            pass
-        print(
-            f"[VAL] running eval on {_eval_tokens} tokens at B={_eval_B} "
-            f"chunk_tokens={eval_diagnostics['eval_chunk_tokens']} attempts={eval_diagnostics['eval_attempt_batches']}...",
-            flush=True,
-        )
-        model.eval()
-        _orig = _prepare_mod.EVAL_TOKENS
-        _orig_chunk = getattr(_prepare_mod, "EVAL_CHUNK_TOKENS", _eval_tokens)
-        _prepare_mod.EVAL_TOKENS = _eval_tokens
-        _prepare_mod.EVAL_CHUNK_TOKENS = int(eval_diagnostics["eval_chunk_tokens"])
-        _successful_batch: int | None = None
-        _attempts: list[int] = []
-        try:
-            for _attempt_batch in eval_diagnostics["eval_attempt_batches"]:
-                _attempts.append(int(_attempt_batch))
-                eval_diagnostics["eval_attempted_batch"] = int(_attempt_batch)
-                try:
-                    with autocast_ctx:
-                        val_bpb = evaluate_bpb(model, tokenizer, int(_attempt_batch))
-                    _successful_batch = int(_attempt_batch)
-                    break
-                except torch.cuda.OutOfMemoryError as _attempt_oom:
-                    eval_error = str(_attempt_oom)
-                    eval_status = "oom"
-                    torch.cuda.empty_cache()
-                    if int(_attempt_batch) == eval_diagnostics["eval_attempt_batches"][-1]:
-                        raise
-        finally:
-            _prepare_mod.EVAL_TOKENS = _orig
-            _prepare_mod.EVAL_CHUNK_TOKENS = _orig_chunk
-        eval_diagnostics["eval_attempt_batches"] = _attempts
-        eval_diagnostics["eval_effective_batch"] = _successful_batch
-        val_ppl = 2 ** val_bpb
-        eval_status = "completed"
-        print(f"[VAL] step={step} val_bpb={val_bpb:.4f} val_ppl={val_ppl:.3f}", flush=True)
-    except torch.cuda.OutOfMemoryError as e:
-        eval_status = "oom"
-        eval_error = str(e)
-        print(f"[VAL] SKIPPED (OOM): {e}", flush=True)
-        torch.cuda.empty_cache()
-        try:
-            eval_diagnostics["eval_free_vram_after_mb"] = float(torch.cuda.mem_get_info()[0] / 1024 / 1024)
-        except Exception:
-            pass
-    except Exception as e:
-        import traceback as _tb
-        eval_status = type(e).__name__
-        eval_error = str(e)
-        print(f"[VAL] SKIPPED ({type(e).__name__}): {e}", flush=True)
-        _tb.print_exc()
-        try:
-            _free = torch.cuda.mem_get_info()[0] / 1024 / 1024
-            eval_diagnostics["eval_free_vram_after_mb"] = float(_free)
-            print(f"[VAL] post-crash free_vram_mb={_free:.0f}", flush=True)
-        except Exception:
-            pass
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+        try:
+            _free_mb = torch.cuda.mem_get_info()[0] / 1024 / 1024
+            eval_diagnostics["eval_free_vram_before_mb"] = float(_free_mb)
+            print(f"[VAL] free_vram_mb={_free_mb:.0f} (cleared optimizer state)", flush=True)
+        except Exception:
+            pass
+        print(
+            f"[VAL] running eval on {_eval_tokens} tokens at B={_eval_B} "
+            f"chunk_tokens={eval_diagnostics['eval_chunk_tokens']} attempts={eval_diagnostics['eval_attempt_batches']}...",
+            flush=True,
+        )
+        model.eval()
+        _orig = _prepare_mod.EVAL_TOKENS
+        _orig_chunk = getattr(_prepare_mod, "EVAL_CHUNK_TOKENS", _eval_tokens)
+        _prepare_mod.EVAL_TOKENS = _eval_tokens
+        _prepare_mod.EVAL_CHUNK_TOKENS = int(eval_diagnostics["eval_chunk_tokens"])
+        _successful_batch: int | None = None
+        _attempts: list[int] = []
+        try:
+            for _attempt_batch in eval_diagnostics["eval_attempt_batches"]:
+                _attempts.append(int(_attempt_batch))
+                eval_diagnostics["eval_attempted_batch"] = int(_attempt_batch)
+                try:
+                    with autocast_ctx:
+                        val_bpb = evaluate_bpb(model, tokenizer, int(_attempt_batch))
+                    _successful_batch = int(_attempt_batch)
+                    break
+                except torch.cuda.OutOfMemoryError as _attempt_oom:
+                    eval_error = str(_attempt_oom)
+                    eval_status = "oom"
+                    torch.cuda.empty_cache()
+                    if int(_attempt_batch) == eval_diagnostics["eval_attempt_batches"][-1]:
+                        raise
+        finally:
+            _prepare_mod.EVAL_TOKENS = _orig
+            _prepare_mod.EVAL_CHUNK_TOKENS = _orig_chunk
+        eval_diagnostics["eval_attempt_batches"] = _attempts
+        eval_diagnostics["eval_effective_batch"] = _successful_batch
+        val_ppl = 2 ** val_bpb
+        eval_status = "completed"
+        print(f"[VAL] step={step} val_bpb={val_bpb:.4f} val_ppl={val_ppl:.3f}", flush=True)
+    except torch.cuda.OutOfMemoryError as e:
+        eval_status = "oom"
+        eval_error = str(e)
+        print(f"[VAL] SKIPPED (OOM): {e}", flush=True)
+        torch.cuda.empty_cache()
+        try:
+            eval_diagnostics["eval_free_vram_after_mb"] = float(torch.cuda.mem_get_info()[0] / 1024 / 1024)
+        except Exception:
+            pass
+    except Exception as e:
+        import traceback as _tb
+        eval_status = type(e).__name__
+        eval_error = str(e)
+        print(f"[VAL] SKIPPED ({type(e).__name__}): {e}", flush=True)
+        _tb.print_exc()
+        try:
+            _free = torch.cuda.mem_get_info()[0] / 1024 / 1024
+            eval_diagnostics["eval_free_vram_after_mb"] = float(_free)
+            print(f"[VAL] post-crash free_vram_mb={_free:.0f}", flush=True)
+        except Exception:
+            pass
 
     # Final ckpts with val_bpb filled in (if eval succeeded).
     save_ckpt(
@@ -1164,13 +1164,13 @@ def main() -> None:
         / total_training_time / GPU_BF16_PEAK_FLOPS
         if total_training_time > 0 else 0
     )
-    peak_vram_mb = torch.cuda.max_memory_allocated() / 1024 / 1024
-    metrics = model.get_secondary_metrics()
-    domain_fingerprints = dataset_domain_payload(
-        env=dict(os.environ),
-        prepare_module=_prepare_mod,
-        nemotron_module=globals().get("_p_nemo"),
-    )
+    peak_vram_mb = torch.cuda.max_memory_allocated() / 1024 / 1024
+    metrics = model.get_secondary_metrics()
+    domain_fingerprints = dataset_domain_payload(
+        env=dict(os.environ),
+        prepare_module=_prepare_mod,
+        nemotron_module=globals().get("_p_nemo"),
+    )
 
     print("---")
     print(f"val_bpb:          {val_bpb:.6f}" if val_bpb is not None else "val_bpb:          SKIPPED")
@@ -1206,28 +1206,28 @@ def main() -> None:
     # Emit full metrics dictionary as JSON for sweep aggregation. Path from
     # HYDRA_METRICS_OUT env var; default=/tmp/hydra_run_metrics.json. Always
     # written (even without diagnostics) so the aggregator can compare runs.
-    _metrics_out = os.environ.get("HYDRA_METRICS_OUT", "/tmp/hydra_run_metrics.json")
-    try:
-        _dump = build_final_metrics_payload(
-            secondary_metrics=metrics,
-            val_bpb=val_bpb,
-            val_ppl=val_ppl,
-            eval_status=eval_status,
-            eval_error=eval_error,
-            n_layer=N_LAYER,
-            d_model=D_MODEL,
-            num_params=num_params,
-            step=step,
-            total_tokens=total_tokens,
-            peak_vram_mb=peak_vram_mb,
-            total_training_time=total_training_time,
-            sdr_target_active=int(os.environ.get("HYDRA_SDR_TARGET_ACTIVE", "327")),
-            architecture_env=dict(os.environ),
-            eval_diagnostics=eval_diagnostics,
-            domain_fingerprints=domain_fingerprints,
-            lineage_payload=lineage_payload,
-        )
-        Path(_metrics_out).parent.mkdir(parents=True, exist_ok=True)
+    _metrics_out = os.environ.get("HYDRA_METRICS_OUT", "/tmp/hydra_run_metrics.json")
+    try:
+        _dump = build_final_metrics_payload(
+            secondary_metrics=metrics,
+            val_bpb=val_bpb,
+            val_ppl=val_ppl,
+            eval_status=eval_status,
+            eval_error=eval_error,
+            n_layer=N_LAYER,
+            d_model=D_MODEL,
+            num_params=num_params,
+            step=step,
+            total_tokens=total_tokens,
+            peak_vram_mb=peak_vram_mb,
+            total_training_time=total_training_time,
+            sdr_target_active=int(os.environ.get("HYDRA_SDR_TARGET_ACTIVE", "327")),
+            architecture_env=dict(os.environ),
+            eval_diagnostics=eval_diagnostics,
+            domain_fingerprints=domain_fingerprints,
+            lineage_payload=lineage_payload,
+        )
+        Path(_metrics_out).parent.mkdir(parents=True, exist_ok=True)
         with open(_metrics_out, 'w') as _f:
             json.dump(_dump, _f, indent=2, sort_keys=True)
         print(f"[METRICS] wrote {_metrics_out}", flush=True)

overlay/prepare.py

@@ -13,10 +13,10 @@ import os
 import sys
 import time
 import math
-import argparse
-import pickle
-from multiprocessing import Pool
-from typing import Any
+import argparse
+import pickle
+from multiprocessing import Pool
+from typing import Any
 
 import requests
 import pyarrow.parquet as pq
@@ -30,8 +30,8 @@ import torch
 
 MAX_SEQ_LEN = int(os.environ.get("HYDRA_SEQ_LEN", "512"))  # context length
 TIME_BUDGET = 300        # training time budget in seconds (5 minutes)
-EVAL_TOKENS = 40 * 524288  # number of tokens for val eval
-EVAL_CHUNK_TOKENS = int(os.environ.get("HYDRA_EVAL_CHUNK_TOKENS", str(EVAL_TOKENS)))
+EVAL_TOKENS = 40 * 524288  # number of tokens for val eval
+EVAL_CHUNK_TOKENS = int(os.environ.get("HYDRA_EVAL_CHUNK_TOKENS", str(EVAL_TOKENS)))
 
 # ---------------------------------------------------------------------------
 # Configuration
@@ -160,8 +160,8 @@ def train_tokenizer():
     print("Tokenizer: training BPE tokenizer...")
     t0 = time.time()
 
-    tokenizer_cls = getattr(rustbpe, "Tokenizer")
-    tokenizer: Any = tokenizer_cls()
+    tokenizer_cls = getattr(rustbpe, "Tokenizer")
+    tokenizer: Any = tokenizer_cls()
     vocab_size_no_special = VOCAB_SIZE - len(SPECIAL_TOKENS)
     tokenizer.train_from_iterator(text_iterator(), vocab_size_no_special, pattern=SPLIT_PATTERN)
 
@@ -228,10 +228,10 @@ class Tokenizer:
     def get_bos_token_id(self):
         return self.bos_token_id
 
-    def encode(self, text, prepend=None, num_threads=8):
-        prepend_id = None
-        if prepend is not None:
-            prepend_id = prepend if isinstance(prepend, int) else self.enc.encode_single_token(prepend)
+    def encode(self, text, prepend=None, num_threads=8):
+        prepend_id = None
+        if prepend is not None:
+            prepend_id = prepend if isinstance(prepend, int) else self.enc.encode_single_token(prepend)
         if isinstance(text, str):
             ids = self.enc.encode_ordinary(text)
             if prepend is not None:
@@ -249,7 +249,7 @@ class Tokenizer:
         return self.enc.decode(ids)
 
 
-_TOKEN_BYTES_CACHE: dict[str, torch.Tensor] = {}
+_TOKEN_BYTES_CACHE: dict[str, torch.Tensor] = {}
 
 def get_token_bytes(device="cpu"):
     key = str(device)
@@ -345,30 +345,30 @@ def make_dataloader(tokenizer, B, T, split, buffer_size=1000):
         gpu_buffer.copy_(cpu_buffer, non_blocking=True)
         yield inputs, targets, epoch
 
-# ---------------------------------------------------------------------------
-# Evaluation (DO NOT CHANGE — this is the fixed metric)
-# ---------------------------------------------------------------------------
-
-def compute_bpb_from_totals(total_nats: torch.Tensor, total_bytes: torch.Tensor) -> torch.Tensor:
-    if int(total_bytes.item()) <= 0:
-        raise ValueError("BPB normalization requires at least one non-special token")
-    return total_nats.to(dtype=torch.float64) / (math.log(2) * total_bytes.to(dtype=torch.float64))
-
-
-def compute_bpb_from_losses(loss_flat: torch.Tensor, nbytes: torch.Tensor) -> torch.Tensor:
-    """Convert per-token losses and token byte lengths into bits-per-byte.
-
-    Tokens with zero byte length (special tokens) are excluded from both the
-    numerator and denominator so BPB remains comparable across tokenizer
-    special-token conventions.
-    """
-    mask = nbytes > 0
-    total_nats = (loss_flat * mask).sum(dtype=torch.float64)
-    total_bytes = nbytes[mask].sum(dtype=torch.int64)
-    return compute_bpb_from_totals(total_nats, total_bytes)
-
-@torch.no_grad()
-def evaluate_bpb(model, tokenizer, batch_size):
+# ---------------------------------------------------------------------------
+# Evaluation (DO NOT CHANGE — this is the fixed metric)
+# ---------------------------------------------------------------------------
+
+def compute_bpb_from_totals(total_nats: torch.Tensor, total_bytes: torch.Tensor) -> torch.Tensor:
+    if int(total_bytes.item()) <= 0:
+        raise ValueError("BPB normalization requires at least one non-special token")
+    return total_nats.to(dtype=torch.float64) / (math.log(2) * total_bytes.to(dtype=torch.float64))
+
+
+def compute_bpb_from_losses(loss_flat: torch.Tensor, nbytes: torch.Tensor) -> torch.Tensor:
+    """Convert per-token losses and token byte lengths into bits-per-byte.
+
+    Tokens with zero byte length (special tokens) are excluded from both the
+    numerator and denominator so BPB remains comparable across tokenizer
+    special-token conventions.
+    """
+    mask = nbytes > 0
+    total_nats = (loss_flat * mask).sum(dtype=torch.float64)
+    total_bytes = nbytes[mask].sum(dtype=torch.int64)
+    return compute_bpb_from_totals(total_nats, total_bytes)
+
+@torch.no_grad()
+def evaluate_bpb(model, tokenizer, batch_size):
     """
     Bits per byte (BPB): vocab size-independent evaluation metric.
     Sums per-token cross-entropy (in nats), sums target byte lengths,
@@ -379,35 +379,35 @@ def evaluate_bpb(model, tokenizer, batch_size):
     Perf: accumulates on GPU (single sync at end), prefetches next batch
     while current forward runs.
     """
-    token_bytes = get_token_bytes(device="cuda")
-    val_loader = make_dataloader(tokenizer, batch_size, MAX_SEQ_LEN, "val")
-    steps = EVAL_TOKENS // (batch_size * MAX_SEQ_LEN)
-    chunk_steps = max(1, EVAL_CHUNK_TOKENS // (batch_size * MAX_SEQ_LEN))
+    token_bytes = get_token_bytes(device="cuda")
+    val_loader = make_dataloader(tokenizer, batch_size, MAX_SEQ_LEN, "val")
+    steps = EVAL_TOKENS // (batch_size * MAX_SEQ_LEN)
+    chunk_steps = max(1, EVAL_CHUNK_TOKENS // (batch_size * MAX_SEQ_LEN))
 
     # GPU-resident accumulators — avoid per-batch .item() sync
     total_nats_t = torch.zeros(1, device="cuda", dtype=torch.float64)
     total_bytes_t = torch.zeros(1, device="cuda", dtype=torch.int64)
 
     # Prefetch first batch
-    next_batch = next(val_loader)
-    steps_done = 0
-    while steps_done < steps:
-        this_chunk = min(chunk_steps, steps - steps_done)
-        for _ in range(this_chunk):
-            x, y, _epoch = next_batch
-            # Prefetch NEXT batch while GPU computes current forward
-            next_batch = next(val_loader)
-            loss_flat = model(x, y, reduction='none').view(-1)
-            y_flat = y.view(-1)
-            nbytes = token_bytes[y_flat]
-            total_nats_t += (loss_flat * (nbytes > 0)).sum(dtype=torch.float64)
-            total_bytes_t += nbytes[nbytes > 0].sum(dtype=torch.int64)
-        steps_done += this_chunk
-        if steps_done < steps:
-            torch.cuda.empty_cache()
-
-    # Single GPU→CPU sync at end
-    return float(compute_bpb_from_totals(total_nats_t, total_bytes_t).item())
+    next_batch = next(val_loader)
+    steps_done = 0
+    while steps_done < steps:
+        this_chunk = min(chunk_steps, steps - steps_done)
+        for _ in range(this_chunk):
+            x, y, _epoch = next_batch
+            # Prefetch NEXT batch while GPU computes current forward
+            next_batch = next(val_loader)
+            loss_flat = model(x, y, reduction='none').view(-1)
+            y_flat = y.view(-1)
+            nbytes = token_bytes[y_flat]
+            total_nats_t += (loss_flat * (nbytes > 0)).sum(dtype=torch.float64)
+            total_bytes_t += nbytes[nbytes > 0].sum(dtype=torch.int64)
+        steps_done += this_chunk
+        if steps_done < steps:
+            torch.cuda.empty_cache()
+
+    # Single GPU→CPU sync at end
+    return float(compute_bpb_from_totals(total_nats_t, total_bytes_t).item())
 
 # ---------------------------------------------------------------------------
 # Main

overlay/prepare_nemotron.py

@@ -20,16 +20,15 @@ Full blend mode (env HYDRA_USE_FULL_BLEND=1):
 """
 from __future__ import annotations
 
-import os
-import random
-import importlib
-import shutil
-from itertools import cycle
-from typing import Any, Iterator, cast
+import os
+import random
+import importlib
+from itertools import cycle
+from typing import Any, Iterator, cast
 
 import torch
 
-import prepare as _p  # reuse tokenizer, BOS, byte-length helpers
+import prepare as _p  # reuse tokenizer, BOS, byte-length helpers
 
 NEMOTRON_REPO = "nvidia/Nemotron-Pretraining-Specialized-v1.1"
 
@@ -66,96 +65,94 @@ PHASE1_WEIGHTS = {
     "Nemotron-Pretraining-Formal-Logic":                0.20,
     "Nemotron-Pretraining-Multiple-Choice":             0.20,
 }
-PHASE2_WEIGHTS = {
+PHASE2_WEIGHTS = {
     "Nemotron-Pretraining-Multiple-Choice":             0.45,
     "Nemotron-Pretraining-Economics":                   0.20,
     "Nemotron-Pretraining-Formal-Logic":                0.15,
     "Nemotron-Pretraining-Code-Concepts":               0.10,
     "Nemotron-Pretraining-Unconditional-Algorithmic":   0.10,
-}
-
-StreamBatch = tuple[list[str], int]
-TokenBatch = tuple[list[list[int]], int]
-
-
-def _tokenizer_cache_repo() -> str:
-    return (
-        os.environ.get("HYDRA_TOKENIZER_CACHE_REPO")
-        or os.environ.get("FEATHER_HF_OUTPUT_REPO")
-        or os.environ.get("HF_REPO_ID")
-        or os.environ.get("HYDRA_RETINA_CACHE_REPO")
-        or os.environ.get("FEATHER_HF_RETINA_CACHE_REPO")
-        or ""
-    )
-
-
-def _tokenizer_cache_prefix() -> str:
-    return f"tokenizer/vocab{_p.VOCAB_SIZE}"
-
-
-def maybe_hydrate_tokenizer_cache() -> bool:
-    """Try to download tokenizer artifacts from HF cache storage."""
-    repo_id = _tokenizer_cache_repo()
-    token = os.environ.get("HF_TOKEN")
-    if not repo_id or not token:
-        return False
-
-    try:
-        from huggingface_hub import hf_hub_download
-    except Exception as e:  # noqa: BLE001
-        print(f"[nemotron] tokenizer cache unavailable: {type(e).__name__}: {e}", flush=True)
-        return False
-
-    os.makedirs(_p.TOKENIZER_DIR, exist_ok=True)
-    prefix = _tokenizer_cache_prefix()
-    try:
-        tok_src = hf_hub_download(
-            repo_id=repo_id,
-            repo_type="model",
-            subfolder=prefix,
-            filename="tokenizer.pkl",
-            token=token,
-            local_dir=_p.TOKENIZER_DIR,
-        )
-        token_bytes_src = hf_hub_download(
-            repo_id=repo_id,
-            repo_type="model",
-            subfolder=prefix,
-            filename="token_bytes.pt",
-            token=token,
-            local_dir=_p.TOKENIZER_DIR,
-        )
-        shutil.copy2(tok_src, os.path.join(_p.TOKENIZER_DIR, "tokenizer.pkl"))
-        shutil.copy2(token_bytes_src, os.path.join(_p.TOKENIZER_DIR, "token_bytes.pt"))
-    except Exception as e:  # noqa: BLE001
-        print(f"[nemotron] tokenizer cache miss in {repo_id}/{prefix}: {type(e).__name__}: {e}", flush=True)
-        return False
-
-    print(f"[nemotron] hydrated tokenizer cache from {repo_id}/{prefix}", flush=True)
-    return True
-
-
-def upload_tokenizer_cache() -> None:
-    """Upload tokenizer artifacts for reuse by future jobs."""
-    repo_id = _tokenizer_cache_repo()
-    token = os.environ.get("HF_TOKEN")
-    if not repo_id or not token:
-        return
-
-    path = os.path.join(_p.TOKENIZER_DIR, "tokenizer.pkl")
-    token_bytes_path = os.path.join(_p.TOKENIZER_DIR, "token_bytes.pt")
-    if not (os.path.exists(path) and os.path.exists(token_bytes_path)):
-        return
-
-    try:
-        from huggingface_hub import HfApi
-        api = HfApi(token=token)
-        prefix = _tokenizer_cache_prefix()
-        api.upload_file(path_or_fileobj=path, path_in_repo=f"{prefix}/tokenizer.pkl", repo_id=repo_id, repo_type="model")
-        api.upload_file(path_or_fileobj=token_bytes_path, path_in_repo=f"{prefix}/token_bytes.pt", repo_id=repo_id, repo_type="model")
-        print(f"[nemotron] uploaded tokenizer cache to {repo_id}/{prefix}", flush=True)
-    except Exception as e:  # noqa: BLE001
-        print(f"[nemotron] tokenizer cache upload skipped: {type(e).__name__}: {e}", flush=True)
+}
+
+type StreamBatch = tuple[list[str], int]
+type TokenBatch = tuple[list[list[int]], int]
+
+
+def _tokenizer_cache_repo() -> str:
+    return (
+        os.environ.get("HYDRA_TOKENIZER_CACHE_REPO")
+        or os.environ.get("FEATHER_HF_OUTPUT_REPO")
+        or os.environ.get("HF_REPO_ID")
+        or os.environ.get("HYDRA_RETINA_CACHE_REPO")
+        or os.environ.get("FEATHER_HF_RETINA_CACHE_REPO")
+        or ""
+    )
+
+
+def _tokenizer_cache_prefix() -> str:
+    return f"tokenizer/vocab{_p.VOCAB_SIZE}"
+
+
+def maybe_hydrate_tokenizer_cache() -> bool:
+    """Try to download tokenizer artifacts from HF cache storage."""
+    repo_id = _tokenizer_cache_repo()
+    token = os.environ.get("HF_TOKEN")
+    if not repo_id or not token:
+        return False
+
+    try:
+        from huggingface_hub import hf_hub_download
+    except Exception as e:  # noqa: BLE001
+        print(f"[nemotron] tokenizer cache unavailable: {type(e).__name__}: {e}", flush=True)
+        return False
+
+    os.makedirs(_p.TOKENIZER_DIR, exist_ok=True)
+    prefix = _tokenizer_cache_prefix()
+    try:
+        hf_hub_download(
+            repo_id=repo_id,
+            repo_type="model",
+            subfolder=prefix,
+            filename="tokenizer.pkl",
+            token=token,
+            local_dir=_p.TOKENIZER_DIR,
+        )
+        hf_hub_download(
+            repo_id=repo_id,
+            repo_type="model",
+            subfolder=prefix,
+            filename="token_bytes.pt",
+            token=token,
+            local_dir=_p.TOKENIZER_DIR,
+        )
+    except Exception as e:  # noqa: BLE001
+        print(f"[nemotron] tokenizer cache miss in {repo_id}/{prefix}: {type(e).__name__}: {e}", flush=True)
+        return False
+
+    print(f"[nemotron] hydrated tokenizer cache from {repo_id}/{prefix}", flush=True)
+    return True
+
+
+def upload_tokenizer_cache() -> None:
+    """Upload tokenizer artifacts for reuse by future jobs."""
+    repo_id = _tokenizer_cache_repo()
+    token = os.environ.get("HF_TOKEN")
+    if not repo_id or not token:
+        return
+
+    path = os.path.join(_p.TOKENIZER_DIR, "tokenizer.pkl")
+    token_bytes_path = os.path.join(_p.TOKENIZER_DIR, "token_bytes.pt")
+    if not (os.path.exists(path) and os.path.exists(token_bytes_path)):
+        return
+
+    try:
+        from huggingface_hub import HfApi
+        api = HfApi(token=token)
+        prefix = _tokenizer_cache_prefix()
+        api.upload_file(path_or_fileobj=path, path_in_repo=f"{prefix}/tokenizer.pkl", repo_id=repo_id, repo_type="model")
+        api.upload_file(path_or_fileobj=token_bytes_path, path_in_repo=f"{prefix}/token_bytes.pt", repo_id=repo_id, repo_type="model")
+        print(f"[nemotron] uploaded tokenizer cache to {repo_id}/{prefix}", flush=True)
+    except Exception as e:  # noqa: BLE001
+        print(f"[nemotron] tokenizer cache upload skipped: {type(e).__name__}: {e}", flush=True)
 
 
 def _phase_weights() -> dict[str, float]:
@@ -166,7 +163,7 @@ def _phase_weights() -> dict[str, float]:
     return PHASE2_WEIGHTS if phase == "phase2" else PHASE1_WEIGHTS
 
 
-def _open_stream(config: str, split: str):
+def _open_stream(config: str, split: str):
     """Open a streaming iterator over one dataset config.
 
     Handles two modes:
@@ -177,17 +174,17 @@ def _open_stream(config: str, split: str):
 
     Yields dicts; text extraction handled downstream by _extract_text.
     """
-    load_dataset = importlib.import_module("datasets").load_dataset
-    token = os.environ.get("HF_TOKEN")
-    shuffle_buf = int(os.environ.get("HYDRA_STREAM_SHUFFLE_BUFFER", "2048"))
-
-    if config in _BLEND_REGISTRY:
-        repo, name, _text_col = _BLEND_REGISTRY[config]
-        kwargs: dict[str, object] = dict(
-            split="train",
-            streaming=True,
-            token=token,
-        )
+    load_dataset = importlib.import_module("datasets").load_dataset
+    token = os.environ.get("HF_TOKEN")
+    shuffle_buf = int(os.environ.get("HYDRA_STREAM_SHUFFLE_BUFFER", "2048"))
+
+    if config in _BLEND_REGISTRY:
+        repo, name, _text_col = _BLEND_REGISTRY[config]
+        kwargs: dict[str, object] = dict(
+            split="train",
+            streaming=True,
+            token=token,
+        )
         if name is not None:
             kwargs["name"] = name
         # nemotron-specialized has multiple sub-configs; pick the first one
@@ -209,18 +206,18 @@ def _open_stream(config: str, split: str):
     return iter(ds)
 
 
-def _extract_text(row: dict[str, object]) -> str:
+def _extract_text(row: dict[str, object]) -> str:
     """Pick the right text column — datasets have different column names.
 
     Priority order: text, content, prompt_completion, question, body.
     For math datasets that split into problem+solution, concatenate both.
     Fallback: concatenate all string-valued fields.
     """
-    # Fast path: most datasets use "text" or "content".
-    for k in ("text", "content", "prompt_completion", "question", "body"):
-        value = row.get(k)
-        if isinstance(value, str) and value:
-            return value
+    # Fast path: most datasets use "text" or "content".
+    for k in ("text", "content", "prompt_completion", "question", "body"):
+        value = row.get(k)
+        if isinstance(value, str) and value:
+            return value
     # Math datasets may have problem + solution as separate fields.
     if "problem" in row and "solution" in row:
         p = row["problem"] or ""
@@ -236,20 +233,20 @@ def _extract_text(row: dict[str, object]) -> str:
     return "\n".join(parts)
 
 
-class _WeightedStream:
+class _WeightedStream:
     """Infinite weighted-round-robin over configs' streaming iterators."""
 
-    def __init__(self, weights: dict[str, float], seed: int = 0):
-        self.configs = list(weights.keys())
-        self.weights = [weights[c] for c in self.configs]
-        self.streams: dict[str, Iterator[dict[str, object]]] = {
-            c: _open_stream(c, "train") for c in self.configs
-        }
-        self.rng = random.Random(seed)
-        self.epoch = 1
-        self._factual_docs: list[str] | None = None
-        self._factual_idx = 0
-        self._inject_counter = 0
+    def __init__(self, weights: dict[str, float], seed: int = 0):
+        self.configs = list(weights.keys())
+        self.weights = [weights[c] for c in self.configs]
+        self.streams: dict[str, Iterator[dict[str, object]]] = {
+            c: _open_stream(c, "train") for c in self.configs
+        }
+        self.rng = random.Random(seed)
+        self.epoch = 1
+        self._factual_docs: list[str] | None = None
+        self._factual_idx = 0
+        self._inject_counter = 0
 
     def _reopen(self, config: str):
         # stream exhausted — reopen (HF streaming typically infinite but restart on edge)
@@ -265,20 +262,20 @@ class _WeightedStream:
         # exist in the Nemotron configs. Controlled by HYDRA_FACTUAL_INJECT_RATE
         # (default 50 = inject one factual doc every 50 Nemotron docs = ~2%).
         inject_rate = int(os.environ.get("HYDRA_FACTUAL_INJECT_RATE", "50"))
-        if inject_rate > 0 and self._factual_docs is None:
-            factual_path = os.path.join(
-                os.path.dirname(os.path.abspath(__file__)), "data", "factual", "facts.txt")
-            if os.path.exists(factual_path):
-                self._factual_docs = open(factual_path).read().strip().split('\n')
-                self._factual_idx = 0
-                self._inject_counter = 0
-        if inject_rate > 0 and self._factual_docs:
-            self._inject_counter += 1
-            if self._inject_counter >= inject_rate:
-                self._inject_counter = 0
-                doc = self._factual_docs[self._factual_idx % len(self._factual_docs)]
-                self._factual_idx += 1
-                return doc, self.epoch
+        if inject_rate > 0 and self._factual_docs is None:
+            factual_path = os.path.join(
+                os.path.dirname(os.path.abspath(__file__)), "data", "factual", "facts.txt")
+            if os.path.exists(factual_path):
+                self._factual_docs = open(factual_path).read().strip().split('\n')
+                self._factual_idx = 0
+                self._inject_counter = 0
+        if inject_rate > 0 and self._factual_docs:
+            self._inject_counter += 1
+            if self._inject_counter >= inject_rate:
+                self._inject_counter = 0
+                doc = self._factual_docs[self._factual_idx % len(self._factual_docs)]
+                self._factual_idx += 1
+                return doc, self.epoch
 
         config = self.rng.choices(self.configs, weights=self.weights, k=1)[0]
         try:
@@ -311,9 +308,9 @@ def _document_batches(split: str, tokenizer_batch_size: int = 128) -> Iterator[t
         stream = _WeightedStream(_phase_weights(), seed=0)
 
     prefetch_depth = int(os.environ.get("HYDRA_STREAM_PREFETCH", "32"))
-    q: queue.Queue[StreamBatch | object] = queue.Queue(maxsize=prefetch_depth)
-    sentinel_stop = object()
-    error_box: list[BaseException] = []
+    q: queue.Queue[StreamBatch | object] = queue.Queue(maxsize=prefetch_depth)
+    sentinel_stop = object()
+    error_box: list[BaseException] = []
 
     def producer():
         try:
@@ -338,7 +335,7 @@ def _document_batches(split: str, tokenizer_batch_size: int = 128) -> Iterator[t
             if error_box:
                 raise error_box[0]
             return
-        yield cast(StreamBatch, item)
+        yield cast(StreamBatch, item)
 
 
 def make_dataloader(tokenizer, B: int, T: int, split: str, buffer_size: int = 1000):
@@ -364,9 +361,9 @@ def make_dataloader(tokenizer, B: int, T: int, split: str, buffer_size: int = 10
     # Stage 2: tokenization prefetch thread. Each queue element is a list of
     # token-id lists (pre-tokenized docs). HYDRA_TOKEN_PREFETCH controls depth.
     tok_prefetch = int(os.environ.get("HYDRA_TOKEN_PREFETCH", "8"))
-    tok_q: queue.Queue[TokenBatch | object] = queue.Queue(maxsize=tok_prefetch)
-    tok_sentinel = object()
-    tok_err_box: list[BaseException] = []
+    tok_q: queue.Queue[TokenBatch | object] = queue.Queue(maxsize=tok_prefetch)
+    tok_sentinel = object()
+    tok_err_box: list[BaseException] = []
 
     def tokenizer_producer():
         try:
@@ -390,8 +387,8 @@ def make_dataloader(tokenizer, B: int, T: int, split: str, buffer_size: int = 10
             if tok_err_box:
                 raise tok_err_box[0]
             raise StopIteration
-        token_lists, epoch = cast(TokenBatch, item)
-        doc_buffer.extend(token_lists)
+        token_lists, epoch = cast(TokenBatch, item)
+        doc_buffer.extend(token_lists)
 
     row_buffer = torch.empty((B, row_capacity), dtype=torch.long)
     cpu_buffer = torch.empty(2 * B * T, dtype=torch.long, pin_memory=True)
@@ -465,24 +462,24 @@ def evaluate_bpb(model, tokenizer, B: int) -> float:
     return total_nats / (math.log(2) * max(total_bytes, 1))
 
 
-def ensure_tokenizer():
+def ensure_tokenizer():
     """Ensure rustbpe tokenizer exists. If absent, train on a Nemotron stream
     sample using the same rustbpe.train_from_iterator API that prepare.py uses
     (production path — don't fork tokenizer training logic).
     """
     import pickle
     import torch
-    path = os.path.join(_p.TOKENIZER_DIR, "tokenizer.pkl")
-    token_bytes_path = os.path.join(_p.TOKENIZER_DIR, "token_bytes.pt")
-    if os.path.exists(path) and os.path.exists(token_bytes_path):
-        print(f"[nemotron] tokenizer + token_bytes already trained at {_p.TOKENIZER_DIR}", flush=True)
-        return
-    if maybe_hydrate_tokenizer_cache() and os.path.exists(path) and os.path.exists(token_bytes_path):
-        return
-    os.makedirs(_p.TOKENIZER_DIR, exist_ok=True)
+    path = os.path.join(_p.TOKENIZER_DIR, "tokenizer.pkl")
+    token_bytes_path = os.path.join(_p.TOKENIZER_DIR, "token_bytes.pt")
+    if os.path.exists(path) and os.path.exists(token_bytes_path):
+        print(f"[nemotron] tokenizer + token_bytes already trained at {_p.TOKENIZER_DIR}", flush=True)
+        return
+    if maybe_hydrate_tokenizer_cache() and os.path.exists(path) and os.path.exists(token_bytes_path):
+        return
+    os.makedirs(_p.TOKENIZER_DIR, exist_ok=True)
     print(f"[nemotron] training BPE (vocab_size={_p.VOCAB_SIZE}) on stream sample…", flush=True)
-    import rustbpe
-    import tiktoken
+    import rustbpe
+    import tiktoken
 
     # Pull a sample of docs — use full blend if active so BPE covers all 7 sources.
     n_docs = int(os.environ.get("HYDRA_BPE_TRAIN_DOCS", "20000"))
@@ -498,8 +495,8 @@ def ensure_tokenizer():
     print(f"[nemotron] collected {len(sample_texts)} sample docs; training BPE…", flush=True)
 
     # Train rustbpe — identical API to prepare.py's train_tokenizer().
-    tokenizer_cls = getattr(rustbpe, "Tokenizer")
-    tokenizer: Any = tokenizer_cls()
+    tokenizer_cls = getattr(rustbpe, "Tokenizer")
+    tokenizer: Any = tokenizer_cls()
     vocab_size_no_special = _p.VOCAB_SIZE - len(_p.SPECIAL_TOKENS)
     tokenizer.train_from_iterator(iter(sample_texts), vocab_size_no_special, pattern=_p.SPLIT_PATTERN)
 
@@ -524,7 +521,7 @@ def ensure_tokenizer():
     for token_id in range(enc.n_vocab):
         tstr = enc.decode([token_id])
         token_bytes_list.append(0 if tstr in special_set else len(tstr.encode("utf-8")))
-    token_bytes_tensor = torch.tensor(token_bytes_list, dtype=torch.int32)
-    torch.save(token_bytes_tensor, token_bytes_path)
-    print(f"[nemotron] BPE + token_bytes saved to {_p.TOKENIZER_DIR}", flush=True)
-    upload_tokenizer_cache()
+    token_bytes_tensor = torch.tensor(token_bytes_list, dtype=torch.int32)
+    torch.save(token_bytes_tensor, token_bytes_path)
+    print(f"[nemotron] BPE + token_bytes saved to {_p.TOKENIZER_DIR}", flush=True)
+    upload_tokenizer_cache()

overlay/scripts/__init__.py

@@ -1 +1 @@
-# Package marker for script-level shared utilities.
+# Package marker for script-level shared utilities.

overlay/scripts/audit_overlay_sync.py

@@ -1,100 +1,100 @@
-#!/usr/bin/env python3
-from __future__ import annotations
-
-import argparse
-import json
-from pathlib import Path
-
-
-DEFAULT_INCLUDE_PATHS = [
-    "hydra",
-    "subsystems",
-    "scripts",
-    "htm_rust",
-    "harness",
-    "configs",
-    "prepare.py",
-    "prepare_nemotron.py",
-    "train.py",
-    "pyproject.toml",
-    "uv.lock",
-]
-
-
-def _iter_files(path: Path) -> list[Path]:
-    if not path.exists():
-        return []
-    if path.is_file():
-        return [path]
-    return sorted(p for p in path.rglob("*") if p.is_file())
-
-
-def classify_overlay_pairs(*, repo_root: Path, include_paths: list[str]) -> dict[str, list[str]]:
-    overlay_root = repo_root / "hf_jobs" / "feather_h200_image" / "overlay"
-    identical: list[str] = []
-    root_ahead: list[str] = []
-    overlay_only: list[str] = []
-    missing_overlay: list[str] = []
-
-    for rel in include_paths:
-        root_path = repo_root / rel
-        overlay_path = overlay_root / rel
-
-        root_files = {p.relative_to(root_path).as_posix(): p for p in _iter_files(root_path)} if root_path.exists() and root_path.is_dir() else {}
-        overlay_files = {p.relative_to(overlay_path).as_posix(): p for p in _iter_files(overlay_path)} if overlay_path.exists() and overlay_path.is_dir() else {}
-
-        if root_path.is_file() or overlay_path.is_file():
-            rel_name = rel.replace("\\", "/")
-            if root_path.exists() and overlay_path.exists():
-                if root_path.read_bytes() == overlay_path.read_bytes():
-                    identical.append(rel_name)
-                else:
-                    root_ahead.append(rel_name)
-            elif root_path.exists():
-                missing_overlay.append(rel_name)
-            elif overlay_path.exists():
-                overlay_only.append(rel_name)
-            continue
-
-        for subrel, root_file in root_files.items():
-            rel_name = f"{rel}/{subrel}".replace("\\", "/")
-            overlay_file = overlay_files.get(subrel)
-            if overlay_file is None:
-                missing_overlay.append(rel_name)
-            elif root_file.read_bytes() == overlay_file.read_bytes():
-                identical.append(rel_name)
-            else:
-                root_ahead.append(rel_name)
-
-        for subrel in overlay_files:
-            if subrel not in root_files:
-                overlay_only.append(f"{rel}/{subrel}".replace("\\", "/"))
-
-    for bucket in (identical, root_ahead, overlay_only, missing_overlay):
-        bucket.sort()
-
-    return {
-        "identical": identical,
-        "root_ahead": root_ahead,
-        "overlay_only": overlay_only,
-        "missing_overlay": missing_overlay,
-    }
-
-
-def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Audit mirrored H200 overlay files against root source-of-truth paths")
-    parser.add_argument("--repo-root", type=Path, default=Path(__file__).resolve().parents[1])
-    parser.add_argument("--include-path", action="append", default=[])
-    return parser.parse_args(argv)
-
-
-def main(argv: list[str] | None = None) -> int:
-    args = parse_args(argv)
-    include_paths = args.include_path or DEFAULT_INCLUDE_PATHS
-    payload = classify_overlay_pairs(repo_root=args.repo_root, include_paths=include_paths)
-    print(json.dumps(payload, indent=2, sort_keys=True))
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+
+DEFAULT_INCLUDE_PATHS = [
+    "hydra",
+    "subsystems",
+    "scripts",
+    "htm_rust",
+    "harness",
+    "configs",
+    "prepare.py",
+    "prepare_nemotron.py",
+    "train.py",
+    "pyproject.toml",
+    "uv.lock",
+]
+
+
+def _iter_files(path: Path) -> list[Path]:
+    if not path.exists():
+        return []
+    if path.is_file():
+        return [path]
+    return sorted(p for p in path.rglob("*") if p.is_file())
+
+
+def classify_overlay_pairs(*, repo_root: Path, include_paths: list[str]) -> dict[str, list[str]]:
+    overlay_root = repo_root / "hf_jobs" / "feather_h200_image" / "overlay"
+    identical: list[str] = []
+    root_ahead: list[str] = []
+    overlay_only: list[str] = []
+    missing_overlay: list[str] = []
+
+    for rel in include_paths:
+        root_path = repo_root / rel
+        overlay_path = overlay_root / rel
+
+        root_files = {p.relative_to(root_path).as_posix(): p for p in _iter_files(root_path)} if root_path.exists() and root_path.is_dir() else {}
+        overlay_files = {p.relative_to(overlay_path).as_posix(): p for p in _iter_files(overlay_path)} if overlay_path.exists() and overlay_path.is_dir() else {}
+
+        if root_path.is_file() or overlay_path.is_file():
+            rel_name = rel.replace("\\", "/")
+            if root_path.exists() and overlay_path.exists():
+                if root_path.read_bytes() == overlay_path.read_bytes():
+                    identical.append(rel_name)
+                else:
+                    root_ahead.append(rel_name)
+            elif root_path.exists():
+                missing_overlay.append(rel_name)
+            elif overlay_path.exists():
+                overlay_only.append(rel_name)
+            continue
+
+        for subrel, root_file in root_files.items():
+            rel_name = f"{rel}/{subrel}".replace("\\", "/")
+            overlay_file = overlay_files.get(subrel)
+            if overlay_file is None:
+                missing_overlay.append(rel_name)
+            elif root_file.read_bytes() == overlay_file.read_bytes():
+                identical.append(rel_name)
+            else:
+                root_ahead.append(rel_name)
+
+        for subrel in overlay_files:
+            if subrel not in root_files:
+                overlay_only.append(f"{rel}/{subrel}".replace("\\", "/"))
+
+    for bucket in (identical, root_ahead, overlay_only, missing_overlay):
+        bucket.sort()
+
+    return {
+        "identical": identical,
+        "root_ahead": root_ahead,
+        "overlay_only": overlay_only,
+        "missing_overlay": missing_overlay,
+    }
+
+
+def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Audit mirrored H200 overlay files against root source-of-truth paths")
+    parser.add_argument("--repo-root", type=Path, default=Path(__file__).resolve().parents[1])
+    parser.add_argument("--include-path", action="append", default=[])
+    return parser.parse_args(argv)
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = parse_args(argv)
+    include_paths = args.include_path or DEFAULT_INCLUDE_PATHS
+    payload = classify_overlay_pairs(repo_root=args.repo_root, include_paths=include_paths)
+    print(json.dumps(payload, indent=2, sort_keys=True))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

overlay/scripts/benchmark_assets.py

@@ -1,124 +1,62 @@
-#!/usr/bin/env python3
-from __future__ import annotations
-
-import os
-import shutil
-from pathlib import Path
-
-from scripts.benchmark_checkpoint import checkpoint_candidates
-
-try:
-    from huggingface_hub import HfApi
-except Exception:  # pragma: no cover - optional import for offline test envs
-    HfApi = None
-
-
-def _download_file(*, repo_id: str, filename: str, local_dir: str, token: str | None, subfolder: str | None = None) -> Path:
-    from huggingface_hub import hf_hub_download
-
-    path = hf_hub_download(
-        repo_id=repo_id,
-        repo_type="model",
-        filename=filename,
-        subfolder=subfolder,
-        token=token,
-        local_dir=local_dir,
-        local_dir_use_symlinks=False,
-    )
-    return Path(path)
-
-
-def resolve_tokenizer_cache_repo(*, output_repo: str, retina_cache_repo: str) -> str:
-    return (
-        os.environ.get("HYDRA_TOKENIZER_CACHE_REPO")
-        or os.environ.get("FEATHER_HF_OUTPUT_REPO")
-        or os.environ.get("HF_REPO_ID")
-        or os.environ.get("HYDRA_RETINA_CACHE_REPO")
-        or os.environ.get("FEATHER_HF_RETINA_CACHE_REPO")
-        or output_repo
-        or retina_cache_repo
-    )
-
-
-def tokenizer_cache_prefix() -> str:
-    vocab_size = int(os.environ.get("HYDRA_VOCAB_SIZE", "65536"))
-    return f"tokenizer/vocab{vocab_size}"
-
-
-def choose_remote_checkpoint_path(files: list[str]) -> str | None:
-    preferred = [
-        path for path in files
-        if path.endswith("/pretrain_final.pt") or path.endswith("/best_bpb.pt") or path.endswith("/latest.pt")
-    ]
-    if not preferred:
-        return None
-    pretrain = sorted([p for p in preferred if p.endswith("/pretrain_final.pt")])
-    if pretrain:
-        return pretrain[-1]
-    best = sorted([p for p in preferred if p.endswith("/best_bpb.pt")])
-    if best:
-        return best[-1]
-    latest = sorted([p for p in preferred if p.endswith("/latest.pt")])
-    if latest:
-        return latest[-1]
-    return None
-
-
-def hydrate_benchmark_assets(*, cache_dir: Path, output_repo: str, tokenizer_repo: str, token: str | None) -> dict[str, str]:
-    cache_dir.mkdir(parents=True, exist_ok=True)
-    tok_dir = cache_dir / "tokenizer"
-    tok_dir.mkdir(parents=True, exist_ok=True)
-    tok_repo = resolve_tokenizer_cache_repo(output_repo=tokenizer_repo, retina_cache_repo=tokenizer_repo)
-    tok_prefix = tokenizer_cache_prefix()
-
-    ckpt_path = None
-    for candidate in checkpoint_candidates(cache_dir):
-        if candidate.exists():
-            ckpt_path = candidate
-            break
-        try:
-            ckpt_path = _download_file(repo_id=output_repo, filename=candidate.name, local_dir=str(cache_dir), token=token)
-            break
-        except Exception:
-            continue
-    if ckpt_path is None:
-        try:
-            if HfApi is None:
-                raise RuntimeError("huggingface_hub unavailable")
-            files = HfApi(token=token).list_repo_files(repo_id=output_repo, repo_type="model", token=token)
-            remote_path = choose_remote_checkpoint_path(files)
-            if remote_path is not None:
-                parent, filename = remote_path.rsplit("/", 1)
-                downloaded_path = _download_file(
-                    repo_id=output_repo,
-                    filename=filename,
-                    local_dir=str(cache_dir),
-                    token=token,
-                    subfolder=parent,
-                )
-                canonical_path = cache_dir / filename
-                if downloaded_path != canonical_path:
-                    canonical_path.parent.mkdir(parents=True, exist_ok=True)
-                    shutil.copy2(downloaded_path, canonical_path)
-                ckpt_path = canonical_path
-        except Exception:
-            pass
-    if ckpt_path is None:
-        raise FileNotFoundError(f"No benchmark checkpoint found in cache or repo {output_repo}")
-
-    tok_path = tok_dir / "tokenizer.pkl"
-    if not tok_path.exists():
-        downloaded_tok = _download_file(repo_id=tok_repo, filename="tokenizer.pkl", local_dir=str(tok_dir), token=token, subfolder=tok_prefix)
-        if downloaded_tok != tok_path:
-            shutil.copy2(downloaded_tok, tok_path)
-
-    token_bytes_path = tok_dir / "token_bytes.pt"
-    if not token_bytes_path.exists():
-        downloaded_token_bytes = _download_file(repo_id=tok_repo, filename="token_bytes.pt", local_dir=str(tok_dir), token=token, subfolder=tok_prefix)
-        if downloaded_token_bytes != token_bytes_path:
-            shutil.copy2(downloaded_token_bytes, token_bytes_path)
-
-    return {
-        "checkpoint_path": str(ckpt_path),
-        "tokenizer_dir": str(tok_dir),
-    }
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+
+def _download_file(*, repo_id: str, filename: str, local_dir: str, token: str | None, subfolder: str | None = None) -> Path:
+    from huggingface_hub import hf_hub_download
+
+    path = hf_hub_download(
+        repo_id=repo_id,
+        repo_type="model",
+        filename=filename,
+        subfolder=subfolder,
+        token=token,
+        local_dir=local_dir,
+        local_dir_use_symlinks=False,
+    )
+    return Path(path)
+
+
+def resolve_tokenizer_cache_repo(*, output_repo: str, retina_cache_repo: str) -> str:
+    return (
+        os.environ.get("HYDRA_TOKENIZER_CACHE_REPO")
+        or os.environ.get("FEATHER_HF_OUTPUT_REPO")
+        or os.environ.get("HF_REPO_ID")
+        or os.environ.get("HYDRA_RETINA_CACHE_REPO")
+        or os.environ.get("FEATHER_HF_RETINA_CACHE_REPO")
+        or output_repo
+        or retina_cache_repo
+    )
+
+
+def tokenizer_cache_prefix() -> str:
+    vocab_size = int(os.environ.get("HYDRA_VOCAB_SIZE", "65536"))
+    return f"tokenizer/vocab{vocab_size}"
+
+
+def hydrate_benchmark_assets(*, cache_dir: Path, output_repo: str, tokenizer_repo: str, token: str | None) -> dict[str, str]:
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    tok_dir = cache_dir / "tokenizer"
+    tok_dir.mkdir(parents=True, exist_ok=True)
+    tok_repo = resolve_tokenizer_cache_repo(output_repo=tokenizer_repo, retina_cache_repo=tokenizer_repo)
+    tok_prefix = tokenizer_cache_prefix()
+
+    ckpt_path = cache_dir / "best_bpb.pt"
+    if not ckpt_path.exists():
+        ckpt_path = _download_file(repo_id=output_repo, filename="best_bpb.pt", local_dir=str(cache_dir), token=token)
+
+    tok_path = tok_dir / "tokenizer.pkl"
+    if not tok_path.exists():
+        tok_path = _download_file(repo_id=tok_repo, filename="tokenizer.pkl", local_dir=str(tok_dir), token=token, subfolder=tok_prefix)
+
+    token_bytes_path = tok_dir / "token_bytes.pt"
+    if not token_bytes_path.exists():
+        token_bytes_path = _download_file(repo_id=tok_repo, filename="token_bytes.pt", local_dir=str(tok_dir), token=token, subfolder=tok_prefix)
+
+    return {
+        "checkpoint_path": str(ckpt_path),
+        "tokenizer_dir": str(tok_dir),
+    }

overlay/scripts/benchmark_checkpoint.py

@@ -1,118 +1,19 @@
-#!/usr/bin/env python3
-from __future__ import annotations
-
-import shutil
-from pathlib import Path
-
-from scripts.hf_routing import resolve_routing
-
-try:
-    from huggingface_hub import HfApi
-except Exception:  # pragma: no cover
-    HfApi = None
-
-
-def choose_remote_checkpoint_path(files: list[str]) -> str | None:
-    preferred = [
-        path for path in files
-        if path.endswith("/pretrain_final.pt") or path.endswith("/best_bpb.pt") or path.endswith("/latest.pt")
-    ]
-    if not preferred:
-        return None
-    pretrain = sorted([p for p in preferred if p.endswith("/pretrain_final.pt")])
-    if pretrain:
-        return pretrain[-1]
-    best = sorted([p for p in preferred if p.endswith("/best_bpb.pt")])
-    if best:
-        return best[-1]
-    latest = sorted([p for p in preferred if p.endswith("/latest.pt")])
-    if latest:
-        return latest[-1]
-    return None
-
-
-def checkpoint_candidates(cache_dir: Path) -> list[Path]:
-    return [
-        cache_dir / "best_bpb.pt",
-        cache_dir / "pretrain_final.pt",
-        cache_dir / "latest.pt",
-    ]
-
-
-def choose_checkpoint_candidate(cache_dir: Path) -> Path | None:
-    for path in checkpoint_candidates(cache_dir):
-        if path.exists():
-            return path
-    return None
-
-
-def resolve_checkpoint_source(*, cache_dir: Path, output_repo: str | None) -> dict[str, str]:
-    local = choose_checkpoint_candidate(cache_dir)
-    if local is not None:
-        return {"mode": "local", "path": str(local)}
-    if output_repo:
-        return {"mode": "remote", "repo_id": output_repo}
-    routing = resolve_routing(token=None)
-    return {"mode": "remote", "repo_id": routing.output_repo}
-
-
-def _download_checkpoint_file(*, repo_id: str, filename: str, local_dir: str, token: str | None, subfolder: str | None = None) -> str:
-    from huggingface_hub import hf_hub_download
-
-    return hf_hub_download(
-        repo_id=repo_id,
-        repo_type="model",
-        filename=filename,
-        subfolder=subfolder,
-        token=token,
-        local_dir=local_dir,
-        local_dir_use_symlinks=False,
-    )
-
-
-def hydrate_checkpoint(*, cache_dir: Path, output_repo: str | None, token: str | None) -> Path | None:
-    local = choose_checkpoint_candidate(cache_dir)
-    if local is not None:
-        return local
-    source = resolve_checkpoint_source(cache_dir=cache_dir, output_repo=output_repo)
-    if source["mode"] != "remote":
-        return None
-    cache_dir.mkdir(parents=True, exist_ok=True)
-    for filename in ("best_bpb.pt", "pretrain_final.pt", "latest.pt"):
-        try:
-            path = Path(
-                _download_checkpoint_file(
-                    repo_id=source["repo_id"],
-                    filename=filename,
-                    local_dir=str(cache_dir),
-                    token=token,
-                )
-            )
-            if path.exists():
-                return path
-        except Exception:
-            continue
-    try:
-        if HfApi is None:
-            raise RuntimeError("huggingface_hub unavailable")
-        files = HfApi(token=token).list_repo_files(repo_id=source["repo_id"], repo_type="model", token=token)
-        remote_path = choose_remote_checkpoint_path(files)
-        if remote_path is not None:
-            parent, filename = remote_path.rsplit("/", 1)
-            downloaded = Path(
-                _download_checkpoint_file(
-                    repo_id=source["repo_id"],
-                    filename=filename,
-                    local_dir=str(cache_dir),
-                    token=token,
-                    subfolder=parent,
-                )
-            )
-            canonical = cache_dir / filename
-            if downloaded != canonical:
-                shutil.copy2(downloaded, canonical)
-            if canonical.exists():
-                return canonical
-    except Exception:
-        pass
-    return None
+#!/usr/bin/env python3
+from __future__ import annotations
+
+from pathlib import Path
+
+
+def checkpoint_candidates(cache_dir: Path) -> list[Path]:
+    return [
+        cache_dir / "best_bpb.pt",
+        cache_dir / "pretrain_final.pt",
+        cache_dir / "latest.pt",
+    ]
+
+
+def choose_checkpoint_candidate(cache_dir: Path) -> Path | None:
+    for path in checkpoint_candidates(cache_dir):
+        if path.exists():
+            return path
+    return None

overlay/scripts/benchmark_checkpoint_report.py

@@ -1,50 +1,50 @@
-#!/usr/bin/env python3
-from __future__ import annotations
-
-import json
-
-
-def build_checkpoint_report(files: list[str]) -> dict[str, object]:
-    by_job: dict[str, dict[str, object]] = {}
-    for path in files:
-        parts = path.split("/")
-        if len(parts) < 3 or parts[0] != "jobs":
-            continue
-        job_id = parts[1]
-        filename = parts[-1]
-        if filename not in {"best_bpb.pt", "pretrain_final.pt", "latest.pt"}:
-            continue
-        row = by_job.setdefault(job_id, {"job_id": job_id, "paths": []})
-        row["paths"].append(path)
-
-    candidates = []
-    for job_id, row in by_job.items():
-        paths = list(row["paths"])
-        preferred = None
-        for suffix in ("pretrain_final.pt", "best_bpb.pt", "latest.pt"):
-            for path in paths:
-                if path.endswith(suffix):
-                    preferred = path
-                    break
-            if preferred is not None:
-                break
-        candidates.append({
-            "job_id": job_id,
-            "preferred_path": preferred,
-            "available_paths": sorted(paths),
-        })
-
-    candidates.sort(key=lambda row: row["job_id"], reverse=True)
-    return {
-        "n_candidates": len(candidates),
-        "candidates": candidates,
-    }
-
-
-def main() -> int:
-    print(json.dumps(build_checkpoint_report([]), indent=2, sort_keys=True))
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import json
+
+
+def build_checkpoint_report(files: list[str]) -> dict[str, object]:
+    by_job: dict[str, dict[str, object]] = {}
+    for path in files:
+        parts = path.split("/")
+        if len(parts) < 3 or parts[0] != "jobs":
+            continue
+        job_id = parts[1]
+        filename = parts[-1]
+        if filename not in {"best_bpb.pt", "pretrain_final.pt", "latest.pt"}:
+            continue
+        row = by_job.setdefault(job_id, {"job_id": job_id, "paths": []})
+        row["paths"].append(path)
+
+    candidates = []
+    for job_id, row in by_job.items():
+        paths = list(row["paths"])
+        preferred = None
+        for suffix in ("pretrain_final.pt", "best_bpb.pt", "latest.pt"):
+            for path in paths:
+                if path.endswith(suffix):
+                    preferred = path
+                    break
+            if preferred is not None:
+                break
+        candidates.append({
+            "job_id": job_id,
+            "preferred_path": preferred,
+            "available_paths": sorted(paths),
+        })
+
+    candidates.sort(key=lambda row: row["job_id"], reverse=True)
+    return {
+        "n_candidates": len(candidates),
+        "candidates": candidates,
+    }
+
+
+def main() -> int:
+    print(json.dumps(build_checkpoint_report([]), indent=2, sort_keys=True))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

overlay/scripts/benchmark_contract.py

@@ -1,67 +1,67 @@
-#!/usr/bin/env python3
-from __future__ import annotations
-
-import json
-from pathlib import Path
-from typing import Any
-
-
-def _require_path(payload: dict[str, Any], path: str) -> None:
-    current: Any = payload
-    for part in path.split('.'):
-        if not isinstance(current, dict) or part not in current:
-            raise ValueError(f"missing required field: {path}")
-        current = current[part]
-
-
-def validate_benchmark_contract(payload: dict[str, Any]) -> None:
-    for field in [
-        "cycle_id",
-        "hardware_class",
-        "seeds",
-        "budget_modes",
-        "coding_benchmarks.fast_iteration",
-        "coding_benchmarks.milestone",
-        "reasoning_benchmarks.fast_iteration",
-        "reasoning_benchmarks.milestone",
-        "variants.hydra_full",
-        "variants.baseline_mamba_matched",
-    ]:
-        _require_path(payload, field)
-
-    for section in [
-        payload["coding_benchmarks"]["fast_iteration"],
-        payload["coding_benchmarks"]["milestone"],
-        payload["reasoning_benchmarks"]["fast_iteration"],
-        payload["reasoning_benchmarks"]["milestone"],
-    ]:
-        if "name" not in section or "primary_metric" not in section or "decode" not in section:
-            raise ValueError("benchmark sections require name, primary_metric, and decode")
-
-    if not isinstance(payload["seeds"], list) or len(payload["seeds"]) < 3:
-        raise ValueError("seeds must contain at least three values")
-
-    if payload["variants"]["hydra_full"].get("status") != "runnable_now":
-        raise ValueError("hydra_full must be runnable_now")
-
-    if payload["variants"]["baseline_mamba_matched"].get("status") != "runnable_now":
-        raise ValueError("baseline_mamba_matched must be runnable_now")
-
-
-def load_benchmark_contract(path: Path) -> dict[str, Any]:
-    payload = json.loads(path.read_text(encoding="utf-8"))
-    if not isinstance(payload, dict):
-        raise ValueError("benchmark contract must be a JSON object")
-    validate_benchmark_contract(payload)
-    return payload
-
-
-def main() -> int:
-    path = Path("artifacts/cycle_1_execution_freeze.json")
-    payload = load_benchmark_contract(path)
-    print(json.dumps(payload, indent=2, sort_keys=True))
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+
+def _require_path(payload: dict[str, Any], path: str) -> None:
+    current: Any = payload
+    for part in path.split('.'):
+        if not isinstance(current, dict) or part not in current:
+            raise ValueError(f"missing required field: {path}")
+        current = current[part]
+
+
+def validate_benchmark_contract(payload: dict[str, Any]) -> None:
+    for field in [
+        "cycle_id",
+        "hardware_class",
+        "seeds",
+        "budget_modes",
+        "coding_benchmarks.fast_iteration",
+        "coding_benchmarks.milestone",
+        "reasoning_benchmarks.fast_iteration",
+        "reasoning_benchmarks.milestone",
+        "variants.hydra_full",
+        "variants.baseline_mamba_matched",
+    ]:
+        _require_path(payload, field)
+
+    for section in [
+        payload["coding_benchmarks"]["fast_iteration"],
+        payload["coding_benchmarks"]["milestone"],
+        payload["reasoning_benchmarks"]["fast_iteration"],
+        payload["reasoning_benchmarks"]["milestone"],
+    ]:
+        if "name" not in section or "primary_metric" not in section or "decode" not in section:
+            raise ValueError("benchmark sections require name, primary_metric, and decode")
+
+    if not isinstance(payload["seeds"], list) or len(payload["seeds"]) < 3:
+        raise ValueError("seeds must contain at least three values")
+
+    if payload["variants"]["hydra_full"].get("status") != "runnable_now":
+        raise ValueError("hydra_full must be runnable_now")
+
+    if payload["variants"]["baseline_mamba_matched"].get("status") != "runnable_now":
+        raise ValueError("baseline_mamba_matched must be runnable_now")
+
+
+def load_benchmark_contract(path: Path) -> dict[str, Any]:
+    payload = json.loads(path.read_text(encoding="utf-8"))
+    if not isinstance(payload, dict):
+        raise ValueError("benchmark contract must be a JSON object")
+    validate_benchmark_contract(payload)
+    return payload
+
+
+def main() -> int:
+    path = Path("artifacts/cycle_1_execution_freeze.json")
+    payload = load_benchmark_contract(path)
+    print(json.dumps(payload, indent=2, sort_keys=True))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

overlay/scripts/benchmark_datasets.py

@@ -1,190 +1,18 @@
-#!/usr/bin/env python3
-from __future__ import annotations
-
-import argparse
-import json
-from pathlib import Path
-from typing import Any
-
-import pyarrow as pa
-import pyarrow.parquet as pq
-
-try:
-    from huggingface_hub import HfApi, hf_hub_download
-except Exception:  # pragma: no cover
-    HfApi = None
-    hf_hub_download = None
-
-
-CANONICAL_SUBSETS = {
-    "MBPP": Path("data/benchmarks/mbpp.cycle1.jsonl"),
-    "GSM8K": Path("data/benchmarks/gsm8k.cycle1.jsonl"),
-    "HumanEval": Path("data/benchmarks/humaneval.cycle1.jsonl"),
-    "ARC-Challenge": Path("data/benchmarks/arc_challenge.cycle1.jsonl"),
-}
-
-
-DATASET_SOURCES: dict[str, dict[str, str]] = {
-    "MBPP": {"repo_id": "Muennighoff/mbpp", "subset": "full", "split": "test", "raw_path": "data/mbpp.jsonl"},
-    "GSM8K": {"repo_id": "openai/gsm8k", "subset": "main", "split": "test"},
-    "HumanEval": {"repo_id": "openai/openai_humaneval", "subset": "default", "split": "test"},
-    "ARC-Challenge": {"repo_id": "allenai/ai2_arc", "subset": "ARC-Challenge", "split": "validation"},
-}
-
-
-def resolve_benchmark_dataset(benchmark_name: str, explicit_path: Path | None) -> Path:
-    if explicit_path is not None:
-        return explicit_path
-    if benchmark_name not in CANONICAL_SUBSETS:
-        raise ValueError(f"Unsupported benchmark dataset: {benchmark_name}")
-    return Path.cwd() / CANONICAL_SUBSETS[benchmark_name]
-
-
-def _normalize_gsm8k_answer(answer: str) -> str:
-    if "####" in answer:
-        return answer.split("####")[-1].strip()
-    return answer.strip()
-
-
-def transform_dataset_row(benchmark_name: str, row: dict[str, Any], *, row_id: int) -> dict[str, Any]:
-    if benchmark_name == "GSM8K":
-        return {
-            "question": str(row["question"]),
-            "answer": _normalize_gsm8k_answer(str(row["answer"])),
-        }
-    if benchmark_name == "ARC-Challenge":
-        choices = row["choices"]
-        labels = list(choices["label"])
-        texts = list(choices["text"])
-        answer_key = str(row["answerKey"])
-        answer_index = labels.index(answer_key)
-        return {
-            "question": str(row["question"]),
-            "choices": [str(choice) for choice in texts],
-            "answer": str(texts[answer_index]),
-        }
-    if benchmark_name == "MBPP":
-        task_id = row.get("task_id", row_id)
-        return {
-            "task_id": str(task_id),
-            "prompt": str(row["text"]),
-            "tests": [str(test) for test in row["test_list"]],
-        }
-    if benchmark_name == "HumanEval":
-        task_id = row.get("task_id", row_id)
-        return {
-            "task_id": str(task_id),
-            "prompt": str(row["prompt"]),
-            "test": str(row["test"]),
-        }
-    raise ValueError(f"Unsupported benchmark dataset: {benchmark_name}")
-
-
-def write_canonical_dataset(*, benchmark_name: str, rows: list[dict[str, Any]], out_path: Path, limit: int) -> int:
-    out_path.parent.mkdir(parents=True, exist_ok=True)
-    transformed = [transform_dataset_row(benchmark_name, row, row_id=index) for index, row in enumerate(rows[:limit])]
-    out_path.write_text("".join(json.dumps(row) + "\n" for row in transformed), encoding="utf-8")
-    return len(transformed)
-
-
-def choose_dataset_parquet_path(benchmark_name: str, files: list[str]) -> str | None:
-    source = DATASET_SOURCES[benchmark_name]
-    subset = source["subset"].lower()
-    split = source["split"].lower()
-    candidates = [path for path in files if path.endswith(".parquet")]
-    preferred = [path for path in candidates if subset in path.lower() and split in path.lower()]
-    if preferred:
-        return sorted(preferred)[0]
-    split_only = [path for path in candidates if split in path.lower()]
-    if split_only:
-        return sorted(split_only)[0]
-    return sorted(candidates)[0] if candidates else None
-
-
-def download_dataset_snapshot(benchmark_name: str, *, cache_dir: Path, token: str | None) -> Path:
-    source = DATASET_SOURCES[benchmark_name]
-    if HfApi is None or hf_hub_download is None:
-        raise RuntimeError("huggingface_hub unavailable")
-    raw_path = source.get("raw_path")
-    if raw_path:
-        if "/" in raw_path:
-            subfolder, filename = raw_path.rsplit("/", 1)
-        else:
-            subfolder, filename = None, raw_path
-        downloaded = hf_hub_download(
-            repo_id=source["repo_id"],
-            repo_type="dataset",
-            filename=filename,
-            subfolder=subfolder,
-            token=token,
-            local_dir=str(cache_dir),
-        )
-        return Path(downloaded)
-    files = HfApi(token=token).list_repo_files(repo_id=source["repo_id"], repo_type="dataset", token=token)
-    parquet_path = choose_dataset_parquet_path(benchmark_name, files)
-    if parquet_path is None:
-        raise FileNotFoundError(f"No parquet dataset file found for {benchmark_name}")
-    if "/" in parquet_path:
-        subfolder, filename = parquet_path.rsplit("/", 1)
-    else:
-        subfolder, filename = None, parquet_path
-    downloaded = hf_hub_download(
-        repo_id=source["repo_id"],
-        repo_type="dataset",
-        filename=filename,
-        subfolder=subfolder,
-        token=token,
-        local_dir=str(cache_dir),
-    )
-    return Path(downloaded)
-
-
-def hydrate_canonical_dataset(
-    *,
-    benchmark_name: str,
-    out_path: Path,
-    limit: int,
-    cache_dir: Path,
-    token: str | None,
-) -> int:
-    source_path = download_dataset_snapshot(benchmark_name, cache_dir=cache_dir, token=token)
-    if source_path.suffix == ".jsonl":
-        rows = [json.loads(line) for line in source_path.read_text(encoding="utf-8").splitlines() if line.strip()]
-    else:
-        table = pq.read_table(source_path)
-        rows = table.to_pylist()
-    return write_canonical_dataset(benchmark_name=benchmark_name, rows=rows, out_path=out_path, limit=limit)
-
-
-def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Hydrate a canonical benchmark dataset JSONL from a public source")
-    parser.add_argument("--benchmark", required=True, choices=list(CANONICAL_SUBSETS))
-    parser.add_argument("--samples", type=Path)
-    parser.add_argument("--out", type=Path)
-    parser.add_argument("--limit", type=int, default=20)
-    parser.add_argument("--cache-dir", type=Path, default=Path(".cache/benchmarks"))
-    parser.add_argument("--token")
-    return parser.parse_args(argv)
-
-
-def main(argv: list[str] | None = None) -> int:
-    args = parse_args(argv)
-    if args.samples is not None:
-        rows = [json.loads(line) for line in args.samples.read_text(encoding="utf-8").splitlines() if line.strip()]
-        out_path = args.out or resolve_benchmark_dataset(args.benchmark, None)
-        write_canonical_dataset(benchmark_name=args.benchmark, rows=rows, out_path=out_path, limit=args.limit)
-        return 0
-
-    out_path = args.out or resolve_benchmark_dataset(args.benchmark, None)
-    hydrate_canonical_dataset(
-        benchmark_name=args.benchmark,
-        out_path=out_path,
-        limit=args.limit,
-        cache_dir=args.cache_dir,
-        token=args.token,
-    )
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())
+#!/usr/bin/env python3
+from __future__ import annotations
+
+from pathlib import Path
+
+
+CANONICAL_SUBSETS = {
+    "MBPP": Path("data/benchmarks/mbpp.cycle1.jsonl"),
+    "GSM8K": Path("data/benchmarks/gsm8k.cycle1.jsonl"),
+}
+
+
+def resolve_benchmark_dataset(benchmark_name: str, explicit_path: Path | None) -> Path:
+    if explicit_path is not None:
+        return explicit_path
+    if benchmark_name not in CANONICAL_SUBSETS:
+        raise ValueError(f"Unsupported benchmark dataset: {benchmark_name}")
+    return Path.cwd() / CANONICAL_SUBSETS[benchmark_name]

overlay/scripts/benchmark_hyena_stack.py

@@ -26,11 +26,8 @@ Invocation:
     # On A100/A10G (production cloud hardware), use time=900 (15 min) for
     # stable steady-state numbers.
 
-After each run the script prints:
-    BENCHMARK config=<name> tps_steady=<avg> bpb_at_500=<val> vram_peak=<MiB>
-
-If `--min-tps` is set (>0), the script exits non-zero when steady-state TPS
-falls below the threshold.
+After each run the script prints:
+    BENCHMARK config=<name> tps_steady=<avg> bpb_at_500=<val> vram_peak=<MiB>
 
 Collate those lines into the matrix table manually, then pick the winner
 for the 6-hour production run (HYDRA_TIME_BUDGET=21600).
@@ -50,34 +47,30 @@ REPO = Path(__file__).resolve().parents[1]
 
 CONFIGS = {
     # Baseline: B=8, no flash, no train-cache. Current reference point.
-    "baseline": {
-        "HYDRA_BATCH_SIZE": "8",
-        "HYDRA_THROUGHPUT_MODE": "1",
-        "HYDRA_HYENA_LAYERS": "3,7",
+    "baseline": {
+        "HYDRA_BATCH_SIZE": "8",
+        "HYDRA_HYENA_LAYERS": "3,7",
         "HYDRA_HYENA_FLASH_FFT": "0",
         "HYDRA_HYENA_TRAIN_CACHE": "0",
         "HYDRA_HYENA_FILTER_CACHE": "0",
     },
-    "b16": {
-        "HYDRA_BATCH_SIZE": "16",
-        "HYDRA_THROUGHPUT_MODE": "1",
-        "HYDRA_HYENA_LAYERS": "3,7",
+    "b16": {
+        "HYDRA_BATCH_SIZE": "16",
+        "HYDRA_HYENA_LAYERS": "3,7",
         "HYDRA_HYENA_FLASH_FFT": "0",
         "HYDRA_HYENA_TRAIN_CACHE": "0",
         "HYDRA_HYENA_FILTER_CACHE": "0",
     },
-    "cache": {
-        "HYDRA_BATCH_SIZE": "16",
-        "HYDRA_THROUGHPUT_MODE": "1",
-        "HYDRA_HYENA_LAYERS": "3,7",
+    "cache": {
+        "HYDRA_BATCH_SIZE": "16",
+        "HYDRA_HYENA_LAYERS": "3,7",
         "HYDRA_HYENA_FLASH_FFT": "0",
         "HYDRA_HYENA_TRAIN_CACHE": "1",
         "HYDRA_HYENA_FILTER_CACHE": "1",
     },
-    "kernel": {
-        "HYDRA_BATCH_SIZE": "16",
-        "HYDRA_THROUGHPUT_MODE": "1",
-        "HYDRA_HYENA_LAYERS": "3,7",
+    "kernel": {
+        "HYDRA_BATCH_SIZE": "16",
+        "HYDRA_HYENA_LAYERS": "3,7",
         "HYDRA_HYENA_FLASH_FFT": "1",
         "HYDRA_HYENA_TRAIN_CACHE": "1",
         "HYDRA_HYENA_FILTER_CACHE": "1",
@@ -88,7 +81,7 @@ CONFIGS = {
 }
 
 
-def build_env(cfg_overrides: dict[str, str]) -> dict[str, str]:
+def build_env(cfg_overrides: dict) -> dict:
     """Compose a full env dict from the inherited env + config overrides."""
     env = os.environ.copy()
     # Ensure the Hyena layer selection is always present (defaults to off).
@@ -98,7 +91,7 @@ def build_env(cfg_overrides: dict[str, str]) -> dict[str, str]:
     return env
 
 
-def parse_step_line(line: str) -> dict[str, float] | None:
+def parse_step_line(line: str) -> dict | None:
     """Parse a single step=... line into a dict of metrics, or None."""
     if not line.startswith("step="):
         return None
@@ -109,7 +102,7 @@ def parse_step_line(line: str) -> dict[str, float] | None:
         return None
 
 
-def summarize(log_path: Path, warmup_steps: int = 50) -> dict[str, float]:
+def summarize(log_path: Path, warmup_steps: int = 50) -> dict:
     """Tail log_path, compute steady-state TPS / BPB@500 / VRAM peak.
 
     Skips the first `warmup_steps` to discard CUDA graph capture / autotune
@@ -145,29 +138,20 @@ def summarize(log_path: Path, warmup_steps: int = 50) -> dict[str, float]:
     tps_sorted = sorted(tps_vals)
     tps_steady = tps_sorted[len(tps_sorted) // 2]  # median
 
-    return {
-        "tps_steady": tps_steady,
-        "bpb_at_500": bpb_at_500 or (bpbs[-1] if bpbs else 0.0),
-        "vram_peak": vram_peak,
-        "steps": len(tps_vals) + warmup_steps,
-    }
-
-
-def fails_tps_floor(summary: dict[str, float], min_tps: float) -> bool:
-    if min_tps <= 0:
-        return False
-    tps_steady = float(summary.get("tps_steady", 0.0))
-    return tps_steady < float(min_tps)
-
-
-def main() -> int:
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--config", required=True, choices=list(CONFIGS))
-    ap.add_argument("--time", type=int, default=300, help="training seconds")
-    ap.add_argument("--log", default=None, help="output log path (default: run_bench_<cfg>.log)")
-    ap.add_argument("--min-tps", type=float, default=50000.0, help="Required steady-state TPS floor (set 0 to disable)")
-    ap.add_argument("--warmup-steps", type=int, default=50, help="Number of initial steps to skip before TPS median")
-    args = ap.parse_args()
+    return {
+        "tps_steady": tps_steady,
+        "bpb_at_500": bpb_at_500 or (bpbs[-1] if bpbs else 0.0),
+        "vram_peak": vram_peak,
+        "steps": len(tps_vals) + warmup_steps,
+    }
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--config", required=True, choices=list(CONFIGS))
+    ap.add_argument("--time", type=int, default=300, help="training seconds")
+    ap.add_argument("--log", default=None, help="output log path (default: run_bench_<cfg>.log)")
+    args = ap.parse_args()
 
     cfg = CONFIGS[args.config]
     log_path = Path(args.log or (REPO / f"run_bench_{args.config}.log"))
@@ -194,25 +178,16 @@ def main() -> int:
         print(f"BENCH FAIL config={args.config}", flush=True)
         return proc.returncode
 
-    summary = summarize(log_path, warmup_steps=max(0, int(args.warmup_steps)))
-    print(
-        f"BENCHMARK config={args.config} "
-        f"tps_steady={summary['tps_steady']:.0f} "
-        f"bpb_at_500={summary['bpb_at_500']:.4f} "
-        f"vram_peak={summary['vram_peak']:.0f}MiB "
-        f"steps={summary['steps']}",
-        flush=True,
-    )
-
-    if fails_tps_floor(summary, args.min_tps):
-        print(
-            f"BENCH FAIL config={args.config} tps_steady={summary['tps_steady']:.0f} < min_tps={args.min_tps:.0f}",
-            flush=True,
-        )
-        return 2
-
-    print(f"BENCH PASS config={args.config} min_tps={args.min_tps:.0f}", flush=True)
-    return 0
+    summary = summarize(log_path)
+    print(
+        f"BENCHMARK config={args.config} "
+        f"tps_steady={summary['tps_steady']:.0f} "
+        f"bpb_at_500={summary['bpb_at_500']:.4f} "
+        f"vram_peak={summary['vram_peak']:.0f}MiB "
+        f"steps={summary['steps']}",
+        flush=True,
+    )
+    return 0
 
 
 if __name__ == "__main__":

overlay/scripts/benchmark_preflight.py

@@ -1,35 +1,31 @@
-#!/usr/bin/env python3
-from __future__ import annotations
-
-from pathlib import Path
-
-from scripts.bootstrap_benchmark_env import build_bootstrap_report
-from scripts.benchmark_checkpoint import choose_checkpoint_candidate
-
-
-def build_readiness_report(*, cache_dir: Path, hf_token_present: bool, dependencies_present: bool = True, missing_dependencies: list[str] | None = None, output_repo: str | None = None, tokenizer_repo: str | None = None) -> dict[str, object]:
-    checkpoint = choose_checkpoint_candidate(cache_dir)
-    tokenizer_dir = cache_dir / "tokenizer"
-    retina_path = cache_dir / "retina.npz"
-    tokenizer_ready = (tokenizer_dir / "tokenizer.pkl").exists() and (tokenizer_dir / "token_bytes.pt").exists()
-    retina_ready = retina_path.exists()
-    checkpoint_present = checkpoint is not None
-    runtime = build_bootstrap_report(missing_dependencies=list(missing_dependencies or []))
-    return {
-        "cache_dir": str(cache_dir),
-        "checkpoint_present": checkpoint_present,
-        "checkpoint_path": str(checkpoint) if checkpoint is not None else None,
-        "tokenizer_ready": tokenizer_ready,
-        "retina_ready": retina_ready,
-        "retina_path": str(retina_path),
-        "hf_token_present": hf_token_present,
-        "dependencies_present": dependencies_present,
-        "missing_dependencies": list(missing_dependencies or []),
-        "install_hint": runtime["install_hint"],
-        "install_command": runtime["install_command"],
-        "install_blockers": runtime["install_blockers"],
-        "output_repo": output_repo,
-        "tokenizer_repo": tokenizer_repo,
-        "hydration_possible": bool(hf_token_present and output_repo and tokenizer_repo),
-        "ready_for_hydra_benchmarks": checkpoint_present and tokenizer_ready and retina_ready and dependencies_present,
-    }
+#!/usr/bin/env python3
+from __future__ import annotations
+
+from pathlib import Path
+
+from scripts.bootstrap_benchmark_env import build_bootstrap_report
+from scripts.benchmark_checkpoint import choose_checkpoint_candidate
+
+
+def build_readiness_report(*, cache_dir: Path, hf_token_present: bool, dependencies_present: bool = True, missing_dependencies: list[str] | None = None, output_repo: str | None = None, tokenizer_repo: str | None = None) -> dict[str, object]:
+    checkpoint = choose_checkpoint_candidate(cache_dir)
+    tokenizer_dir = cache_dir / "tokenizer"
+    tokenizer_ready = (tokenizer_dir / "tokenizer.pkl").exists() and (tokenizer_dir / "token_bytes.pt").exists()
+    checkpoint_present = checkpoint is not None
+    runtime = build_bootstrap_report(missing_dependencies=list(missing_dependencies or []))
+    return {
+        "cache_dir": str(cache_dir),
+        "checkpoint_present": checkpoint_present,
+        "checkpoint_path": str(checkpoint) if checkpoint is not None else None,
+        "tokenizer_ready": tokenizer_ready,
+        "hf_token_present": hf_token_present,
+        "dependencies_present": dependencies_present,
+        "missing_dependencies": list(missing_dependencies or []),
+        "install_hint": runtime["install_hint"],
+        "install_command": runtime["install_command"],
+        "install_blockers": runtime["install_blockers"],
+        "output_repo": output_repo,
+        "tokenizer_repo": tokenizer_repo,
+        "hydration_possible": bool(hf_token_present and output_repo and tokenizer_repo),
+        "ready_for_hydra_benchmarks": checkpoint_present and tokenizer_ready and dependencies_present,
+    }

overlay/scripts/benchmark_runner.py

@@ -1,327 +1,248 @@
-#!/usr/bin/env python3
-from __future__ import annotations
-
-import argparse
-import json
-import re
-import sys
-from pathlib import Path
-from typing import Any, Callable
-
-REPO_ROOT = Path(__file__).resolve().parents[1]
-if str(REPO_ROOT) not in sys.path:
-    sys.path.insert(0, str(REPO_ROOT))
-
-LEDGER_TEMPLATE_PATH = REPO_ROOT / "artifacts" / "benchmark_ledger.template.json"
-
-from scripts.hydra_generation import build_hydra_generator
-from scripts.benchmark_datasets import resolve_benchmark_dataset as resolve_canonical_dataset
-from scripts.benchmark_suite import build_prompt, validate_sample
-
-
-class BenchmarkExecutionError(RuntimeError):
-    def __init__(
-        self,
-        *,
-        benchmark: str,
-        sample: dict[str, Any],
-        generated_output: str,
-        cause: BaseException,
-        extracted_output: str | None = None,
-    ):
-        super().__init__(str(cause))
-        self.benchmark = benchmark
-        self.sample = sample
-        self.generated_output = generated_output
-        self.cause = cause
-        self.extracted_output = extracted_output
-
-
-def load_jsonl_samples(path: Path) -> list[dict[str, Any]]:
-    rows: list[dict[str, Any]] = []
-    for line in path.read_text(encoding="utf-8").splitlines():
-        if line.strip():
-            rows.append(json.loads(line))
-    return rows
-
-
-def _normalize_samples_path(path: Path) -> Path:
-    return path if path.is_absolute() else REPO_ROOT / path
-
-
-def _preview_text(text: str, *, limit: int = 1000) -> str:
-    if len(text) <= limit:
-        return text
-    return text[:limit] + "\n...[truncated]"
-
-
-def extract_python_code(text: str) -> str:
-    fenced = re.search(r"```python\s*(.*?)```", text, flags=re.IGNORECASE | re.DOTALL)
-    if fenced:
-        extracted = fenced.group(1).strip("\n")
-        return extracted + "\n"
-
-    lines = text.splitlines()
-    for index, line in enumerate(lines):
-        if line.startswith("def "):
-            extracted = "\n".join(lines[index:]).strip("\n")
-            return extracted + "\n"
-    return text
-
-
-def _score_mbpp(samples: list[dict[str, Any]], generate_fn: Callable[[str], str]) -> float:
-    passed = 0
-    for sample in samples:
-        validate_sample("MBPP", sample)
-        raw_output = generate_fn(build_prompt("MBPP", sample))
-        code = extract_python_code(raw_output)
-        namespace: dict[str, Any] = {}
-        try:
-            exec(code, namespace, namespace)
-            for test in sample["tests"]:
-                exec(test, namespace, namespace)
-        except Exception as exc:
-            raise BenchmarkExecutionError(
-                benchmark="MBPP",
-                sample=sample,
-                generated_output=raw_output,
-                cause=exc,
-                extracted_output=code,
-            ) from exc
-        passed += 1
-    return passed / len(samples) if samples else 0.0
-
-
-def _extract_last_number(text: str) -> str | None:
-    matches = re.findall(r"-?\d+(?:\.\d+)?", text)
-    return matches[-1] if matches else None
-
-
-def _score_gsm8k(samples: list[dict[str, Any]], generate_fn: Callable[[str], str]) -> float:
-    passed = 0
-    for sample in samples:
-        validate_sample("GSM8K", sample)
-        output = generate_fn(build_prompt("GSM8K", sample))
-        pred = _extract_last_number(output)
-        if pred is not None and pred == str(sample["answer"]):
-            passed += 1
-    return passed / len(samples) if samples else 0.0
-
-
-def _score_humaneval(samples: list[dict[str, Any]], generate_fn: Callable[[str], str]) -> float:
-    passed = 0
-    for sample in samples:
-        validate_sample("HumanEval", sample)
-        raw_output = generate_fn(build_prompt("HumanEval", sample))
-        code = extract_python_code(raw_output)
-        namespace: dict[str, Any] = {}
-        try:
-            exec(code, namespace, namespace)
-            exec(sample["test"], namespace, namespace)
-        except Exception as exc:
-            raise BenchmarkExecutionError(
-                benchmark="HumanEval",
-                sample=sample,
-                generated_output=raw_output,
-                cause=exc,
-                extracted_output=code,
-            ) from exc
-        passed += 1
-    return passed / len(samples) if samples else 0.0
-
-
-def _score_arc(samples: list[dict[str, Any]], generate_fn: Callable[[str], str]) -> float:
-    passed = 0
-    for sample in samples:
-        validate_sample("ARC-Challenge", sample)
-        output = generate_fn(build_prompt("ARC-Challenge", sample)).strip()
-        if output == str(sample["answer"]):
-            passed += 1
-    return passed / len(samples) if samples else 0.0
-
-
-def run_benchmark(benchmark_name: str, path: Path, generate_fn: Callable[[str], str]) -> dict[str, Any]:
-    samples = load_jsonl_samples(path)
-    if benchmark_name == "MBPP":
-        return {
-            "benchmark": "MBPP",
-            "primary_metric": "pass_at_1",
-            "score": _score_mbpp(samples, generate_fn),
-            "n_samples": len(samples),
-        }
-    if benchmark_name == "GSM8K":
-        return {
-            "benchmark": "GSM8K",
-            "primary_metric": "exact_match",
-            "score": _score_gsm8k(samples, generate_fn),
-            "n_samples": len(samples),
-        }
-    if benchmark_name == "HumanEval":
-        return {
-            "benchmark": "HumanEval",
-            "primary_metric": "pass_at_1",
-            "score": _score_humaneval(samples, generate_fn),
-            "n_samples": len(samples),
-        }
-    if benchmark_name == "ARC-Challenge":
-        return {
-            "benchmark": "ARC-Challenge",
-            "primary_metric": "accuracy",
-            "score": _score_arc(samples, generate_fn),
-            "n_samples": len(samples),
-        }
-    raise ValueError(f"Unsupported runnable benchmark: {benchmark_name}")
-
-
-def write_benchmark_result(path: Path, payload: dict[str, Any]) -> None:
-    path.parent.mkdir(parents=True, exist_ok=True)
-    path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
-
-
-def append_benchmark_run_record(
-    ledger_path: Path,
-    result: dict[str, Any],
-    *,
-    benchmark_name: str,
-    variant: str,
-    seed: int,
-    samples_path: Path,
-) -> None:
-    if not ledger_path.exists():
-        ledger_path.parent.mkdir(parents=True, exist_ok=True)
-        ledger_path.write_text(LEDGER_TEMPLATE_PATH.read_text(encoding="utf-8"), encoding="utf-8")
-    payload = json.loads(ledger_path.read_text(encoding="utf-8"))
-    run_records = payload.setdefault("run_records", [])
-    if len(run_records) == 1 and run_records[0].get("run_id") == "example-run-0001":
-        run_records.clear()
-    run_records.append(
-        {
-            "run_id": result.get("run_id", f"{benchmark_name.lower()}-{seed}"),
-            "commit": "HEAD",
-            "model_family": "hydra",
-            "variant": variant,
-            "seed": seed,
-            "hardware": {
-                "hardware_class": payload.get("benchmark_cycle", {}).get("hardware_class", "unknown"),
-            },
-            "budget": {
-                "budget_mode": payload.get("benchmark_cycle", {}).get("budget_modes", [None])[0],
-            },
-            "capability": {
-                "coding_score": result["score"] if benchmark_name in {"MBPP", "HumanEval"} else None,
-                "reasoning_score": result["score"] if benchmark_name in {"GSM8K", "ARC-Challenge"} else None,
-            },
-            "artifacts": {
-                "samples_path": str(samples_path),
-            },
-        }
-    )
-    ledger_path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
-
-
-def resolve_samples_path(benchmark_name: str, samples: Path | None, suite_path: Path) -> Path:
-    if samples is not None:
-        return _normalize_samples_path(samples)
-    payload = json.loads(suite_path.read_text(encoding="utf-8"))
-    for section in ("coding_benchmarks", "reasoning_benchmarks"):
-        if section not in payload:
-            continue
-        for slot in ("fast_iteration", "milestone"):
-            entry = payload[section].get(slot)
-            if isinstance(entry, dict) and entry.get("name") == benchmark_name and "sample_path" in entry:
-                return _normalize_samples_path(Path(entry["sample_path"]))
-    try:
-        return _normalize_samples_path(resolve_canonical_dataset(benchmark_name, None))
-    except ValueError:
-        raise ValueError(f"No sample path found for benchmark: {benchmark_name}")
-
-
-def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Run a local benchmark against JSONL samples")
-    parser.add_argument("--benchmark", required=True, choices=["MBPP", "GSM8K", "HumanEval", "ARC-Challenge"])
-    parser.add_argument("--samples", type=Path)
-    parser.add_argument("--suite", type=Path, default=REPO_ROOT / "artifacts" / "benchmark_suite.cycle1.json")
-    parser.add_argument("--out", type=Path)
-    parser.add_argument("--ledger", type=Path)
-    parser.add_argument("--variant", default="hydra_full")
-    parser.add_argument("--seed", type=int, default=42)
-    parser.add_argument("--generator-mode", choices=["stub", "hydra"], default="stub")
-    parser.add_argument("--checkpoint", type=Path)
-    parser.add_argument("--device")
-    parser.add_argument("--max-new-tokens", type=int, default=256)
-    parser.add_argument("--temperature", type=float, default=0.2)
-    parser.add_argument("--top-p", type=float, default=0.95)
-    return parser.parse_args(argv)
-
-
-def main(argv: list[str] | None = None) -> int:
-    args = parse_args(argv)
-    sample_path = resolve_samples_path(args.benchmark, args.samples, args.suite)
-    try:
-        if args.generator_mode == "hydra":
-            generator = build_hydra_generator(
-                checkpoint_path=args.checkpoint,
-                device=args.device,
-                max_new_tokens=args.max_new_tokens,
-                temperature=args.temperature,
-                top_p=args.top_p,
-            )
-        else:
-            def generator(prompt: str) -> str:
-                return prompt
-
-        result = run_benchmark(args.benchmark, sample_path, generator)
-        exit_code = 0
-    except FileNotFoundError as exc:
-        result = {
-            "benchmark": args.benchmark,
-            "status": "failed",
-            "failure_type": "missing_checkpoint",
-            "error": str(exc),
-            "n_samples": 0,
-        }
-        exit_code = 1
-    except BenchmarkExecutionError as exc:
-        result = {
-            "benchmark": args.benchmark,
-            "status": "failed",
-            "failure_type": type(exc.cause).__name__,
-            "error": str(exc.cause),
-            "n_samples": 0,
-            "debug": {
-                "sample": {
-                    "task_id": exc.sample.get("task_id"),
-                    "question": exc.sample.get("question"),
-                },
-                "generated_output_preview": _preview_text(exc.generated_output),
-                "extracted_code_preview": _preview_text(exc.extracted_output) if exc.extracted_output is not None else None,
-            },
-        }
-        exit_code = 1
-    except Exception as exc:  # noqa: BLE001
-        result = {
-            "benchmark": args.benchmark,
-            "status": "failed",
-            "failure_type": type(exc).__name__,
-            "error": str(exc),
-            "n_samples": 0,
-        }
-        exit_code = 1
-
-    if args.out is not None:
-        write_benchmark_result(args.out, result)
-    if args.ledger is not None and exit_code == 0:
-        append_benchmark_run_record(
-            args.ledger,
-            result,
-            benchmark_name=args.benchmark,
-            variant=args.variant,
-            seed=args.seed,
-            samples_path=sample_path,
-        )
-    print(json.dumps(result, indent=2, sort_keys=True))
-    return exit_code
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import json
+import re
+import sys
+from pathlib import Path
+from typing import Any, Callable
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+LEDGER_TEMPLATE_PATH = REPO_ROOT / "artifacts" / "benchmark_ledger.template.json"
+
+from scripts.hydra_generation import build_hydra_generator
+from scripts.benchmark_datasets import resolve_benchmark_dataset as resolve_canonical_dataset
+from scripts.benchmark_suite import build_prompt, validate_sample
+
+
+def load_jsonl_samples(path: Path) -> list[dict[str, Any]]:
+    rows: list[dict[str, Any]] = []
+    for line in path.read_text(encoding="utf-8").splitlines():
+        if line.strip():
+            rows.append(json.loads(line))
+    return rows
+
+
+def _score_mbpp(samples: list[dict[str, Any]], generate_fn: Callable[[str], str]) -> float:
+    passed = 0
+    for sample in samples:
+        validate_sample("MBPP", sample)
+        code = generate_fn(build_prompt("MBPP", sample))
+        namespace: dict[str, Any] = {}
+        exec(code, namespace, namespace)
+        for test in sample["tests"]:
+            exec(test, namespace, namespace)
+        passed += 1
+    return passed / len(samples) if samples else 0.0
+
+
+def _extract_last_number(text: str) -> str | None:
+    matches = re.findall(r"-?\d+(?:\.\d+)?", text)
+    return matches[-1] if matches else None
+
+
+def _score_gsm8k(samples: list[dict[str, Any]], generate_fn: Callable[[str], str]) -> float:
+    passed = 0
+    for sample in samples:
+        validate_sample("GSM8K", sample)
+        output = generate_fn(build_prompt("GSM8K", sample))
+        pred = _extract_last_number(output)
+        if pred is not None and pred == str(sample["answer"]):
+            passed += 1
+    return passed / len(samples) if samples else 0.0
+
+
+def _score_humaneval(samples: list[dict[str, Any]], generate_fn: Callable[[str], str]) -> float:
+    passed = 0
+    for sample in samples:
+        validate_sample("HumanEval", sample)
+        code = generate_fn(build_prompt("HumanEval", sample))
+        namespace: dict[str, Any] = {}
+        exec(code, namespace, namespace)
+        exec(sample["test"], namespace, namespace)
+        passed += 1
+    return passed / len(samples) if samples else 0.0
+
+
+def _score_arc(samples: list[dict[str, Any]], generate_fn: Callable[[str], str]) -> float:
+    passed = 0
+    for sample in samples:
+        validate_sample("ARC-Challenge", sample)
+        output = generate_fn(build_prompt("ARC-Challenge", sample)).strip()
+        if output == str(sample["answer"]):
+            passed += 1
+    return passed / len(samples) if samples else 0.0
+
+
+def run_benchmark(benchmark_name: str, path: Path, generate_fn: Callable[[str], str]) -> dict[str, Any]:
+    samples = load_jsonl_samples(path)
+    if benchmark_name == "MBPP":
+        return {
+            "benchmark": "MBPP",
+            "primary_metric": "pass_at_1",
+            "score": _score_mbpp(samples, generate_fn),
+            "n_samples": len(samples),
+        }
+    if benchmark_name == "GSM8K":
+        return {
+            "benchmark": "GSM8K",
+            "primary_metric": "exact_match",
+            "score": _score_gsm8k(samples, generate_fn),
+            "n_samples": len(samples),
+        }
+    if benchmark_name == "HumanEval":
+        return {
+            "benchmark": "HumanEval",
+            "primary_metric": "pass_at_1",
+            "score": _score_humaneval(samples, generate_fn),
+            "n_samples": len(samples),
+        }
+    if benchmark_name == "ARC-Challenge":
+        return {
+            "benchmark": "ARC-Challenge",
+            "primary_metric": "accuracy",
+            "score": _score_arc(samples, generate_fn),
+            "n_samples": len(samples),
+        }
+    raise ValueError(f"Unsupported runnable benchmark: {benchmark_name}")
+
+
+def write_benchmark_result(path: Path, payload: dict[str, Any]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
+
+
+def append_benchmark_run_record(
+    ledger_path: Path,
+    result: dict[str, Any],
+    *,
+    benchmark_name: str,
+    variant: str,
+    seed: int,
+    samples_path: Path,
+) -> None:
+    if not ledger_path.exists():
+        ledger_path.parent.mkdir(parents=True, exist_ok=True)
+        ledger_path.write_text(LEDGER_TEMPLATE_PATH.read_text(encoding="utf-8"), encoding="utf-8")
+    payload = json.loads(ledger_path.read_text(encoding="utf-8"))
+    run_records = payload.setdefault("run_records", [])
+    if len(run_records) == 1 and run_records[0].get("run_id") == "example-run-0001":
+        run_records.clear()
+    run_records.append(
+        {
+            "run_id": result.get("run_id", f"{benchmark_name.lower()}-{seed}"),
+            "commit": "HEAD",
+            "model_family": "hydra",
+            "variant": variant,
+            "seed": seed,
+            "hardware": {
+                "hardware_class": payload.get("benchmark_cycle", {}).get("hardware_class", "unknown"),
+            },
+            "budget": {
+                "budget_mode": payload.get("benchmark_cycle", {}).get("budget_modes", [None])[0],
+            },
+            "capability": {
+                "coding_score": result["score"] if benchmark_name in {"MBPP", "HumanEval"} else None,
+                "reasoning_score": result["score"] if benchmark_name in {"GSM8K", "ARC-Challenge"} else None,
+            },
+            "artifacts": {
+                "samples_path": str(samples_path),
+            },
+        }
+    )
+    ledger_path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
+
+
+def resolve_samples_path(benchmark_name: str, samples: Path | None, suite_path: Path) -> Path:
+    if samples is not None:
+        return samples
+    payload = json.loads(suite_path.read_text(encoding="utf-8"))
+    for section in ("coding_benchmarks", "reasoning_benchmarks"):
+        if section not in payload:
+            continue
+        for slot in ("fast_iteration", "milestone"):
+            entry = payload[section].get(slot)
+            if isinstance(entry, dict) and entry.get("name") == benchmark_name and "sample_path" in entry:
+                return Path(entry["sample_path"])
+    try:
+        return resolve_canonical_dataset(benchmark_name, None)
+    except ValueError:
+        raise ValueError(f"No sample path found for benchmark: {benchmark_name}")
+
+
+def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run a local benchmark against JSONL samples")
+    parser.add_argument("--benchmark", required=True, choices=["MBPP", "GSM8K", "HumanEval", "ARC-Challenge"])
+    parser.add_argument("--samples", type=Path)
+    parser.add_argument("--suite", type=Path, default=REPO_ROOT / "artifacts" / "benchmark_suite.cycle1.json")
+    parser.add_argument("--out", type=Path)
+    parser.add_argument("--ledger", type=Path)
+    parser.add_argument("--variant", default="hydra_full")
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--generator-mode", choices=["stub", "hydra"], default="stub")
+    parser.add_argument("--checkpoint", type=Path)
+    parser.add_argument("--device")
+    parser.add_argument("--max-new-tokens", type=int, default=256)
+    parser.add_argument("--temperature", type=float, default=0.2)
+    parser.add_argument("--top-p", type=float, default=0.95)
+    return parser.parse_args(argv)
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = parse_args(argv)
+    sample_path = resolve_samples_path(args.benchmark, args.samples, args.suite)
+    try:
+        if args.generator_mode == "hydra":
+            generator = build_hydra_generator(
+                checkpoint_path=args.checkpoint,
+                device=args.device,
+                max_new_tokens=args.max_new_tokens,
+                temperature=args.temperature,
+                top_p=args.top_p,
+            )
+        else:
+            def generator(prompt: str) -> str:
+                return prompt
+
+        result = run_benchmark(args.benchmark, sample_path, generator)
+        exit_code = 0
+    except FileNotFoundError as exc:
+        result = {
+            "benchmark": args.benchmark,
+            "status": "failed",
+            "failure_type": "missing_checkpoint",
+            "error": str(exc),
+            "n_samples": 0,
+        }
+        exit_code = 1
+    except Exception as exc:  # noqa: BLE001
+        result = {
+            "benchmark": args.benchmark,
+            "status": "failed",
+            "failure_type": type(exc).__name__,
+            "error": str(exc),
+            "n_samples": 0,
+        }
+        exit_code = 1
+
+    if args.out is not None:
+        write_benchmark_result(args.out, result)
+    if args.ledger is not None and exit_code == 0:
+        append_benchmark_run_record(
+            args.ledger,
+            result,
+            benchmark_name=args.benchmark,
+            variant=args.variant,
+            seed=args.seed,
+            samples_path=sample_path,
+        )
+    print(json.dumps(result, indent=2, sort_keys=True))
+    return exit_code
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

overlay/scripts/benchmark_suite.py

@@ -1,84 +1,84 @@
-#!/usr/bin/env python3
-from __future__ import annotations
-
-import json
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Any
-
-
-@dataclass(frozen=True)
-class BenchmarkSpec:
-    name: str
-    family: str
-    required_fields: tuple[str, ...]
-
-
-REGISTRY: dict[str, BenchmarkSpec] = {
-    "MBPP": BenchmarkSpec("MBPP", "coding", ("task_id", "prompt", "tests")),
-    "HumanEval": BenchmarkSpec("HumanEval", "coding", ("task_id", "prompt", "test")),
-    "GSM8K": BenchmarkSpec("GSM8K", "reasoning", ("question", "answer")),
-    "ARC-Challenge": BenchmarkSpec("ARC-Challenge", "reasoning", ("question", "choices", "answer")),
-}
-
-
-def validate_sample(benchmark_name: str, sample: dict[str, Any]) -> None:
-    spec = REGISTRY[benchmark_name]
-    for field in spec.required_fields:
-        if field not in sample:
-            raise ValueError(f"{benchmark_name} sample missing required field: {field}")
-
-
-def build_prompt(benchmark_name: str, sample: dict[str, Any]) -> str:
-    validate_sample(benchmark_name, sample)
-    if benchmark_name == "MBPP":
-        tests = sample["tests"]
-        rendered_tests = "\n".join(str(t) for t in tests)
-        return (
-            "Write a Python function that solves the task below.\n\n"
-            f"Task:\n{sample['prompt']}\n\n"
-            f"Tests:\n{rendered_tests}\n"
-        )
-    if benchmark_name == "HumanEval":
-        return (
-            "Complete the following Python function exactly as specified.\n\n"
-            f"Prompt:\n{sample['prompt']}\n\n"
-            f"Reference test:\n{sample['test']}\n"
-        )
-    if benchmark_name == "GSM8K":
-        return f"Solve the following math word problem. Return only the final answer.\n\nQuestion: {sample['question']}\n"
-    if benchmark_name == "ARC-Challenge":
-        choices = sample["choices"]
-        rendered_choices = "\n".join(f"- {choice}" for choice in choices)
-        return (
-            "Answer the following multiple-choice science question. Return only the correct option text or label.\n\n"
-            f"Question: {sample['question']}\nChoices:\n{rendered_choices}\n"
-        )
-    raise ValueError(f"Unknown benchmark: {benchmark_name}")
-
-
-def load_cycle_benchmark_suite(path: Path) -> dict[str, dict[str, BenchmarkSpec]]:
-    payload = json.loads(path.read_text(encoding="utf-8"))
-    out: dict[str, dict[str, BenchmarkSpec]] = {"coding_benchmarks": {}, "reasoning_benchmarks": {}}
-    for section in ("coding_benchmarks", "reasoning_benchmarks"):
-        if section not in payload:
-            raise ValueError(f"missing benchmark section: {section}")
-        for slot in ("fast_iteration", "milestone"):
-            if slot not in payload[section]:
-                raise ValueError(f"missing benchmark slot: {section}.{slot}")
-            name = payload[section][slot]["name"]
-            if name not in REGISTRY:
-                raise ValueError(f"unsupported benchmark: {name}")
-            out[section][slot] = REGISTRY[name]
-    return out
-
-
-def main() -> int:
-    path = Path("artifacts/benchmark_suite.cycle1.json")
-    suite = load_cycle_benchmark_suite(path)
-    print(json.dumps({k: {slot: spec.name for slot, spec in section.items()} for k, section in suite.items()}, indent=2))
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+
+@dataclass(frozen=True)
+class BenchmarkSpec:
+    name: str
+    family: str
+    required_fields: tuple[str, ...]
+
+
+REGISTRY: dict[str, BenchmarkSpec] = {
+    "MBPP": BenchmarkSpec("MBPP", "coding", ("task_id", "prompt", "tests")),
+    "HumanEval": BenchmarkSpec("HumanEval", "coding", ("task_id", "prompt", "test")),
+    "GSM8K": BenchmarkSpec("GSM8K", "reasoning", ("question", "answer")),
+    "ARC-Challenge": BenchmarkSpec("ARC-Challenge", "reasoning", ("question", "choices", "answer")),
+}
+
+
+def validate_sample(benchmark_name: str, sample: dict[str, Any]) -> None:
+    spec = REGISTRY[benchmark_name]
+    for field in spec.required_fields:
+        if field not in sample:
+            raise ValueError(f"{benchmark_name} sample missing required field: {field}")
+
+
+def build_prompt(benchmark_name: str, sample: dict[str, Any]) -> str:
+    validate_sample(benchmark_name, sample)
+    if benchmark_name == "MBPP":
+        tests = sample["tests"]
+        rendered_tests = "\n".join(str(t) for t in tests)
+        return (
+            "Write a Python function that solves the task below.\n\n"
+            f"Task:\n{sample['prompt']}\n\n"
+            f"Tests:\n{rendered_tests}\n"
+        )
+    if benchmark_name == "HumanEval":
+        return (
+            "Complete the following Python function exactly as specified.\n\n"
+            f"Prompt:\n{sample['prompt']}\n\n"
+            f"Reference test:\n{sample['test']}\n"
+        )
+    if benchmark_name == "GSM8K":
+        return f"Solve the following math word problem. Return only the final answer.\n\nQuestion: {sample['question']}\n"
+    if benchmark_name == "ARC-Challenge":
+        choices = sample["choices"]
+        rendered_choices = "\n".join(f"- {choice}" for choice in choices)
+        return (
+            "Answer the following multiple-choice science question. Return only the correct option text or label.\n\n"
+            f"Question: {sample['question']}\nChoices:\n{rendered_choices}\n"
+        )
+    raise ValueError(f"Unknown benchmark: {benchmark_name}")
+
+
+def load_cycle_benchmark_suite(path: Path) -> dict[str, dict[str, BenchmarkSpec]]:
+    payload = json.loads(path.read_text(encoding="utf-8"))
+    out: dict[str, dict[str, BenchmarkSpec]] = {"coding_benchmarks": {}, "reasoning_benchmarks": {}}
+    for section in ("coding_benchmarks", "reasoning_benchmarks"):
+        if section not in payload:
+            raise ValueError(f"missing benchmark section: {section}")
+        for slot in ("fast_iteration", "milestone"):
+            if slot not in payload[section]:
+                raise ValueError(f"missing benchmark slot: {section}.{slot}")
+            name = payload[section][slot]["name"]
+            if name not in REGISTRY:
+                raise ValueError(f"unsupported benchmark: {name}")
+            out[section][slot] = REGISTRY[name]
+    return out
+
+
+def main() -> int:
+    path = Path("artifacts/benchmark_suite.cycle1.json")
+    suite = load_cycle_benchmark_suite(path)
+    print(json.dumps({k: {slot: spec.name for slot, spec in section.items()} for k, section in suite.items()}, indent=2))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

overlay/scripts/bootstrap_benchmark_env.py

@@ -1,63 +1,63 @@
-#!/usr/bin/env python3
-from __future__ import annotations
-
-import json
-import shutil
-
-import torch
-
-
-PACKAGE_MAP = {
-    "mamba_ssm": "mamba-ssm",
-    "transformers": "transformers",
-}
-
-
-def build_install_command(*, missing_dependencies: list[str]) -> list[str]:
-    packages = [PACKAGE_MAP.get(name, name) for name in missing_dependencies]
-    return [] if not packages else ["python", "-m", "pip", "install", *packages]
-
-
-def diagnose_install_blockers(
-    *,
-    missing_dependencies: list[str],
-    torch_version: str,
-    cuda_available: bool,
-    nvcc_present: bool,
-) -> list[str]:
-    blockers: list[str] = []
-    if "mamba_ssm" in missing_dependencies:
-        if "+cpu" in torch_version or not cuda_available:
-            blockers.append("mamba_ssm install likely blocked by CPU-only torch runtime")
-        if not nvcc_present:
-            blockers.append("mamba_ssm install likely blocked because nvcc is unavailable")
-    return blockers
-
-
-def build_bootstrap_report(*, missing_dependencies: list[str]) -> dict[str, object]:
-    ready = len(missing_dependencies) == 0
-    packages = [PACKAGE_MAP.get(name, name) for name in missing_dependencies]
-    install_hint = "" if ready else f"Install missing benchmark dependencies: {', '.join(packages)}"
-    blockers = diagnose_install_blockers(
-        missing_dependencies=missing_dependencies,
-        torch_version=getattr(torch, "__version__", "unknown"),
-        cuda_available=torch.cuda.is_available(),
-        nvcc_present=shutil.which("nvcc") is not None,
-    )
-    return {
-        "ready": ready,
-        "missing_dependencies": list(missing_dependencies),
-        "install_hint": install_hint,
-        "install_command": build_install_command(missing_dependencies=missing_dependencies),
-        "install_blockers": blockers,
-    }
-
-
-def main() -> int:
-    report = build_bootstrap_report(missing_dependencies=["mamba_ssm"])
-    print(json.dumps(report, indent=2, sort_keys=True))
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import json
+import shutil
+
+import torch
+
+
+PACKAGE_MAP = {
+    "mamba_ssm": "mamba-ssm",
+    "transformers": "transformers",
+}
+
+
+def build_install_command(*, missing_dependencies: list[str]) -> list[str]:
+    packages = [PACKAGE_MAP.get(name, name) for name in missing_dependencies]
+    return [] if not packages else ["python", "-m", "pip", "install", *packages]
+
+
+def diagnose_install_blockers(
+    *,
+    missing_dependencies: list[str],
+    torch_version: str,
+    cuda_available: bool,
+    nvcc_present: bool,
+) -> list[str]:
+    blockers: list[str] = []
+    if "mamba_ssm" in missing_dependencies:
+        if "+cpu" in torch_version or not cuda_available:
+            blockers.append("mamba_ssm install likely blocked by CPU-only torch runtime")
+        if not nvcc_present:
+            blockers.append("mamba_ssm install likely blocked because nvcc is unavailable")
+    return blockers
+
+
+def build_bootstrap_report(*, missing_dependencies: list[str]) -> dict[str, object]:
+    ready = len(missing_dependencies) == 0
+    packages = [PACKAGE_MAP.get(name, name) for name in missing_dependencies]
+    install_hint = "" if ready else f"Install missing benchmark dependencies: {', '.join(packages)}"
+    blockers = diagnose_install_blockers(
+        missing_dependencies=missing_dependencies,
+        torch_version=getattr(torch, "__version__", "unknown"),
+        cuda_available=torch.cuda.is_available(),
+        nvcc_present=shutil.which("nvcc") is not None,
+    )
+    return {
+        "ready": ready,
+        "missing_dependencies": list(missing_dependencies),
+        "install_hint": install_hint,
+        "install_command": build_install_command(missing_dependencies=missing_dependencies),
+        "install_blockers": blockers,
+    }
+
+
+def main() -> int:
+    report = build_bootstrap_report(missing_dependencies=["mamba_ssm"])
+    print(json.dumps(report, indent=2, sort_keys=True))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

overlay/scripts/cycle1a_report.py

@@ -1,52 +1,52 @@
-#!/usr/bin/env python3
-from __future__ import annotations
-
-import json
-from collections import defaultdict
-from pathlib import Path
-from typing import Any
-
-
-def build_cycle1a_report(run_dir: Path) -> dict[str, Any]:
-    runs = []
-    for path in sorted(run_dir.glob("*.json")):
-        try:
-            payload = json.loads(path.read_text(encoding="utf-8"))
-        except Exception:
-            continue
-        if isinstance(payload, dict) and "benchmark" in payload:
-            runs.append((path.name, payload))
-
-    n_failed = sum(1 for _, payload in runs if payload.get("status") == "failed")
-    by_benchmark: dict[str, dict[str, dict[str, float]]] = defaultdict(dict)
-    for filename, payload in runs:
-        if payload.get("status") == "failed":
-            continue
-        parts = filename.removesuffix('.json').split('_')
-        if len(parts) < 3:
-            continue
-        benchmark = payload["benchmark"]
-        variant = '_'.join(parts[1:-1])
-        score = float(payload.get("score", 0.0))
-        slot = by_benchmark.setdefault(benchmark, {}).setdefault(variant, {"scores": []})
-        slot["scores"].append(score)
-
-    for benchmark, variants in by_benchmark.items():
-        for variant, slot in variants.items():
-            scores = slot.pop("scores")
-            slot["mean_score"] = sum(scores) / len(scores)
-            slot["n_scores"] = len(scores)
-
-    if runs and n_failed == len(runs):
-        panel_status = "blocked"
-    elif n_failed > 0:
-        panel_status = "partial"
-    else:
-        panel_status = "ready"
-
-    return {
-        "n_runs": len(runs),
-        "n_failed": n_failed,
-        "panel_status": panel_status,
-        "by_benchmark": by_benchmark,
-    }
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import json
+from collections import defaultdict
+from pathlib import Path
+from typing import Any
+
+
+def build_cycle1a_report(run_dir: Path) -> dict[str, Any]:
+    runs = []
+    for path in sorted(run_dir.glob("*.json")):
+        try:
+            payload = json.loads(path.read_text(encoding="utf-8"))
+        except Exception:
+            continue
+        if isinstance(payload, dict) and "benchmark" in payload:
+            runs.append((path.name, payload))
+
+    n_failed = sum(1 for _, payload in runs if payload.get("status") == "failed")
+    by_benchmark: dict[str, dict[str, dict[str, float]]] = defaultdict(dict)
+    for filename, payload in runs:
+        if payload.get("status") == "failed":
+            continue
+        parts = filename.removesuffix('.json').split('_')
+        if len(parts) < 3:
+            continue
+        benchmark = payload["benchmark"]
+        variant = '_'.join(parts[1:-1])
+        score = float(payload.get("score", 0.0))
+        slot = by_benchmark.setdefault(benchmark, {}).setdefault(variant, {"scores": []})
+        slot["scores"].append(score)
+
+    for benchmark, variants in by_benchmark.items():
+        for variant, slot in variants.items():
+            scores = slot.pop("scores")
+            slot["mean_score"] = sum(scores) / len(scores)
+            slot["n_scores"] = len(scores)
+
+    if runs and n_failed == len(runs):
+        panel_status = "blocked"
+    elif n_failed > 0:
+        panel_status = "partial"
+    else:
+        panel_status = "ready"
+
+    return {
+        "n_runs": len(runs),
+        "n_failed": n_failed,
+        "panel_status": panel_status,
+        "by_benchmark": by_benchmark,
+    }

overlay/scripts/cycle_executor.py

@@ -1,332 +1,312 @@
-#!/usr/bin/env python3
-from __future__ import annotations
-
-import argparse
-import importlib.util
-import importlib
-import json
-import os
-import subprocess
-import sys
-from pathlib import Path
-from typing import Any
-
-from scripts.benchmark_preflight import build_readiness_report
-from scripts.hf_routing import resolve_routing
-
-
-REPO_ROOT = Path(__file__).resolve().parents[1]
-FREEZE_PATH = REPO_ROOT / "artifacts" / "cycle_1_execution_freeze.json"
-RUNNER_PATH = REPO_ROOT / "scripts" / "benchmark_runner.py"
-
-
-def active_hf_token() -> str | None:
-    token = os.environ.get("HF_TOKEN")
-    if token:
-        return token
-    try:
-        from huggingface_hub.utils import get_token
-        return get_token()
-    except Exception:
-        return None
-
-
-def missing_benchmark_dependencies() -> list[str]:
-    required = ["mamba_ssm", "transformers"]
-    missing: list[str] = []
-    for name in required:
-        try:
-            spec = importlib.util.find_spec(name)
-        except (ImportError, ValueError):
-            spec = None
-        if spec is None:
-            try:
-                importlib.import_module(name)
-            except Exception:
-                missing.append(name)
-    return missing
-
-
-def load_cycle_freeze(path: Path) -> dict[str, Any]:
-    return json.loads(path.read_text(encoding="utf-8"))
-
-
-def load_cycle_benchmarks(path: Path) -> list[str]:
-    payload = json.loads(path.read_text(encoding="utf-8"))
-    out: list[str] = []
-    for section in ("coding_benchmarks", "reasoning_benchmarks"):
-        for slot in ("fast_iteration", "milestone"):
-            entry = payload.get(section, {}).get(slot)
-            if isinstance(entry, dict) and entry.get("name"):
-                out.append(str(entry["name"]))
-    return out
-
-
-def variant_env_for_benchmark(freeze: dict[str, Any], variant: str) -> dict[str, str]:
-    variant_cfg = freeze["variants"][variant]
-    return {str(k): str(v) for k, v in variant_cfg.get("env", {}).items()}
-
-
-def decode_config_for_benchmark(freeze: dict[str, Any], benchmark: str) -> dict[str, Any]:
-    for section in ("coding_benchmarks", "reasoning_benchmarks"):
-        for slot in ("fast_iteration", "milestone"):
-            entry = freeze.get(section, {}).get(slot)
-            if isinstance(entry, dict) and entry.get("name") == benchmark:
-                return dict(entry.get("decode", {}))
-    return {}
-
-
-def build_preflight_report(
-    *,
-    cache_dir: Path,
-    output_repo: str | None = None,
-    tokenizer_repo: str | None = None,
-) -> dict[str, object]:
-    return build_readiness_report(
-        cache_dir=cache_dir,
-        hf_token_present=bool(active_hf_token()),
-        dependencies_present=not bool(missing_benchmark_dependencies()),
-        missing_dependencies=missing_benchmark_dependencies(),
-        output_repo=output_repo,
-        tokenizer_repo=tokenizer_repo,
-    )
-
-
-def write_preflight_report(path: Path, payload: dict[str, object]) -> None:
-    path.parent.mkdir(parents=True, exist_ok=True)
-    path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
-
-
-def write_cycle_summary(path: Path, payload: list[dict[str, Any]]) -> None:
-    path.parent.mkdir(parents=True, exist_ok=True)
-    path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
-
-
-def build_remote_checkpoint_report(output_repo: str, token: str | None) -> dict[str, Any]:
-    from huggingface_hub import HfApi
-
-    from scripts.benchmark_checkpoint_report import build_checkpoint_report
-
-    files = HfApi(token=token).list_repo_files(repo_id=output_repo, repo_type="model", token=token)
-    return build_checkpoint_report(files)
-
-
-def ensure_benchmark_assets(
-    *,
-    cache_dir: Path,
-    output_repo: str,
-    tokenizer_repo: str,
-    token: str | None,
-    hydrate: bool,
-) -> dict[str, str] | None:
-    if not hydrate:
-        return None
-    from scripts.benchmark_assets import hydrate_benchmark_assets
-
-    return hydrate_benchmark_assets(
-        cache_dir=cache_dir,
-        output_repo=output_repo,
-        tokenizer_repo=tokenizer_repo,
-        token=token,
-    )
-
-
-def build_benchmark_command(
-    freeze: dict[str, Any],
-    *,
-    benchmark: str,
-    variant: str,
-    seed: int,
-    out_dir: Path,
-) -> tuple[list[str], dict[str, str]]:
-    env = os.environ.copy()
-    env.update(variant_env_for_benchmark(freeze, variant))
-    env["HYDRA_SEED"] = str(seed)
-    decode_cfg = decode_config_for_benchmark(freeze, benchmark)
-
-    out_dir.mkdir(parents=True, exist_ok=True)
-    result_path = out_dir / f"{benchmark.lower()}_{variant}_seed{seed}.json"
-    ledger_path = out_dir / "benchmark_ledger.json"
-    cmd = [
-        sys.executable,
-        str(RUNNER_PATH),
-        "--benchmark",
-        benchmark,
-        "--generator-mode",
-        "hydra",
-        "--out",
-        str(result_path),
-        "--ledger",
-        str(ledger_path),
-        "--variant",
-        variant,
-        "--seed",
-        str(seed),
-        "--max-new-tokens",
-        str(int(decode_cfg.get("max_tokens", 256))),
-        "--temperature",
-        str(float(decode_cfg.get("temperature", 0.2))),
-        "--top-p",
-        str(float(decode_cfg.get("top_p", 0.95))),
-    ]
-    return cmd, env
-
-
-def build_cycle_plan(freeze: dict[str, Any], *, benchmark: str, out_dir: Path) -> list[dict[str, Any]]:
-    runnable_variants = [
-        name for name, cfg in freeze.get("variants", {}).items()
-        if isinstance(cfg, dict) and cfg.get("status") == "runnable_now"
-    ]
-    seeds = [int(seed) for seed in freeze.get("seeds", [])]
-    plan: list[dict[str, Any]] = []
-    for variant in runnable_variants:
-        for seed in seeds:
-            cmd, env = build_benchmark_command(
-                freeze,
-                benchmark=benchmark,
-                variant=variant,
-                seed=seed,
-                out_dir=out_dir,
-            )
-            plan.append({
-                "benchmark": benchmark,
-                "variant": variant,
-                "seed": seed,
-                "command": cmd,
-                "env": env,
-            })
-    return plan
-
-
-def execute_cycle_plan(plan: list[dict[str, Any]], *, repo_root: Path) -> list[dict[str, Any]]:
-    results: list[dict[str, Any]] = []
-    for item in plan:
-        proc = subprocess.run(item["command"], cwd=str(repo_root), env=item["env"])
-        results.append(
-            {
-                "benchmark": item["benchmark"],
-                "variant": item["variant"],
-                "seed": item["seed"],
-                "returncode": proc.returncode,
-            }
-        )
-    return results
-
-
-def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Execute a frozen Cycle 1 benchmark run")
-    parser.add_argument("--freeze", type=Path, default=FREEZE_PATH)
-    parser.add_argument("--suite", type=Path, default=REPO_ROOT / "artifacts" / "benchmark_suite.cycle1.json")
-    parser.add_argument("--benchmark", required=True)
-    parser.add_argument("--variant", required=True)
-    parser.add_argument("--seed", type=int, required=True)
-    parser.add_argument("--out-dir", type=Path, default=REPO_ROOT / "artifacts" / "runs")
-    parser.add_argument("--preflight-out", type=Path)
-    parser.add_argument("--summary-out", type=Path)
-    parser.add_argument("--hydrate-assets", action="store_true")
-    parser.add_argument("--all-runnable", action="store_true")
-    parser.add_argument("--all-benchmarks", action="store_true")
-    parser.add_argument("--require-ready", action="store_true")
-    parser.add_argument("--output-repo")
-    parser.add_argument("--tokenizer-repo")
-    return parser.parse_args(argv)
-
-
-def main(argv: list[str] | None = None) -> int:
-    args = parse_args(argv)
-    cache_dir = Path(os.path.expanduser("~/.cache/autoresearch"))
-    report = None
-    token = active_hf_token()
-    routing = resolve_routing(token=token)
-    output_repo = args.output_repo or routing.output_repo
-    tokenizer_repo = args.tokenizer_repo or routing.output_repo
-    if args.hydrate_assets:
-        try:
-            ensure_benchmark_assets(
-                cache_dir=cache_dir,
-                output_repo=output_repo,
-                tokenizer_repo=tokenizer_repo,
-                token=token,
-                hydrate=True,
-            )
-        except FileNotFoundError as exc:
-            checkpoint_report = None
-            try:
-                checkpoint_report = build_remote_checkpoint_report(output_repo, token)
-            except Exception:
-                checkpoint_report = None
-            if args.summary_out is not None:
-                write_cycle_summary(
-                    args.summary_out,
-                    [{
-                        "status": "blocked",
-                        "reason": "asset_hydration_failed",
-                        "error": str(exc),
-                        "checkpoint_candidates": checkpoint_report,
-                    }],
-                )
-            return 3
-    if args.preflight_out is not None:
-        report = build_preflight_report(
-            cache_dir=cache_dir,
-            output_repo=output_repo,
-            tokenizer_repo=tokenizer_repo,
-        )
-        write_preflight_report(args.preflight_out, report)
-    if args.require_ready:
-        if report is None:
-            report = build_preflight_report(
-                cache_dir=cache_dir,
-                output_repo=output_repo,
-                tokenizer_repo=tokenizer_repo,
-            )
-        if not bool(report.get("ready_for_hydra_benchmarks")):
-            checkpoint_report = None
-            try:
-                checkpoint_report = build_remote_checkpoint_report(output_repo, token)
-            except Exception:
-                checkpoint_report = None
-            if args.summary_out is not None:
-                write_cycle_summary(
-                    args.summary_out,
-                    [{
-                        "status": "blocked",
-                        "reason": "preflight_not_ready",
-                        "preflight": report,
-                        "checkpoint_candidates": checkpoint_report,
-                    }],
-                )
-            return 2
-    freeze = load_cycle_freeze(args.freeze)
-    if args.all_runnable:
-        benchmarks = load_cycle_benchmarks(args.suite) if args.all_benchmarks else [args.benchmark]
-        plan = []
-        for benchmark in benchmarks:
-            plan.extend(build_cycle_plan(freeze, benchmark=benchmark, out_dir=args.out_dir))
-        results = execute_cycle_plan(plan, repo_root=REPO_ROOT)
-        if args.summary_out is not None:
-            write_cycle_summary(args.summary_out, results)
-        return 0 if all(item["returncode"] == 0 for item in results) else 1
-    cmd, env = build_benchmark_command(
-        freeze,
-        benchmark=args.benchmark,
-        variant=args.variant,
-        seed=args.seed,
-        out_dir=args.out_dir,
-    )
-    proc = subprocess.run(cmd, cwd=str(REPO_ROOT), env=env)
-    if args.summary_out is not None:
-        write_cycle_summary(
-            args.summary_out,
-            [{
-                "benchmark": args.benchmark,
-                "variant": args.variant,
-                "seed": args.seed,
-                "returncode": proc.returncode,
-            }],
-        )
-    return proc.returncode
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import importlib.util
+import importlib
+import json
+import os
+import subprocess
+import sys
+from pathlib import Path
+from typing import Any
+
+from scripts.benchmark_preflight import build_readiness_report
+from scripts.hf_routing import resolve_routing
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+FREEZE_PATH = REPO_ROOT / "artifacts" / "cycle_1_execution_freeze.json"
+RUNNER_PATH = REPO_ROOT / "scripts" / "benchmark_runner.py"
+
+
+def active_hf_token() -> str | None:
+    token = os.environ.get("HF_TOKEN")
+    if token:
+        return token
+    try:
+        from huggingface_hub.utils import get_token
+        return get_token()
+    except Exception:
+        return None
+
+
+def missing_benchmark_dependencies() -> list[str]:
+    required = ["mamba_ssm", "transformers"]
+    missing: list[str] = []
+    for name in required:
+        try:
+            spec = importlib.util.find_spec(name)
+        except (ImportError, ValueError):
+            spec = None
+        if spec is None:
+            try:
+                importlib.import_module(name)
+            except Exception:
+                missing.append(name)
+    return missing
+
+
+def load_cycle_freeze(path: Path) -> dict[str, Any]:
+    return json.loads(path.read_text(encoding="utf-8"))
+
+
+def load_cycle_benchmarks(path: Path) -> list[str]:
+    payload = json.loads(path.read_text(encoding="utf-8"))
+    out: list[str] = []
+    for section in ("coding_benchmarks", "reasoning_benchmarks"):
+        for slot in ("fast_iteration", "milestone"):
+            entry = payload.get(section, {}).get(slot)
+            if isinstance(entry, dict) and entry.get("name"):
+                out.append(str(entry["name"]))
+    return out
+
+
+def build_preflight_report(
+    *,
+    cache_dir: Path,
+    output_repo: str | None = None,
+    tokenizer_repo: str | None = None,
+) -> dict[str, object]:
+    return build_readiness_report(
+        cache_dir=cache_dir,
+        hf_token_present=bool(active_hf_token()),
+        dependencies_present=not bool(missing_benchmark_dependencies()),
+        missing_dependencies=missing_benchmark_dependencies(),
+        output_repo=output_repo,
+        tokenizer_repo=tokenizer_repo,
+    )
+
+
+def write_preflight_report(path: Path, payload: dict[str, object]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
+
+
+def write_cycle_summary(path: Path, payload: list[dict[str, Any]]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
+
+
+def build_remote_checkpoint_report(output_repo: str, token: str | None) -> dict[str, Any]:
+    from huggingface_hub import HfApi
+
+    from scripts.benchmark_checkpoint_report import build_checkpoint_report
+
+    files = HfApi(token=token).list_repo_files(repo_id=output_repo, repo_type="model", token=token)
+    return build_checkpoint_report(files)
+
+
+def ensure_benchmark_assets(
+    *,
+    cache_dir: Path,
+    output_repo: str,
+    tokenizer_repo: str,
+    token: str | None,
+    hydrate: bool,
+) -> dict[str, str] | None:
+    if not hydrate:
+        return None
+    from scripts.benchmark_assets import hydrate_benchmark_assets
+
+    return hydrate_benchmark_assets(
+        cache_dir=cache_dir,
+        output_repo=output_repo,
+        tokenizer_repo=tokenizer_repo,
+        token=token,
+    )
+
+
+def build_benchmark_command(
+    freeze: dict[str, Any],
+    *,
+    benchmark: str,
+    variant: str,
+    seed: int,
+    out_dir: Path,
+) -> tuple[list[str], dict[str, str]]:
+    variant_cfg = freeze["variants"][variant]
+    env = os.environ.copy()
+    env.update({str(k): str(v) for k, v in variant_cfg.get("env", {}).items()})
+    env["HYDRA_SEED"] = str(seed)
+
+    out_dir.mkdir(parents=True, exist_ok=True)
+    result_path = out_dir / f"{benchmark.lower()}_{variant}_seed{seed}.json"
+    ledger_path = out_dir / "benchmark_ledger.json"
+    cmd = [
+        sys.executable,
+        str(RUNNER_PATH),
+        "--benchmark",
+        benchmark,
+        "--generator-mode",
+        "hydra",
+        "--out",
+        str(result_path),
+        "--ledger",
+        str(ledger_path),
+        "--variant",
+        variant,
+        "--seed",
+        str(seed),
+    ]
+    return cmd, env
+
+
+def build_cycle_plan(freeze: dict[str, Any], *, benchmark: str, out_dir: Path) -> list[dict[str, Any]]:
+    runnable_variants = [
+        name for name, cfg in freeze.get("variants", {}).items()
+        if isinstance(cfg, dict) and cfg.get("status") == "runnable_now"
+    ]
+    seeds = [int(seed) for seed in freeze.get("seeds", [])]
+    plan: list[dict[str, Any]] = []
+    for variant in runnable_variants:
+        for seed in seeds:
+            cmd, env = build_benchmark_command(
+                freeze,
+                benchmark=benchmark,
+                variant=variant,
+                seed=seed,
+                out_dir=out_dir,
+            )
+            plan.append({
+                "benchmark": benchmark,
+                "variant": variant,
+                "seed": seed,
+                "command": cmd,
+                "env": env,
+            })
+    return plan
+
+
+def execute_cycle_plan(plan: list[dict[str, Any]], *, repo_root: Path) -> list[dict[str, Any]]:
+    results: list[dict[str, Any]] = []
+    for item in plan:
+        proc = subprocess.run(item["command"], cwd=str(repo_root), env=item["env"])
+        results.append(
+            {
+                "benchmark": item["benchmark"],
+                "variant": item["variant"],
+                "seed": item["seed"],
+                "returncode": proc.returncode,
+            }
+        )
+    return results
+
+
+def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Execute a frozen Cycle 1 benchmark run")
+    parser.add_argument("--freeze", type=Path, default=FREEZE_PATH)
+    parser.add_argument("--suite", type=Path, default=REPO_ROOT / "artifacts" / "benchmark_suite.cycle1.json")
+    parser.add_argument("--benchmark", required=True)
+    parser.add_argument("--variant", required=True)
+    parser.add_argument("--seed", type=int, required=True)
+    parser.add_argument("--out-dir", type=Path, default=REPO_ROOT / "artifacts" / "runs")
+    parser.add_argument("--preflight-out", type=Path)
+    parser.add_argument("--summary-out", type=Path)
+    parser.add_argument("--hydrate-assets", action="store_true")
+    parser.add_argument("--all-runnable", action="store_true")
+    parser.add_argument("--all-benchmarks", action="store_true")
+    parser.add_argument("--require-ready", action="store_true")
+    parser.add_argument("--output-repo")
+    parser.add_argument("--tokenizer-repo")
+    return parser.parse_args(argv)
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = parse_args(argv)
+    cache_dir = Path(os.path.expanduser("~/.cache/autoresearch"))
+    report = None
+    token = active_hf_token()
+    routing = resolve_routing(token=token)
+    output_repo = args.output_repo or routing.output_repo
+    tokenizer_repo = args.tokenizer_repo or routing.output_repo
+    if args.hydrate_assets:
+        try:
+            ensure_benchmark_assets(
+                cache_dir=cache_dir,
+                output_repo=output_repo,
+                tokenizer_repo=tokenizer_repo,
+                token=token,
+                hydrate=True,
+            )
+        except FileNotFoundError as exc:
+            checkpoint_report = None
+            try:
+                checkpoint_report = build_remote_checkpoint_report(output_repo, token)
+            except Exception:
+                checkpoint_report = None
+            if args.summary_out is not None:
+                write_cycle_summary(
+                    args.summary_out,
+                    [{
+                        "status": "blocked",
+                        "reason": "asset_hydration_failed",
+                        "error": str(exc),
+                        "checkpoint_candidates": checkpoint_report,
+                    }],
+                )
+            return 3
+    if args.preflight_out is not None:
+        report = build_preflight_report(
+            cache_dir=cache_dir,
+            output_repo=output_repo,
+            tokenizer_repo=tokenizer_repo,
+        )
+        write_preflight_report(args.preflight_out, report)
+    if args.require_ready:
+        if report is None:
+            report = build_preflight_report(
+                cache_dir=cache_dir,
+                output_repo=output_repo,
+                tokenizer_repo=tokenizer_repo,
+            )
+        if not bool(report.get("ready_for_hydra_benchmarks")):
+            checkpoint_report = None
+            try:
+                checkpoint_report = build_remote_checkpoint_report(output_repo, token)
+            except Exception:
+                checkpoint_report = None
+            if args.summary_out is not None:
+                write_cycle_summary(
+                    args.summary_out,
+                    [{
+                        "status": "blocked",
+                        "reason": "preflight_not_ready",
+                        "preflight": report,
+                        "checkpoint_candidates": checkpoint_report,
+                    }],
+                )
+            return 2
+    freeze = load_cycle_freeze(args.freeze)
+    if args.all_runnable:
+        benchmarks = load_cycle_benchmarks(args.suite) if args.all_benchmarks else [args.benchmark]
+        plan = []
+        for benchmark in benchmarks:
+            plan.extend(build_cycle_plan(freeze, benchmark=benchmark, out_dir=args.out_dir))
+        results = execute_cycle_plan(plan, repo_root=REPO_ROOT)
+        if args.summary_out is not None:
+            write_cycle_summary(args.summary_out, results)
+        return 0 if all(item["returncode"] == 0 for item in results) else 1
+    cmd, env = build_benchmark_command(
+        freeze,
+        benchmark=args.benchmark,
+        variant=args.variant,
+        seed=args.seed,
+        out_dir=args.out_dir,
+    )
+    proc = subprocess.run(cmd, cwd=str(REPO_ROOT), env=env)
+    if args.summary_out is not None:
+        write_cycle_summary(
+            args.summary_out,
+            [{
+                "benchmark": args.benchmark,
+                "variant": args.variant,
+                "seed": args.seed,
+                "returncode": proc.returncode,
+            }],
+        )
+    return proc.returncode
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

overlay/scripts/export_hpo_priors.py

@@ -1,94 +1,94 @@
-#!/usr/bin/env python3
-from __future__ import annotations
-
-import argparse
-import datetime as dt
-import json
-from pathlib import Path
-from typing import Any
-
-import optuna
-
-from scripts.hpo_leaderboard import build_leaderboard
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Export top Optuna trials as transfer-learning priors")
-    parser.add_argument("--study-name", action="append", default=[], help="Repeat to merge multiple studies")
-    parser.add_argument("--storage", default="sqlite:///optuna_hpo.db")
-    parser.add_argument("--top-k", type=int, default=20)
-    parser.add_argument("--out", type=Path, default=Path("docs") / "hpo_transfer_priors.json")
-    parser.add_argument("--metric", default="val_bpb")
-    return parser.parse_args()
-
-
-def _completed_trials(study: optuna.Study) -> list[optuna.trial.FrozenTrial]:
-    trials = [t for t in study.trials if t.value is not None]
-    reverse = study.direction == optuna.study.StudyDirection.MAXIMIZE
-    return sorted(trials, key=lambda t: float(t.value), reverse=reverse)
-
-
-def _serialize_trial(trial: optuna.trial.FrozenTrial) -> dict[str, Any]:
-    return {
-        "trial_number": trial.number,
-        "value": float(trial.value) if trial.value is not None else None,
-        "params": dict(trial.params),
-        "user_attrs": dict(trial.user_attrs),
-    }
-
-
-def collect_prior_trials(*, storage: str, study_names: list[str], top_k: int, metric: str) -> dict[str, Any]:
-    leaderboard = build_leaderboard(storage=storage, study_names=study_names, metric=metric)
-    selected = leaderboard["clean_trials"][: max(0, top_k)]
-    trials = [
-        {
-            "study_name": row["study_name"],
-            "trial_number": row["trial_number"],
-            "value": row["value"],
-            "params": row["params"],
-            "user_attrs": row["user_attrs"],
-        }
-        for row in selected
-    ]
-    quarantined = [
-        {
-            "study_name": row["study_name"],
-            "trial_number": row["trial_number"],
-            "value": row["value"],
-            "params": row["params"],
-            "user_attrs": row["user_attrs"],
-            "contamination_reason": row["contamination_reason"],
-        }
-        for row in leaderboard["contaminated_trials"]
-    ]
-    return {
-        "schema_version": 2,
-        "generated_at": dt.datetime.now(dt.UTC).isoformat(timespec="seconds"),
-        "study_names": study_names,
-        "metric": metric,
-        "n_total_trials": sum(int(s["n_trials"]) for s in leaderboard["studies"]),
-        "n_completed_trials": sum(int(s["n_completed"]) for s in leaderboard["studies"]),
-        "n_exported_trials": len(trials),
-        "n_quarantined_trials": len(quarantined),
-        "top_k": top_k,
-        "trials": trials,
-        "quarantined_trials": quarantined,
-    }
-
-
-def main() -> int:
-    args = parse_args()
-    study_names = args.study_name or ["hydra_hpo"]
-    payload = collect_prior_trials(storage=args.storage, study_names=study_names, top_k=args.top_k, metric=args.metric)
-
-    args.out.parent.mkdir(parents=True, exist_ok=True)
-    args.out.write_text(json.dumps(payload, indent=2), encoding="utf-8")
-    print(
-        f"[hpo-priors] wrote {args.out} with {payload['n_exported_trials']} clean trials "
-        f"({payload['n_quarantined_trials']} quarantined)"
-    )
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import datetime as dt
+import json
+from pathlib import Path
+from typing import Any
+
+import optuna
+
+from scripts.hpo_leaderboard import build_leaderboard
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Export top Optuna trials as transfer-learning priors")
+    parser.add_argument("--study-name", action="append", default=[], help="Repeat to merge multiple studies")
+    parser.add_argument("--storage", default="sqlite:///optuna_hpo.db")
+    parser.add_argument("--top-k", type=int, default=20)
+    parser.add_argument("--out", type=Path, default=Path("docs") / "hpo_transfer_priors.json")
+    parser.add_argument("--metric", default="val_bpb")
+    return parser.parse_args()
+
+
+def _completed_trials(study: optuna.Study) -> list[optuna.trial.FrozenTrial]:
+    trials = [t for t in study.trials if t.value is not None]
+    reverse = study.direction == optuna.study.StudyDirection.MAXIMIZE
+    return sorted(trials, key=lambda t: float(t.value), reverse=reverse)
+
+
+def _serialize_trial(trial: optuna.trial.FrozenTrial) -> dict[str, Any]:
+    return {
+        "trial_number": trial.number,
+        "value": float(trial.value) if trial.value is not None else None,
+        "params": dict(trial.params),
+        "user_attrs": dict(trial.user_attrs),
+    }
+
+
+def collect_prior_trials(*, storage: str, study_names: list[str], top_k: int, metric: str) -> dict[str, Any]:
+    leaderboard = build_leaderboard(storage=storage, study_names=study_names, metric=metric)
+    selected = leaderboard["clean_trials"][: max(0, top_k)]
+    trials = [
+        {
+            "study_name": row["study_name"],
+            "trial_number": row["trial_number"],
+            "value": row["value"],
+            "params": row["params"],
+            "user_attrs": row["user_attrs"],
+        }
+        for row in selected
+    ]
+    quarantined = [
+        {
+            "study_name": row["study_name"],
+            "trial_number": row["trial_number"],
+            "value": row["value"],
+            "params": row["params"],
+            "user_attrs": row["user_attrs"],
+            "contamination_reason": row["contamination_reason"],
+        }
+        for row in leaderboard["contaminated_trials"]
+    ]
+    return {
+        "schema_version": 2,
+        "generated_at": dt.datetime.now(dt.UTC).isoformat(timespec="seconds"),
+        "study_names": study_names,
+        "metric": metric,
+        "n_total_trials": sum(int(s["n_trials"]) for s in leaderboard["studies"]),
+        "n_completed_trials": sum(int(s["n_completed"]) for s in leaderboard["studies"]),
+        "n_exported_trials": len(trials),
+        "n_quarantined_trials": len(quarantined),
+        "top_k": top_k,
+        "trials": trials,
+        "quarantined_trials": quarantined,
+    }
+
+
+def main() -> int:
+    args = parse_args()
+    study_names = args.study_name or ["hydra_hpo"]
+    payload = collect_prior_trials(storage=args.storage, study_names=study_names, top_k=args.top_k, metric=args.metric)
+
+    args.out.parent.mkdir(parents=True, exist_ok=True)
+    args.out.write_text(json.dumps(payload, indent=2), encoding="utf-8")
+    print(
+        f"[hpo-priors] wrote {args.out} with {payload['n_exported_trials']} clean trials "
+        f"({payload['n_quarantined_trials']} quarantined)"
+    )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

overlay/scripts/hf_routing.py

@@ -1,94 +1,94 @@
-from __future__ import annotations
-
-import os
-from dataclasses import dataclass
-
-from huggingface_hub import HfApi
-
-
-_OWNER_ALIASES = {
-    'jack': 'jackoatmon',
-    'jackoatmon': 'jackoatmon',
-    'icarus': 'icarus112',
-    'icarus112': 'icarus112',
-}
-
-
-def _normalize_owner(value: str | None) -> str | None:
-    if not value:
-        return None
-    normalized = value.strip().lower().lstrip('@')
-    if not normalized:
-        return None
-    return _OWNER_ALIASES.get(normalized, normalized)
-
-
-def _owner_from_env() -> str | None:
-    for key in (
-        'FEATHER_HF_OWNER',
-        'FEATHER_HF_NAMESPACE_OWNER',
-        'FEATHER_HF_PROFILE',
-        'FEATHER_HF_NAMESPACE',
-    ):
-        owner = _normalize_owner(os.environ.get(key))
-        if owner:
-            return owner
-    return None
-
-
-def resolve_owner(token: str | None = None) -> str:
-    """Resolve active HF owner in a collaborator-safe way.
-
-    Resolution precedence:
-      1) explicit env owner override (FEATHER_HF_OWNER/...)
-      2) Hugging Face `whoami` from HF_TOKEN (unless disabled)
-      3) default to jackoatmon
-    """
-    owner = _owner_from_env()
-    if owner:
-        return owner
-
-    if os.environ.get('FEATHER_HF_DISABLE_WHOAMI', '0') != '1':
-        active_token = token or os.environ.get('HF_TOKEN')
-        if active_token:
-            try:
-                info = HfApi(token=active_token).whoami(token=active_token)
-                if isinstance(info, dict):
-                    whoami_owner = _normalize_owner(info.get('name'))
-                    if whoami_owner:
-                        return whoami_owner
-            except Exception:
-                # We intentionally fail-open to deterministic defaults.
-                pass
-
-    return 'jackoatmon'
-
-
-@dataclass(frozen=True)
-class HfRouting:
-    owner: str
-    space_repo: str
-    output_repo: str
-    retina_cache_repo: str
-    job_namespace: str
-
-
-def resolve_routing(token: str | None = None) -> HfRouting:
-    owner = resolve_owner(token=token)
-
-    space_name = os.environ.get('FEATHER_HF_SPACE_NAME', 'feather-a10-runtime')
-    output_name = os.environ.get('FEATHER_HF_OUTPUT_REPO_NAME', 'feather-pretrain-checkpoints')
-    retina_name = os.environ.get('FEATHER_HF_RETINA_REPO_NAME', 'feather-retina-cache')
-
-    space_repo = os.environ.get('FEATHER_HF_SPACE_REPO') or f'{owner}/{space_name}'
-    output_repo = os.environ.get('FEATHER_HF_OUTPUT_REPO') or f'{owner}/{output_name}'
-    retina_cache_repo = os.environ.get('FEATHER_HF_RETINA_CACHE_REPO') or f'{owner}/{retina_name}'
-    job_namespace = os.environ.get('FEATHER_HF_JOB_NAMESPACE') or owner
-
-    return HfRouting(
-        owner=owner,
-        space_repo=space_repo,
-        output_repo=output_repo,
-        retina_cache_repo=retina_cache_repo,
-        job_namespace=job_namespace,
-    )
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+
+from huggingface_hub import HfApi
+
+
+_OWNER_ALIASES = {
+    'jack': 'jackoatmon',
+    'jackoatmon': 'jackoatmon',
+    'icarus': 'icarus112',
+    'icarus112': 'icarus112',
+}
+
+
+def _normalize_owner(value: str | None) -> str | None:
+    if not value:
+        return None
+    normalized = value.strip().lower().lstrip('@')
+    if not normalized:
+        return None
+    return _OWNER_ALIASES.get(normalized, normalized)
+
+
+def _owner_from_env() -> str | None:
+    for key in (
+        'FEATHER_HF_OWNER',
+        'FEATHER_HF_NAMESPACE_OWNER',
+        'FEATHER_HF_PROFILE',
+        'FEATHER_HF_NAMESPACE',
+    ):
+        owner = _normalize_owner(os.environ.get(key))
+        if owner:
+            return owner
+    return None
+
+
+def resolve_owner(token: str | None = None) -> str:
+    """Resolve active HF owner in a collaborator-safe way.
+
+    Resolution precedence:
+      1) explicit env owner override (FEATHER_HF_OWNER/...)
+      2) Hugging Face `whoami` from HF_TOKEN (unless disabled)
+      3) default to jackoatmon
+    """
+    owner = _owner_from_env()
+    if owner:
+        return owner
+
+    if os.environ.get('FEATHER_HF_DISABLE_WHOAMI', '0') != '1':
+        active_token = token or os.environ.get('HF_TOKEN')
+        if active_token:
+            try:
+                info = HfApi(token=active_token).whoami(token=active_token)
+                if isinstance(info, dict):
+                    whoami_owner = _normalize_owner(info.get('name'))
+                    if whoami_owner:
+                        return whoami_owner
+            except Exception:
+                # We intentionally fail-open to deterministic defaults.
+                pass
+
+    return 'jackoatmon'
+
+
+@dataclass(frozen=True)
+class HfRouting:
+    owner: str
+    space_repo: str
+    output_repo: str
+    retina_cache_repo: str
+    job_namespace: str
+
+
+def resolve_routing(token: str | None = None) -> HfRouting:
+    owner = resolve_owner(token=token)
+
+    space_name = os.environ.get('FEATHER_HF_SPACE_NAME', 'feather-a10-runtime')
+    output_name = os.environ.get('FEATHER_HF_OUTPUT_REPO_NAME', 'feather-pretrain-checkpoints')
+    retina_name = os.environ.get('FEATHER_HF_RETINA_REPO_NAME', 'feather-retina-cache')
+
+    space_repo = os.environ.get('FEATHER_HF_SPACE_REPO') or f'{owner}/{space_name}'
+    output_repo = os.environ.get('FEATHER_HF_OUTPUT_REPO') or f'{owner}/{output_name}'
+    retina_cache_repo = os.environ.get('FEATHER_HF_RETINA_CACHE_REPO') or f'{owner}/{retina_name}'
+    job_namespace = os.environ.get('FEATHER_HF_JOB_NAMESPACE') or owner
+
+    return HfRouting(
+        owner=owner,
+        space_repo=space_repo,
+        output_repo=output_repo,
+        retina_cache_repo=retina_cache_repo,
+        job_namespace=job_namespace,
+    )

overlay/scripts/hpo_component_report.py

@@ -1,130 +1,130 @@
-#!/usr/bin/env python3
-from __future__ import annotations
-
-import argparse
-import datetime as dt
-import json
-import math
-from collections import defaultdict
-from pathlib import Path
-from typing import Any
-
-from scripts.hpo_leaderboard import build_leaderboard
-
-
-_COMPONENT_KEYS = [
-    "engram_subsample",
-    "htm_subsample",
-    "htm_learn_every",
-    "engram_n_columns",
-    "engram_layer_idx",
-    "sdr_target_active",
-    "mamba3_chunk",
-    "dropout",
-    "hyena_layers",
-]
-
-
-def _recover_params(row: dict[str, Any]) -> dict[str, Any]:
-    params = dict(row.get("params") or {})
-    attrs = row.get("user_attrs") or {}
-    for key, value in attrs.items():
-        if key.startswith("param_"):
-            params.setdefault(key.removeprefix("param_"), value)
-    return params
-
-
-def _pearson(xs: list[float], ys: list[float]) -> float | None:
-    if len(xs) < 2 or len(xs) != len(ys):
-        return None
-    mean_x = sum(xs) / len(xs)
-    mean_y = sum(ys) / len(ys)
-    cov = sum((x - mean_x) * (y - mean_y) for x, y in zip(xs, ys))
-    var_x = sum((x - mean_x) ** 2 for x in xs)
-    var_y = sum((y - mean_y) ** 2 for y in ys)
-    if var_x <= 0 or var_y <= 0:
-        return None
-    return cov / math.sqrt(var_x * var_y)
-
-
-def build_component_report(*, storage: str, study_names: list[str], metric: str = "val_bpb") -> dict[str, Any]:
-    leaderboard = build_leaderboard(storage=storage, study_names=study_names, metric=metric)
-    clean_trials = leaderboard["clean_trials"]
-
-    ablations: dict[str, list[dict[str, Any]]] = {}
-    numeric_correlations: list[dict[str, Any]] = []
-
-    for key in _COMPONENT_KEYS:
-        grouped: dict[str, list[dict[str, Any]]] = defaultdict(list)
-        numeric_x: list[float] = []
-        metric_y: list[float] = []
-        tps_y: list[float] = []
-        for row in clean_trials:
-            params = _recover_params(row)
-            if key not in params:
-                continue
-            value = params[key]
-            grouped[str(value)].append({"value": value, "metric": float(row["value"]), "tps": row.get("tps")})
-            if isinstance(value, (int, float)) and isinstance(row.get("tps"), (int, float)):
-                numeric_x.append(float(value))
-                metric_y.append(float(row["value"]))
-                tps_y.append(float(row["tps"]))
-
-        rows: list[dict[str, Any]] = []
-        for grouped_rows in grouped.values():
-            value = grouped_rows[0]["value"]
-            metric_vals = [r["metric"] for r in grouped_rows]
-            tps_vals = [float(r["tps"]) for r in grouped_rows if isinstance(r["tps"], (int, float))]
-            rows.append({
-                "value": value,
-                "n_trials": len(grouped_rows),
-                "mean_metric": sum(metric_vals) / len(metric_vals),
-                "mean_tps": (sum(tps_vals) / len(tps_vals)) if tps_vals else None,
-            })
-        if rows:
-            rows.sort(key=lambda row: str(row["value"]))
-            ablations[key] = rows
-
-        pearson_metric = _pearson(numeric_x, metric_y)
-        pearson_tps = _pearson(numeric_x, tps_y)
-        if pearson_metric is not None or pearson_tps is not None:
-            numeric_correlations.append({
-                "param": key,
-                "pearson_with_metric": pearson_metric,
-                "pearson_with_tps": pearson_tps,
-                "n_points": len(numeric_x),
-            })
-
-    numeric_correlations.sort(key=lambda row: row["param"])
-    return {
-        "schema_version": 1,
-        "generated_at": dt.datetime.now(dt.UTC).isoformat(timespec="seconds"),
-        "metric": metric,
-        "study_names": study_names,
-        "n_clean_trials": len(clean_trials),
-        "component_ablations": ablations,
-        "numeric_correlations": numeric_correlations,
-    }
-
-
-def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Build component ablation and correlation report from clean HPO trials")
-    parser.add_argument("--storage", default="sqlite:///optuna_hpo.db")
-    parser.add_argument("--study-name", action="append", default=[])
-    parser.add_argument("--metric", default="val_bpb")
-    parser.add_argument("--out", type=Path, default=Path(".tmp") / "optuna" / "component_report.json")
-    return parser.parse_args(argv)
-
-
-def main(argv: list[str] | None = None) -> int:
-    args = parse_args(argv)
-    study_names = args.study_name or ["hydra_hpo"]
-    payload = build_component_report(storage=args.storage, study_names=study_names, metric=args.metric)
-    args.out.parent.mkdir(parents=True, exist_ok=True)
-    args.out.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
-    print(json.dumps(payload, indent=2, sort_keys=True))
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import datetime as dt
+import json
+import math
+from collections import defaultdict
+from pathlib import Path
+from typing import Any
+
+from scripts.hpo_leaderboard import build_leaderboard
+
+
+_COMPONENT_KEYS = [
+    "engram_subsample",
+    "htm_subsample",
+    "htm_learn_every",
+    "engram_n_columns",
+    "engram_layer_idx",
+    "sdr_target_active",
+    "mamba3_chunk",
+    "dropout",
+    "hyena_layers",
+]
+
+
+def _recover_params(row: dict[str, Any]) -> dict[str, Any]:
+    params = dict(row.get("params") or {})
+    attrs = row.get("user_attrs") or {}
+    for key, value in attrs.items():
+        if key.startswith("param_"):
+            params.setdefault(key.removeprefix("param_"), value)
+    return params
+
+
+def _pearson(xs: list[float], ys: list[float]) -> float | None:
+    if len(xs) < 2 or len(xs) != len(ys):
+        return None
+    mean_x = sum(xs) / len(xs)
+    mean_y = sum(ys) / len(ys)
+    cov = sum((x - mean_x) * (y - mean_y) for x, y in zip(xs, ys))
+    var_x = sum((x - mean_x) ** 2 for x in xs)
+    var_y = sum((y - mean_y) ** 2 for y in ys)
+    if var_x <= 0 or var_y <= 0:
+        return None
+    return cov / math.sqrt(var_x * var_y)
+
+
+def build_component_report(*, storage: str, study_names: list[str], metric: str = "val_bpb") -> dict[str, Any]:
+    leaderboard = build_leaderboard(storage=storage, study_names=study_names, metric=metric)
+    clean_trials = leaderboard["clean_trials"]
+
+    ablations: dict[str, list[dict[str, Any]]] = {}
+    numeric_correlations: list[dict[str, Any]] = []
+
+    for key in _COMPONENT_KEYS:
+        grouped: dict[str, list[dict[str, Any]]] = defaultdict(list)
+        numeric_x: list[float] = []
+        metric_y: list[float] = []
+        tps_y: list[float] = []
+        for row in clean_trials:
+            params = _recover_params(row)
+            if key not in params:
+                continue
+            value = params[key]
+            grouped[str(value)].append({"value": value, "metric": float(row["value"]), "tps": row.get("tps")})
+            if isinstance(value, (int, float)) and isinstance(row.get("tps"), (int, float)):
+                numeric_x.append(float(value))
+                metric_y.append(float(row["value"]))
+                tps_y.append(float(row["tps"]))
+
+        rows: list[dict[str, Any]] = []
+        for grouped_rows in grouped.values():
+            value = grouped_rows[0]["value"]
+            metric_vals = [r["metric"] for r in grouped_rows]
+            tps_vals = [float(r["tps"]) for r in grouped_rows if isinstance(r["tps"], (int, float))]
+            rows.append({
+                "value": value,
+                "n_trials": len(grouped_rows),
+                "mean_metric": sum(metric_vals) / len(metric_vals),
+                "mean_tps": (sum(tps_vals) / len(tps_vals)) if tps_vals else None,
+            })
+        if rows:
+            rows.sort(key=lambda row: str(row["value"]))
+            ablations[key] = rows
+
+        pearson_metric = _pearson(numeric_x, metric_y)
+        pearson_tps = _pearson(numeric_x, tps_y)
+        if pearson_metric is not None or pearson_tps is not None:
+            numeric_correlations.append({
+                "param": key,
+                "pearson_with_metric": pearson_metric,
+                "pearson_with_tps": pearson_tps,
+                "n_points": len(numeric_x),
+            })
+
+    numeric_correlations.sort(key=lambda row: row["param"])
+    return {
+        "schema_version": 1,
+        "generated_at": dt.datetime.now(dt.UTC).isoformat(timespec="seconds"),
+        "metric": metric,
+        "study_names": study_names,
+        "n_clean_trials": len(clean_trials),
+        "component_ablations": ablations,
+        "numeric_correlations": numeric_correlations,
+    }
+
+
+def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Build component ablation and correlation report from clean HPO trials")
+    parser.add_argument("--storage", default="sqlite:///optuna_hpo.db")
+    parser.add_argument("--study-name", action="append", default=[])
+    parser.add_argument("--metric", default="val_bpb")
+    parser.add_argument("--out", type=Path, default=Path(".tmp") / "optuna" / "component_report.json")
+    return parser.parse_args(argv)
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = parse_args(argv)
+    study_names = args.study_name or ["hydra_hpo"]
+    payload = build_component_report(storage=args.storage, study_names=study_names, metric=args.metric)
+    args.out.parent.mkdir(parents=True, exist_ok=True)
+    args.out.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
+    print(json.dumps(payload, indent=2, sort_keys=True))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

overlay/scripts/hpo_leaderboard.py

@@ -1,156 +1,156 @@
-#!/usr/bin/env python3
-from __future__ import annotations
-
-import argparse
-import datetime as dt
-import json
-from pathlib import Path
-from typing import Any
-
-import optuna
-
-
-def _trial_direction(study: optuna.Study) -> str:
-    return "maximize" if study.direction == optuna.study.StudyDirection.MAXIMIZE else "minimize"
-
-
-def _contamination_reason(trial: optuna.trial.FrozenTrial, metric: str) -> str | None:
-    if trial.value is None:
-        return "missing_value"
-    attrs = trial.user_attrs
-    source = attrs.get("objective_source")
-    eval_status = attrs.get("eval_status")
-    objective_metric = attrs.get("objective_metric")
-
-    if source in {"train_log_fallback", "missing_metric", "missing_metrics", "missing_final_val"}:
-        return f"objective_source={source}"
-    if eval_status not in {None, "completed"}:
-        return f"eval_status={eval_status}"
-    if objective_metric not in {None, metric}:
-        return f"objective_metric={objective_metric}"
-    return None
-
-
-def _serialize_trial(study_name: str, trial: optuna.trial.FrozenTrial, metric: str) -> dict[str, Any]:
-    attrs = dict(trial.user_attrs)
-    source = attrs.get("objective_source") or "legacy_completed_value"
-    row = {
-        "study_name": study_name,
-        "trial_number": trial.number,
-        "value": float(trial.value) if trial.value is not None else None,
-        "metric": metric,
-        "objective_source": source,
-        "objective_metric": attrs.get("objective_metric", metric),
-        "eval_status": attrs.get("eval_status"),
-        "hf_job_id": attrs.get("hf_job_id"),
-        "tps": attrs.get("tps"),
-        "params": dict(trial.params),
-        "user_attrs": attrs,
-    }
-    reason = _contamination_reason(trial, metric)
-    if reason is not None:
-        row["contamination_reason"] = reason
-    return row
-
-
-def _is_pareto_dominated(candidate: dict[str, Any], peers: list[dict[str, Any]]) -> bool:
-    candidate_value = float(candidate["value"])
-    candidate_tps = float(candidate["tps"])
-    for peer in peers:
-        if peer is candidate or peer.get("tps") is None:
-            continue
-        peer_value = float(peer["value"])
-        peer_tps = float(peer["tps"])
-        no_worse = peer_value <= candidate_value and peer_tps >= candidate_tps
-        strictly_better = peer_value < candidate_value or peer_tps > candidate_tps
-        if no_worse and strictly_better:
-            return True
-    return False
-
-
-def _annotate_pareto(clean_trials: list[dict[str, Any]]) -> list[dict[str, Any]]:
-    pareto_trials: list[dict[str, Any]] = []
-    comparable = [row for row in clean_trials if row.get("tps") is not None]
-    for row in clean_trials:
-        if row.get("tps") is None:
-            row["pareto_frontier"] = False
-            row["pareto_dominated"] = None
-            row["pareto_reason"] = "missing_tps"
-            continue
-        dominated = _is_pareto_dominated(row, comparable)
-        row["pareto_frontier"] = not dominated
-        row["pareto_dominated"] = dominated
-        row["pareto_reason"] = "frontier" if not dominated else "dominated"
-        if not dominated:
-            pareto_trials.append(row)
-    pareto_trials.sort(key=lambda row: (float(row["value"]), -float(row["tps"])))
-    return pareto_trials
-
-
-def build_leaderboard(*, storage: str, study_names: list[str], metric: str = "val_bpb") -> dict[str, Any]:
-    clean_trials: list[dict[str, Any]] = []
-    contaminated_trials: list[dict[str, Any]] = []
-    study_summaries: list[dict[str, Any]] = []
-    direction = "minimize"
-
-    for study_name in study_names:
-        study = optuna.load_study(study_name=study_name, storage=storage)
-        direction = _trial_direction(study)
-        completed = [t for t in study.trials if t.value is not None]
-        study_summaries.append({
-            "study_name": study_name,
-            "direction": direction,
-            "n_trials": len(study.trials),
-            "n_completed": len(completed),
-        })
-        for trial in completed:
-            row = _serialize_trial(study_name, trial, metric)
-            if "contamination_reason" in row:
-                contaminated_trials.append(row)
-            else:
-                clean_trials.append(row)
-
-    reverse = direction == "maximize"
-    clean_trials.sort(key=lambda row: float(row["value"]), reverse=reverse)
-    contaminated_trials.sort(key=lambda row: float(row["value"]), reverse=reverse)
-    pareto_trials = _annotate_pareto(clean_trials)
-
-    return {
-        "schema_version": 1,
-        "generated_at": dt.datetime.now(dt.UTC).isoformat(timespec="seconds"),
-        "metric": metric,
-        "direction": direction,
-        "study_names": study_names,
-        "studies": study_summaries,
-        "n_clean_trials": len(clean_trials),
-        "n_contaminated_trials": len(contaminated_trials),
-        "pareto_metric_x": metric,
-        "pareto_metric_y": "tps",
-        "n_pareto_trials": len(pareto_trials),
-        "clean_trials": clean_trials,
-        "contaminated_trials": contaminated_trials,
-        "pareto_trials": pareto_trials,
-    }
-
-
-def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Build a clean Optuna HPO leaderboard")
-    parser.add_argument("--storage", default="sqlite:///optuna_hpo.db")
-    parser.add_argument("--study-name", action="append", default=[], help="Repeat to merge multiple studies")
-    parser.add_argument("--metric", default="val_bpb")
-    parser.add_argument("--out", type=Path, default=Path(".tmp") / "optuna" / "leaderboard.json")
-    return parser.parse_args(argv)
-
-
-def main(argv: list[str] | None = None) -> int:
-    args = parse_args(argv)
-    study_names = args.study_name or ["hydra_hpo"]
-    payload = build_leaderboard(storage=args.storage, study_names=study_names, metric=args.metric)
-    args.out.parent.mkdir(parents=True, exist_ok=True)
-    args.out.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
-    print(json.dumps(payload, indent=2, sort_keys=True))
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import datetime as dt
+import json
+from pathlib import Path
+from typing import Any
+
+import optuna
+
+
+def _trial_direction(study: optuna.Study) -> str:
+    return "maximize" if study.direction == optuna.study.StudyDirection.MAXIMIZE else "minimize"
+
+
+def _contamination_reason(trial: optuna.trial.FrozenTrial, metric: str) -> str | None:
+    if trial.value is None:
+        return "missing_value"
+    attrs = trial.user_attrs
+    source = attrs.get("objective_source")
+    eval_status = attrs.get("eval_status")
+    objective_metric = attrs.get("objective_metric")
+
+    if source in {"train_log_fallback", "missing_metric", "missing_metrics", "missing_final_val"}:
+        return f"objective_source={source}"
+    if eval_status not in {None, "completed"}:
+        return f"eval_status={eval_status}"
+    if objective_metric not in {None, metric}:
+        return f"objective_metric={objective_metric}"
+    return None
+
+
+def _serialize_trial(study_name: str, trial: optuna.trial.FrozenTrial, metric: str) -> dict[str, Any]:
+    attrs = dict(trial.user_attrs)
+    source = attrs.get("objective_source") or "legacy_completed_value"
+    row = {
+        "study_name": study_name,
+        "trial_number": trial.number,
+        "value": float(trial.value) if trial.value is not None else None,
+        "metric": metric,
+        "objective_source": source,
+        "objective_metric": attrs.get("objective_metric", metric),
+        "eval_status": attrs.get("eval_status"),
+        "hf_job_id": attrs.get("hf_job_id"),
+        "tps": attrs.get("tps"),
+        "params": dict(trial.params),
+        "user_attrs": attrs,
+    }
+    reason = _contamination_reason(trial, metric)
+    if reason is not None:
+        row["contamination_reason"] = reason
+    return row
+
+
+def _is_pareto_dominated(candidate: dict[str, Any], peers: list[dict[str, Any]]) -> bool:
+    candidate_value = float(candidate["value"])
+    candidate_tps = float(candidate["tps"])
+    for peer in peers:
+        if peer is candidate or peer.get("tps") is None:
+            continue
+        peer_value = float(peer["value"])
+        peer_tps = float(peer["tps"])
+        no_worse = peer_value <= candidate_value and peer_tps >= candidate_tps
+        strictly_better = peer_value < candidate_value or peer_tps > candidate_tps
+        if no_worse and strictly_better:
+            return True
+    return False
+
+
+def _annotate_pareto(clean_trials: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    pareto_trials: list[dict[str, Any]] = []
+    comparable = [row for row in clean_trials if row.get("tps") is not None]
+    for row in clean_trials:
+        if row.get("tps") is None:
+            row["pareto_frontier"] = False
+            row["pareto_dominated"] = None
+            row["pareto_reason"] = "missing_tps"
+            continue
+        dominated = _is_pareto_dominated(row, comparable)
+        row["pareto_frontier"] = not dominated
+        row["pareto_dominated"] = dominated
+        row["pareto_reason"] = "frontier" if not dominated else "dominated"
+        if not dominated:
+            pareto_trials.append(row)
+    pareto_trials.sort(key=lambda row: (float(row["value"]), -float(row["tps"])))
+    return pareto_trials
+
+
+def build_leaderboard(*, storage: str, study_names: list[str], metric: str = "val_bpb") -> dict[str, Any]:
+    clean_trials: list[dict[str, Any]] = []
+    contaminated_trials: list[dict[str, Any]] = []
+    study_summaries: list[dict[str, Any]] = []
+    direction = "minimize"
+
+    for study_name in study_names:
+        study = optuna.load_study(study_name=study_name, storage=storage)
+        direction = _trial_direction(study)
+        completed = [t for t in study.trials if t.value is not None]
+        study_summaries.append({
+            "study_name": study_name,
+            "direction": direction,
+            "n_trials": len(study.trials),
+            "n_completed": len(completed),
+        })
+        for trial in completed:
+            row = _serialize_trial(study_name, trial, metric)
+            if "contamination_reason" in row:
+                contaminated_trials.append(row)
+            else:
+                clean_trials.append(row)
+
+    reverse = direction == "maximize"
+    clean_trials.sort(key=lambda row: float(row["value"]), reverse=reverse)
+    contaminated_trials.sort(key=lambda row: float(row["value"]), reverse=reverse)
+    pareto_trials = _annotate_pareto(clean_trials)
+
+    return {
+        "schema_version": 1,
+        "generated_at": dt.datetime.now(dt.UTC).isoformat(timespec="seconds"),
+        "metric": metric,
+        "direction": direction,
+        "study_names": study_names,
+        "studies": study_summaries,
+        "n_clean_trials": len(clean_trials),
+        "n_contaminated_trials": len(contaminated_trials),
+        "pareto_metric_x": metric,
+        "pareto_metric_y": "tps",
+        "n_pareto_trials": len(pareto_trials),
+        "clean_trials": clean_trials,
+        "contaminated_trials": contaminated_trials,
+        "pareto_trials": pareto_trials,
+    }
+
+
+def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Build a clean Optuna HPO leaderboard")
+    parser.add_argument("--storage", default="sqlite:///optuna_hpo.db")
+    parser.add_argument("--study-name", action="append", default=[], help="Repeat to merge multiple studies")
+    parser.add_argument("--metric", default="val_bpb")
+    parser.add_argument("--out", type=Path, default=Path(".tmp") / "optuna" / "leaderboard.json")
+    return parser.parse_args(argv)
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = parse_args(argv)
+    study_names = args.study_name or ["hydra_hpo"]
+    payload = build_leaderboard(storage=args.storage, study_names=study_names, metric=args.metric)
+    args.out.parent.mkdir(parents=True, exist_ok=True)
+    args.out.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
+    print(json.dumps(payload, indent=2, sort_keys=True))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

overlay/scripts/hpo_orchestrator.py

@@ -1,25 +1,25 @@
 #!/usr/bin/env python3
 from __future__ import annotations
 
-import argparse
-import json
-import os
-import subprocess
-import sys
-from pathlib import Path
-from typing import Any
-
-import optuna
-
-
-REPO_ROOT = Path(__file__).resolve().parents[1]
-if str(REPO_ROOT) not in sys.path:
-    sys.path.insert(0, str(REPO_ROOT))
-
-from scripts.hf_routing import resolve_routing
-from scripts.optuna_hpo import _enqueue_transfer_priors
-
-HPO_SCRIPT = REPO_ROOT / "scripts" / "optuna_hpo.py"
+import argparse
+import json
+import os
+import subprocess
+import sys
+from pathlib import Path
+from typing import Any
+
+import optuna
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from scripts.hf_routing import resolve_routing
+from scripts.optuna_hpo import _enqueue_transfer_priors
+
+HPO_SCRIPT = REPO_ROOT / "scripts" / "optuna_hpo.py"
 
 
 def _run_worker(args: list[str]) -> int:
@@ -28,7 +28,7 @@ def _run_worker(args: list[str]) -> int:
     return proc.returncode
 
 
-def _study_stats(storage: str, study_name: str) -> dict[str, Any]:
+def _study_stats(storage: str, study_name: str) -> dict[str, Any]:
     try:
         study = optuna.load_study(study_name=study_name, storage=storage)
     except KeyError:
@@ -62,29 +62,29 @@ def _study_stats(storage: str, study_name: str) -> dict[str, Any]:
                 "best_trial_user_attrs": study.best_trial.user_attrs,
             }
         )
-    return stats
-
-
-def _prime_transfer_priors(storage: str, study_name: str, priors_file: Path, apply_priors: bool) -> int:
-    if not apply_priors:
-        return 0
-    study = optuna.create_study(
-        study_name=study_name,
-        storage=storage,
-        load_if_exists=True,
-        direction="minimize",
-    )
-    return _enqueue_transfer_priors(study, priors_file, apply_priors=True)
-
-
-def _disable_worker_priors(worker_args: list[str]) -> list[str]:
-    cleaned: list[str] = []
-    for item in worker_args:
-        if item in {"--apply-priors", "--no-apply-priors"}:
-            continue
-        cleaned.append(item)
-    cleaned.append("--no-apply-priors")
-    return cleaned
+    return stats
+
+
+def _prime_transfer_priors(storage: str, study_name: str, priors_file: Path, apply_priors: bool) -> int:
+    if not apply_priors:
+        return 0
+    study = optuna.create_study(
+        study_name=study_name,
+        storage=storage,
+        load_if_exists=True,
+        direction="minimize",
+    )
+    return _enqueue_transfer_priors(study, priors_file, apply_priors=True)
+
+
+def _disable_worker_priors(worker_args: list[str]) -> list[str]:
+    cleaned: list[str] = []
+    for item in worker_args:
+        if item in {"--apply-priors", "--no-apply-priors"}:
+            continue
+        cleaned.append(item)
+    cleaned.append("--no-apply-priors")
+    return cleaned
 
 
 def _phase_args(phase: str, base: argparse.Namespace) -> list[str]:
@@ -117,25 +117,25 @@ def _phase_args(phase: str, base: argparse.Namespace) -> list[str]:
         base.hf_command,
         "--hf-token-env",
         base.hf_token_env,
-        "--hf-poll-interval",
-        str(base.hf_poll_interval),
-        "--hf-launcher-script",
-        str(base.hf_launcher_script),
-        "--priors-file",
-        str(base.priors_file),
-    ]
+        "--hf-poll-interval",
+        str(base.hf_poll_interval),
+        "--hf-launcher-script",
+        str(base.hf_launcher_script),
+        "--priors-file",
+        str(base.priors_file),
+    ]
     if base.hf_output_repo:
         common.extend(["--hf-output-repo", base.hf_output_repo])
     if base.hf_use_bash:
         common.append("--hf-use-bash")
-    if base.hf_stop_after_metric:
-        common.append("--hf-stop-after-metric")
-    else:
-        common.append("--no-hf-stop-after-metric")
-    if base.apply_priors:
-        common.append("--apply-priors")
-    else:
-        common.append("--no-apply-priors")
+    if base.hf_stop_after_metric:
+        common.append("--hf-stop-after-metric")
+    else:
+        common.append("--no-hf-stop-after-metric")
+    if base.apply_priors:
+        common.append("--apply-priors")
+    else:
+        common.append("--no-apply-priors")
     if phase == "phase1":
         return [
             *common,
@@ -184,32 +184,32 @@ def cmd_phase(args: argparse.Namespace) -> int:
     return rc
 
 
-def cmd_parallel(args: argparse.Namespace) -> int:
-    enqueued_priors = _prime_transfer_priors(args.storage, args.study_name, args.priors_file, args.apply_priors)
-    worker_args = _disable_worker_priors(_phase_args(args.phase, args))
-    procs: list[subprocess.Popen[str]] = []
-    for _ in range(args.workers):
-        cmd = [sys.executable, str(HPO_SCRIPT), *worker_args]
+def cmd_parallel(args: argparse.Namespace) -> int:
+    enqueued_priors = _prime_transfer_priors(args.storage, args.study_name, args.priors_file, args.apply_priors)
+    worker_args = _disable_worker_priors(_phase_args(args.phase, args))
+    procs: list[subprocess.Popen[str]] = []
+    for _ in range(args.workers):
+        cmd = [sys.executable, str(HPO_SCRIPT), *worker_args]
         procs.append(subprocess.Popen(cmd, cwd=str(REPO_ROOT), text=True))
 
     exit_codes = [p.wait() for p in procs]
     stats = _study_stats(args.storage, args.study_name)
     payload = {
         "phase": args.phase,
-        "workers": args.workers,
-        "exit_codes": exit_codes,
-        "enqueued_priors": enqueued_priors,
-        "stats": stats,
-    }
+        "workers": args.workers,
+        "exit_codes": exit_codes,
+        "enqueued_priors": enqueued_priors,
+        "stats": stats,
+    }
     args.summary_out.parent.mkdir(parents=True, exist_ok=True)
     args.summary_out.write_text(json.dumps(payload, indent=2), encoding="utf-8")
     print(json.dumps(payload, indent=2))
     return 0 if all(code == 0 for code in exit_codes) else 1
 
 
-def cmd_recommend(args: argparse.Namespace) -> int:
-    stats = _study_stats(args.storage, args.study_name)
-    min_tps_floor = float(args.min_tps)
+def cmd_recommend(args: argparse.Namespace) -> int:
+    stats = _study_stats(args.storage, args.study_name)
+    min_tps_floor = float(args.min_tps)
     if stats.get("status") == "missing":
         payload = {
             "stats": stats,
@@ -226,41 +226,41 @@ def cmd_recommend(args: argparse.Namespace) -> int:
 
     n_completed = int(stats.get("n_completed", 0))
 
-    if n_completed < 10:
-        recommendation = {
-            "status": "insufficient_data",
-            "next_step": "Run phase1 with 2-4 parallel workers until >=10 completed trials.",
-            "early_stop_policy": {
-                "patience_trials": 8,
-                "min_improvement": 0.001,
-            },
-            "throughput_guard": {
-                "min_tps": min_tps_floor,
-                "note": "Trials below this TPS floor are pruned.",
-            },
-            "transfer_learning": {
-                "export_priors": f"python scripts/export_hpo_priors.py --storage {args.storage} --study-name {args.study_name} --top-k 10 --out docs/hpo_transfer_priors.json",
-                "use_priors": "Enabled by default in scripts/optuna_hpo.py (override with --no-apply-priors)",
-            },
-        }
-    else:
-        recommendation = {
-            "status": "ready_for_full_optimization",
-            "next_step": "Run phase2 with 3-4 parallel workers.",
+    if n_completed < 10:
+        recommendation = {
+            "status": "insufficient_data",
+            "next_step": "Run phase1 with 2-4 parallel workers until >=10 completed trials.",
+            "early_stop_policy": {
+                "patience_trials": 8,
+                "min_improvement": 0.001,
+            },
+            "throughput_guard": {
+                "min_tps": min_tps_floor,
+                "note": "Trials below this TPS floor are pruned.",
+            },
+            "transfer_learning": {
+                "export_priors": f"python scripts/export_hpo_priors.py --storage {args.storage} --study-name {args.study_name} --top-k 10 --out docs/hpo_transfer_priors.json",
+                "use_priors": "Enabled by default in scripts/optuna_hpo.py (override with --no-apply-priors)",
+            },
+        }
+    else:
+        recommendation = {
+            "status": "ready_for_full_optimization",
+            "next_step": "Run phase2 with 3-4 parallel workers.",
             "suggested_full_run": {
                 "trials": 60,
-                "workers": 4,
-                "trial_time_budget": 300,
-                "trial_timeout": 900,
-                "min_tps": min_tps_floor,
-                "patience_trials": 12,
-                "min_improvement": 0.0005,
-            },
-            "transfer_learning": {
-                "refresh_priors": f"python scripts/export_hpo_priors.py --storage {args.storage} --study-name {args.study_name} --top-k 20 --out docs/hpo_transfer_priors.json",
-                "notes": "Carry priors into new studies unless architecture/objective diverges significantly.",
-            },
-        }
+                "workers": 4,
+                "trial_time_budget": 300,
+                "trial_timeout": 900,
+                "min_tps": min_tps_floor,
+                "patience_trials": 12,
+                "min_improvement": 0.0005,
+            },
+            "transfer_learning": {
+                "refresh_priors": f"python scripts/export_hpo_priors.py --storage {args.storage} --study-name {args.study_name} --top-k 20 --out docs/hpo_transfer_priors.json",
+                "notes": "Carry priors into new studies unless architecture/objective diverges significantly.",
+            },
+        }
 
     payload = {"stats": stats, "recommendation": recommendation}
     args.summary_out.parent.mkdir(parents=True, exist_ok=True)
@@ -269,9 +269,9 @@ def cmd_recommend(args: argparse.Namespace) -> int:
     return 0
 
 
-def build_parser() -> argparse.ArgumentParser:
-    routing_defaults = resolve_routing(token=os.environ.get("HF_TOKEN"))
-    parser = argparse.ArgumentParser(description="Phase-oriented orchestration for Optuna HPO")
+def build_parser() -> argparse.ArgumentParser:
+    routing_defaults = resolve_routing(token=os.environ.get("HF_TOKEN"))
+    parser = argparse.ArgumentParser(description="Phase-oriented orchestration for Optuna HPO")
     sub = parser.add_subparsers(dest="cmd", required=True)
 
     def add_common(p: argparse.ArgumentParser) -> None:
@@ -283,21 +283,21 @@ def build_parser() -> argparse.ArgumentParser:
         p.add_argument("--min-tps", type=float, default=50000.0)
         p.add_argument("--summary-out", type=Path, default=REPO_ROOT / ".tmp" / "optuna" / "orchestrator_summary.json")
         p.add_argument("--runner", choices=["local", "hf-job", "hf-launcher"], default="local")
-        p.add_argument("--hf-namespace", default=routing_defaults.job_namespace)
-        p.add_argument("--hf-image", default=f"hf.co/spaces/{routing_defaults.space_repo}")
+        p.add_argument("--hf-namespace", default=routing_defaults.job_namespace)
+        p.add_argument("--hf-image", default=f"hf.co/spaces/{routing_defaults.space_repo}")
         p.add_argument("--hf-flavor", default="a10g-large")
         p.add_argument("--hf-timeout", default="25m")
         p.add_argument("--hf-command", default="/app/entrypoint.py")
         p.add_argument("--hf-use-bash", action="store_true")
         p.add_argument("--hf-token-env", default="HF_TOKEN")
-        p.add_argument("--hf-poll-interval", type=int, default=12)
-        p.add_argument("--hf-launcher-script", type=Path, default=REPO_ROOT / "scripts" / "launch_feather_hf_job.py")
-        p.add_argument("--hf-output-repo", default=routing_defaults.output_repo)
-        p.add_argument("--priors-file", type=Path, default=REPO_ROOT / "docs" / "hpo_transfer_priors.json")
-        p.add_argument("--apply-priors", action="store_true", default=True)
-        p.add_argument("--no-apply-priors", action="store_false", dest="apply_priors")
-        p.add_argument("--hf-stop-after-metric", action="store_true", default=True)
-        p.add_argument("--no-hf-stop-after-metric", action="store_false", dest="hf_stop_after_metric")
+        p.add_argument("--hf-poll-interval", type=int, default=12)
+        p.add_argument("--hf-launcher-script", type=Path, default=REPO_ROOT / "scripts" / "launch_feather_hf_job.py")
+        p.add_argument("--hf-output-repo", default=routing_defaults.output_repo)
+        p.add_argument("--priors-file", type=Path, default=REPO_ROOT / "docs" / "hpo_transfer_priors.json")
+        p.add_argument("--apply-priors", action="store_true", default=True)
+        p.add_argument("--no-apply-priors", action="store_false", dest="apply_priors")
+        p.add_argument("--hf-stop-after-metric", action="store_true", default=True)
+        p.add_argument("--no-hf-stop-after-metric", action="store_false", dest="hf_stop_after_metric")
 
         # Phase-1 defaults
         p.add_argument("--phase1-trials", type=int, default=30)

overlay/scripts/hpo_retest.py

@@ -1,151 +1,151 @@
-#!/usr/bin/env python3
-from __future__ import annotations
-
-import argparse
-import datetime as dt
-import json
-from pathlib import Path
-from typing import Any
-
-import optuna
-
-from scripts.hpo_leaderboard import build_leaderboard
-
-
-_PARAM_TO_ENV = {
-    "d_model": "HYDRA_D_MODEL",
-    "n_layer": "HYDRA_N_LAYER",
-    "d_state": "HYDRA_D_STATE",
-    "headdim": "HYDRA_HEADDIM",
-    "expand": "HYDRA_EXPAND",
-    "seq_len": "HYDRA_SEQ_LEN",
-    "batch_size": "HYDRA_BATCH_SIZE",
-    "matrix_lr": "HYDRA_MATRIX_LR",
-    "embed_lr": "HYDRA_EMBED_LR",
-    "unembed_lr": "HYDRA_UNEMBED_LR",
-    "engram_n_columns": "HYDRA_ENGRAM_N_COLUMNS",
-    "engram_layer_idx": "HYDRA_ENGRAM_LAYER_IDX",
-    "sdr_target_active": "HYDRA_SDR_TARGET_ACTIVE",
-    "htm_learn_every": "HYDRA_HTM_LEARN_EVERY",
-    "htm_subsample": "HYDRA_HTM_SUBSAMPLE",
-    "engram_subsample": "HYDRA_ENGRAM_SUBSAMPLE",
-    "mamba3_chunk": "HYDRA_MAMBA3_CHUNK",
-    "dropout": "HYDRA_DROPOUT",
-}
-
-_DEFAULT_ENV = {
-    "HYDRA_USE_NEMOTRON": "1",
-    "HYDRA_LOCAL_SHARDS_ONLY": "0",
-    "HYDRA_THROUGHPUT_MODE": "0",
-    "HYDRA_FASTPATH": "0",
-    "HYDRA_FORCE_HTM_CPU": "0",
-    "HYDRA_INERT_MAMBA": "0",
-    "HYDRA_ALLOW_SYNTHETIC_RETINA": "0",
-    "HYDRA_HTM_FUSED": "1",
-    "HYDRA_HYENA_LAYERS": "",
-    "HYDRA_CKPT_INTERVAL": "0",
-    "HYDRA_ENGRAM_SUBSAMPLE": "1",
-    "HYDRA_HTM_SUBSAMPLE": "2",
-    "HYDRA_HTM_LEARN_EVERY": "8",
-}
-
-
-def _recover_params(row: dict[str, Any]) -> dict[str, Any]:
-    params = dict(row.get("params") or {})
-    attrs = row.get("user_attrs") or {}
-    for key, value in attrs.items():
-        if key.startswith("param_"):
-            params.setdefault(key.removeprefix("param_"), value)
-    return params
-
-
-def _candidate_env(params: dict[str, Any], *, eval_tokens: int, eval_batch: int, time_budget: int) -> dict[str, str]:
-    env = dict(_DEFAULT_ENV)
-    env["HYDRA_EVAL_TOKENS"] = str(eval_tokens)
-    env["HYDRA_EVAL_BATCH"] = str(eval_batch)
-    env["HYDRA_TIME_BUDGET"] = str(time_budget)
-    for key, value in params.items():
-        env_key = _PARAM_TO_ENV.get(key)
-        if env_key is not None:
-            env[env_key] = str(value)
-    if "HYDRA_BATCH_SIZE" in env and "HYDRA_SEQ_LEN" in env:
-        grad_accum = int(params.get("grad_accum", 16))
-        env["HYDRA_TOTAL_BATCH"] = str(int(env["HYDRA_BATCH_SIZE"]) * int(env["HYDRA_SEQ_LEN"]) * grad_accum)
-    return env
-
-
-def build_retest_plan(
-    *,
-    storage: str,
-    study_names: list[str],
-    top_k: int,
-    metric: str = "val_bpb",
-    eval_tokens: int = 16384,
-    eval_batch: int = 2,
-    time_budget: int = 420,
-) -> dict[str, Any]:
-    leaderboard = build_leaderboard(storage=storage, study_names=study_names, metric=metric)
-    rows = [*leaderboard["contaminated_trials"], *leaderboard["clean_trials"]]
-    reverse = leaderboard["direction"] == "maximize"
-    rows.sort(key=lambda row: float(row["value"]), reverse=reverse)
-    candidates = []
-    for row in rows[: max(0, top_k)]:
-        params = _recover_params(row)
-        env = _candidate_env(params, eval_tokens=eval_tokens, eval_batch=eval_batch, time_budget=time_budget)
-        reason = row.get("contamination_reason") or "canonical_truth_eval_retest"
-        candidates.append({
-            "study_name": row["study_name"],
-            "trial_number": row["trial_number"],
-            "source_value": row["value"],
-            "source_objective": row["objective_source"],
-            "source_job_id": row.get("hf_job_id"),
-            "needs_retest_reason": reason,
-            "params": params,
-            "env": env,
-        })
-    return {
-        "schema_version": 1,
-        "generated_at": dt.datetime.now(dt.UTC).isoformat(timespec="seconds"),
-        "metric": metric,
-        "study_names": study_names,
-        "eval_tokens": eval_tokens,
-        "eval_batch": eval_batch,
-        "time_budget": time_budget,
-        "n_candidates": len(candidates),
-        "candidates": candidates,
-    }
-
-
-def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Plan canonical-eval retests for historical HPO configs")
-    parser.add_argument("--storage", default="sqlite:///optuna_hpo.db")
-    parser.add_argument("--study-name", action="append", default=[])
-    parser.add_argument("--metric", default="val_bpb")
-    parser.add_argument("--top-k", type=int, default=10)
-    parser.add_argument("--eval-tokens", type=int, default=16384)
-    parser.add_argument("--eval-batch", type=int, default=2)
-    parser.add_argument("--time-budget", type=int, default=420)
-    parser.add_argument("--out", type=Path, default=Path(".tmp") / "optuna" / "retest_plan.json")
-    return parser.parse_args(argv)
-
-
-def main(argv: list[str] | None = None) -> int:
-    args = parse_args(argv)
-    study_names = args.study_name or ["hydra_hpo"]
-    payload = build_retest_plan(
-        storage=args.storage,
-        study_names=study_names,
-        top_k=args.top_k,
-        metric=args.metric,
-        eval_tokens=args.eval_tokens,
-        eval_batch=args.eval_batch,
-        time_budget=args.time_budget,
-    )
-    args.out.parent.mkdir(parents=True, exist_ok=True)
-    args.out.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
-    print(json.dumps(payload, indent=2, sort_keys=True))
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import datetime as dt
+import json
+from pathlib import Path
+from typing import Any
+
+import optuna
+
+from scripts.hpo_leaderboard import build_leaderboard
+
+
+_PARAM_TO_ENV = {
+    "d_model": "HYDRA_D_MODEL",
+    "n_layer": "HYDRA_N_LAYER",
+    "d_state": "HYDRA_D_STATE",
+    "headdim": "HYDRA_HEADDIM",
+    "expand": "HYDRA_EXPAND",
+    "seq_len": "HYDRA_SEQ_LEN",
+    "batch_size": "HYDRA_BATCH_SIZE",
+    "matrix_lr": "HYDRA_MATRIX_LR",
+    "embed_lr": "HYDRA_EMBED_LR",
+    "unembed_lr": "HYDRA_UNEMBED_LR",
+    "engram_n_columns": "HYDRA_ENGRAM_N_COLUMNS",
+    "engram_layer_idx": "HYDRA_ENGRAM_LAYER_IDX",
+    "sdr_target_active": "HYDRA_SDR_TARGET_ACTIVE",
+    "htm_learn_every": "HYDRA_HTM_LEARN_EVERY",
+    "htm_subsample": "HYDRA_HTM_SUBSAMPLE",
+    "engram_subsample": "HYDRA_ENGRAM_SUBSAMPLE",
+    "mamba3_chunk": "HYDRA_MAMBA3_CHUNK",
+    "dropout": "HYDRA_DROPOUT",
+}
+
+_DEFAULT_ENV = {
+    "HYDRA_USE_NEMOTRON": "1",
+    "HYDRA_LOCAL_SHARDS_ONLY": "0",
+    "HYDRA_THROUGHPUT_MODE": "0",
+    "HYDRA_FASTPATH": "0",
+    "HYDRA_FORCE_HTM_CPU": "0",
+    "HYDRA_INERT_MAMBA": "0",
+    "HYDRA_ALLOW_SYNTHETIC_RETINA": "0",
+    "HYDRA_HTM_FUSED": "1",
+    "HYDRA_HYENA_LAYERS": "",
+    "HYDRA_CKPT_INTERVAL": "0",
+    "HYDRA_ENGRAM_SUBSAMPLE": "1",
+    "HYDRA_HTM_SUBSAMPLE": "2",
+    "HYDRA_HTM_LEARN_EVERY": "8",
+}
+
+
+def _recover_params(row: dict[str, Any]) -> dict[str, Any]:
+    params = dict(row.get("params") or {})
+    attrs = row.get("user_attrs") or {}
+    for key, value in attrs.items():
+        if key.startswith("param_"):
+            params.setdefault(key.removeprefix("param_"), value)
+    return params
+
+
+def _candidate_env(params: dict[str, Any], *, eval_tokens: int, eval_batch: int, time_budget: int) -> dict[str, str]:
+    env = dict(_DEFAULT_ENV)
+    env["HYDRA_EVAL_TOKENS"] = str(eval_tokens)
+    env["HYDRA_EVAL_BATCH"] = str(eval_batch)
+    env["HYDRA_TIME_BUDGET"] = str(time_budget)
+    for key, value in params.items():
+        env_key = _PARAM_TO_ENV.get(key)
+        if env_key is not None:
+            env[env_key] = str(value)
+    if "HYDRA_BATCH_SIZE" in env and "HYDRA_SEQ_LEN" in env:
+        grad_accum = int(params.get("grad_accum", 16))
+        env["HYDRA_TOTAL_BATCH"] = str(int(env["HYDRA_BATCH_SIZE"]) * int(env["HYDRA_SEQ_LEN"]) * grad_accum)
+    return env
+
+
+def build_retest_plan(
+    *,
+    storage: str,
+    study_names: list[str],
+    top_k: int,
+    metric: str = "val_bpb",
+    eval_tokens: int = 16384,
+    eval_batch: int = 2,
+    time_budget: int = 420,
+) -> dict[str, Any]:
+    leaderboard = build_leaderboard(storage=storage, study_names=study_names, metric=metric)
+    rows = [*leaderboard["contaminated_trials"], *leaderboard["clean_trials"]]
+    reverse = leaderboard["direction"] == "maximize"
+    rows.sort(key=lambda row: float(row["value"]), reverse=reverse)
+    candidates = []
+    for row in rows[: max(0, top_k)]:
+        params = _recover_params(row)
+        env = _candidate_env(params, eval_tokens=eval_tokens, eval_batch=eval_batch, time_budget=time_budget)
+        reason = row.get("contamination_reason") or "canonical_truth_eval_retest"
+        candidates.append({
+            "study_name": row["study_name"],
+            "trial_number": row["trial_number"],
+            "source_value": row["value"],
+            "source_objective": row["objective_source"],
+            "source_job_id": row.get("hf_job_id"),
+            "needs_retest_reason": reason,
+            "params": params,
+            "env": env,
+        })
+    return {
+        "schema_version": 1,
+        "generated_at": dt.datetime.now(dt.UTC).isoformat(timespec="seconds"),
+        "metric": metric,
+        "study_names": study_names,
+        "eval_tokens": eval_tokens,
+        "eval_batch": eval_batch,
+        "time_budget": time_budget,
+        "n_candidates": len(candidates),
+        "candidates": candidates,
+    }
+
+
+def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Plan canonical-eval retests for historical HPO configs")
+    parser.add_argument("--storage", default="sqlite:///optuna_hpo.db")
+    parser.add_argument("--study-name", action="append", default=[])
+    parser.add_argument("--metric", default="val_bpb")
+    parser.add_argument("--top-k", type=int, default=10)
+    parser.add_argument("--eval-tokens", type=int, default=16384)
+    parser.add_argument("--eval-batch", type=int, default=2)
+    parser.add_argument("--time-budget", type=int, default=420)
+    parser.add_argument("--out", type=Path, default=Path(".tmp") / "optuna" / "retest_plan.json")
+    return parser.parse_args(argv)
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = parse_args(argv)
+    study_names = args.study_name or ["hydra_hpo"]
+    payload = build_retest_plan(
+        storage=args.storage,
+        study_names=study_names,
+        top_k=args.top_k,
+        metric=args.metric,
+        eval_tokens=args.eval_tokens,
+        eval_batch=args.eval_batch,
+        time_budget=args.time_budget,
+    )
+    args.out.parent.mkdir(parents=True, exist_ok=True)
+    args.out.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
+    print(json.dumps(payload, indent=2, sort_keys=True))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

overlay/scripts/hydra_generation.py

@@ -1,183 +1,180 @@
-#!/usr/bin/env python3
-from __future__ import annotations
-
-import os
-from pathlib import Path
-from typing import Callable
-
-import torch
-
-from scripts.benchmark_checkpoint import hydrate_checkpoint
-from scripts.hf_routing import resolve_routing
-
-
-def default_checkpoint_path() -> Path:
-    return Path(os.path.expanduser("~/.cache/autoresearch/latest.pt"))
-
-
-def checkpoint_candidates(*, cache_dir: Path | None = None) -> list[Path]:
-    base = cache_dir or Path(os.path.expanduser("~/.cache/autoresearch"))
-    return [
-        base / "best_bpb.pt",
-        base / "pretrain_final.pt",
-        base / "latest.pt",
-    ]
-
-
-def resolve_checkpoint_path(explicit_path: Path | None, *, cache_dir: Path | None = None) -> Path:
-    if explicit_path is not None:
-        return explicit_path
-    env_checkpoint = os.environ.get("HYDRA_HF_CHECKPOINT_PATH")
-    if env_checkpoint:
-        return Path(env_checkpoint).expanduser()
-    for candidate in checkpoint_candidates(cache_dir=cache_dir):
-        if candidate.exists():
-            return candidate
-    return default_checkpoint_path()
-
-
-def validate_checkpoint_compatibility(
-    *,
-    baseline_arch: str,
-    missing_keys: list[str],
-    unexpected_keys: list[str],
-    total_model_keys: int,
-) -> None:
-    if baseline_arch == "transformer" and (missing_keys or unexpected_keys):
-        raise RuntimeError(
-            "checkpoint incompatible with transformer baseline architecture; "
-            "use a transformer-trained checkpoint or keep HYDRA_BASELINE_ARCH=mamba3"
-        )
-    mismatch_count = len(missing_keys) + len(unexpected_keys)
-    if total_model_keys > 0 and mismatch_count > max(8, total_model_keys // 2):
-        raise RuntimeError("checkpoint incompatible with requested model architecture")
-
-
-def generate_from_callable(
-    generator: Callable[[str], str] | Callable[..., str],
-    prompt: str,
-    *,
-    max_new_tokens: int,
-    temperature: float,
-    top_p: float,
-) -> str:
-    text = generator(
-        prompt,
-        max_new_tokens=max_new_tokens,
-        temperature=temperature,
-        top_p=top_p,
-    )
-    return str(text).strip()
-
-
-def load_hydra_causal_lm(checkpoint_path: Path | None = None, device: str | None = None):
-    ckpt_path = resolve_checkpoint_path(checkpoint_path)
-    if not ckpt_path.exists():
-        hydrated = hydrate_checkpoint(
-            cache_dir=ckpt_path.parent,
-            output_repo=resolve_routing(token=os.environ.get("HF_TOKEN")).output_repo,
-            token=os.environ.get("HF_TOKEN"),
-        )
-        if hydrated is not None:
-            ckpt_path = hydrated
-        if not ckpt_path.exists():
-            raise FileNotFoundError(f"Checkpoint not found: {ckpt_path}")
-
-    from transformers import GenerationConfig, GenerationMixin, PretrainedConfig, PreTrainedModel
-    from transformers.modeling_outputs import CausalLMOutputWithPast
-
-    from hydra.config import PostSemClawConfig
-    from hydra.model import PostSemClawModel
-    from prepare import Tokenizer
-
-    resolved_device = device or ("cuda" if torch.cuda.is_available() else "cpu")
-
-    class _HydraGenConfig(PretrainedConfig):
-        model_type = "hydra"
-
-        def __init__(self, vocab_size: int = 65536, **kw):
-            super().__init__(**kw)
-            self.vocab_size = vocab_size
-
-    class HydraForCausalLM(PreTrainedModel, GenerationMixin):
-        config_class = _HydraGenConfig
-
-        def __init__(self, gen_config, inner_model):
-            super().__init__(gen_config)
-            self.inner = inner_model
-            self.config.vocab_size = gen_config.vocab_size
-
-        def forward(self, input_ids, attention_mask=None, **kw):
-            logits = self.inner(input_ids)
-            return CausalLMOutputWithPast(loss=None, logits=logits, past_key_values=None)
-
-        def prepare_inputs_for_generation(self, input_ids, **kw):
-            return {"input_ids": input_ids}
-
-        def get_input_embeddings(self):
-            return self.inner.wte
-
-        def can_generate(self) -> bool:
-            return True
-
-        @property
-        def _supports_cache_class(self):
-            return False
-
-    tokenizer = Tokenizer.from_directory()
-    vocab_size = tokenizer.get_vocab_size()
-    bos = tokenizer.get_bos_token_id()
-    ckpt = torch.load(str(ckpt_path), map_location="cpu", weights_only=False)
-    cfg = PostSemClawConfig(**ckpt["config"])
-    with torch.device("meta"):
-        inner = PostSemClawModel(cfg)
-    inner.to_empty(device=resolved_device)
-    missing, unexpected = inner.load_state_dict(ckpt["model_state_dict"], strict=False)
-    validate_checkpoint_compatibility(
-        baseline_arch=os.environ.get("HYDRA_BASELINE_ARCH", "mamba3").strip().lower(),
-        missing_keys=list(missing),
-        unexpected_keys=list(unexpected),
-        total_model_keys=len(inner.state_dict()),
-    )
-    inner.eval()
-
-    gen_cfg = _HydraGenConfig(vocab_size=vocab_size)
-    gen_cfg.bos_token_id = bos
-    gen_cfg.eos_token_id = bos
-    gen_cfg.pad_token_id = bos
-    model = HydraForCausalLM(gen_cfg, inner).to(resolved_device)
-    model.eval()
-    return tokenizer, model, bos, resolved_device, GenerationConfig
-
-
-def build_hydra_generator(
-    *,
-    checkpoint_path: Path | None = None,
-    device: str | None = None,
-    max_new_tokens: int,
-    temperature: float,
-    top_p: float,
-):
-    tokenizer, model, bos, resolved_device, GenerationConfig = load_hydra_causal_lm(checkpoint_path=checkpoint_path, device=device)
-
-    def _generate(prompt: str) -> str:
-        ids = torch.tensor([tokenizer.encode(prompt)], dtype=torch.long, device=resolved_device)
-        gen_config = GenerationConfig(
-            max_new_tokens=max_new_tokens,
-            use_cache=False,
-            do_sample=temperature > 0.0,
-            temperature=temperature,
-            top_p=top_p,
-            bos_token_id=bos,
-            eos_token_id=bos,
-            pad_token_id=bos,
-        )
-        if str(resolved_device).startswith("cuda"):
-            with torch.no_grad(), torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
-                out = model.generate(ids, generation_config=gen_config)
-        else:
-            with torch.no_grad():
-                out = model.generate(ids, generation_config=gen_config)
-        return tokenizer.decode(out[0].tolist())
-
-    return _generate
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Callable
+
+import torch
+
+from scripts.benchmark_checkpoint import hydrate_checkpoint
+from scripts.hf_routing import resolve_routing
+
+
+def default_checkpoint_path() -> Path:
+    return Path(os.path.expanduser("~/.cache/autoresearch/latest.pt"))
+
+
+def checkpoint_candidates(*, cache_dir: Path | None = None) -> list[Path]:
+    base = cache_dir or Path(os.path.expanduser("~/.cache/autoresearch"))
+    return [
+        base / "best_bpb.pt",
+        base / "pretrain_final.pt",
+        base / "latest.pt",
+    ]
+
+
+def resolve_checkpoint_path(explicit_path: Path | None, *, cache_dir: Path | None = None) -> Path:
+    if explicit_path is not None:
+        return explicit_path
+    for candidate in checkpoint_candidates(cache_dir=cache_dir):
+        if candidate.exists():
+            return candidate
+    return default_checkpoint_path()
+
+
+def validate_checkpoint_compatibility(
+    *,
+    baseline_arch: str,
+    missing_keys: list[str],
+    unexpected_keys: list[str],
+    total_model_keys: int,
+) -> None:
+    if baseline_arch == "transformer" and (missing_keys or unexpected_keys):
+        raise RuntimeError(
+            "checkpoint incompatible with transformer baseline architecture; "
+            "use a transformer-trained checkpoint or keep HYDRA_BASELINE_ARCH=mamba3"
+        )
+    mismatch_count = len(missing_keys) + len(unexpected_keys)
+    if total_model_keys > 0 and mismatch_count > max(8, total_model_keys // 2):
+        raise RuntimeError("checkpoint incompatible with requested model architecture")
+
+
+def generate_from_callable(
+    generator: Callable[[str], str] | Callable[..., str],
+    prompt: str,
+    *,
+    max_new_tokens: int,
+    temperature: float,
+    top_p: float,
+) -> str:
+    text = generator(
+        prompt,
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+        top_p=top_p,
+    )
+    return str(text).strip()
+
+
+def load_hydra_causal_lm(checkpoint_path: Path | None = None, device: str | None = None):
+    ckpt_path = resolve_checkpoint_path(checkpoint_path)
+    if not ckpt_path.exists():
+        hydrated = hydrate_checkpoint(
+            cache_dir=ckpt_path.parent,
+            output_repo=resolve_routing(token=os.environ.get("HF_TOKEN")).output_repo,
+            token=os.environ.get("HF_TOKEN"),
+        )
+        if hydrated is not None:
+            ckpt_path = hydrated
+        if not ckpt_path.exists():
+            raise FileNotFoundError(f"Checkpoint not found: {ckpt_path}")
+
+    from transformers import GenerationConfig, GenerationMixin, PretrainedConfig, PreTrainedModel
+    from transformers.modeling_outputs import CausalLMOutputWithPast
+
+    from hydra.config import PostSemClawConfig
+    from hydra.model import PostSemClawModel
+    from prepare import Tokenizer
+
+    resolved_device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+
+    class _HydraGenConfig(PretrainedConfig):
+        model_type = "hydra"
+
+        def __init__(self, vocab_size: int = 65536, **kw):
+            super().__init__(**kw)
+            self.vocab_size = vocab_size
+
+    class HydraForCausalLM(PreTrainedModel, GenerationMixin):
+        config_class = _HydraGenConfig
+
+        def __init__(self, gen_config, inner_model):
+            super().__init__(gen_config)
+            self.inner = inner_model
+            self.config.vocab_size = gen_config.vocab_size
+
+        def forward(self, input_ids, attention_mask=None, **kw):
+            logits = self.inner(input_ids)
+            return CausalLMOutputWithPast(loss=None, logits=logits, past_key_values=None)
+
+        def prepare_inputs_for_generation(self, input_ids, **kw):
+            return {"input_ids": input_ids}
+
+        def get_input_embeddings(self):
+            return self.inner.wte
+
+        def can_generate(self) -> bool:
+            return True
+
+        @property
+        def _supports_cache_class(self):
+            return False
+
+    tokenizer = Tokenizer.from_directory()
+    vocab_size = tokenizer.get_vocab_size()
+    bos = tokenizer.get_bos_token_id()
+    ckpt = torch.load(str(ckpt_path), map_location="cpu", weights_only=False)
+    cfg = PostSemClawConfig(**ckpt["config"])
+    with torch.device("meta"):
+        inner = PostSemClawModel(cfg)
+    inner.to_empty(device=resolved_device)
+    missing, unexpected = inner.load_state_dict(ckpt["model_state_dict"], strict=False)
+    validate_checkpoint_compatibility(
+        baseline_arch=os.environ.get("HYDRA_BASELINE_ARCH", "mamba3").strip().lower(),
+        missing_keys=list(missing),
+        unexpected_keys=list(unexpected),
+        total_model_keys=len(inner.state_dict()),
+    )
+    inner.eval()
+
+    gen_cfg = _HydraGenConfig(vocab_size=vocab_size)
+    gen_cfg.bos_token_id = bos
+    gen_cfg.eos_token_id = bos
+    gen_cfg.pad_token_id = bos
+    model = HydraForCausalLM(gen_cfg, inner).to(resolved_device)
+    model.eval()
+    return tokenizer, model, bos, resolved_device, GenerationConfig
+
+
+def build_hydra_generator(
+    *,
+    checkpoint_path: Path | None = None,
+    device: str | None = None,
+    max_new_tokens: int,
+    temperature: float,
+    top_p: float,
+):
+    tokenizer, model, bos, resolved_device, GenerationConfig = load_hydra_causal_lm(checkpoint_path=checkpoint_path, device=device)
+
+    def _generate(prompt: str) -> str:
+        ids = torch.tensor([tokenizer.encode(prompt)], dtype=torch.long, device=resolved_device)
+        gen_config = GenerationConfig(
+            max_new_tokens=max_new_tokens,
+            use_cache=False,
+            do_sample=temperature > 0.0,
+            temperature=temperature,
+            top_p=top_p,
+            bos_token_id=bos,
+            eos_token_id=bos,
+            pad_token_id=bos,
+        )
+        if str(resolved_device).startswith("cuda"):
+            with torch.no_grad(), torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
+                out = model.generate(ids, generation_config=gen_config)
+        else:
+            with torch.no_grad():
+                out = model.generate(ids, generation_config=gen_config)
+        return tokenizer.decode(out[0].tolist())
+
+    return _generate

overlay/scripts/launch_benchmark_hf_job.py

@@ -1,222 +1,157 @@
-#!/usr/bin/env python3
-from __future__ import annotations
-
-import argparse
-import json
-import os
-import sys
-from pathlib import Path
-
-REPO_ROOT = Path(__file__).resolve().parents[1]
-if str(REPO_ROOT) not in sys.path:
-    sys.path.insert(0, str(REPO_ROOT))
-
-from huggingface_hub import HfApi
-from huggingface_hub.utils import get_token
-
-from scripts.cycle_executor import decode_config_for_benchmark, load_cycle_freeze, variant_env_for_benchmark
-from scripts.hf_routing import resolve_routing
-from scripts.launch_feather_hf_job import IMAGE_DIR, sync_overlay_from_repo, wait_for_space
-
-
-FREEZE_PATH = REPO_ROOT / "artifacts" / "cycle_1_execution_freeze.json"
-
-
-def resolve_variant_checkpoint_override(*, variant: str) -> str | None:
-    if variant == "baseline_transformer_matched":
-        return os.environ.get("HYDRA_TRANSFORMER_CHECKPOINT_PATH") or os.environ.get("HYDRA_HF_CHECKPOINT_PATH")
-    if variant == "baseline_mamba_matched":
-        return os.environ.get("HYDRA_MAMBA_CHECKPOINT_PATH") or os.environ.get("HYDRA_HF_CHECKPOINT_PATH")
-    return os.environ.get("HYDRA_HF_CHECKPOINT_PATH")
-
-
-def validate_variant_ready_for_submission(*, variant: str, freeze: dict[str, object]) -> None:
-    variants = freeze.get("variants")
-    variant_cfg: object = variants.get(variant, {}) if isinstance(variants, dict) else {}
-    status = variant_cfg.get("status") if isinstance(variant_cfg, dict) else None
-    checkpoint_override = resolve_variant_checkpoint_override(variant=variant)
-    if variant == "baseline_transformer_matched" and status == "blocked_checkpoint_incompatible" and not checkpoint_override:
-        raise SystemExit(
-            "baseline_transformer_matched is blocked by checkpoint incompatibility; set HYDRA_TRANSFORMER_CHECKPOINT_PATH to a transformer-compatible checkpoint before submission"
-        )
-
-
-def build_benchmark_job_env(
-    *,
-    benchmark: str,
-    variant: str,
-    seed: int,
-    output_repo: str,
-    tokenizer_repo: str,
-    retina_repo: str,
-    freeze: dict[str, object],
-) -> dict[str, str]:
-    env = {
-        "FEATHER_HF_OUTPUT_REPO": output_repo,
-        "FEATHER_HF_RETINA_CACHE_REPO": retina_repo,
-        "HF_REPO_ID": output_repo,
-        "FEATHER_RUNTIME_MODE": "benchmark",
-        "HYDRA_TOKENIZER_CACHE_REPO": tokenizer_repo,
-        "HYDRA_RETINA_CACHE_REPO": retina_repo,
-        "HYDRA_BENCHMARK_NAME": benchmark,
-        "HYDRA_BENCHMARK_VARIANT": variant,
-        "HYDRA_SEED": str(seed),
-        "PYTHONUNBUFFERED": "1",
-    }
-    env.update(variant_env_for_benchmark(freeze, variant))
-    checkpoint_override = resolve_variant_checkpoint_override(variant=variant)
-    if checkpoint_override:
-        env["HYDRA_HF_CHECKPOINT_PATH"] = checkpoint_override
-    for key, value in os.environ.items():
-        if key.startswith("HYDRA_") and key not in env:
-            env[key] = value
-    return env
-
-
-def build_benchmark_job_command(
-    *,
-    benchmark: str,
-    variant: str,
-    seed: int,
-    suite_path: Path | None,
-    freeze: dict[str, object],
-) -> list[str]:
-    decode_cfg = decode_config_for_benchmark(freeze, benchmark)
-    command = [
-        "python",
-        "/app/entrypoint.py",
-        "--max-new-tokens",
-        str(int(decode_cfg.get("max_tokens", 256))),
-        "--temperature",
-        str(float(decode_cfg.get("temperature", 0.2))),
-        "--top-p",
-        str(float(decode_cfg.get("top_p", 0.95))),
-    ]
-    if suite_path is not None:
-        command.extend(["--suite", str(suite_path)])
-    return command
-
-
-def submit_benchmark_job(
-    *,
-    api,
-    image: str,
-    command: list[str],
-    env: dict[str, str],
-    token: str,
-    namespace: str,
-    flavor: str,
-    timeout: str,
-) -> dict[str, str]:
-    job = api.run_job(
-        image=image,
-        command=command,
-        env=env,
-        secrets={"HF_TOKEN": token},
-        flavor=flavor,
-        timeout=timeout,
-        namespace=namespace,
-        token=token,
-    )
-    return {
-        "job_id": job.id,
-        "job_url": job.url,
-        "job_stage": str(job.status.stage),
-    }
-
-
-def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
-    routing = resolve_routing(token=os.environ.get("HF_TOKEN"))
-    parser = argparse.ArgumentParser(description="Prepare or submit a remote HF benchmark job")
-    parser.add_argument("--benchmark", required=True)
-    parser.add_argument("--variant", required=True)
-    parser.add_argument("--seed", type=int, required=True)
-    parser.add_argument("--output-repo", default=routing.output_repo)
-    parser.add_argument("--tokenizer-repo", default=routing.output_repo)
-    parser.add_argument("--retina-repo", default=routing.retina_cache_repo)
-    parser.add_argument("--freeze", type=Path, default=FREEZE_PATH)
-    parser.add_argument("--suite", type=Path)
-    parser.add_argument("--image", default=f"hf.co/spaces/{routing.space_repo}")
-    parser.add_argument("--namespace", default=routing.job_namespace)
-    parser.add_argument("--flavor", default="a10g-small")
-    parser.add_argument("--timeout", default="30m")
-    parser.add_argument("--summary-out", type=Path)
-    parser.add_argument("--dry-run", action="store_true")
-    parser.add_argument("--refresh-image", action="store_true")
-    parser.add_argument("--sync-overlay", action="store_true")
-    return parser.parse_args(argv)
-
-
-def main(argv: list[str] | None = None) -> int:
-    args = parse_args(argv)
-    freeze = load_cycle_freeze(args.freeze)
-    validate_variant_ready_for_submission(variant=args.variant, freeze=freeze)
-    env = build_benchmark_job_env(
-        benchmark=args.benchmark,
-        variant=args.variant,
-        seed=args.seed,
-        output_repo=args.output_repo,
-        tokenizer_repo=args.tokenizer_repo,
-        retina_repo=args.retina_repo,
-        freeze=freeze,
-    )
-    command = build_benchmark_job_command(
-        benchmark=args.benchmark,
-        variant=args.variant,
-        seed=args.seed,
-        suite_path=args.suite,
-        freeze=freeze,
-    )
-    payload = {
-        "benchmark": args.benchmark,
-        "variant": args.variant,
-        "seed": args.seed,
-        "output_repo": args.output_repo,
-        "tokenizer_repo": args.tokenizer_repo,
-        "retina_repo": args.retina_repo,
-        "freeze": str(args.freeze),
-        "suite": str(args.suite) if args.suite is not None else None,
-        "image": args.image,
-        "namespace": args.namespace,
-        "command": command,
-        "env": env,
-        "dry_run": args.dry_run,
-    }
-    if not args.dry_run:
-        token = os.environ.get("HF_TOKEN") or get_token()
-        if not token:
-            raise SystemExit("HF_TOKEN must be set or cached via huggingface-cli login")
-        api = HfApi(token=token)
-        if args.refresh_image:
-            space_repo = args.image.removeprefix("hf.co/spaces/")
-            if args.sync_overlay:
-                sync_overlay_from_repo()
-            api.upload_folder(
-                repo_id=space_repo,
-                repo_type="space",
-                folder_path=str(IMAGE_DIR),
-                commit_message="Update benchmark runtime image",
-                token=token,
-            )
-            wait_for_space(api, space_repo, token=token)
-        payload.update(
-            submit_benchmark_job(
-                api=api,
-                image=args.image,
-                command=command,
-                env=env,
-                token=token,
-                namespace=args.namespace,
-                flavor=args.flavor,
-                timeout=args.timeout,
-            )
-        )
-    if args.summary_out is not None:
-        args.summary_out.parent.mkdir(parents=True, exist_ok=True)
-        args.summary_out.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
-    print(json.dumps(payload, indent=2, sort_keys=True))
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from huggingface_hub import HfApi
+from huggingface_hub.utils import get_token
+
+from scripts.hf_routing import resolve_routing
+from scripts.launch_feather_hf_job import IMAGE_DIR, sync_overlay_from_repo, wait_for_space
+
+
+def build_benchmark_job_env(
+    *,
+    benchmark: str,
+    variant: str,
+    seed: int,
+    output_repo: str,
+    tokenizer_repo: str,
+) -> dict[str, str]:
+    env = {
+        "FEATHER_HF_OUTPUT_REPO": output_repo,
+        "FEATHER_RUNTIME_MODE": "benchmark",
+        "HYDRA_TOKENIZER_CACHE_REPO": tokenizer_repo,
+        "HYDRA_BENCHMARK_NAME": benchmark,
+        "HYDRA_BENCHMARK_VARIANT": variant,
+        "HYDRA_SEED": str(seed),
+        "PYTHONUNBUFFERED": "1",
+    }
+    for key, value in os.environ.items():
+        if key.startswith("HYDRA_") and key not in env:
+            env[key] = value
+    return env
+
+
+def build_benchmark_job_command(*, benchmark: str, variant: str, seed: int) -> list[str]:
+    return [
+        "python",
+        "/app/entrypoint.py",
+    ]
+
+
+def submit_benchmark_job(
+    *,
+    api,
+    image: str,
+    command: list[str],
+    env: dict[str, str],
+    token: str,
+    namespace: str,
+    flavor: str,
+    timeout: str,
+) -> dict[str, str]:
+    job = api.run_job(
+        image=image,
+        command=command,
+        env=env,
+        secrets={"HF_TOKEN": token},
+        flavor=flavor,
+        timeout=timeout,
+        namespace=namespace,
+        token=token,
+    )
+    return {
+        "job_id": job.id,
+        "job_url": job.url,
+        "job_stage": str(job.status.stage),
+    }
+
+
+def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
+    routing = resolve_routing(token=os.environ.get("HF_TOKEN"))
+    parser = argparse.ArgumentParser(description="Prepare or submit a remote HF benchmark job")
+    parser.add_argument("--benchmark", required=True)
+    parser.add_argument("--variant", required=True)
+    parser.add_argument("--seed", type=int, required=True)
+    parser.add_argument("--output-repo", default=routing.output_repo)
+    parser.add_argument("--tokenizer-repo", default=routing.output_repo)
+    parser.add_argument("--image", default=f"hf.co/spaces/{routing.space_repo}")
+    parser.add_argument("--namespace", default=routing.job_namespace)
+    parser.add_argument("--flavor", default="a10g-small")
+    parser.add_argument("--timeout", default="30m")
+    parser.add_argument("--summary-out", type=Path)
+    parser.add_argument("--dry-run", action="store_true")
+    parser.add_argument("--refresh-image", action="store_true")
+    parser.add_argument("--sync-overlay", action="store_true")
+    return parser.parse_args(argv)
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = parse_args(argv)
+    env = build_benchmark_job_env(
+        benchmark=args.benchmark,
+        variant=args.variant,
+        seed=args.seed,
+        output_repo=args.output_repo,
+        tokenizer_repo=args.tokenizer_repo,
+    )
+    command = build_benchmark_job_command(benchmark=args.benchmark, variant=args.variant, seed=args.seed)
+    payload = {
+        "benchmark": args.benchmark,
+        "variant": args.variant,
+        "seed": args.seed,
+        "output_repo": args.output_repo,
+        "tokenizer_repo": args.tokenizer_repo,
+        "image": args.image,
+        "namespace": args.namespace,
+        "command": command,
+        "env": env,
+        "dry_run": args.dry_run,
+    }
+    if not args.dry_run:
+        token = os.environ.get("HF_TOKEN") or get_token()
+        if not token:
+            raise SystemExit("HF_TOKEN must be set or cached via huggingface-cli login")
+        api = HfApi(token=token)
+        if args.refresh_image:
+            space_repo = args.image.removeprefix("hf.co/spaces/")
+            if args.sync_overlay:
+                sync_overlay_from_repo()
+            api.upload_folder(
+                repo_id=space_repo,
+                repo_type="space",
+                folder_path=str(IMAGE_DIR),
+                commit_message="Update benchmark runtime image",
+                token=token,
+            )
+            wait_for_space(api, space_repo, token=token)
+        payload.update(
+            submit_benchmark_job(
+                api=api,
+                image=args.image,
+                command=command,
+                env=env,
+                token=token,
+                namespace=args.namespace,
+                flavor=args.flavor,
+                timeout=args.timeout,
+            )
+        )
+    if args.summary_out is not None:
+        args.summary_out.parent.mkdir(parents=True, exist_ok=True)
+        args.summary_out.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
+    print(json.dumps(payload, indent=2, sort_keys=True))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

overlay/scripts/launch_feather_hf_job.py

@@ -1,224 +1,218 @@
 #!/usr/bin/env python3
 from __future__ import annotations
 
-import os
-import shutil
-import sys
-import time
-import json
-from collections.abc import Mapping, Sequence
-from typing import Any, cast
-from pathlib import Path
-
-import httpx
-from huggingface_hub import HfApi
-from huggingface_hub.utils import HfHubHTTPError
-from huggingface_hub.utils import get_token
-
-REPO_ROOT = Path(__file__).resolve().parents[1]
-if str(REPO_ROOT) not in sys.path:
-    sys.path.insert(0, str(REPO_ROOT))
-
-from scripts.hf_routing import resolve_routing
-from configs.harness_config import HarnessConfig
-
-DEFAULT_IMAGE = os.environ.get('FEATHER_HF_IMAGE', 'ghcr.io/slapglif/feather-hf-runtime:latest')
-IMAGE_DIR = Path(__file__).resolve().parents[1] / 'hf_jobs' / 'feather_h200_image'
-TIMEOUT = os.environ.get('FEATHER_HF_JOB_TIMEOUT', '12h')
-TARGET_SHARDS = os.environ.get('HYDRA_TARGET_SHARDS', '2048')
-TIME_BUDGET = os.environ.get('HYDRA_TIME_BUDGET', '43200')
-DOWNLOAD_WORKERS = os.environ.get('HYDRA_DOWNLOAD_WORKERS', '16')
-CKPT_INTERVAL = os.environ.get('HYDRA_CKPT_INTERVAL', '1000')
-JOB_FLAVOR = os.environ.get('FEATHER_HF_FLAVOR', 'a10g-small')
-DRY_RUN = os.environ.get('FEATHER_HF_DRY_RUN', '0') == '1'
-USE_SPACE_IMAGE = os.environ.get('FEATHER_HF_USE_SPACE_IMAGE', '0') == '1'
+import os
+import shutil
+import sys
+import time
+import json
+from collections.abc import Mapping, Sequence
+from typing import Any, cast
+from pathlib import Path
+
+import httpx
+from huggingface_hub import HfApi
+from huggingface_hub.utils import HfHubHTTPError
+from huggingface_hub.utils import get_token
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from scripts.hf_routing import resolve_routing
+from configs.harness_config import HarnessConfig
+
+DEFAULT_IMAGE = os.environ.get('FEATHER_HF_IMAGE', 'ghcr.io/slapglif/feather-hf-runtime:latest')
+IMAGE_DIR = Path(__file__).resolve().parents[1] / 'hf_jobs' / 'feather_h200_image'
+TIMEOUT = os.environ.get('FEATHER_HF_JOB_TIMEOUT', '12h')
+TARGET_SHARDS = os.environ.get('HYDRA_TARGET_SHARDS', '2048')
+TIME_BUDGET = os.environ.get('HYDRA_TIME_BUDGET', '43200')
+DOWNLOAD_WORKERS = os.environ.get('HYDRA_DOWNLOAD_WORKERS', '16')
+CKPT_INTERVAL = os.environ.get('HYDRA_CKPT_INTERVAL', '1000')
+JOB_FLAVOR = os.environ.get('FEATHER_HF_FLAVOR', 'a10g-small')
+DRY_RUN = os.environ.get('FEATHER_HF_DRY_RUN', '0') == '1'
+USE_SPACE_IMAGE = os.environ.get('FEATHER_HF_USE_SPACE_IMAGE', '0') == '1'
 # When true, assume the Space image has already been built by a previous
 # invocation and skip the upload+build wait. Used by sweep drivers that fan
 # out many jobs against a single pre-uploaded image.
-SKIP_UPLOAD = os.environ.get('FEATHER_HF_SKIP_UPLOAD', '0') == '1'
-SYNC_OVERLAY = os.environ.get('FEATHER_HF_SYNC_OVERLAY', '1') == '1'
-JOB_SUBMIT_RETRIES = max(1, int(os.environ.get('FEATHER_HF_JOB_SUBMIT_RETRIES', '3')))
-JOB_SUBMIT_RETRY_BASE_S = float(os.environ.get('FEATHER_HF_JOB_SUBMIT_RETRY_BASE_S', '5'))
-BUILD_LOG_TAIL_LINES = max(1, int(os.environ.get('FEATHER_HF_BUILD_LOG_TAIL_LINES', '120')))
-
-
-def should_enable_fast_start_streaming(target_shards: str, time_budget: str) -> bool:
-    """Use streaming data path for short-budget launch profiles."""
-    try:
-        shards = int(target_shards)
-        budget = int(time_budget)
-    except ValueError:
-        return False
-    return shards > 0 and shards <= 256 and budget > 0 and budget <= 1800
-
-
-def apply_a10_env_profile(
-    env: dict[str, str],
-    *,
-    job_flavor: str,
-    parent_env: Mapping[str, str] = os.environ,
-) -> str | None:
-    if not job_flavor.startswith('a10'):
-        return None
-
-    full_arch = parent_env.get('HYDRA_THROUGHPUT_MODE') == '0' or env.get('HYDRA_THROUGHPUT_MODE') == '0'
-    if full_arch:
-        defaults = {
-            'HYDRA_THROUGHPUT_MODE': '0',
-            'HYDRA_MUON_COMPILE': '0',
-            'HYDRA_FORCE_HTM_CPU': '0',
-            'HYDRA_INERT_MAMBA': '0',
-            'HYDRA_ALLOW_SYNTHETIC_RETINA': '0',
-            'HYDRA_FASTPATH': '0',
-        }
-        profile = 'full-architecture'
-    else:
-        defaults = {
-            'HYDRA_MUON_COMPILE': '0',
-            'HYDRA_FORCE_HTM_CPU': '1',
-            'HYDRA_INERT_MAMBA': '1',
-            'HYDRA_ALLOW_SYNTHETIC_RETINA': '1',
-            'HYDRA_FASTPATH': '1',
-        }
-        profile = 'compatibility'
-
-    for key, default in defaults.items():
-        env[key] = parent_env.get(key, env.get(key, default))
-    return profile
-
-
-def _http_status(exc: BaseException) -> int | None:
-    response = getattr(exc, 'response', None)
-    status = getattr(response, 'status_code', None)
-    return int(status) if isinstance(status, int) else None
-
-
-def submit_job_with_retry(
-    api: HfApi,
-    *,
-    image: str,
-    command: Sequence[str],
-    env: dict[str, str],
-    secrets: dict[str, str],
-    flavor: Any,
-    timeout: str,
-    token: str,
-    namespace: str,
-):
-    for attempt in range(1, JOB_SUBMIT_RETRIES + 1):
-        try:
-            return api.run_job(
-                image=image,
-                command=list(command),
-                env=env,
-                secrets=secrets,
-                flavor=flavor,
-                timeout=timeout,
-                namespace=namespace,
-                token=token,
-            )
-        except HfHubHTTPError as exc:
-            status = _http_status(exc)
-            if status is None or status < 500 or status >= 600:
-                raise
-            if attempt >= JOB_SUBMIT_RETRIES:
-                raise SystemExit(
-                    f'HF job submit failed after {JOB_SUBMIT_RETRIES} attempts '
-                    f'(http {status}); failing fast'
-                ) from exc
-            time.sleep(JOB_SUBMIT_RETRY_BASE_S * attempt)
-    raise SystemExit('HF job submit failed unexpectedly; failing fast')
-
-
-def fetch_space_build_log_tail(
-    api: HfApi,
-    repo_id: str,
-    token: str,
-    *,
-    limit: int = BUILD_LOG_TAIL_LINES,
-) -> str:
-    try:
-        lines = list(api.fetch_space_logs(repo_id, build=True, follow=False, token=token))
-    except Exception as exc:
-        return f'[space-build-log] failed to fetch build logs: {exc!r}'
-    tail = lines[-limit:]
-    text = ''.join(tail)
-    if text and not text.endswith('\n'):
-        text += '\n'
-    return text or '[space-build-log] no buffered build logs available\n'
-
-
-def sync_overlay_from_repo() -> None:
-    """Refresh Space overlay with required project files."""
-    overlay = IMAGE_DIR / 'overlay'
-    overlay.mkdir(parents=True, exist_ok=True)
-
-    for child in overlay.iterdir():
-        if child.is_dir():
-            shutil.rmtree(child)
-        else:
-            child.unlink()
-
-    include_paths = [
-        'hydra',
-        'subsystems',
-        'scripts',
-        'data',
-        'htm_rust',
-        'harness',
-        'configs',
-        'artifacts',
-        'prepare.py',
-        'prepare_nemotron.py',
-        'train.py',
-        'pyproject.toml',
-        'uv.lock',
-    ]
-    ignore = shutil.ignore_patterns(
-        '__pycache__',
-        '.pytest_cache',
-        '.ruff_cache',
-        '.venv',
-        '.git',
-        'target',
-        '*.pyc',
-        '_tmp*',
-        'cycle1a_runs',
-        'cycle1a_probe',
-        'remote_benchmark_submission*.json',
-    )
-
-    copied: list[str] = []
-    for rel in include_paths:
-        src = REPO_ROOT / rel
-        dst = overlay / rel
-        if not src.exists():
-            continue
-        if src.is_dir():
-            shutil.copytree(src, dst, dirs_exist_ok=True, ignore=ignore)
-        else:
-            dst.parent.mkdir(parents=True, exist_ok=True)
-            shutil.copy2(src, dst)
-        copied.append(rel)
-
-    scripts_dir = overlay / 'scripts'
-    if scripts_dir.exists():
-        for sh_path in scripts_dir.rglob('*.sh'):
-            data = sh_path.read_bytes()
-            data = data.replace(b'\r\n', b'\n').replace(b'\r', b'\n')
-            sh_path.write_bytes(data)
-
-    print(f'[launch] overlay synced from repo ({len(copied)} paths): {copied}', flush=True)
-
-
-def require_token() -> str:
-    token = os.environ.get('HF_TOKEN') or get_token()
-    if not token:
-        raise SystemExit('HF_TOKEN must be set or cached via huggingface-cli login for launch_feather_hf_job.py')
-    return token
-
-
-def wait_for_space(api: HfApi, repo_id: str, token: str, timeout_s: int = 1800) -> None:
+SKIP_UPLOAD = os.environ.get('FEATHER_HF_SKIP_UPLOAD', '0') == '1'
+SYNC_OVERLAY = os.environ.get('FEATHER_HF_SYNC_OVERLAY', '1') == '1'
+JOB_SUBMIT_RETRIES = max(1, int(os.environ.get('FEATHER_HF_JOB_SUBMIT_RETRIES', '3')))
+JOB_SUBMIT_RETRY_BASE_S = float(os.environ.get('FEATHER_HF_JOB_SUBMIT_RETRY_BASE_S', '5'))
+BUILD_LOG_TAIL_LINES = max(1, int(os.environ.get('FEATHER_HF_BUILD_LOG_TAIL_LINES', '120')))
+
+
+def should_enable_fast_start_streaming(target_shards: str, time_budget: str) -> bool:
+    """Use streaming data path for short-budget launch profiles."""
+    try:
+        shards = int(target_shards)
+        budget = int(time_budget)
+    except ValueError:
+        return False
+    return shards > 0 and shards <= 256 and budget > 0 and budget <= 1800
+
+
+def apply_a10_env_profile(
+    env: dict[str, str],
+    *,
+    job_flavor: str,
+    parent_env: Mapping[str, str] = os.environ,
+) -> str | None:
+    if not job_flavor.startswith('a10'):
+        return None
+
+    full_arch = parent_env.get('HYDRA_THROUGHPUT_MODE') == '0' or env.get('HYDRA_THROUGHPUT_MODE') == '0'
+    if full_arch:
+        defaults = {
+            'HYDRA_THROUGHPUT_MODE': '0',
+            'HYDRA_MUON_COMPILE': '0',
+            'HYDRA_FORCE_HTM_CPU': '0',
+            'HYDRA_INERT_MAMBA': '0',
+            'HYDRA_ALLOW_SYNTHETIC_RETINA': '0',
+            'HYDRA_FASTPATH': '0',
+        }
+        profile = 'full-architecture'
+    else:
+        defaults = {
+            'HYDRA_MUON_COMPILE': '0',
+            'HYDRA_FORCE_HTM_CPU': '1',
+            'HYDRA_INERT_MAMBA': '1',
+            'HYDRA_ALLOW_SYNTHETIC_RETINA': '1',
+            'HYDRA_FASTPATH': '1',
+        }
+        profile = 'compatibility'
+
+    for key, default in defaults.items():
+        env[key] = parent_env.get(key, env.get(key, default))
+    return profile
+
+
+def _http_status(exc: BaseException) -> int | None:
+    response = getattr(exc, 'response', None)
+    status = getattr(response, 'status_code', None)
+    return int(status) if isinstance(status, int) else None
+
+
+def submit_job_with_retry(
+    api: HfApi,
+    *,
+    image: str,
+    command: Sequence[str],
+    env: dict[str, str],
+    secrets: dict[str, str],
+    flavor: Any,
+    timeout: str,
+    token: str,
+    namespace: str,
+):
+    for attempt in range(1, JOB_SUBMIT_RETRIES + 1):
+        try:
+            return api.run_job(
+                image=image,
+                command=list(command),
+                env=env,
+                secrets=secrets,
+                flavor=flavor,
+                timeout=timeout,
+                namespace=namespace,
+                token=token,
+            )
+        except HfHubHTTPError as exc:
+            status = _http_status(exc)
+            if status is None or status < 500 or status >= 600:
+                raise
+            if attempt >= JOB_SUBMIT_RETRIES:
+                raise SystemExit(
+                    f'HF job submit failed after {JOB_SUBMIT_RETRIES} attempts '
+                    f'(http {status}); failing fast'
+                ) from exc
+            time.sleep(JOB_SUBMIT_RETRY_BASE_S * attempt)
+    raise SystemExit('HF job submit failed unexpectedly; failing fast')
+
+
+def fetch_space_build_log_tail(
+    api: HfApi,
+    repo_id: str,
+    token: str,
+    *,
+    limit: int = BUILD_LOG_TAIL_LINES,
+) -> str:
+    try:
+        lines = list(api.fetch_space_logs(repo_id, build=True, follow=False, token=token))
+    except Exception as exc:
+        return f'[space-build-log] failed to fetch build logs: {exc!r}'
+    tail = lines[-limit:]
+    text = ''.join(tail)
+    if text and not text.endswith('\n'):
+        text += '\n'
+    return text or '[space-build-log] no buffered build logs available\n'
+
+
+def sync_overlay_from_repo() -> None:
+    """Refresh Space overlay with required project files."""
+    overlay = IMAGE_DIR / 'overlay'
+    overlay.mkdir(parents=True, exist_ok=True)
+
+    for child in overlay.iterdir():
+        if child.is_dir():
+            shutil.rmtree(child)
+        else:
+            child.unlink()
+
+    include_paths = [
+        'hydra',
+        'subsystems',
+        'scripts',
+        'htm_rust',
+        'harness',
+        'configs',
+        'prepare.py',
+        'prepare_nemotron.py',
+        'train.py',
+        'pyproject.toml',
+        'uv.lock',
+    ]
+    ignore = shutil.ignore_patterns(
+        '__pycache__',
+        '.pytest_cache',
+        '.ruff_cache',
+        '.venv',
+        '.git',
+        'target',
+        '*.pyc',
+    )
+
+    copied: list[str] = []
+    for rel in include_paths:
+        src = REPO_ROOT / rel
+        dst = overlay / rel
+        if not src.exists():
+            continue
+        if src.is_dir():
+            shutil.copytree(src, dst, dirs_exist_ok=True, ignore=ignore)
+        else:
+            dst.parent.mkdir(parents=True, exist_ok=True)
+            shutil.copy2(src, dst)
+        copied.append(rel)
+
+    scripts_dir = overlay / 'scripts'
+    if scripts_dir.exists():
+        for sh_path in scripts_dir.rglob('*.sh'):
+            data = sh_path.read_bytes()
+            data = data.replace(b'\r\n', b'\n').replace(b'\r', b'\n')
+            sh_path.write_bytes(data)
+
+    print(f'[launch] overlay synced from repo ({len(copied)} paths): {copied}', flush=True)
+
+
+def require_token() -> str:
+    token = os.environ.get('HF_TOKEN') or get_token()
+    if not token:
+        raise SystemExit('HF_TOKEN must be set or cached via huggingface-cli login for launch_feather_hf_job.py')
+    return token
+
+
+def wait_for_space(api: HfApi, repo_id: str, token: str, timeout_s: int = 1800) -> None:
     """Wait until the Space image has been built.
 
     We use the Space purely as a container-image builder for HF Jobs. The Space
@@ -233,134 +227,134 @@ def wait_for_space(api: HfApi, repo_id: str, token: str, timeout_s: int = 1800)
     and APP_STARTING_ERROR after a successful BUILDING→APP_STARTING transition
     are acceptable — the image exists in the registry and Jobs can use it.
     """
-    start = time.time()
-    seen_build_completion = False
-    seen_building = False
-    while True:
-        try:
-            runtime = api.get_space_runtime(repo_id, token=token)
-        except httpx.TransportError as exc:
-            if time.time() - start > timeout_s:
-                raise TimeoutError(f'Space {repo_id} runtime endpoint kept failing with network errors') from exc
-            print(f'[space] transient runtime endpoint network error: {exc!r}; retrying', flush=True)
-            time.sleep(20)
-            continue
-        except HfHubHTTPError as exc:
-            status = _http_status(exc)
-            if status is not None and 500 <= status < 600:
-                if time.time() - start > timeout_s:
-                    raise TimeoutError(f'Space {repo_id} runtime endpoint kept failing with HTTP {status}') from exc
-                time.sleep(20)
-                continue
-            raise
-        stage = getattr(runtime, 'stage', None)
-        hardware = getattr(runtime, 'hardware', None)
-        err = getattr(runtime, 'errorMessage', None) or getattr(runtime, 'error_message', None)
-        print(f'[space] stage={stage} hardware={hardware}', flush=True)
-        if stage == 'BUILDING':
-            seen_building = True
-        if stage in {'APP_STARTING', 'RUNNING', 'PAUSED', 'SLEEPING'}:
-            seen_build_completion = True
-        if stage in {'RUNNING', 'PAUSED', 'SLEEPING'}:
-            return
-        # Image is built — Jobs can use it regardless of Space boot outcome.
-        if (seen_build_completion or seen_building) and stage in {'RUNTIME_ERROR', 'APP_STARTING_ERROR'}:
-            print(f'[space] Space boot failed with {stage} but built image is '
-                  f'available in the Space registry and is usable by HF Jobs.',
-                  flush=True)
-            return
-        # Hard build failures — no image was produced.
-        if stage in {'BUILD_ERROR', 'CONFIG_ERROR', 'NO_APP_FILE'}:
-            build_log_tail = fetch_space_build_log_tail(api, repo_id, token)
-            raise RuntimeError(
-                f'Space {repo_id} build failed: stage={stage} error={err!r}\n'
-                f'--- Space build log tail ---\n{build_log_tail}'
-            )
+    start = time.time()
+    seen_build_completion = False
+    seen_building = False
+    while True:
+        try:
+            runtime = api.get_space_runtime(repo_id, token=token)
+        except httpx.TransportError as exc:
+            if time.time() - start > timeout_s:
+                raise TimeoutError(f'Space {repo_id} runtime endpoint kept failing with network errors') from exc
+            print(f'[space] transient runtime endpoint network error: {exc!r}; retrying', flush=True)
+            time.sleep(20)
+            continue
+        except HfHubHTTPError as exc:
+            status = _http_status(exc)
+            if status is not None and 500 <= status < 600:
+                if time.time() - start > timeout_s:
+                    raise TimeoutError(f'Space {repo_id} runtime endpoint kept failing with HTTP {status}') from exc
+                time.sleep(20)
+                continue
+            raise
+        stage = getattr(runtime, 'stage', None)
+        hardware = getattr(runtime, 'hardware', None)
+        err = getattr(runtime, 'errorMessage', None) or getattr(runtime, 'error_message', None)
+        print(f'[space] stage={stage} hardware={hardware}', flush=True)
+        if stage == 'BUILDING':
+            seen_building = True
+        if stage in {'APP_STARTING', 'RUNNING', 'PAUSED', 'SLEEPING'}:
+            seen_build_completion = True
+        if stage in {'RUNNING', 'PAUSED', 'SLEEPING'}:
+            return
+        # Image is built — Jobs can use it regardless of Space boot outcome.
+        if (seen_build_completion or seen_building) and stage in {'RUNTIME_ERROR', 'APP_STARTING_ERROR'}:
+            print(f'[space] Space boot failed with {stage} but built image is '
+                  f'available in the Space registry and is usable by HF Jobs.',
+                  flush=True)
+            return
+        # Hard build failures — no image was produced.
+        if stage in {'BUILD_ERROR', 'CONFIG_ERROR', 'NO_APP_FILE'}:
+            build_log_tail = fetch_space_build_log_tail(api, repo_id, token)
+            raise RuntimeError(
+                f'Space {repo_id} build failed: stage={stage} error={err!r}\n'
+                f'--- Space build log tail ---\n{build_log_tail}'
+            )
         if time.time() - start > timeout_s:
             raise TimeoutError(f'Space {repo_id} did not become ready in {timeout_s}s (last stage={stage})')
         time.sleep(20)
 
 
-def main() -> int:
-    token = require_token()
-    routing = resolve_routing(token=token)
-    api = HfApi(token=token)
-    secondary_gates = HarnessConfig().to_secondary_gates()
-
-    print(f'[launch] image_dir={IMAGE_DIR}', flush=True)
-    print(f'[launch] owner={routing.owner}', flush=True)
-    print(f'[launch] space_repo={routing.space_repo}', flush=True)
-    print(f'[launch] output_repo={routing.output_repo}', flush=True)
-    print(f'[launch] retina_cache_repo={routing.retina_cache_repo}', flush=True)
-    print(f'[launch] target_shards={TARGET_SHARDS} time_budget={TIME_BUDGET} timeout={TIMEOUT}', flush=True)
-    print(f'[launch] flavor={JOB_FLAVOR}', flush=True)
-    print(f'[launch] namespace={routing.job_namespace}', flush=True)
-    print(f'[launch] image_mode={"space" if USE_SPACE_IMAGE else "ghcr"}', flush=True)
-    print(f'[launch] secondary_gates={json.dumps(secondary_gates, sort_keys=True)}', flush=True)
-    if not USE_SPACE_IMAGE:
-        print(f'[launch] image={DEFAULT_IMAGE}', flush=True)
-
-    api.create_repo(repo_id=routing.space_repo, repo_type='space', space_sdk='docker', private=True, exist_ok=True, token=token)
-    api.create_repo(repo_id=routing.output_repo, repo_type='model', private=True, exist_ok=True, token=token)
+def main() -> int:
+    token = require_token()
+    routing = resolve_routing(token=token)
+    api = HfApi(token=token)
+    secondary_gates = HarnessConfig().to_secondary_gates()
+
+    print(f'[launch] image_dir={IMAGE_DIR}', flush=True)
+    print(f'[launch] owner={routing.owner}', flush=True)
+    print(f'[launch] space_repo={routing.space_repo}', flush=True)
+    print(f'[launch] output_repo={routing.output_repo}', flush=True)
+    print(f'[launch] retina_cache_repo={routing.retina_cache_repo}', flush=True)
+    print(f'[launch] target_shards={TARGET_SHARDS} time_budget={TIME_BUDGET} timeout={TIMEOUT}', flush=True)
+    print(f'[launch] flavor={JOB_FLAVOR}', flush=True)
+    print(f'[launch] namespace={routing.job_namespace}', flush=True)
+    print(f'[launch] image_mode={"space" if USE_SPACE_IMAGE else "ghcr"}', flush=True)
+    print(f'[launch] secondary_gates={json.dumps(secondary_gates, sort_keys=True)}', flush=True)
+    if not USE_SPACE_IMAGE:
+        print(f'[launch] image={DEFAULT_IMAGE}', flush=True)
+
+    api.create_repo(repo_id=routing.space_repo, repo_type='space', space_sdk='docker', private=True, exist_ok=True, token=token)
+    api.create_repo(repo_id=routing.output_repo, repo_type='model', private=True, exist_ok=True, token=token)
 
     if DRY_RUN:
         print('[launch] dry-run mode; skipping upload and job submission', flush=True)
         return 0
 
-    image_ref = DEFAULT_IMAGE
-    if USE_SPACE_IMAGE:
-        if SKIP_UPLOAD:
-            print('[launch] FEATHER_HF_SKIP_UPLOAD=1; reusing existing Space image', flush=True)
-        else:
-            if SYNC_OVERLAY:
-                sync_overlay_from_repo()
-            print('[launch] uploading custom Docker Space image context...', flush=True)
-            api.upload_folder(
-                repo_id=routing.space_repo,
-                repo_type='space',
-                folder_path=str(IMAGE_DIR),
-                commit_message='Update Feather training runtime image',
-                token=token,
-            )
-
-            print('[launch] waiting for Space image build to become ready...', flush=True)
-            wait_for_space(api, routing.space_repo, token=token)
-        image_ref = f'hf.co/spaces/{routing.space_repo}'
-
-    env = {
-        'HF_REPO_ID': routing.output_repo,
-        'FEATHER_HF_OWNER': routing.owner,
-        'FEATHER_HF_SPACE_REPO': routing.space_repo,
-        'FEATHER_HF_OUTPUT_REPO': routing.output_repo,
-        'FEATHER_HF_RETINA_CACHE_REPO': routing.retina_cache_repo,
-        'HYDRA_RETINA_CACHE_REPO': routing.retina_cache_repo,
-        'HYDRA_TARGET_SHARDS': TARGET_SHARDS,
-        'HYDRA_TIME_BUDGET': TIME_BUDGET,
-        'HYDRA_DOWNLOAD_WORKERS': DOWNLOAD_WORKERS,
+    image_ref = DEFAULT_IMAGE
+    if USE_SPACE_IMAGE:
+        if SKIP_UPLOAD:
+            print('[launch] FEATHER_HF_SKIP_UPLOAD=1; reusing existing Space image', flush=True)
+        else:
+            if SYNC_OVERLAY:
+                sync_overlay_from_repo()
+            print('[launch] uploading custom Docker Space image context...', flush=True)
+            api.upload_folder(
+                repo_id=routing.space_repo,
+                repo_type='space',
+                folder_path=str(IMAGE_DIR),
+                commit_message='Update Feather training runtime image',
+                token=token,
+            )
+
+            print('[launch] waiting for Space image build to become ready...', flush=True)
+            wait_for_space(api, routing.space_repo, token=token)
+        image_ref = f'hf.co/spaces/{routing.space_repo}'
+
+    env = {
+        'HF_REPO_ID': routing.output_repo,
+        'FEATHER_HF_OWNER': routing.owner,
+        'FEATHER_HF_SPACE_REPO': routing.space_repo,
+        'FEATHER_HF_OUTPUT_REPO': routing.output_repo,
+        'FEATHER_HF_RETINA_CACHE_REPO': routing.retina_cache_repo,
+        'HYDRA_RETINA_CACHE_REPO': routing.retina_cache_repo,
+        'HYDRA_TARGET_SHARDS': TARGET_SHARDS,
+        'HYDRA_TIME_BUDGET': TIME_BUDGET,
+        'HYDRA_DOWNLOAD_WORKERS': DOWNLOAD_WORKERS,
         'HYDRA_CKPT_INTERVAL': CKPT_INTERVAL,
         'PYTHONUNBUFFERED': '1',
-        'FEATHER_RUNTIME_MODE': 'job',
-    }
-    if 'HYDRA_USE_NEMOTRON' not in os.environ and should_enable_fast_start_streaming(TARGET_SHARDS, TIME_BUDGET):
-        env['HYDRA_USE_NEMOTRON'] = '1'
-        print('[launch] auto-enabled HYDRA_USE_NEMOTRON=1 for short-budget fast-start profile', flush=True)
-    # A10 profile: default compatibility mode avoids known PTX/compile runtime
-    # pitfalls, but HYDRA_THROUGHPUT_MODE=0 explicitly selects the full
-    # SDR/HTM/Engram architecture instead of silently inheriting bypass defaults.
-    _a10_profile = apply_a10_env_profile(env, job_flavor=JOB_FLAVOR)
-    if _a10_profile is not None:
-        if env.get('HYDRA_INERT_MAMBA') == '0' and 'HYDRA_FASTPATH' not in os.environ:
-            env['HYDRA_FASTPATH'] = '0'
-        print(
-            f'[launch] applied A10 {_a10_profile} env profile '
-            f"(HYDRA_MUON_COMPILE={env['HYDRA_MUON_COMPILE']}, "
-            f"HYDRA_THROUGHPUT_MODE={env.get('HYDRA_THROUGHPUT_MODE', 'unset')}, "
-            f"HYDRA_FORCE_HTM_CPU={env['HYDRA_FORCE_HTM_CPU']}, "
-            f"HYDRA_INERT_MAMBA={env['HYDRA_INERT_MAMBA']}, "
-            f"HYDRA_ALLOW_SYNTHETIC_RETINA={env['HYDRA_ALLOW_SYNTHETIC_RETINA']}, "
-            f"HYDRA_FASTPATH={env['HYDRA_FASTPATH']})",
-            flush=True,
-        )
+        'FEATHER_RUNTIME_MODE': 'job',
+    }
+    if 'HYDRA_USE_NEMOTRON' not in os.environ and should_enable_fast_start_streaming(TARGET_SHARDS, TIME_BUDGET):
+        env['HYDRA_USE_NEMOTRON'] = '1'
+        print('[launch] auto-enabled HYDRA_USE_NEMOTRON=1 for short-budget fast-start profile', flush=True)
+    # A10 profile: default compatibility mode avoids known PTX/compile runtime
+    # pitfalls, but HYDRA_THROUGHPUT_MODE=0 explicitly selects the full
+    # SDR/HTM/Engram architecture instead of silently inheriting bypass defaults.
+    _a10_profile = apply_a10_env_profile(env, job_flavor=JOB_FLAVOR)
+    if _a10_profile is not None:
+        if env.get('HYDRA_INERT_MAMBA') == '0' and 'HYDRA_FASTPATH' not in os.environ:
+            env['HYDRA_FASTPATH'] = '0'
+        print(
+            f'[launch] applied A10 {_a10_profile} env profile '
+            f"(HYDRA_MUON_COMPILE={env['HYDRA_MUON_COMPILE']}, "
+            f"HYDRA_THROUGHPUT_MODE={env.get('HYDRA_THROUGHPUT_MODE', 'unset')}, "
+            f"HYDRA_FORCE_HTM_CPU={env['HYDRA_FORCE_HTM_CPU']}, "
+            f"HYDRA_INERT_MAMBA={env['HYDRA_INERT_MAMBA']}, "
+            f"HYDRA_ALLOW_SYNTHETIC_RETINA={env['HYDRA_ALLOW_SYNTHETIC_RETINA']}, "
+            f"HYDRA_FASTPATH={env['HYDRA_FASTPATH']})",
+            flush=True,
+        )
     # Pass through any HYDRA_* / FEATHER_* overrides from the caller's env so
     # sweep drivers can set HYDRA_N_LAYER, HYDRA_SDR_TARGET_ACTIVE,
     # HYDRA_LAYER_DIAGNOSTICS, HYDRA_METRICS_OUT, HYDRA_MID_VAL_INTERVAL, etc.
@@ -370,18 +364,18 @@ def main() -> int:
             env[_k] = _v
     secrets = {'HF_TOKEN': token}
 
-    print(f'[launch] submitting HF Job on flavor={JOB_FLAVOR}...', flush=True)
-    job = submit_job_with_retry(
-        api,
-        image=image_ref,
-        command=['python', '/app/entrypoint.py'],
-        env=env,
-        secrets=secrets,
-        flavor=cast(Any, JOB_FLAVOR),
-        timeout=TIMEOUT,
-        namespace=routing.job_namespace,
-        token=token,
-    )
+    print(f'[launch] submitting HF Job on flavor={JOB_FLAVOR}...', flush=True)
+    job = submit_job_with_retry(
+        api,
+        image=image_ref,
+        command=['python', '/app/entrypoint.py'],
+        env=env,
+        secrets=secrets,
+        flavor=cast(Any, JOB_FLAVOR),
+        timeout=TIMEOUT,
+        namespace=routing.job_namespace,
+        token=token,
+    )
     print(f'[launch] submitted job_id={job.id} status={job.status.stage} url={job.url}', flush=True)
     return 0
 

overlay/scripts/optuna_hpo.py

@@ -5,131 +5,131 @@ import argparse
 import json
 import os
 import re
-import subprocess
-import sys
-import time
-import tempfile
-from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeoutError
-from pathlib import Path
-from typing import Any
-
-import optuna
-
-
-_HF_ENV_KEY_RE = re.compile(r"^[A-Z][A-Z0-9_]*$")
-
-
-REPO_ROOT = Path(__file__).resolve().parents[1]
-if str(REPO_ROOT) not in sys.path:
-    sys.path.insert(0, str(REPO_ROOT))
-
-from scripts.hf_routing import resolve_routing
-
-TRAIN_ENTRYPOINT = REPO_ROOT / "train.py"
-SEARCH_SPACE_KEYS = {
-    "d_model",
-    "n_layer",
-    "d_state",
-    "headdim",
-    "expand",
-    "seq_len",
-    "batch_size",
-    "grad_accum",
-    "matrix_lr",
-    "embed_lr",
-    "unembed_lr",
-    "hyena_layers",
-    "engram_n_columns",
-    "engram_layer_idx",
-    "sdr_target_active",
-    "htm_learn_every",
-    "htm_subsample",
-    "engram_subsample",
-    "mamba3_chunk",
-    "dropout",
-}
-
-
-def _filter_prior_params(raw: dict[str, Any]) -> dict[str, Any]:
-    return {k: v for k, v in raw.items() if k in SEARCH_SPACE_KEYS}
-
-
-def _load_prior_param_sets(path: Path) -> list[dict[str, Any]]:
-    if not path.exists():
-        return []
-
-    payload = json.loads(path.read_text(encoding="utf-8"))
-    if isinstance(payload, dict):
-        rows = payload.get("trials", [])
-    elif isinstance(payload, list):
-        rows = payload
-    else:
-        rows = []
-
-    out: list[dict[str, Any]] = []
-    for item in rows:
-        if not isinstance(item, dict):
-            continue
-        params_obj = item.get("params", item)
-        if not isinstance(params_obj, dict):
-            continue
-        filtered = _filter_prior_params(params_obj)
-        if filtered:
-            out.append(filtered)
-    return out
-
-
-def _enqueue_transfer_priors(study: optuna.Study, priors_file: Path, apply_priors: bool) -> int:
-    if not apply_priors:
-        return 0
-
-    priors_raw = _load_prior_param_sets(priors_file)
-    if not priors_raw:
-        return 0
-
-    # Deduplicate param sets across merged studies.
-    priors: list[dict[str, Any]] = []
-    seen: set[str] = set()
-    for params in priors_raw:
-        key = json.dumps(params, sort_keys=True)
-        if key in seen:
-            continue
-        seen.add(key)
-        priors.append(params)
-
-    enqueued = 0
-    for params in priors:
-        before = len(study.get_trials(deepcopy=False))
-        try:
-            study.enqueue_trial(params, user_attrs={"seed_source": "transfer_priors"}, skip_if_exists=True)
-        except TypeError:
-            study.enqueue_trial(params, user_attrs={"seed_source": "transfer_priors"})
-        after = len(study.get_trials(deepcopy=False))
-        if after > before:
-            enqueued += 1
-    return enqueued
-
-
-def _enqueue_quality_anchors(study: optuna.Study, priors_file: Path, quality_mode_local: bool, top_k: int) -> int:
-    if not quality_mode_local or top_k <= 0:
-        return 0
-
-    priors = _load_prior_param_sets(priors_file)[:top_k]
-    enqueued = 0
-    for params in priors:
-        before = len(study.get_trials(deepcopy=False))
-        try:
-            study.enqueue_trial(
-                params,
-                user_attrs={"seed_source": "quality_anchor"},
-                skip_if_exists=True,
-            )
-        except TypeError:
-            study.enqueue_trial(params, user_attrs={"seed_source": "quality_anchor"})
-        after = len(study.get_trials(deepcopy=False))
-        if after > before:
-            enqueued += 1
-    return enqueued
+import subprocess
+import sys
+import time
+import tempfile
+from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeoutError
+from pathlib import Path
+from typing import Any
+
+import optuna
+
+
+_HF_ENV_KEY_RE = re.compile(r"^[A-Z][A-Z0-9_]*$")
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from scripts.hf_routing import resolve_routing
+
+TRAIN_ENTRYPOINT = REPO_ROOT / "train.py"
+SEARCH_SPACE_KEYS = {
+    "d_model",
+    "n_layer",
+    "d_state",
+    "headdim",
+    "expand",
+    "seq_len",
+    "batch_size",
+    "grad_accum",
+    "matrix_lr",
+    "embed_lr",
+    "unembed_lr",
+    "hyena_layers",
+    "engram_n_columns",
+    "engram_layer_idx",
+    "sdr_target_active",
+    "htm_learn_every",
+    "htm_subsample",
+    "engram_subsample",
+    "mamba3_chunk",
+    "dropout",
+}
+
+
+def _filter_prior_params(raw: dict[str, Any]) -> dict[str, Any]:
+    return {k: v for k, v in raw.items() if k in SEARCH_SPACE_KEYS}
+
+
+def _load_prior_param_sets(path: Path) -> list[dict[str, Any]]:
+    if not path.exists():
+        return []
+
+    payload = json.loads(path.read_text(encoding="utf-8"))
+    if isinstance(payload, dict):
+        rows = payload.get("trials", [])
+    elif isinstance(payload, list):
+        rows = payload
+    else:
+        rows = []
+
+    out: list[dict[str, Any]] = []
+    for item in rows:
+        if not isinstance(item, dict):
+            continue
+        params_obj = item.get("params", item)
+        if not isinstance(params_obj, dict):
+            continue
+        filtered = _filter_prior_params(params_obj)
+        if filtered:
+            out.append(filtered)
+    return out
+
+
+def _enqueue_transfer_priors(study: optuna.Study, priors_file: Path, apply_priors: bool) -> int:
+    if not apply_priors:
+        return 0
+
+    priors_raw = _load_prior_param_sets(priors_file)
+    if not priors_raw:
+        return 0
+
+    # Deduplicate param sets across merged studies.
+    priors: list[dict[str, Any]] = []
+    seen: set[str] = set()
+    for params in priors_raw:
+        key = json.dumps(params, sort_keys=True)
+        if key in seen:
+            continue
+        seen.add(key)
+        priors.append(params)
+
+    enqueued = 0
+    for params in priors:
+        before = len(study.get_trials(deepcopy=False))
+        try:
+            study.enqueue_trial(params, user_attrs={"seed_source": "transfer_priors"}, skip_if_exists=True)
+        except TypeError:
+            study.enqueue_trial(params, user_attrs={"seed_source": "transfer_priors"})
+        after = len(study.get_trials(deepcopy=False))
+        if after > before:
+            enqueued += 1
+    return enqueued
+
+
+def _enqueue_quality_anchors(study: optuna.Study, priors_file: Path, quality_mode_local: bool, top_k: int) -> int:
+    if not quality_mode_local or top_k <= 0:
+        return 0
+
+    priors = _load_prior_param_sets(priors_file)[:top_k]
+    enqueued = 0
+    for params in priors:
+        before = len(study.get_trials(deepcopy=False))
+        try:
+            study.enqueue_trial(
+                params,
+                user_attrs={"seed_source": "quality_anchor"},
+                skip_if_exists=True,
+            )
+        except TypeError:
+            study.enqueue_trial(params, user_attrs={"seed_source": "quality_anchor"})
+        after = len(study.get_trials(deepcopy=False))
+        if after > before:
+            enqueued += 1
+    return enqueued
 
 
 def _parse_metrics_from_stdout(stdout: str) -> dict[str, Any] | None:
@@ -164,241 +164,241 @@ def _parse_metrics_from_log_lines(lines: list[str]) -> dict[str, Any] | None:
         return None
 
 
-def _parse_last_train_bpb_from_logs(lines: list[str]) -> float | None:
-    """Best-effort fallback when final eval crashes before metrics JSON write."""
-    last: float | None = None
-    for line in lines:
-        m = re.search(r"\bbpb=([0-9]+(?:\.[0-9]+)?)", line)
+def _parse_last_train_bpb_from_logs(lines: list[str]) -> float | None:
+    """Best-effort fallback when final eval crashes before metrics JSON write."""
+    last: float | None = None
+    for line in lines:
+        m = re.search(r"\bbpb=([0-9]+(?:\.[0-9]+)?)", line)
         if m:
-            last = float(m.group(1))
-    return last
-
-
-def _persist_trial_artifacts(
-    *,
-    trial_dir: Path,
-    metrics: dict[str, Any] | None,
-    log_lines: list[str] | None,
-    log_name: str,
-    metadata: dict[str, Any],
-) -> dict[str, str | None]:
-    trial_dir.mkdir(parents=True, exist_ok=True)
-    metrics_path = trial_dir / "metrics.json"
-    log_path = trial_dir / log_name
-    manifest_path = trial_dir / "trial_artifacts.json"
-
-    if metrics is not None:
-        metrics_path.write_text(json.dumps(metrics, indent=2, sort_keys=True), encoding="utf-8")
-    if log_lines is not None:
-        log_path.write_text("\n".join(log_lines), encoding="utf-8")
-
-    manifest = {
-        **metadata,
-        "metrics_path": str(metrics_path) if metrics is not None else None,
-        "log_path": str(log_path) if log_lines is not None else None,
-    }
-    manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True), encoding="utf-8")
-    return {
-        "metrics_path": str(metrics_path) if metrics is not None else None,
-        "log_path": str(log_path) if log_lines is not None else None,
-        "manifest_path": str(manifest_path),
-    }
-
-
-def _resolve_objective_metric(
-    trial: optuna.Trial,
-    *,
-    metric_key: str,
-    metrics: dict[str, Any] | None,
-    allow_log_metric_fallback: bool,
-    fallback_bpb: float | None,
-    tps_seen: float | None,
-) -> float:
-    """Resolve the objective value while labeling where it came from.
-
-    Validation metrics and live training-log fallbacks are intentionally
-    different sources. Keeping that distinction in trial attrs prevents a
-    skipped/OOM eval from being mistaken for a real validation result.
-    """
-    if metrics is None:
-        if allow_log_metric_fallback and metric_key == "val_bpb" and fallback_bpb is not None:
-            trial.set_user_attr("objective_source", "train_log_fallback")
-            trial.set_user_attr("objective_metric", "train_bpb")
-            trial.set_user_attr("eval_status", "missing_metrics")
-            trial.set_user_attr("train_bpb_fallback", float(fallback_bpb))
-            if tps_seen is not None:
-                trial.set_user_attr("tps", float(tps_seen))
-            return float(fallback_bpb)
-        trial.set_user_attr("objective_source", "missing_metrics")
-        raise optuna.TrialPruned("No metrics payload found")
-
-    eval_status = str(
-        metrics.get(
-            "eval_status",
-            "completed" if metrics.get("val_bpb") is not None else "unknown",
-        )
-    )
-    trial.set_user_attr("eval_status", eval_status)
-
-    if fallback_bpb is not None:
-        trial.set_user_attr("train_bpb_fallback", float(fallback_bpb))
-
-    if metric_key not in metrics or metrics[metric_key] is None:
-        trial.set_user_attr("objective_source", "missing_metric")
-        trial.set_user_attr("objective_metric", metric_key)
-        raise optuna.TrialPruned(f"Metric '{metric_key}' missing in metrics payload")
-
-    value = float(metrics[metric_key])
-    trial.set_user_attr("objective_metric", metric_key)
-    if metric_key == "val_bpb":
-        trial.set_user_attr("objective_source", "final_val")
-        trial.set_user_attr("final_val_bpb", value)
-    else:
-        trial.set_user_attr("objective_source", "metrics_json")
-    return value
-
-
-def _fetch_job_logs_safe(
-    api,
-    *,
-    job_id: str,
-    token: str,
-    namespace: str,
-    retries: int = 3,
-    sleep_s: float = 2.0,
-    timeout_s: float = 20.0,
-) -> list[str]:
-    last_exc: Exception | None = None
-    for attempt in range(1, retries + 1):
-        try:
-            with ThreadPoolExecutor(max_workers=1) as executor:
-                future = executor.submit(
-                    lambda: list(api.fetch_job_logs(job_id=job_id, follow=False, token=token, namespace=namespace))
-                )
-                return future.result(timeout=timeout_s)
-        except FuturesTimeoutError as exc:
-            last_exc = TimeoutError(f"Timed out fetching HF job logs for {job_id} after {timeout_s:.1f}s")
-        except Exception as exc:  # noqa: BLE001
-            last_exc = exc
-            if attempt >= retries:
-                raise
-            time.sleep(sleep_s)
-    if last_exc is not None:
-        raise last_exc
-    return []
-
-
-def _effective_min_tps(args: argparse.Namespace) -> float | None:
-    min_tps = args.min_tps
-    if getattr(args, "quality_mode_local", False) and min_tps == 50000.0:
-        return 0.0
-    return min_tps
-
-
-def _trial_env(trial: optuna.Trial, args: argparse.Namespace, metrics_path: Path) -> dict[str, str]:
-    env = os.environ.copy()
-    full_arch_hpo = env.get("HYDRA_HPO_FULL_ARCH", "0") == "1"
-    speed_arch_hpo = full_arch_hpo and env.get("HYDRA_HPO_SPEED_ARCH", "0") == "1"
-    quality_mode_local = bool(getattr(args, "quality_mode_local", False))
+            last = float(m.group(1))
+    return last
+
+
+def _persist_trial_artifacts(
+    *,
+    trial_dir: Path,
+    metrics: dict[str, Any] | None,
+    log_lines: list[str] | None,
+    log_name: str,
+    metadata: dict[str, Any],
+) -> dict[str, str | None]:
+    trial_dir.mkdir(parents=True, exist_ok=True)
+    metrics_path = trial_dir / "metrics.json"
+    log_path = trial_dir / log_name
+    manifest_path = trial_dir / "trial_artifacts.json"
+
+    if metrics is not None:
+        metrics_path.write_text(json.dumps(metrics, indent=2, sort_keys=True), encoding="utf-8")
+    if log_lines is not None:
+        log_path.write_text("\n".join(log_lines), encoding="utf-8")
+
+    manifest = {
+        **metadata,
+        "metrics_path": str(metrics_path) if metrics is not None else None,
+        "log_path": str(log_path) if log_lines is not None else None,
+    }
+    manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True), encoding="utf-8")
+    return {
+        "metrics_path": str(metrics_path) if metrics is not None else None,
+        "log_path": str(log_path) if log_lines is not None else None,
+        "manifest_path": str(manifest_path),
+    }
+
+
+def _resolve_objective_metric(
+    trial: optuna.Trial,
+    *,
+    metric_key: str,
+    metrics: dict[str, Any] | None,
+    allow_log_metric_fallback: bool,
+    fallback_bpb: float | None,
+    tps_seen: float | None,
+) -> float:
+    """Resolve the objective value while labeling where it came from.
+
+    Validation metrics and live training-log fallbacks are intentionally
+    different sources. Keeping that distinction in trial attrs prevents a
+    skipped/OOM eval from being mistaken for a real validation result.
+    """
+    if metrics is None:
+        if allow_log_metric_fallback and metric_key == "val_bpb" and fallback_bpb is not None:
+            trial.set_user_attr("objective_source", "train_log_fallback")
+            trial.set_user_attr("objective_metric", "train_bpb")
+            trial.set_user_attr("eval_status", "missing_metrics")
+            trial.set_user_attr("train_bpb_fallback", float(fallback_bpb))
+            if tps_seen is not None:
+                trial.set_user_attr("tps", float(tps_seen))
+            return float(fallback_bpb)
+        trial.set_user_attr("objective_source", "missing_metrics")
+        raise optuna.TrialPruned("No metrics payload found")
+
+    eval_status = str(
+        metrics.get(
+            "eval_status",
+            "completed" if metrics.get("val_bpb") is not None else "unknown",
+        )
+    )
+    trial.set_user_attr("eval_status", eval_status)
+
+    if fallback_bpb is not None:
+        trial.set_user_attr("train_bpb_fallback", float(fallback_bpb))
+
+    if metric_key not in metrics or metrics[metric_key] is None:
+        trial.set_user_attr("objective_source", "missing_metric")
+        trial.set_user_attr("objective_metric", metric_key)
+        raise optuna.TrialPruned(f"Metric '{metric_key}' missing in metrics payload")
+
+    value = float(metrics[metric_key])
+    trial.set_user_attr("objective_metric", metric_key)
+    if metric_key == "val_bpb":
+        trial.set_user_attr("objective_source", "final_val")
+        trial.set_user_attr("final_val_bpb", value)
+    else:
+        trial.set_user_attr("objective_source", "metrics_json")
+    return value
+
+
+def _fetch_job_logs_safe(
+    api,
+    *,
+    job_id: str,
+    token: str,
+    namespace: str,
+    retries: int = 3,
+    sleep_s: float = 2.0,
+    timeout_s: float = 20.0,
+) -> list[str]:
+    last_exc: Exception | None = None
+    for attempt in range(1, retries + 1):
+        try:
+            with ThreadPoolExecutor(max_workers=1) as executor:
+                future = executor.submit(
+                    lambda: list(api.fetch_job_logs(job_id=job_id, follow=False, token=token, namespace=namespace))
+                )
+                return future.result(timeout=timeout_s)
+        except FuturesTimeoutError as exc:
+            last_exc = TimeoutError(f"Timed out fetching HF job logs for {job_id} after {timeout_s:.1f}s")
+        except Exception as exc:  # noqa: BLE001
+            last_exc = exc
+            if attempt >= retries:
+                raise
+            time.sleep(sleep_s)
+    if last_exc is not None:
+        raise last_exc
+    return []
+
+
+def _effective_min_tps(args: argparse.Namespace) -> float | None:
+    min_tps = args.min_tps
+    if getattr(args, "quality_mode_local", False) and min_tps == 50000.0:
+        return 0.0
+    return min_tps
+
+
+def _trial_env(trial: optuna.Trial, args: argparse.Namespace, metrics_path: Path) -> dict[str, str]:
+    env = os.environ.copy()
+    full_arch_hpo = env.get("HYDRA_HPO_FULL_ARCH", "0") == "1"
+    speed_arch_hpo = full_arch_hpo and env.get("HYDRA_HPO_SPEED_ARCH", "0") == "1"
+    quality_mode_local = bool(getattr(args, "quality_mode_local", False))
 
     # Runtime and reporting
     env["HYDRA_METRICS_OUT"] = str(metrics_path)
     env["HYDRA_TIME_BUDGET"] = str(args.trial_time_budget)
     env["PYTHONUNBUFFERED"] = "1"
 
-    # Search space — fully env-driven to match existing training stack.
-    if speed_arch_hpo:
-        # Full-arch speed mode targets A10 underutilization observed in HPO:
-        # low VRAM/MFU, strong BPB from shallow models, and fixed SDR/HTM
-        # overhead dominating small microbatches. Keep all components enabled
-        # while amortizing overhead over more tokens.
-        env["HYDRA_D_MODEL"] = str(trial.suggest_categorical("d_model", [64, 96]))
-        env["HYDRA_N_LAYER"] = str(trial.suggest_categorical("n_layer", [2]))
-        env["HYDRA_D_STATE"] = str(trial.suggest_categorical("d_state", [16, 32]))
-        env["HYDRA_HEADDIM"] = str(trial.suggest_categorical("headdim", [16, 32]))
-        env["HYDRA_EXPAND"] = str(trial.suggest_categorical("expand", [1, 2]))
-    elif quality_mode_local and full_arch_hpo:
-        env["HYDRA_D_MODEL"] = str(trial.suggest_categorical("d_model", [64, 96, 128]))
-        env["HYDRA_N_LAYER"] = str(trial.suggest_int("n_layer", 2, 3))
-        env["HYDRA_D_STATE"] = str(trial.suggest_categorical("d_state", [16, 32]))
-        env["HYDRA_HEADDIM"] = str(trial.suggest_categorical("headdim", [16, 32]))
-        env["HYDRA_EXPAND"] = str(trial.suggest_categorical("expand", [1, 2]))
-    else:
-        env["HYDRA_D_MODEL"] = str(trial.suggest_categorical("d_model", [64, 96, 128, 160, 192]))
-        env["HYDRA_N_LAYER"] = str(trial.suggest_int("n_layer", 1, 4))
-        env["HYDRA_D_STATE"] = str(trial.suggest_categorical("d_state", [16, 32, 48]))
-        env["HYDRA_HEADDIM"] = str(trial.suggest_categorical("headdim", [8, 16, 32]))
-        env["HYDRA_EXPAND"] = str(trial.suggest_categorical("expand", [1, 2]))
-
-    if speed_arch_hpo:
-        seq_len = trial.suggest_categorical("seq_len", [64, 128])
-        batch_size = trial.suggest_categorical("batch_size", [8, 16, 32])
-        grad_accum = trial.suggest_categorical("grad_accum", [4, 8, 16])
-    elif quality_mode_local and full_arch_hpo:
-        seq_len = trial.suggest_categorical("seq_len", [64])
-        batch_size = trial.suggest_categorical("batch_size", [4, 8])
-        grad_accum = trial.suggest_categorical("grad_accum", [4, 8, 16])
-    else:
-        seq_len = trial.suggest_categorical("seq_len", [32, 64])
-        batch_size = trial.suggest_categorical("batch_size", [4, 8] if full_arch_hpo else [4, 8, 16])
-        grad_accum = trial.suggest_categorical("grad_accum", [1, 4, 8, 16] if full_arch_hpo else [8, 16, 32, 64])
+    # Search space — fully env-driven to match existing training stack.
+    if speed_arch_hpo:
+        # Full-arch speed mode targets A10 underutilization observed in HPO:
+        # low VRAM/MFU, strong BPB from shallow models, and fixed SDR/HTM
+        # overhead dominating small microbatches. Keep all components enabled
+        # while amortizing overhead over more tokens.
+        env["HYDRA_D_MODEL"] = str(trial.suggest_categorical("d_model", [64, 96]))
+        env["HYDRA_N_LAYER"] = str(trial.suggest_categorical("n_layer", [2]))
+        env["HYDRA_D_STATE"] = str(trial.suggest_categorical("d_state", [16, 32]))
+        env["HYDRA_HEADDIM"] = str(trial.suggest_categorical("headdim", [16, 32]))
+        env["HYDRA_EXPAND"] = str(trial.suggest_categorical("expand", [1, 2]))
+    elif quality_mode_local and full_arch_hpo:
+        env["HYDRA_D_MODEL"] = str(trial.suggest_categorical("d_model", [64, 96, 128]))
+        env["HYDRA_N_LAYER"] = str(trial.suggest_int("n_layer", 2, 3))
+        env["HYDRA_D_STATE"] = str(trial.suggest_categorical("d_state", [16, 32]))
+        env["HYDRA_HEADDIM"] = str(trial.suggest_categorical("headdim", [16, 32]))
+        env["HYDRA_EXPAND"] = str(trial.suggest_categorical("expand", [1, 2]))
+    else:
+        env["HYDRA_D_MODEL"] = str(trial.suggest_categorical("d_model", [64, 96, 128, 160, 192]))
+        env["HYDRA_N_LAYER"] = str(trial.suggest_int("n_layer", 1, 4))
+        env["HYDRA_D_STATE"] = str(trial.suggest_categorical("d_state", [16, 32, 48]))
+        env["HYDRA_HEADDIM"] = str(trial.suggest_categorical("headdim", [8, 16, 32]))
+        env["HYDRA_EXPAND"] = str(trial.suggest_categorical("expand", [1, 2]))
+
+    if speed_arch_hpo:
+        seq_len = trial.suggest_categorical("seq_len", [64, 128])
+        batch_size = trial.suggest_categorical("batch_size", [8, 16, 32])
+        grad_accum = trial.suggest_categorical("grad_accum", [4, 8, 16])
+    elif quality_mode_local and full_arch_hpo:
+        seq_len = trial.suggest_categorical("seq_len", [64])
+        batch_size = trial.suggest_categorical("batch_size", [4, 8])
+        grad_accum = trial.suggest_categorical("grad_accum", [4, 8, 16])
+    else:
+        seq_len = trial.suggest_categorical("seq_len", [32, 64])
+        batch_size = trial.suggest_categorical("batch_size", [4, 8] if full_arch_hpo else [4, 8, 16])
+        grad_accum = trial.suggest_categorical("grad_accum", [1, 4, 8, 16] if full_arch_hpo else [8, 16, 32, 64])
     # Keep TOTAL_BATCH_SIZE divisible by DEVICE_BATCH_SIZE * MAX_SEQ_LEN.
     total_batch = batch_size * seq_len * grad_accum
     env["HYDRA_SEQ_LEN"] = str(seq_len)
     env["HYDRA_BATCH_SIZE"] = str(batch_size)
     env["HYDRA_TOTAL_BATCH"] = str(total_batch)
 
-    if quality_mode_local and full_arch_hpo:
-        env["HYDRA_MATRIX_LR"] = str(trial.suggest_float("matrix_lr", 0.008, 0.03, log=True))
-        env["HYDRA_EMBED_LR"] = str(trial.suggest_float("embed_lr", 0.15, 0.6, log=True))
-        env["HYDRA_UNEMBED_LR"] = str(trial.suggest_float("unembed_lr", 0.001, 0.01, log=True))
-    else:
-        env["HYDRA_MATRIX_LR"] = str(trial.suggest_float("matrix_lr", 0.005, 0.2, log=True))
-        env["HYDRA_EMBED_LR"] = str(trial.suggest_float("embed_lr", 0.05, 1.0, log=True))
-        env["HYDRA_UNEMBED_LR"] = str(trial.suggest_float("unembed_lr", 0.0005, 0.02, log=True))
-
-    if full_arch_hpo:
-        env["HYDRA_HYENA_LAYERS"] = ""
-        env["HYDRA_ENGRAM_N_COLUMNS"] = str(
-            trial.suggest_categorical(
-                "engram_n_columns",
-                [512, 1024] if (speed_arch_hpo or quality_mode_local) else [512, 1024, 2048],
-            )
-        )
-        env["HYDRA_ENGRAM_LAYER_IDX"] = str(trial.suggest_int("engram_layer_idx", 0, max(0, int(env["HYDRA_N_LAYER"]) - 1)))
-        env["HYDRA_SDR_TARGET_ACTIVE"] = str(
-            trial.suggest_categorical(
-                "sdr_target_active",
-                [327] if quality_mode_local else ([164, 327] if speed_arch_hpo else [164, 327, 512]),
-            )
-        )
-        env["HYDRA_HTM_LEARN_EVERY"] = str(
-            trial.suggest_categorical("htm_learn_every", [8, 16] if (speed_arch_hpo or quality_mode_local) else [4, 8, 16])
-        )
-        env["HYDRA_HTM_SUBSAMPLE"] = str(
-            trial.suggest_categorical("htm_subsample", [1, 2] if quality_mode_local else ([4, 8, 16] if speed_arch_hpo else [1, 2, 4, 8]))
-        )
-        env["HYDRA_ENGRAM_SUBSAMPLE"] = str(
-            trial.suggest_categorical("engram_subsample", [1, 2] if quality_mode_local else ([1, 2, 4] if speed_arch_hpo else [1]))
-        )
-        env["HYDRA_MAMBA3_CHUNK"] = str(trial.suggest_categorical("mamba3_chunk", [32, 64]))
-        env["HYDRA_DROPOUT"] = str(trial.suggest_categorical("dropout", [0.0, 0.1] if (speed_arch_hpo or quality_mode_local) else [0.0, 0.1, 0.2]))
-    else:
-        env["HYDRA_HYENA_LAYERS"] = trial.suggest_categorical("hyena_layers", ["", "0", "1", "0,1"])
+    if quality_mode_local and full_arch_hpo:
+        env["HYDRA_MATRIX_LR"] = str(trial.suggest_float("matrix_lr", 0.008, 0.03, log=True))
+        env["HYDRA_EMBED_LR"] = str(trial.suggest_float("embed_lr", 0.15, 0.6, log=True))
+        env["HYDRA_UNEMBED_LR"] = str(trial.suggest_float("unembed_lr", 0.001, 0.01, log=True))
+    else:
+        env["HYDRA_MATRIX_LR"] = str(trial.suggest_float("matrix_lr", 0.005, 0.2, log=True))
+        env["HYDRA_EMBED_LR"] = str(trial.suggest_float("embed_lr", 0.05, 1.0, log=True))
+        env["HYDRA_UNEMBED_LR"] = str(trial.suggest_float("unembed_lr", 0.0005, 0.02, log=True))
+
+    if full_arch_hpo:
+        env["HYDRA_HYENA_LAYERS"] = ""
+        env["HYDRA_ENGRAM_N_COLUMNS"] = str(
+            trial.suggest_categorical(
+                "engram_n_columns",
+                [512, 1024] if (speed_arch_hpo or quality_mode_local) else [512, 1024, 2048],
+            )
+        )
+        env["HYDRA_ENGRAM_LAYER_IDX"] = str(trial.suggest_int("engram_layer_idx", 0, max(0, int(env["HYDRA_N_LAYER"]) - 1)))
+        env["HYDRA_SDR_TARGET_ACTIVE"] = str(
+            trial.suggest_categorical(
+                "sdr_target_active",
+                [327] if quality_mode_local else ([164, 327] if speed_arch_hpo else [164, 327, 512]),
+            )
+        )
+        env["HYDRA_HTM_LEARN_EVERY"] = str(
+            trial.suggest_categorical("htm_learn_every", [8, 16] if (speed_arch_hpo or quality_mode_local) else [4, 8, 16])
+        )
+        env["HYDRA_HTM_SUBSAMPLE"] = str(
+            trial.suggest_categorical("htm_subsample", [1, 2] if quality_mode_local else ([4, 8, 16] if speed_arch_hpo else [1, 2, 4, 8]))
+        )
+        env["HYDRA_ENGRAM_SUBSAMPLE"] = str(
+            trial.suggest_categorical("engram_subsample", [1, 2] if quality_mode_local else ([1, 2, 4] if speed_arch_hpo else [1]))
+        )
+        env["HYDRA_MAMBA3_CHUNK"] = str(trial.suggest_categorical("mamba3_chunk", [32, 64]))
+        env["HYDRA_DROPOUT"] = str(trial.suggest_categorical("dropout", [0.0, 0.1] if (speed_arch_hpo or quality_mode_local) else [0.0, 0.1, 0.2]))
+    else:
+        env["HYDRA_HYENA_LAYERS"] = trial.suggest_categorical("hyena_layers", ["", "0", "1", "0,1"])
 
     # Keep trials alive long enough to emit metrics.
     env["HYDRA_FAIL_LOSS_THRESHOLD"] = "1000000"
     env["HYDRA_USE_NEMOTRON"] = os.environ.get("HYDRA_USE_NEMOTRON", "1")
     env["HYDRA_LOCAL_SHARDS_ONLY"] = os.environ.get("HYDRA_LOCAL_SHARDS_ONLY", "0")
     # Strict optimal-path defaults (no forced fallback profile).
-    env["HYDRA_MUON_COMPILE"] = os.environ.get("HYDRA_MUON_COMPILE", "1")
-    env["HYDRA_THROUGHPUT_MODE"] = os.environ.get("HYDRA_THROUGHPUT_MODE", "0" if full_arch_hpo else "1")
-    env["HYDRA_FORCE_HTM_CPU"] = os.environ.get("HYDRA_FORCE_HTM_CPU", "0")
-    env["HYDRA_ALLOW_SYNTHETIC_RETINA"] = os.environ.get("HYDRA_ALLOW_SYNTHETIC_RETINA", "0")
-    env["HYDRA_INERT_MAMBA"] = os.environ.get("HYDRA_INERT_MAMBA", "0")
-    env["HYDRA_FASTPATH"] = os.environ.get("HYDRA_FASTPATH", "0" if full_arch_hpo else "1")
-
-    return env
+    env["HYDRA_MUON_COMPILE"] = os.environ.get("HYDRA_MUON_COMPILE", "1")
+    env["HYDRA_THROUGHPUT_MODE"] = os.environ.get("HYDRA_THROUGHPUT_MODE", "0" if full_arch_hpo else "1")
+    env["HYDRA_FORCE_HTM_CPU"] = os.environ.get("HYDRA_FORCE_HTM_CPU", "0")
+    env["HYDRA_ALLOW_SYNTHETIC_RETINA"] = os.environ.get("HYDRA_ALLOW_SYNTHETIC_RETINA", "0")
+    env["HYDRA_INERT_MAMBA"] = os.environ.get("HYDRA_INERT_MAMBA", "0")
+    env["HYDRA_FASTPATH"] = os.environ.get("HYDRA_FASTPATH", "0" if full_arch_hpo else "1")
+
+    return env
 
 
 def _sanitize_hf_env(env: dict[str, str]) -> dict[str, str]:
@@ -410,7 +410,7 @@ def _sanitize_hf_env(env: dict[str, str]) -> dict[str, str]:
     return sanitized
 
 
-def _hf_command_candidates(args: argparse.Namespace) -> list[list[str]]:
+def _hf_command_candidates(args: argparse.Namespace) -> list[list[str]]:
     if args.hf_use_bash:
         return [["bash", "-lc", args.hf_command]]
 
@@ -432,20 +432,20 @@ def _hf_command_candidates(args: argparse.Namespace) -> list[list[str]]:
                 uniq.append(c)
         return uniq
 
-    return [raw.split()]
-
-
-def _space_repo_from_hf_image(image: str, namespace: str) -> str:
-    prefix = "hf.co/spaces/"
-    if image.startswith(prefix):
-        return image[len(prefix):]
-    return os.environ.get("FEATHER_HF_SPACE_REPO", f"{namespace}/feather-a10-runtime")
+    return [raw.split()]
+
 
+def _space_repo_from_hf_image(image: str, namespace: str) -> str:
+    prefix = "hf.co/spaces/"
+    if image.startswith(prefix):
+        return image[len(prefix):]
+    return os.environ.get("FEATHER_HF_SPACE_REPO", f"{namespace}/feather-a10-runtime")
 
-def _objective_local(args: argparse.Namespace):
-    effective_min_tps = _effective_min_tps(args)
-
-    def objective(trial: optuna.Trial) -> float:
+
+def _objective_local(args: argparse.Namespace):
+    effective_min_tps = _effective_min_tps(args)
+
+    def objective(trial: optuna.Trial) -> float:
         trial_dir = Path(tempfile.mkdtemp(prefix=f"optuna_trial_{trial.number}_", dir=str(args.work_dir)))
         metrics_path = trial_dir / "metrics.json"
 
@@ -460,67 +460,67 @@ def _objective_local(args: argparse.Namespace):
             timeout=args.trial_timeout,
         )
 
-        metrics: dict[str, Any] | None = None
+        metrics: dict[str, Any] | None = None
         if metrics_path.exists():
             try:
                 metrics = json.loads(metrics_path.read_text(encoding="utf-8"))
             except json.JSONDecodeError:
                 metrics = None
-        if metrics is None:
-            metrics = _parse_metrics_from_stdout(proc.stdout)
-
-        artifact_paths = _persist_trial_artifacts(
-            trial_dir=trial_dir,
-            metrics=metrics,
-            log_lines=(proc.stdout or "").splitlines(),
-            log_name="train_stdout.log",
-            metadata={"runner": "local", "returncode": proc.returncode},
-        )
-        (trial_dir / "train_stderr.log").write_text(proc.stderr or "", encoding="utf-8")
-
-        fallback_bpb = _parse_last_train_bpb_from_logs(proc.stdout.splitlines())
-        if metrics is None:
-            _resolve_objective_metric(
-                trial,
-                metric_key=args.metric,
-                metrics=None,
-                allow_log_metric_fallback=args.allow_log_metric_fallback,
-                fallback_bpb=fallback_bpb,
-                tps_seen=None,
-            )
-            raise optuna.TrialPruned("No metrics found (HYDRA_METRICS_OUT/[METRICS_JSON])")
+        if metrics is None:
+            metrics = _parse_metrics_from_stdout(proc.stdout)
+
+        artifact_paths = _persist_trial_artifacts(
+            trial_dir=trial_dir,
+            metrics=metrics,
+            log_lines=(proc.stdout or "").splitlines(),
+            log_name="train_stdout.log",
+            metadata={"runner": "local", "returncode": proc.returncode},
+        )
+        (trial_dir / "train_stderr.log").write_text(proc.stderr or "", encoding="utf-8")
+
+        fallback_bpb = _parse_last_train_bpb_from_logs(proc.stdout.splitlines())
+        if metrics is None:
+            _resolve_objective_metric(
+                trial,
+                metric_key=args.metric,
+                metrics=None,
+                allow_log_metric_fallback=args.allow_log_metric_fallback,
+                fallback_bpb=fallback_bpb,
+                tps_seen=None,
+            )
+            raise optuna.TrialPruned("No metrics found (HYDRA_METRICS_OUT/[METRICS_JSON])")
 
         if proc.returncode != 0:
             raise optuna.TrialPruned(f"Training failed rc={proc.returncode}")
 
-        metric_key = args.metric
+        metric_key = args.metric
 
         tps_val = metrics.get("tps")
         if tps_val is not None:
             tps_f = float(tps_val)
             trial.set_user_attr("tps", tps_f)
-            if effective_min_tps is not None and tps_f < effective_min_tps:
-                raise optuna.TrialPruned(f"TPS below floor: {tps_f} < {effective_min_tps}")
-
-        value = _resolve_objective_metric(
-            trial,
-            metric_key=metric_key,
-            metrics=metrics,
-            allow_log_metric_fallback=args.allow_log_metric_fallback,
-            fallback_bpb=fallback_bpb,
-            tps_seen=None,
-        )
-
-        # Keep useful context on trial
-        trial.set_user_attr("summary_path", metrics.get("summary_path") or artifact_paths["manifest_path"])
-        trial.set_user_attr("run_log_path", metrics.get("run_log_path") or artifact_paths["log_path"])
-
-        return value
+            if effective_min_tps is not None and tps_f < effective_min_tps:
+                raise optuna.TrialPruned(f"TPS below floor: {tps_f} < {effective_min_tps}")
+
+        value = _resolve_objective_metric(
+            trial,
+            metric_key=metric_key,
+            metrics=metrics,
+            allow_log_metric_fallback=args.allow_log_metric_fallback,
+            fallback_bpb=fallback_bpb,
+            tps_seen=None,
+        )
+
+        # Keep useful context on trial
+        trial.set_user_attr("summary_path", metrics.get("summary_path") or artifact_paths["manifest_path"])
+        trial.set_user_attr("run_log_path", metrics.get("run_log_path") or artifact_paths["log_path"])
+
+        return value
 
     return objective
 
 
-def _objective_hf_job(args: argparse.Namespace):
+def _objective_hf_job(args: argparse.Namespace):
     from huggingface_hub import HfApi
     from huggingface_hub.utils import get_token
 
@@ -530,9 +530,9 @@ def _objective_hf_job(args: argparse.Namespace):
             f"No Hugging Face token found. Set {args.hf_token_env} or run huggingface-cli login."
         )
 
-    api = HfApi(token=token)
-    terminal_states = {"ERROR", "COMPLETED", "CANCELLED", "TIMEOUT", "FAILED", "CANCELED"}
-    effective_min_tps = _effective_min_tps(args)
+    api = HfApi(token=token)
+    terminal_states = {"ERROR", "COMPLETED", "CANCELLED", "TIMEOUT", "FAILED", "CANCELED"}
+    effective_min_tps = _effective_min_tps(args)
 
     def objective(trial: optuna.Trial) -> float:
         trial_dir = Path(tempfile.mkdtemp(prefix=f"optuna_trial_{trial.number}_", dir=str(args.work_dir)))
@@ -568,14 +568,14 @@ def _objective_hf_job(args: argparse.Namespace):
                 info = api.inspect_job(job_id=job.id, token=token, namespace=args.hf_namespace)
                 bootstrap_stage = str(info.status.stage)
                 bootstrap_msg = str(getattr(info.status, "message", "") or "")
-                bootstrap_logs = _fetch_job_logs_safe(
-                    api,
-                    job_id=job.id,
-                    token=token,
-                    namespace=args.hf_namespace,
-                    retries=2,
-                    sleep_s=1.0,
-                )
+                bootstrap_logs = _fetch_job_logs_safe(
+                    api,
+                    job_id=job.id,
+                    token=token,
+                    namespace=args.hf_namespace,
+                    retries=2,
+                    sleep_s=1.0,
+                )
                 if bootstrap_stage in {"RUNNING", "COMPLETED"} or bootstrap_logs:
                     break
                 if bootstrap_stage in {"ERROR", "FAILED", "CANCELLED", "CANCELED", "TIMEOUT"}:
@@ -611,12 +611,12 @@ def _objective_hf_job(args: argparse.Namespace):
             info = api.inspect_job(job_id=job_id, token=token, namespace=args.hf_namespace)
             stage = str(info.status.stage)
             terminal_detail = str(getattr(info.status, "message", "")) or terminal_detail
-            log_lines = _fetch_job_logs_safe(
-                api,
-                job_id=job_id,
-                token=token,
-                namespace=args.hf_namespace,
-            )
+            log_lines = _fetch_job_logs_safe(
+                api,
+                job_id=job_id,
+                token=token,
+                namespace=args.hf_namespace,
+            )
 
             m = _parse_metrics_from_log_lines(log_lines)
             if m is not None:
@@ -643,66 +643,66 @@ def _objective_hf_job(args: argparse.Namespace):
         except Exception:
             pass
 
-        artifact_paths = _persist_trial_artifacts(
-            trial_dir=trial_dir,
-            metrics=metrics,
-            log_lines=log_lines,
-            log_name="hf_job.log",
-            metadata={"runner": "hf-job", "hf_job_id": job_id, "hf_stage": stage},
-        )
-        trial.set_user_attr("hf_stage", stage)
-        trial.set_user_attr("hf_log_lines", len(log_lines))
+        artifact_paths = _persist_trial_artifacts(
+            trial_dir=trial_dir,
+            metrics=metrics,
+            log_lines=log_lines,
+            log_name="hf_job.log",
+            metadata={"runner": "hf-job", "hf_job_id": job_id, "hf_stage": stage},
+        )
+        trial.set_user_attr("hf_stage", stage)
+        trial.set_user_attr("hf_log_lines", len(log_lines))
         if terminal_detail:
             trial.set_user_attr("hf_status_message", terminal_detail)
 
-        fallback_bpb = _parse_last_train_bpb_from_logs(log_lines)
-        if metrics is None:
-            try:
-                value = _resolve_objective_metric(
-                    trial,
-                    metric_key=args.metric,
-                    metrics=None,
-                    allow_log_metric_fallback=args.allow_log_metric_fallback,
-                    fallback_bpb=fallback_bpb,
-                    tps_seen=tps_seen,
-                )
-                if tps_seen is not None and effective_min_tps is not None and tps_seen < effective_min_tps:
-                    raise optuna.TrialPruned(f"TPS below floor: {tps_seen} < {effective_min_tps}")
-                return value
-            except optuna.TrialPruned:
-                pass
-            if tps_seen is not None:
-                trial.set_user_attr("tps", tps_seen)
-            detail = f"stage={stage}, logs={len(log_lines)}"
-            if terminal_detail:
-                detail = f"{detail}, message={terminal_detail}"
+        fallback_bpb = _parse_last_train_bpb_from_logs(log_lines)
+        if metrics is None:
+            try:
+                value = _resolve_objective_metric(
+                    trial,
+                    metric_key=args.metric,
+                    metrics=None,
+                    allow_log_metric_fallback=args.allow_log_metric_fallback,
+                    fallback_bpb=fallback_bpb,
+                    tps_seen=tps_seen,
+                )
+                if tps_seen is not None and effective_min_tps is not None and tps_seen < effective_min_tps:
+                    raise optuna.TrialPruned(f"TPS below floor: {tps_seen} < {effective_min_tps}")
+                return value
+            except optuna.TrialPruned:
+                pass
+            if tps_seen is not None:
+                trial.set_user_attr("tps", tps_seen)
+            detail = f"stage={stage}, logs={len(log_lines)}"
+            if terminal_detail:
+                detail = f"{detail}, message={terminal_detail}"
             raise optuna.TrialPruned(f"No metrics found from HF job ({detail})")
 
-        metric_key = args.metric
+        metric_key = args.metric
 
         tps_val = metrics.get("tps")
         if tps_val is not None:
             tps_f = float(tps_val)
             trial.set_user_attr("tps", tps_f)
-            if effective_min_tps is not None and tps_f < effective_min_tps:
-                raise optuna.TrialPruned(f"TPS below floor: {tps_f} < {effective_min_tps}")
-
-        value = _resolve_objective_metric(
-            trial,
-            metric_key=metric_key,
-            metrics=metrics,
-            allow_log_metric_fallback=args.allow_log_metric_fallback,
-            fallback_bpb=fallback_bpb,
-            tps_seen=tps_seen,
-        )
-        trial.set_user_attr("summary_path", metrics.get("summary_path") or artifact_paths["manifest_path"])
-        trial.set_user_attr("run_log_path", metrics.get("run_log_path") or artifact_paths["log_path"])
-        return value
+            if effective_min_tps is not None and tps_f < effective_min_tps:
+                raise optuna.TrialPruned(f"TPS below floor: {tps_f} < {effective_min_tps}")
+
+        value = _resolve_objective_metric(
+            trial,
+            metric_key=metric_key,
+            metrics=metrics,
+            allow_log_metric_fallback=args.allow_log_metric_fallback,
+            fallback_bpb=fallback_bpb,
+            tps_seen=tps_seen,
+        )
+        trial.set_user_attr("summary_path", metrics.get("summary_path") or artifact_paths["manifest_path"])
+        trial.set_user_attr("run_log_path", metrics.get("run_log_path") or artifact_paths["log_path"])
+        return value
 
     return objective
 
 
-def _objective_hf_launcher(args: argparse.Namespace):
+def _objective_hf_launcher(args: argparse.Namespace):
     from huggingface_hub import HfApi
     from huggingface_hub.utils import get_token
 
@@ -712,9 +712,9 @@ def _objective_hf_launcher(args: argparse.Namespace):
             f"No Hugging Face token found. Set {args.hf_token_env} or run huggingface-cli login."
         )
 
-    api = HfApi(token=token)
-    terminal_states = {"ERROR", "COMPLETED", "CANCELLED", "TIMEOUT", "FAILED", "CANCELED"}
-    effective_min_tps = _effective_min_tps(args)
+    api = HfApi(token=token)
+    terminal_states = {"ERROR", "COMPLETED", "CANCELLED", "TIMEOUT", "FAILED", "CANCELED"}
+    effective_min_tps = _effective_min_tps(args)
 
     def objective(trial: optuna.Trial) -> float:
         trial_dir = Path(tempfile.mkdtemp(prefix=f"optuna_trial_{trial.number}_", dir=str(args.work_dir)))
@@ -725,11 +725,11 @@ def _objective_hf_launcher(args: argparse.Namespace):
         local_env = os.environ.copy()
         local_env.update(env)
         local_env[args.hf_token_env] = token
-        local_env["FEATHER_HF_NAMESPACE"] = args.hf_namespace
-        local_env["FEATHER_HF_FLAVOR"] = args.hf_flavor
-        local_env["FEATHER_HF_JOB_TIMEOUT"] = args.hf_timeout
-        local_env["FEATHER_HF_IMAGE"] = args.hf_image
-        local_env["FEATHER_HF_SPACE_REPO"] = _space_repo_from_hf_image(args.hf_image, args.hf_namespace)
+        local_env["FEATHER_HF_NAMESPACE"] = args.hf_namespace
+        local_env["FEATHER_HF_FLAVOR"] = args.hf_flavor
+        local_env["FEATHER_HF_JOB_TIMEOUT"] = args.hf_timeout
+        local_env["FEATHER_HF_IMAGE"] = args.hf_image
+        local_env["FEATHER_HF_SPACE_REPO"] = _space_repo_from_hf_image(args.hf_image, args.hf_namespace)
         if args.hf_output_repo:
             local_env["FEATHER_HF_OUTPUT_REPO"] = args.hf_output_repo
         else:
@@ -766,12 +766,12 @@ def _objective_hf_launcher(args: argparse.Namespace):
             info = api.inspect_job(job_id=job_id, token=token, namespace=args.hf_namespace)
             stage = str(info.status.stage)
             terminal_detail = str(getattr(info.status, "message", "") or "") or terminal_detail
-            log_lines = _fetch_job_logs_safe(
-                api,
-                job_id=job_id,
-                token=token,
-                namespace=args.hf_namespace,
-            )
+            log_lines = _fetch_job_logs_safe(
+                api,
+                job_id=job_id,
+                token=token,
+                namespace=args.hf_namespace,
+            )
 
             mtr = _parse_metrics_from_log_lines(log_lines)
             if mtr is not None:
@@ -796,85 +796,85 @@ def _objective_hf_launcher(args: argparse.Namespace):
         except Exception:
             pass
 
-        artifact_paths = _persist_trial_artifacts(
-            trial_dir=trial_dir,
-            metrics=metrics,
-            log_lines=log_lines,
-            log_name="hf_job.log",
-            metadata={"runner": "hf-launcher", "hf_job_id": job_id, "hf_stage": stage},
-        )
-        trial.set_user_attr("hf_stage", stage)
-        trial.set_user_attr("hf_log_lines", len(log_lines))
+        artifact_paths = _persist_trial_artifacts(
+            trial_dir=trial_dir,
+            metrics=metrics,
+            log_lines=log_lines,
+            log_name="hf_job.log",
+            metadata={"runner": "hf-launcher", "hf_job_id": job_id, "hf_stage": stage},
+        )
+        trial.set_user_attr("hf_stage", stage)
+        trial.set_user_attr("hf_log_lines", len(log_lines))
         if terminal_detail:
             trial.set_user_attr("hf_status_message", terminal_detail)
 
-        fallback_bpb = _parse_last_train_bpb_from_logs(log_lines)
-        if metrics is None:
-            try:
-                value = _resolve_objective_metric(
-                    trial,
-                    metric_key=args.metric,
-                    metrics=None,
-                    allow_log_metric_fallback=args.allow_log_metric_fallback,
-                    fallback_bpb=fallback_bpb,
-                    tps_seen=tps_seen,
-                )
-                if tps_seen is not None and effective_min_tps is not None and tps_seen < effective_min_tps:
-                    raise optuna.TrialPruned(f"TPS below floor: {tps_seen} < {effective_min_tps}")
-                return value
-            except optuna.TrialPruned:
-                pass
-            if tps_seen is not None:
-                trial.set_user_attr("tps", tps_seen)
-            detail = f"stage={stage}, logs={len(log_lines)}"
-            if terminal_detail:
-                detail = f"{detail}, message={terminal_detail}"
+        fallback_bpb = _parse_last_train_bpb_from_logs(log_lines)
+        if metrics is None:
+            try:
+                value = _resolve_objective_metric(
+                    trial,
+                    metric_key=args.metric,
+                    metrics=None,
+                    allow_log_metric_fallback=args.allow_log_metric_fallback,
+                    fallback_bpb=fallback_bpb,
+                    tps_seen=tps_seen,
+                )
+                if tps_seen is not None and effective_min_tps is not None and tps_seen < effective_min_tps:
+                    raise optuna.TrialPruned(f"TPS below floor: {tps_seen} < {effective_min_tps}")
+                return value
+            except optuna.TrialPruned:
+                pass
+            if tps_seen is not None:
+                trial.set_user_attr("tps", tps_seen)
+            detail = f"stage={stage}, logs={len(log_lines)}"
+            if terminal_detail:
+                detail = f"{detail}, message={terminal_detail}"
             raise optuna.TrialPruned(f"No metrics found from HF launcher job ({detail})")
 
-        metric_key = args.metric
+        metric_key = args.metric
 
         tps_val = metrics.get("tps")
         if tps_val is not None:
             tps_f = float(tps_val)
             trial.set_user_attr("tps", tps_f)
-            if effective_min_tps is not None and tps_f < effective_min_tps:
-                raise optuna.TrialPruned(f"TPS below floor: {tps_f} < {effective_min_tps}")
-
-        value = _resolve_objective_metric(
-            trial,
-            metric_key=metric_key,
-            metrics=metrics,
-            allow_log_metric_fallback=args.allow_log_metric_fallback,
-            fallback_bpb=fallback_bpb,
-            tps_seen=tps_seen,
-        )
-        trial.set_user_attr("summary_path", metrics.get("summary_path") or artifact_paths["manifest_path"])
-        trial.set_user_attr("run_log_path", metrics.get("run_log_path") or artifact_paths["log_path"])
-        return value
+            if effective_min_tps is not None and tps_f < effective_min_tps:
+                raise optuna.TrialPruned(f"TPS below floor: {tps_f} < {effective_min_tps}")
+
+        value = _resolve_objective_metric(
+            trial,
+            metric_key=metric_key,
+            metrics=metrics,
+            allow_log_metric_fallback=args.allow_log_metric_fallback,
+            fallback_bpb=fallback_bpb,
+            tps_seen=tps_seen,
+        )
+        trial.set_user_attr("summary_path", metrics.get("summary_path") or artifact_paths["manifest_path"])
+        trial.set_user_attr("run_log_path", metrics.get("run_log_path") or artifact_paths["log_path"])
+        return value
 
     return objective
 
 
-def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
-    routing_defaults = resolve_routing(token=os.environ.get("HF_TOKEN"))
-    parser = argparse.ArgumentParser(description="Optuna HPO runner for HYDRA train.py")
+def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
+    routing_defaults = resolve_routing(token=os.environ.get("HF_TOKEN"))
+    parser = argparse.ArgumentParser(description="Optuna HPO runner for HYDRA train.py")
     parser.add_argument("--study-name", default="hydra_hpo", help="Optuna study name")
     parser.add_argument("--storage", default="sqlite:///optuna_hpo.db", help="Optuna storage URL")
     parser.add_argument("--direction", choices=["minimize", "maximize"], default="minimize")
     parser.add_argument("--metric", default="val_bpb", help="Metric key to optimize from HYDRA metrics")
-    parser.add_argument(
-        "--min-tps",
-        type=float,
-        default=50000.0,
-        help="TPS floor; prune trials under this value (set 0 to disable)",
-    )
+    parser.add_argument(
+        "--min-tps",
+        type=float,
+        default=50000.0,
+        help="TPS floor; prune trials under this value (set 0 to disable)",
+    )
     parser.add_argument("--trials", type=int, default=20, help="Number of Optuna trials")
     parser.add_argument("--study-timeout", type=int, default=None, help="Study timeout in seconds")
     parser.add_argument("--trial-time-budget", type=int, default=300, help="HYDRA_TIME_BUDGET passed to each trial")
     parser.add_argument("--trial-timeout", type=int, default=900, help="Subprocess timeout per trial in seconds")
     parser.add_argument("--runner", choices=["local", "hf-job", "hf-launcher"], default="local", help="Trial execution backend")
-    parser.add_argument("--hf-namespace", default=routing_defaults.job_namespace, help="HF namespace for jobs")
-    parser.add_argument("--hf-image", default=f"hf.co/spaces/{routing_defaults.space_repo}", help="HF jobs image")
+    parser.add_argument("--hf-namespace", default=routing_defaults.job_namespace, help="HF namespace for jobs")
+    parser.add_argument("--hf-image", default=f"hf.co/spaces/{routing_defaults.space_repo}", help="HF jobs image")
     parser.add_argument("--hf-flavor", default="a10g-large", help="HF jobs hardware flavor")
     parser.add_argument("--hf-timeout", default="25m", help="HF job timeout string")
     parser.add_argument("--hf-command", default="/app/entrypoint.py", help="Command executed inside HF job")
@@ -886,23 +886,23 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
     parser.add_argument("--hf-token-env", default="HF_TOKEN", help="Token env key passed as HF job secret")
     parser.add_argument("--hf-stop-after-metric", action="store_true", default=True, help="Cancel running job after metrics captured")
     parser.add_argument("--no-hf-stop-after-metric", action="store_false", dest="hf_stop_after_metric")
-    parser.add_argument("--hf-launcher-script", type=Path, default=REPO_ROOT / "scripts" / "launch_feather_hf_job.py", help="Local launcher script for hf-launcher runner")
-    parser.add_argument("--hf-output-repo", default=routing_defaults.output_repo, help="Optional FEATHER_HF_OUTPUT_REPO override for launcher runner")
-    parser.add_argument("--allow-log-metric-fallback", action="store_true", default=False, help="When metrics JSON is missing, allow val_bpb fallback from latest logged train bpb")
-    parser.add_argument("--no-allow-log-metric-fallback", action="store_false", dest="allow_log_metric_fallback")
-    parser.add_argument("--priors-file", type=Path, default=REPO_ROOT / "docs" / "hpo_transfer_priors.json", help="Path to transfer-learning prior trials JSON")
-    parser.add_argument("--apply-priors", action="store_true", default=True, help="Enqueue transfer-learning prior trials before optimize")
-    parser.add_argument("--no-apply-priors", action="store_false", dest="apply_priors")
-    parser.add_argument("--quality-mode-local", action="store_true", default=False, help="Narrow local full-architecture search around the proven quality-winning region")
-    parser.add_argument("--quality-anchor-top-k", type=int, default=3, help="Number of top clean priors to enqueue as deterministic local quality anchors")
-    parser.add_argument("--seed", type=int, default=42, help="Seed for sampler")
+    parser.add_argument("--hf-launcher-script", type=Path, default=REPO_ROOT / "scripts" / "launch_feather_hf_job.py", help="Local launcher script for hf-launcher runner")
+    parser.add_argument("--hf-output-repo", default=routing_defaults.output_repo, help="Optional FEATHER_HF_OUTPUT_REPO override for launcher runner")
+    parser.add_argument("--allow-log-metric-fallback", action="store_true", default=False, help="When metrics JSON is missing, allow val_bpb fallback from latest logged train bpb")
+    parser.add_argument("--no-allow-log-metric-fallback", action="store_false", dest="allow_log_metric_fallback")
+    parser.add_argument("--priors-file", type=Path, default=REPO_ROOT / "docs" / "hpo_transfer_priors.json", help="Path to transfer-learning prior trials JSON")
+    parser.add_argument("--apply-priors", action="store_true", default=True, help="Enqueue transfer-learning prior trials before optimize")
+    parser.add_argument("--no-apply-priors", action="store_false", dest="apply_priors")
+    parser.add_argument("--quality-mode-local", action="store_true", default=False, help="Narrow local full-architecture search around the proven quality-winning region")
+    parser.add_argument("--quality-anchor-top-k", type=int, default=3, help="Number of top clean priors to enqueue as deterministic local quality anchors")
+    parser.add_argument("--seed", type=int, default=42, help="Seed for sampler")
     parser.add_argument("--n-startup-trials", type=int, default=5, help="Pruner startup trials before pruning")
     parser.add_argument("--n-warmup-steps", type=int, default=0, help="Pruner warmup steps")
     parser.add_argument("--patience-trials", type=int, default=None, help="Stop study after this many completed trials without meaningful improvement")
     parser.add_argument("--min-improvement", type=float, default=0.0, help="Minimum best-value improvement to reset patience")
     parser.add_argument("--work-dir", type=Path, default=REPO_ROOT / ".tmp" / "optuna", help="Directory for trial artifacts")
     parser.add_argument("--summary-out", type=Path, default=REPO_ROOT / ".tmp" / "optuna" / "best_summary.json")
-    return parser.parse_args(argv)
+    return parser.parse_args(argv)
 
 
 def main() -> int:
@@ -916,22 +916,22 @@ def main() -> int:
         n_warmup_steps=args.n_warmup_steps,
     )
 
-    study = optuna.create_study(
-        study_name=args.study_name,
-        storage=args.storage,
-        load_if_exists=True,
-        direction=args.direction,
-        sampler=sampler,
-        pruner=pruner,
-    )
-
-    enqueued_quality_anchors = _enqueue_quality_anchors(study, args.priors_file, args.quality_mode_local, args.quality_anchor_top_k)
-    if enqueued_quality_anchors:
-        print(f"[hpo] enqueued {enqueued_quality_anchors} local quality anchors from {args.priors_file}")
-
-    enqueued_priors = _enqueue_transfer_priors(study, args.priors_file, args.apply_priors)
-    if enqueued_priors:
-        print(f"[hpo] enqueued {enqueued_priors} transfer priors from {args.priors_file}")
+    study = optuna.create_study(
+        study_name=args.study_name,
+        storage=args.storage,
+        load_if_exists=True,
+        direction=args.direction,
+        sampler=sampler,
+        pruner=pruner,
+    )
+
+    enqueued_quality_anchors = _enqueue_quality_anchors(study, args.priors_file, args.quality_mode_local, args.quality_anchor_top_k)
+    if enqueued_quality_anchors:
+        print(f"[hpo] enqueued {enqueued_quality_anchors} local quality anchors from {args.priors_file}")
+
+    enqueued_priors = _enqueue_transfer_priors(study, args.priors_file, args.apply_priors)
+    if enqueued_priors:
+        print(f"[hpo] enqueued {enqueued_priors} transfer priors from {args.priors_file}")
 
     state: dict[str, Any] = {
         "best": None,
@@ -990,29 +990,29 @@ def main() -> int:
             "best_trial_number": study.best_trial.number,
             "best_trial_user_attrs": study.best_trial.user_attrs,
             "n_trials": len(study.trials),
-            "n_completed": len(completed),
-            "patience_trials": args.patience_trials,
-            "min_improvement": args.min_improvement,
-            "quality_mode_local": args.quality_mode_local,
-            "enqueued_quality_anchors": enqueued_quality_anchors,
-            "enqueued_priors": enqueued_priors,
-        }
-    else:
-        best = {
+            "n_completed": len(completed),
+            "patience_trials": args.patience_trials,
+            "min_improvement": args.min_improvement,
+            "quality_mode_local": args.quality_mode_local,
+            "enqueued_quality_anchors": enqueued_quality_anchors,
+            "enqueued_priors": enqueued_priors,
+        }
+    else:
+        best = {
             "study_name": study.study_name,
             "direction": args.direction,
             "metric": args.metric,
             "best_value": None,
             "best_params": {},
-            "best_trial_number": None,
-            "best_trial_user_attrs": {},
-            "n_trials": len(study.trials),
-            "n_completed": 0,
-            "quality_mode_local": args.quality_mode_local,
-            "enqueued_quality_anchors": enqueued_quality_anchors,
-            "enqueued_priors": enqueued_priors,
-            "note": "No completed trials with metrics found.",
-        }
+            "best_trial_number": None,
+            "best_trial_user_attrs": {},
+            "n_trials": len(study.trials),
+            "n_completed": 0,
+            "quality_mode_local": args.quality_mode_local,
+            "enqueued_quality_anchors": enqueued_quality_anchors,
+            "enqueued_priors": enqueued_priors,
+            "note": "No completed trials with metrics found.",
+        }
     args.summary_out.write_text(json.dumps(best, indent=2), encoding="utf-8")
     print(json.dumps(best, indent=2))
     return 0

overlay/scripts/run_cycle1a.py

@@ -1,46 +1,45 @@
-#!/usr/bin/env python3
-from __future__ import annotations
-
-import argparse
-import sys
-from pathlib import Path
-
-REPO_ROOT = Path(__file__).resolve().parents[1]
-if str(REPO_ROOT) not in sys.path:
-    sys.path.insert(0, str(REPO_ROOT))
-
-from scripts import cycle_executor
-
-
-def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Run the full local Cycle 1a benchmark suite")
-    parser.add_argument("--out-dir", type=Path, default=REPO_ROOT / "artifacts" / "cycle1a_runs")
-    parser.add_argument("--preflight-out", type=Path, default=REPO_ROOT / "artifacts" / "cycle1a_preflight.json")
-    parser.add_argument("--summary-out", type=Path, default=REPO_ROOT / "artifacts" / "cycle1a_summary.json")
-    parser.add_argument("--hydrate-assets", action="store_true")
-    parser.add_argument("--require-ready", action="store_true")
-    parser.add_argument("--output-repo")
-    parser.add_argument("--tokenizer-repo")
-    return parser.parse_args(argv)
-
-
-def main(argv: list[str] | None = None) -> int:
-    args = parse_args(argv)
-    return cycle_executor.main([
-        "--benchmark", "GSM8K",
-        "--variant", "hydra_full",
-        "--seed", "42",
-        "--out-dir", str(args.out_dir),
-        "--preflight-out", str(args.preflight_out),
-        "--summary-out", str(args.summary_out),
-        "--all-runnable",
-        "--all-benchmarks",
-        *( ["--hydrate-assets"] if args.hydrate_assets else [] ),
-        *( ["--require-ready"] if args.require_ready else [] ),
-        *( ["--output-repo", args.output_repo] if args.output_repo else [] ),
-        *( ["--tokenizer-repo", args.tokenizer_repo] if args.tokenizer_repo else [] ),
-    ])
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+from scripts import cycle_executor
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+
+
+def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run the full local Cycle 1a benchmark suite")
+    parser.add_argument("--out-dir", type=Path, default=REPO_ROOT / "artifacts" / "cycle1a_runs")
+    parser.add_argument("--preflight-out", type=Path, default=REPO_ROOT / "artifacts" / "cycle1a_preflight.json")
+    parser.add_argument("--summary-out", type=Path, default=REPO_ROOT / "artifacts" / "cycle1a_summary.json")
+    parser.add_argument("--hydrate-assets", action="store_true")
+    parser.add_argument("--require-ready", action="store_true")
+    parser.add_argument("--output-repo")
+    parser.add_argument("--tokenizer-repo")
+    return parser.parse_args(argv)
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = parse_args(argv)
+    return cycle_executor.main([
+        "--benchmark", "GSM8K",
+        "--variant", "hydra_full",
+        "--seed", "42",
+        "--out-dir", str(args.out_dir),
+        "--preflight-out", str(args.preflight_out),
+        "--summary-out", str(args.summary_out),
+        "--all-runnable",
+        "--all-benchmarks",
+        *( ["--hydrate-assets"] if args.hydrate_assets else [] ),
+        *( ["--require-ready"] if args.require_ready else [] ),
+        *( ["--output-repo", args.output_repo] if args.output_repo else [] ),
+        *( ["--tokenizer-repo", args.tokenizer_repo] if args.tokenizer_repo else [] ),
+    ])
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

overlay/scripts/setup.sh

@@ -25,4 +25,3 @@ echo "=== Setup complete ==="
 echo "Run experiments with: uv run train.py"
 echo "Run orchestrator with: uv run -m harness.orchestrator"
 echo "Run Phase 1 subsystems with: bash scripts/run_phase1.sh"
-echo "For WSL/CUDA throughput gate: see docs/WSL_TPS_RUNBOOK.md"

overlay/scripts/sweep_depth_aggregate.py

@@ -11,77 +11,77 @@ Usage:
 """
 from __future__ import annotations
 
-import json
-import os
-import statistics
-import re
-import sys
-from pathlib import Path
-
-from configs.harness_config import HarnessConfig
-
-type MetricValue = float | int | str | bool | None
-type MetricsDict = dict[str, MetricValue]
-
-MANIFEST = Path(sys.argv[1] if len(sys.argv) > 1 else '/tmp/sweep_depth_manifest.txt')
-STEP_TPS_PATTERN = re.compile(r"step=(\d+).*?\btps=(\d+)\b")
-MIN_TPS = float(os.environ.get('SWEEP_MIN_TPS', '0'))
-TARGET_TOKENS_M = float(os.environ.get('SWEEP_TARGET_TOKENS_M', '0'))
-TARGET_SECONDS = float(os.environ.get('SWEEP_TARGET_SECONDS', '0'))
-
-
-def _zero_shot_score(result: MetricsDict) -> float:
-    """Composite quality score for tie-breaking among BPB-near runs."""
-    factual = float(result.get('factual_english_score', 0.0) or 0.0)
-    instruction = float(result.get('instruction_following_score', 0.0) or 0.0)
-    distinct_2 = float(result.get('distinct_2', 0.0) or 0.0)
-    repetition = float(result.get('repetition_rate', 0.0) or 0.0)
-    return factual + instruction + distinct_2 - repetition
-
-
-def _metric_float(result: MetricsDict, key: str, default: float = 0.0) -> float:
-    value = result.get(key, default)
-    return float(value) if isinstance(value, (int, float)) else default
-
-
-def _metric_int(result: MetricsDict, key: str, default: int = 0) -> int:
-    value = result.get(key, default)
-    return int(value) if isinstance(value, int) else default
-
-
-def _fixed_budget_ranking(results: dict[int, MetricsDict], *, metric_key: str, target: float) -> list[tuple[int, MetricsDict, float]]:
-    ranked: list[tuple[int, MetricsDict, float]] = []
-    for n_layer, row in results.items():
-        budget_val = row.get(metric_key)
-        if not isinstance(budget_val, (int, float)):
-            continue
-        gap = abs(float(budget_val) - target)
-        ranked.append((n_layer, row, gap))
-    ranked.sort(
-        key=lambda item: (
-            item[2],
-            _metric_float(item[1], 'val_bpb', float('inf')),
-            -_zero_shot_score(item[1]),
-            -_metric_float(item[1], 'tps_median', 0.0),
-        )
-    )
-    return ranked
-
-
-def _percentile_linear(sorted_values: list[float], pct: float) -> float:
-    if not sorted_values:
-        return 0.0
-    if len(sorted_values) == 1:
-        return sorted_values[0]
-    rank = (len(sorted_values) - 1) * (pct / 100.0)
-    lo = int(rank)
-    hi = min(lo + 1, len(sorted_values) - 1)
-    frac = rank - lo
-    return sorted_values[lo] * (1.0 - frac) + sorted_values[hi] * frac
-
-
-def fetch_metrics_from_job(job_id: str) -> MetricsDict | None:
-    """Fetch HF Job stdout and parse the [METRICS_JSON] line."""
+import json
+import os
+import statistics
+import re
+import sys
+from pathlib import Path
+
+from configs.harness_config import HarnessConfig
+
+type MetricValue = float | int | str | bool | None
+type MetricsDict = dict[str, MetricValue]
+
+MANIFEST = Path(sys.argv[1] if len(sys.argv) > 1 else '/tmp/sweep_depth_manifest.txt')
+STEP_TPS_PATTERN = re.compile(r"step=(\d+).*?\btps=(\d+)\b")
+MIN_TPS = float(os.environ.get('SWEEP_MIN_TPS', '0'))
+TARGET_TOKENS_M = float(os.environ.get('SWEEP_TARGET_TOKENS_M', '0'))
+TARGET_SECONDS = float(os.environ.get('SWEEP_TARGET_SECONDS', '0'))
+
+
+def _zero_shot_score(result: MetricsDict) -> float:
+    """Composite quality score for tie-breaking among BPB-near runs."""
+    factual = float(result.get('factual_english_score', 0.0) or 0.0)
+    instruction = float(result.get('instruction_following_score', 0.0) or 0.0)
+    distinct_2 = float(result.get('distinct_2', 0.0) or 0.0)
+    repetition = float(result.get('repetition_rate', 0.0) or 0.0)
+    return factual + instruction + distinct_2 - repetition
+
+
+def _metric_float(result: MetricsDict, key: str, default: float = 0.0) -> float:
+    value = result.get(key, default)
+    return float(value) if isinstance(value, (int, float)) else default
+
+
+def _metric_int(result: MetricsDict, key: str, default: int = 0) -> int:
+    value = result.get(key, default)
+    return int(value) if isinstance(value, int) else default
+
+
+def _fixed_budget_ranking(results: dict[int, MetricsDict], *, metric_key: str, target: float) -> list[tuple[int, MetricsDict, float]]:
+    ranked: list[tuple[int, MetricsDict, float]] = []
+    for n_layer, row in results.items():
+        budget_val = row.get(metric_key)
+        if not isinstance(budget_val, (int, float)):
+            continue
+        gap = abs(float(budget_val) - target)
+        ranked.append((n_layer, row, gap))
+    ranked.sort(
+        key=lambda item: (
+            item[2],
+            _metric_float(item[1], 'val_bpb', float('inf')),
+            -_zero_shot_score(item[1]),
+            -_metric_float(item[1], 'tps_median', 0.0),
+        )
+    )
+    return ranked
+
+
+def _percentile_linear(sorted_values: list[float], pct: float) -> float:
+    if not sorted_values:
+        return 0.0
+    if len(sorted_values) == 1:
+        return sorted_values[0]
+    rank = (len(sorted_values) - 1) * (pct / 100.0)
+    lo = int(rank)
+    hi = min(lo + 1, len(sorted_values) - 1)
+    frac = rank - lo
+    return sorted_values[lo] * (1.0 - frac) + sorted_values[hi] * frac
+
+
+def fetch_metrics_from_job(job_id: str) -> MetricsDict | None:
+    """Fetch HF Job stdout and parse the [METRICS_JSON] line."""
     try:
         from huggingface_hub import HfApi  # type: ignore
     except Exception as e:
@@ -94,73 +94,73 @@ def fetch_metrics_from_job(job_id: str) -> MetricsDict | None:
         print(f'[agg] could not fetch logs for job={job_id}: {e}', file=sys.stderr)
         return None
 
-    last_json = None
-    tps_samples: list[tuple[int, int]] = []
-    warmup_steps = 25
-    for line in logs_stream:
-        # HfApi returns strings or JobLogEntry-like objects depending on version.
-        text = getattr(line, 'data', None) or str(line)
-
-        wm = re.search(r"\[TPS_GUARD\] enabled .*?warmup_steps=(\d+)", text)
-        if wm:
-            warmup_steps = int(wm.group(1))
-
-        sm = STEP_TPS_PATTERN.search(text)
-        if sm:
-            tps_samples.append((int(sm.group(1)), int(sm.group(2))))
-
-        if '[METRICS_JSON]' in text:
-            payload = text.split('[METRICS_JSON]', 1)[1].strip()
-            try:
-                last_json = json.loads(payload)
-            except Exception:
-                # Might be truncated on a line boundary — keep looking.
-                pass
-    if last_json is None:
-        return None
-
-    steady_tps = [float(tps) for step, tps in tps_samples if step >= warmup_steps]
-    if not steady_tps:
-        steady_tps = [float(tps) for _, tps in tps_samples]
-    if steady_tps:
-        sorted_tps = sorted(steady_tps)
-        last_json['tps_samples'] = len(steady_tps)
-        last_json['tps_median'] = float(statistics.median(steady_tps))
-        last_json['tps_p10'] = float(_percentile_linear(sorted_tps, 10.0))
-        last_json['tps_min'] = float(sorted_tps[0])
-        last_json['tps_max'] = float(sorted_tps[-1])
-        last_json['tps_warmup_steps'] = int(warmup_steps)
-
-    return last_json
-
-
-def compare(results: dict[int, MetricsDict]) -> None:
-    """Pretty-print comparison across n_layer values."""
-    if not results:
-        print('[agg] no results')
-        return
-    sorted_n = sorted(results.keys())
-    secondary_gates = HarnessConfig().to_secondary_gates()
-
-    print('\n=== Active secondary gates ===')
-    for metric, thresholds in sorted(secondary_gates.items()):
-        print(f'  {metric}: {json.dumps(thresholds, sort_keys=True)}')
-
-    # Top-level scalars
-    print('\n=== Top-level scalars ===')
+    last_json = None
+    tps_samples: list[tuple[int, int]] = []
+    warmup_steps = 25
+    for line in logs_stream:
+        # HfApi returns strings or JobLogEntry-like objects depending on version.
+        text = getattr(line, 'data', None) or str(line)
+
+        wm = re.search(r"\[TPS_GUARD\] enabled .*?warmup_steps=(\d+)", text)
+        if wm:
+            warmup_steps = int(wm.group(1))
+
+        sm = STEP_TPS_PATTERN.search(text)
+        if sm:
+            tps_samples.append((int(sm.group(1)), int(sm.group(2))))
+
+        if '[METRICS_JSON]' in text:
+            payload = text.split('[METRICS_JSON]', 1)[1].strip()
+            try:
+                last_json = json.loads(payload)
+            except Exception:
+                # Might be truncated on a line boundary — keep looking.
+                pass
+    if last_json is None:
+        return None
+
+    steady_tps = [float(tps) for step, tps in tps_samples if step >= warmup_steps]
+    if not steady_tps:
+        steady_tps = [float(tps) for _, tps in tps_samples]
+    if steady_tps:
+        sorted_tps = sorted(steady_tps)
+        last_json['tps_samples'] = len(steady_tps)
+        last_json['tps_median'] = float(statistics.median(steady_tps))
+        last_json['tps_p10'] = float(_percentile_linear(sorted_tps, 10.0))
+        last_json['tps_min'] = float(sorted_tps[0])
+        last_json['tps_max'] = float(sorted_tps[-1])
+        last_json['tps_warmup_steps'] = int(warmup_steps)
+
+    return last_json
+
+
+def compare(results: dict[int, MetricsDict]) -> None:
+    """Pretty-print comparison across n_layer values."""
+    if not results:
+        print('[agg] no results')
+        return
+    sorted_n = sorted(results.keys())
+    secondary_gates = HarnessConfig().to_secondary_gates()
+
+    print('\n=== Active secondary gates ===')
+    for metric, thresholds in sorted(secondary_gates.items()):
+        print(f'  {metric}: {json.dumps(thresholds, sort_keys=True)}')
+
+    # Top-level scalars
+    print('\n=== Top-level scalars ===')
     hdr = ['metric'] + [f'L={n}' for n in sorted_n]
     print('  '.join(f'{h:>14}' for h in hdr))
-    for key in ('val_bpb', 'val_ppl', 'num_params_M', 'total_tokens_M',
-                'training_seconds', 'peak_vram_mb', 'sdr_target_active',
-                'htm_anomaly', 'engram_hit_rate', 'sdr_active_bits',
-                'tps_median', 'tps_p10', 'tps_min', 'tps_max', 'tps_samples'):
-        row = [key] + [f'{results[n].get(key, float("nan")):.4f}' if isinstance(results[n].get(key), (int, float)) else 'n/a' for n in sorted_n]
-        print('  '.join(f'{c:>14}' for c in row))
+    for key in ('val_bpb', 'val_ppl', 'num_params_M', 'total_tokens_M',
+                'training_seconds', 'peak_vram_mb', 'sdr_target_active',
+                'htm_anomaly', 'engram_hit_rate', 'sdr_active_bits',
+                'tps_median', 'tps_p10', 'tps_min', 'tps_max', 'tps_samples'):
+        row = [key] + [f'{results[n].get(key, float("nan")):.4f}' if isinstance(results[n].get(key), (int, float)) else 'n/a' for n in sorted_n]
+        print('  '.join(f'{c:>14}' for c in row))
 
     # Per-layer panel — one table per metric.
     print('\n=== Per-layer: delta_ratio (residual contribution) ===')
     print('  '.join(['layer'] + [f'L={n:>2}' for n in sorted_n]))
-    max_depth = max(_metric_int(results[n], 'n_layer', 0) for n in sorted_n)
+    max_depth = max(_metric_int(results[n], 'n_layer', 0) for n in sorted_n)
     for li in range(max_depth):
         row = [f'L{li:02d}']
         for n in sorted_n:
@@ -197,62 +197,62 @@ def compare(results: dict[int, MetricsDict]) -> None:
 
     # Dead-layer detection
     print('\n=== Dead-layer detection (delta_ratio < 0.02) ===')
-    for n in sorted_n:
-        r = results[n]
-        n_layer = _metric_int(r, 'n_layer', 0)
+    for n in sorted_n:
+        r = results[n]
+        n_layer = _metric_int(r, 'n_layer', 0)
         dead = []
         for li in range(n_layer):
             v = r.get(f'layer_{li}_delta_ratio')
             if isinstance(v, (int, float)) and v < 0.02:
                 dead.append(li)
-        status = 'ALL LIVE' if not dead else f'DEAD LAYERS: {dead}'
-        print(f'  n_layer={n:2d}  val_bpb={r.get("val_bpb", float("nan")):.4f}  {status}')
-
-    print('\n=== Throughput-constrained ranking ===')
-    ranked = sorted(
-        ((n, r) for n, r in results.items() if isinstance(r.get('val_bpb'), (int, float))),
-        key=lambda x: (
-            (MIN_TPS > 0) and (_metric_float(x[1], 'tps_median', 0.0) < MIN_TPS),
-            _metric_float(x[1], 'val_bpb', float('inf')),
-            -_zero_shot_score(x[1]),
-        ),
-    )
-    feasible_count = 0
-    for n, r in ranked:
-        tps_median = _metric_float(r, 'tps_median', 0.0)
-        feasible = (MIN_TPS <= 0) or (tps_median >= MIN_TPS)
-        zero_shot_score = _zero_shot_score(r)
-        if feasible:
-            feasible_count += 1
-        print(
-            f"  n_layer={n:2d} val_bpb={_metric_float(r, 'val_bpb', float('nan')):.4f} "
-            f"tps_median={tps_median:.0f} zero_shot_score={zero_shot_score:.4f} feasible={feasible}",
-            flush=True,
-        )
-    if MIN_TPS > 0:
-        print(f"[agg] throughput gate: tps_median >= {MIN_TPS:.0f}; feasible={feasible_count}/{len(ranked)}")
-
-    if TARGET_TOKENS_M > 0:
-        print('\n=== Fixed-token champion comparison ===')
-        print(f'  target_tokens_M={TARGET_TOKENS_M:.4f}')
-        for n, r, gap in _fixed_budget_ranking(results, metric_key='total_tokens_M', target=TARGET_TOKENS_M):
-            print(
-                f"  n_layer={n:2d} val_bpb={_metric_float(r, 'val_bpb', float('nan')):.4f} "
-                f"total_tokens_M={_metric_float(r, 'total_tokens_M', float('nan')):.4f} "
-                f"token_gap_M={gap:.4f} tps_median={_metric_float(r, 'tps_median', 0.0):.0f}",
-                flush=True,
-            )
-
-    if TARGET_SECONDS > 0:
-        print('\n=== Fixed-time champion comparison ===')
-        print(f'  target_seconds={TARGET_SECONDS:.1f}')
-        for n, r, gap in _fixed_budget_ranking(results, metric_key='training_seconds', target=TARGET_SECONDS):
-            print(
-                f"  n_layer={n:2d} val_bpb={_metric_float(r, 'val_bpb', float('nan')):.4f} "
-                f"training_seconds={_metric_float(r, 'training_seconds', float('nan')):.1f} "
-                f"time_gap_s={gap:.1f} tps_median={_metric_float(r, 'tps_median', 0.0):.0f}",
-                flush=True,
-            )
+        status = 'ALL LIVE' if not dead else f'DEAD LAYERS: {dead}'
+        print(f'  n_layer={n:2d}  val_bpb={r.get("val_bpb", float("nan")):.4f}  {status}')
+
+    print('\n=== Throughput-constrained ranking ===')
+    ranked = sorted(
+        ((n, r) for n, r in results.items() if isinstance(r.get('val_bpb'), (int, float))),
+        key=lambda x: (
+            (MIN_TPS > 0) and (_metric_float(x[1], 'tps_median', 0.0) < MIN_TPS),
+            _metric_float(x[1], 'val_bpb', float('inf')),
+            -_zero_shot_score(x[1]),
+        ),
+    )
+    feasible_count = 0
+    for n, r in ranked:
+        tps_median = _metric_float(r, 'tps_median', 0.0)
+        feasible = (MIN_TPS <= 0) or (tps_median >= MIN_TPS)
+        zero_shot_score = _zero_shot_score(r)
+        if feasible:
+            feasible_count += 1
+        print(
+            f"  n_layer={n:2d} val_bpb={_metric_float(r, 'val_bpb', float('nan')):.4f} "
+            f"tps_median={tps_median:.0f} zero_shot_score={zero_shot_score:.4f} feasible={feasible}",
+            flush=True,
+        )
+    if MIN_TPS > 0:
+        print(f"[agg] throughput gate: tps_median >= {MIN_TPS:.0f}; feasible={feasible_count}/{len(ranked)}")
+
+    if TARGET_TOKENS_M > 0:
+        print('\n=== Fixed-token champion comparison ===')
+        print(f'  target_tokens_M={TARGET_TOKENS_M:.4f}')
+        for n, r, gap in _fixed_budget_ranking(results, metric_key='total_tokens_M', target=TARGET_TOKENS_M):
+            print(
+                f"  n_layer={n:2d} val_bpb={_metric_float(r, 'val_bpb', float('nan')):.4f} "
+                f"total_tokens_M={_metric_float(r, 'total_tokens_M', float('nan')):.4f} "
+                f"token_gap_M={gap:.4f} tps_median={_metric_float(r, 'tps_median', 0.0):.0f}",
+                flush=True,
+            )
+
+    if TARGET_SECONDS > 0:
+        print('\n=== Fixed-time champion comparison ===')
+        print(f'  target_seconds={TARGET_SECONDS:.1f}')
+        for n, r, gap in _fixed_budget_ranking(results, metric_key='training_seconds', target=TARGET_SECONDS):
+            print(
+                f"  n_layer={n:2d} val_bpb={_metric_float(r, 'val_bpb', float('nan')):.4f} "
+                f"training_seconds={_metric_float(r, 'training_seconds', float('nan')):.1f} "
+                f"time_gap_s={gap:.1f} tps_median={_metric_float(r, 'tps_median', 0.0):.0f}",
+                flush=True,
+            )
 
 
 def main() -> int:
@@ -273,7 +273,7 @@ def main() -> int:
         jobs[n_layer] = job_id
 
     print(f'[agg] reading {len(jobs)} jobs from {MANIFEST}')
-    results: dict[int, MetricsDict] = {}
+    results: dict[int, MetricsDict] = {}
     for n, jid in jobs.items():
         print(f'[agg] fetching job={jid} (n_layer={n}) ...')
         m = fetch_metrics_from_job(jid)

overlay/scripts/watch_benchmark_hf_job.py

@@ -1,81 +1,33 @@
-#!/usr/bin/env python3
-from __future__ import annotations
-
-import argparse
-import json
-import time
-from pathlib import Path
-
-from huggingface_hub import HfApi
-from huggingface_hub.utils import get_token
-
-
-def parse_benchmark_result_from_logs(lines: list[str]):
-    for line in reversed(lines):
-        text = line.strip()
-        if not text.startswith("{"):
-            continue
-        try:
-            payload = json.loads(text)
-        except json.JSONDecodeError:
-            continue
-        if isinstance(payload, dict) and "benchmark" in payload:
-            return payload
-    return None
-
-
-def collect_job_snapshot(api, *, job_id: str, token: str, namespace: str) -> dict[str, object]:
-    info = api.inspect_job(job_id=job_id, token=token, namespace=namespace)
-    stage = str(info.status.stage)
-    message = str(getattr(info.status, "message", "") or "")
-    logs = list(api.fetch_job_logs(job_id=job_id, follow=False, token=token, namespace=namespace))
-    texts = [(getattr(line, "data", None) or str(line)) for line in logs]
-    return {
-        "job_id": job_id,
-        "stage": stage,
-        "message": message,
-        "log_lines": len(texts),
-        "result": parse_benchmark_result_from_logs(texts),
-    }
-
-
-def wait_for_terminal_snapshot(api, *, job_id: str, token: str, namespace: str, poll_interval: float = 10.0, timeout_s: float = 1800.0) -> dict[str, object]:
-    deadline = time.time() + timeout_s
-    terminal = {"COMPLETED", "ERROR", "FAILED", "CANCELLED", "CANCELED", "TIMEOUT"}
-    while True:
-        payload = collect_job_snapshot(api, job_id=job_id, token=token, namespace=namespace)
-        if payload["stage"] in terminal:
-            return payload
-        if time.time() >= deadline:
-            return payload
-        time.sleep(poll_interval)
-
-
-def write_watch_summary(path: Path, payload: dict[str, object]) -> None:
-    path.parent.mkdir(parents=True, exist_ok=True)
-    path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
-
-
-def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Watch or snapshot a remote benchmark job")
-    parser.add_argument("--job-id", required=True)
-    parser.add_argument("--namespace", default="jackoatmon")
-    parser.add_argument("--summary-out", type=Path)
-    return parser.parse_args(argv)
-
-
-def main(argv: list[str] | None = None) -> int:
-    args = parse_args(argv)
-    token = get_token()
-    if not token:
-        raise SystemExit("HF_TOKEN must be set or cached via huggingface-cli login")
-    api = HfApi(token=token)
-    payload = collect_job_snapshot(api, job_id=args.job_id, token=token, namespace=args.namespace)
-    if args.summary_out is not None:
-        write_watch_summary(args.summary_out, payload)
-    print(json.dumps(payload, indent=2, sort_keys=True))
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+
+def parse_benchmark_result_from_logs(lines: list[str]):
+    for line in reversed(lines):
+        text = line.strip()
+        if not text.startswith("{"):
+            continue
+        try:
+            payload = json.loads(text)
+        except json.JSONDecodeError:
+            continue
+        if isinstance(payload, dict) and "benchmark" in payload:
+            return payload
+    return None
+
+
+def write_watch_summary(path: Path, payload: dict[str, object]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
+
+
+def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Watch or snapshot a remote benchmark job")
+    parser.add_argument("--job-id", required=True)
+    parser.add_argument("--namespace", default="jackoatmon")
+    parser.add_argument("--summary-out", type=Path)
+    return parser.parse_args(argv)

overlay/subsystems/htm.py

@@ -29,46 +29,46 @@ copy is small compared to the SP/TM compute.
 from __future__ import annotations
 
 import time
-from concurrent.futures import ThreadPoolExecutor
-from typing import Any
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any
 
 import numpy as np
 import torch
 import torch.nn as nn
 
-import htm_rust
-
-_HTM_REGION: Any = getattr(htm_rust, "HTMRegion", None)
-_HTM_REGION_GPU: Any = getattr(htm_rust, "HTMRegionGpu", None)
-_HTM_STEP_BATCH_FUSED_CUDA: Any = getattr(htm_rust, "step_batch_fused_cuda", None)
-
-# step_many releases the GIL for the whole pass, so multiple threads can
-# truly run regions in parallel — wall-clock scales with B up to CPU cores.
-_HTM_HAS_STEP_MANY = hasattr(_HTM_REGION, "step_many")
+import htm_rust
+
+_HTM_REGION: Any = getattr(htm_rust, "HTMRegion", None)
+_HTM_REGION_GPU: Any = getattr(htm_rust, "HTMRegionGpu", None)
+_HTM_STEP_BATCH_FUSED_CUDA: Any = getattr(htm_rust, "step_batch_fused_cuda", None)
+
+# step_many releases the GIL for the whole pass, so multiple threads can
+# truly run regions in parallel — wall-clock scales with B up to CPU cores.
+_HTM_HAS_STEP_MANY = hasattr(_HTM_REGION, "step_many")
 # GPU backend: built with `maturin develop --features gpu`. One CUDA region
 # per batch slot, persistent device state for SP synapses. Transparent
 # fallback to CPU when not available.
-_HTM_HAS_GPU = hasattr(htm_rust, "HTMRegionGpu")
+_HTM_HAS_GPU = hasattr(htm_rust, "HTMRegionGpu")
 # Zero-copy CUDA path: consumes torch CUDA tensors directly via the
 # __cuda_array_interface__ protocol, skipping the sdr.cpu()/numpy round-trip
 # and the D2H of outputs. Huge win when the input SDR already lives on GPU
 # (which is the train.py hot path — retina is a device buffer).
-_HTM_HAS_CAI = _HTM_HAS_GPU and hasattr(_HTM_REGION_GPU, "step_many_cuda")
+_HTM_HAS_CAI = _HTM_HAS_GPU and hasattr(_HTM_REGION_GPU, "step_many_cuda")
 # Fused megakernel path: collapses all T timesteps + SP + TM into a single
 # CUDA launch per forward. Replaces global top-K with per-column threshold
 # inhibition (see htm_rust/docs/GPU_HTM.md §Fused Kernel).
 # Opt-in via env var (default on when available).
 import os as _os_fused
-_HTM_HAS_FUSED = _HTM_HAS_GPU and hasattr(_HTM_REGION_GPU, "step_many_fused_cuda")
-_HTM_USE_FUSED = _HTM_HAS_FUSED and bool(int(_os_fused.environ.get("HYDRA_HTM_FUSED", "1")))
-
-
-def _is_fused_unavailable_error(exc: RuntimeError) -> bool:
-    message = str(exc)
-    return (
-        "Fused HTM kernel is unavailable" in message
-        or "fused HTM kernel disabled for this CUDA arch" in message
-    )
+_HTM_HAS_FUSED = _HTM_HAS_GPU and hasattr(_HTM_REGION_GPU, "step_many_fused_cuda")
+_HTM_USE_FUSED = _HTM_HAS_FUSED and bool(int(_os_fused.environ.get("HYDRA_HTM_FUSED", "1")))
+
+
+def _is_fused_unavailable_error(exc: RuntimeError) -> bool:
+    message = str(exc)
+    return (
+        "Fused HTM kernel is unavailable" in message
+        or "fused HTM kernel disabled for this CUDA arch" in message
+    )
 
 
 class HTMLayer(nn.Module):
@@ -93,11 +93,11 @@ class HTMLayer(nn.Module):
         learn: bool = True,
         reset_each_forward: bool = True,
         use_gpu: bool | None = None,
-    ) -> None:
-        super().__init__()
-        self.input_bits = input_bits
-        self.n_columns = n_columns
-        self.cells_per_column = cells_per_column
+    ) -> None:
+        super().__init__()
+        self.input_bits = input_bits
+        self.n_columns = n_columns
+        self.cells_per_column = cells_per_column
         self.learn = learn
         self.reset_each_forward = reset_each_forward
         self._seed_base = seed
@@ -107,23 +107,23 @@ class HTMLayer(nn.Module):
         # converges since the EMA accumulates over many calls. Env:
         # HYDRA_HTM_LEARN_EVERY=N (default 1 = every forward, 0 = disabled).
         import os as _os
-        self._learn_every = int(_os.environ.get("HYDRA_HTM_LEARN_EVERY", "1"))
-        self._forward_counter = 0
-        force_cpu = _os.environ.get("HYDRA_FORCE_HTM_CPU", "0") == "1"
-        # GPU backend gate. Default: auto-detect — use GPU when the pyo3
-        # module was built with --features gpu AND CUDA is actually usable.
-        if use_gpu is None:
-            use_gpu = (not force_cpu) and _HTM_HAS_GPU and torch.cuda.is_available()
-        elif use_gpu and not _HTM_HAS_GPU:
-            raise RuntimeError(
-                "HTMLayer(use_gpu=True) but htm_rust was not built with "
-                "--features gpu. Re-run `maturin develop --features gpu`."
-            )
-        elif use_gpu and force_cpu:
-            use_gpu = False
-        self._use_gpu = bool(use_gpu)
-        cls = _HTM_REGION_GPU if self._use_gpu else _HTM_REGION
-        self._region_cls = cls
+        self._learn_every = int(_os.environ.get("HYDRA_HTM_LEARN_EVERY", "1"))
+        self._forward_counter = 0
+        force_cpu = _os.environ.get("HYDRA_FORCE_HTM_CPU", "0") == "1"
+        # GPU backend gate. Default: auto-detect — use GPU when the pyo3
+        # module was built with --features gpu AND CUDA is actually usable.
+        if use_gpu is None:
+            use_gpu = (not force_cpu) and _HTM_HAS_GPU and torch.cuda.is_available()
+        elif use_gpu and not _HTM_HAS_GPU:
+            raise RuntimeError(
+                "HTMLayer(use_gpu=True) but htm_rust was not built with "
+                "--features gpu. Re-run `maturin develop --features gpu`."
+            )
+        elif use_gpu and force_cpu:
+            use_gpu = False
+        self._use_gpu = bool(use_gpu)
+        cls = _HTM_REGION_GPU if self._use_gpu else _HTM_REGION
+        self._region_cls = cls
         self._regions = [
             cls(input_bits, n_columns, cells_per_column, seed + i)
             for i in range(batch_size)
@@ -144,19 +144,19 @@ class HTMLayer(nn.Module):
                 )
             )
 
-    def reset(self) -> None:
-        """Clear TM predictive state on every region (keeps SP synapses)."""
-        for r in self._regions:
-            r.reset()
-
-    def _next_learn_flag(self) -> bool:
-        self._forward_counter += 1
-        return bool(
-            self.learn
-            and self.training
-            and self._learn_every > 0
-            and (self._forward_counter % self._learn_every == 0)
-        )
+    def reset(self) -> None:
+        """Clear TM predictive state on every region (keeps SP synapses)."""
+        for r in self._regions:
+            r.reset()
+
+    def _next_learn_flag(self) -> bool:
+        self._forward_counter += 1
+        return bool(
+            self.learn
+            and self.training
+            and self._learn_every > 0
+            and (self._forward_counter % self._learn_every == 0)
+        )
 
     @torch.no_grad()
     def forward(self, sdr: torch.Tensor) -> torch.Tensor:
@@ -167,9 +167,9 @@ class HTMLayer(nn.Module):
         if self.reset_each_forward:
             self.reset()
 
-        # Learn-gate: run learn kernels only every N forwards (skips 56% of
-        # HTM CUDA time on skip-forwards; Hebbian EMA still converges).
-        learn = self._next_learn_flag()
+        # Learn-gate: run learn kernels only every N forwards (skips 56% of
+        # HTM CUDA time on skip-forwards; Hebbian EMA still converges).
+        learn = self._next_learn_flag()
 
         # Zero-copy CUDA hot path. SDR already lives on GPU (retina buffer),
         # so we skip sdr.cpu()/numpy round-trip AND the output D2H. The Rust
@@ -178,30 +178,30 @@ class HTMLayer(nn.Module):
         if _HTM_HAS_CAI and self._use_gpu and sdr.is_cuda:
             sdr_u8 = sdr.contiguous().to(torch.uint8) if sdr.dtype != torch.uint8 else sdr.contiguous()
             cols_out = torch.empty((B, T, self.n_columns), dtype=torch.uint8, device=sdr.device)
-            anom_out = torch.empty((B, T), dtype=torch.float32, device=sdr.device)
-            # Pick fused (1 launch) or legacy (12*T launches) path.
-            if _HTM_USE_FUSED:
-                try:
-                    for b in range(B):
-                        self._regions[b].step_many_fused_cuda(
-                            sdr_u8[b].__cuda_array_interface__,
-                            cols_out[b].__cuda_array_interface__,
-                            anom_out[b].__cuda_array_interface__,
-                            learn,
-                        )
-                except RuntimeError as exc:
-                    if not _is_fused_unavailable_error(exc):
-                        raise
-                    for b in range(B):
-                        self._regions[b].step_many_cuda(
-                            sdr_u8[b].__cuda_array_interface__,
-                            cols_out[b].__cuda_array_interface__,
-                            anom_out[b].__cuda_array_interface__,
-                            learn,
-                        )
-            else:
-                for b in range(B):
-                    self._regions[b].step_many_cuda(
+            anom_out = torch.empty((B, T), dtype=torch.float32, device=sdr.device)
+            # Pick fused (1 launch) or legacy (12*T launches) path.
+            if _HTM_USE_FUSED:
+                try:
+                    for b in range(B):
+                        self._regions[b].step_many_fused_cuda(
+                            sdr_u8[b].__cuda_array_interface__,
+                            cols_out[b].__cuda_array_interface__,
+                            anom_out[b].__cuda_array_interface__,
+                            learn,
+                        )
+                except RuntimeError as exc:
+                    if not _is_fused_unavailable_error(exc):
+                        raise
+                    for b in range(B):
+                        self._regions[b].step_many_cuda(
+                            sdr_u8[b].__cuda_array_interface__,
+                            cols_out[b].__cuda_array_interface__,
+                            anom_out[b].__cuda_array_interface__,
+                            learn,
+                        )
+            else:
+                for b in range(B):
+                    self._regions[b].step_many_cuda(
                         sdr_u8[b].__cuda_array_interface__,
                         cols_out[b].__cuda_array_interface__,
                         anom_out[b].__cuda_array_interface__,
@@ -275,7 +275,7 @@ class HTMLayer(nn.Module):
         self._ensure_regions(B)
         if self.reset_each_forward:
             self.reset()
-        learn = self._next_learn_flag()
+        learn = self._next_learn_flag()
 
         if _HTM_HAS_CAI and self._use_gpu and sdr.is_cuda:
             sdr_u8 = sdr.contiguous().to(torch.uint8) if sdr.dtype != torch.uint8 else sdr.contiguous()
@@ -287,61 +287,61 @@ class HTMLayer(nn.Module):
             # grid.y = B processes all regions concurrently — ~B× speedup.
             # Falls back to sequential dispatch if the batched entry isn't
             # available (older htm_rust wheel).
-            if _HTM_USE_FUSED and _HTM_STEP_BATCH_FUSED_CUDA is not None:
+            if _HTM_USE_FUSED and _HTM_STEP_BATCH_FUSED_CUDA is not None:
                 # Slice self._regions to match B: _ensure_regions may have
                 # allocated more regions than the current batch size needs
                 # (e.g. factual eval uses smaller batches than training).
                 try:
-                    _HTM_STEP_BATCH_FUSED_CUDA(
+                    _HTM_STEP_BATCH_FUSED_CUDA(
                         self._regions[:B],
                         [sdr_u8[b].__cuda_array_interface__ for b in range(B)],
                         [cols_out[b].__cuda_array_interface__ for b in range(B)],
                         [anom_out[b].__cuda_array_interface__ for b in range(B)],
                         learn,
                     )
-                except RuntimeError as _e:
-                    if "COOPERATIVE_LAUNCH_TOO_LARGE" in str(_e):
-                        # Batch too large for cooperative grid. Fall back to
-                        # sequential per-region fused launches (each B=1).
-                        for b in range(B):
+                except RuntimeError as _e:
+                    if "COOPERATIVE_LAUNCH_TOO_LARGE" in str(_e):
+                        # Batch too large for cooperative grid. Fall back to
+                        # sequential per-region fused launches (each B=1).
+                        for b in range(B):
                             self._regions[b].step_many_fused_cuda(
                                 sdr_u8[b].__cuda_array_interface__,
                                 cols_out[b].__cuda_array_interface__,
-                                anom_out[b].__cuda_array_interface__,
-                                learn,
-                            )
-                    elif _is_fused_unavailable_error(_e):
-                        for b in range(B):
-                            self._regions[b].step_many_cuda(
-                                sdr_u8[b].__cuda_array_interface__,
-                                cols_out[b].__cuda_array_interface__,
-                                anom_out[b].__cuda_array_interface__,
-                                learn,
-                            )
-                    else:
-                        raise
-            elif _HTM_USE_FUSED:
-                try:
-                    for b in range(B):
-                        self._regions[b].step_many_fused_cuda(
-                            sdr_u8[b].__cuda_array_interface__,
-                            cols_out[b].__cuda_array_interface__,
-                            anom_out[b].__cuda_array_interface__,
-                            learn,
-                        )
-                except RuntimeError as exc:
-                    if not _is_fused_unavailable_error(exc):
-                        raise
-                    for b in range(B):
-                        self._regions[b].step_many_cuda(
-                            sdr_u8[b].__cuda_array_interface__,
-                            cols_out[b].__cuda_array_interface__,
-                            anom_out[b].__cuda_array_interface__,
-                            learn,
-                        )
-            else:
-                for b in range(B):
-                    self._regions[b].step_many_cuda(
+                                anom_out[b].__cuda_array_interface__,
+                                learn,
+                            )
+                    elif _is_fused_unavailable_error(_e):
+                        for b in range(B):
+                            self._regions[b].step_many_cuda(
+                                sdr_u8[b].__cuda_array_interface__,
+                                cols_out[b].__cuda_array_interface__,
+                                anom_out[b].__cuda_array_interface__,
+                                learn,
+                            )
+                    else:
+                        raise
+            elif _HTM_USE_FUSED:
+                try:
+                    for b in range(B):
+                        self._regions[b].step_many_fused_cuda(
+                            sdr_u8[b].__cuda_array_interface__,
+                            cols_out[b].__cuda_array_interface__,
+                            anom_out[b].__cuda_array_interface__,
+                            learn,
+                        )
+                except RuntimeError as exc:
+                    if not _is_fused_unavailable_error(exc):
+                        raise
+                    for b in range(B):
+                        self._regions[b].step_many_cuda(
+                            sdr_u8[b].__cuda_array_interface__,
+                            cols_out[b].__cuda_array_interface__,
+                            anom_out[b].__cuda_array_interface__,
+                            learn,
+                        )
+            else:
+                for b in range(B):
+                    self._regions[b].step_many_cuda(
                         sdr_u8[b].__cuda_array_interface__,
                         cols_out[b].__cuda_array_interface__,
                         anom_out[b].__cuda_array_interface__,