Update Feather H200 training runtime image

overlay/hydra/config.py

@@ -48,7 +48,10 @@ class PostSemClawConfig:
 
     # SemanticFoldingSDR (offline retina with STE; no-bypass, runs every step)
     sdr_n_bits: int = 16384          # retina width
-    sdr_target_active: int = 327     # exact active bits (2% sparsity)
+    # Default 327 = 2% sparsity (Webber/Numenta canonical). Override with
+    # HYDRA_SDR_TARGET_ACTIVE env var; value MUST match subsystems/sdr_retina.py
+    # TARGET_ACTIVE (same env var is read there, so just setting it once works).
+    sdr_target_active: int = int(os.environ.get("HYDRA_SDR_TARGET_ACTIVE", "327"))
     sdr_delta_rank: int = 32         # low-rank STE delta rank
     sdr_som_warmup: int = 500
     sdr_som_interval: int = 100

overlay/hydra/model.py

@@ -141,6 +141,29 @@ class PostSemClawModel(nn.Module):
         # Secondary metrics storage
         self._metrics = {}
 
+        # Per-layer diagnostic panel. Env-gated; zero overhead when off.
+        # Emits residual-contribution (delta_ratio), feature std, effective rank,
+        # gradient norm per layer; used to identify minimum viable n_layer + find
+        # entropy leakage / dead layers. See docs/depth-sweep.md.
+        self._diag_enabled = os.environ.get("HYDRA_LAYER_DIAGNOSTICS", "0") == "1"
+        self._diag_step = 0
+        self._diag_svd_every = int(os.environ.get("HYDRA_LAYER_DIAG_SVD_EVERY", "100"))
+        if self._diag_enabled:
+            # Gradient-norm backward hooks on each Mamba3 block output.
+            # grad_output[0] is dL/d(block_out) — measures how much learning
+            # signal survives to reach this layer. Used to detect vanishing or
+            # exploding gradients per-depth.
+            for _i, _block in enumerate(self.blocks):
+                def _mk_grad_hook(_layer_idx):
+                    def _hook(module, grad_input, grad_output):
+                        if grad_output and grad_output[0] is not None:
+                            g = grad_output[0].detach()
+                            self._metrics[f'layer_{_layer_idx}_grad_norm'] = float(
+                                g.pow(2).mean().sqrt().item()
+                            )
+                    return _hook
+                _block.register_full_backward_hook(_mk_grad_hook(_i))
+
         # Triton kernel integration gates (Phase 2 — deferred, see module docstring).
         self._fused_bcnorm = os.environ.get("HYDRA_FUSED_BCNORM", "0") == "1"
         self._fused_ssd = os.environ.get("HYDRA_FUSED_SSD", "0") == "1"
@@ -385,6 +408,16 @@ class PostSemClawModel(nn.Module):
         # mHC-routed Mamba-3 stack with Engram injection at configured layer.
         streams = self.mhc[0].init_streams(x)
         _engram_ev = None
+
+        # 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.
+        # All reads are detached no-grad to avoid autograd graph pollution.
+        _diag = self._diag_enabled
+        if _diag:
+            with torch.no_grad():
+                h_pre = self.mhc[0].merge_streams(streams).detach()
+            _run_svd = (self._diag_step % self._diag_svd_every) == 0
+
         for i, (block, mhc_layer) in enumerate(zip(self.blocks, self.mhc)):
             def _block_fn(h, _block=block):
                 return self.drop(_block(norm(h)))
@@ -399,6 +432,32 @@ class PostSemClawModel(nn.Module):
                 self._metrics['engram_hit_rate'] = hit_rate
                 if _profile: _engram_ev = _ev()
 
+            if _diag:
+                with torch.no_grad():
+                    h_post = mhc_layer.merge_streams(streams).detach()
+                    in_n  = h_pre.pow(2).mean().sqrt()
+                    out_n = h_post.pow(2).mean().sqrt()
+                    d_n   = (h_post - h_pre).pow(2).mean().sqrt()
+                    self._metrics[f'layer_{i}_in_norm']     = float(in_n.item())
+                    self._metrics[f'layer_{i}_out_norm']    = float(out_n.item())
+                    self._metrics[f'layer_{i}_delta_ratio'] = float((d_n / (in_n + 1e-6)).item())
+                    self._metrics[f'layer_{i}_feat_std']    = float(h_post.std(dim=-1).mean().item())
+                    if _run_svd:
+                        # Effective rank via participation ratio of singular values.
+                        # eff_rank = (Σσ)^2 / Σσ² — smooth rank proxy, bounded by d_model.
+                        # Sampled to keep overhead low (SVD is O(min(B*T, D)^2·D)).
+                        flat = h_post.reshape(-1, h_post.shape[-1])[:512].float()
+                        try:
+                            s = torch.linalg.svdvals(flat)
+                            eff_rank = float(((s.sum() ** 2) / (s.pow(2).sum() + 1e-6)).item())
+                            self._metrics[f'layer_{i}_eff_rank'] = eff_rank
+                        except Exception:
+                            pass
+                    h_pre = h_post
+
+        if _diag:
+            self._diag_step += 1
+
         if _profile: _t_blocks = _ev()
 
         self._metrics['sdr_active_bits'] = sdr_active_bits

overlay/hydra/training.py

@@ -7,6 +7,7 @@ preserved. Public entrypoint: `main()`.
 from __future__ import annotations
 
 import gc
+import json
 import math
 import os
 import sys
@@ -427,6 +428,30 @@ def main() -> None:
                 _prepare_mod.EVAL_TOKENS = _orig_mid
                 mid_ppl = 2.0 ** mid_bpb
                 print(f"[MID_VAL] step={step} val_bpb={mid_bpb:.4f} val_ppl={mid_ppl:.3f}", flush=True)
+
+                # Per-layer diagnostic panel. Only printed when HYDRA_LAYER_DIAGNOSTICS=1
+                # is set (otherwise the layer_* keys are absent from _metrics).
+                _diag_metrics = model.get_secondary_metrics()
+                _layer_keys = sorted([k for k in _diag_metrics.keys() if k.startswith('layer_')])
+                if _layer_keys:
+                    # Condense: one row per layer showing the four core signals.
+                    n_layers = len(model.blocks)
+                    print(f"[LAYER_DIAG] step={step}", flush=True)
+                    for li in range(n_layers):
+                        d_ratio = _diag_metrics.get(f'layer_{li}_delta_ratio', float('nan'))
+                        out_n   = _diag_metrics.get(f'layer_{li}_out_norm',    float('nan'))
+                        g_norm  = _diag_metrics.get(f'layer_{li}_grad_norm',   float('nan'))
+                        eff_r   = _diag_metrics.get(f'layer_{li}_eff_rank',    float('nan'))
+                        f_std   = _diag_metrics.get(f'layer_{li}_feat_std',    float('nan'))
+                        print(
+                            f"[LAYER_DIAG]   L{li:02d}  delta_ratio={d_ratio:.4f}  "
+                            f"out_norm={out_n:.4f}  grad_norm={g_norm:.3e}  "
+                            f"eff_rank={eff_r:.1f}  feat_std={f_std:.4f}",
+                            flush=True,
+                        )
+                    htm_proj_g = _diag_metrics.get('htm_proj_grad_norm', None)
+                    if htm_proj_g is not None:
+                        print(f"[LAYER_DIAG]   htm_proj grad_norm={htm_proj_g:.3e}", flush=True)
             except Exception as e:
                 print(f"[MID_VAL] failed: {e}", flush=True)
             model.train()
@@ -509,6 +534,50 @@ def main() -> None:
     print(f"sdr_active_bits:  {metrics.get('sdr_active_bits', 0):.1f}")
     print(f"htm_anomaly:      {metrics.get('htm_anomaly', 0):.4f}")
 
+    # Per-layer summary panel — only printed when diagnostics were active.
+    _layer_keys = sorted([k for k in metrics.keys() if k.startswith('layer_')])
+    if _layer_keys:
+        n_layers = len(model.blocks)
+        print("--- per-layer diagnostic panel ---")
+        for li in range(n_layers):
+            d_ratio = metrics.get(f'layer_{li}_delta_ratio', float('nan'))
+            out_n   = metrics.get(f'layer_{li}_out_norm',    float('nan'))
+            g_norm  = metrics.get(f'layer_{li}_grad_norm',   float('nan'))
+            eff_r   = metrics.get(f'layer_{li}_eff_rank',    float('nan'))
+            f_std   = metrics.get(f'layer_{li}_feat_std',    float('nan'))
+            print(
+                f"L{li:02d}  delta_ratio={d_ratio:.4f}  out_norm={out_n:.4f}  "
+                f"grad_norm={g_norm:.3e}  eff_rank={eff_r:.1f}  feat_std={f_std:.4f}"
+            )
+
+    # 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 = dict(metrics)
+        _dump.update({
+            'val_bpb': float(val_bpb),
+            'val_ppl': float(val_ppl),
+            '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(os.environ.get("HYDRA_SDR_TARGET_ACTIVE", "327")),
+        })
+        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)
+        # Also emit a single-line JSON to stdout so the sweep aggregator can
+        # scrape it from HF Jobs logs without pulling files out of the container.
+        print("[METRICS_JSON] " + json.dumps(_dump, sort_keys=True), flush=True)
+    except Exception as _e:
+        print(f"[METRICS] write failed: {_e}", flush=True)
+
     run_factual_english(model, tokenizer, MAX_SEQ_LEN)
     # startup_time is informative but not printed (preserve historical output)
     _ = startup_time

overlay/prep_nemotron.py

@@ -0,0 +1,281 @@
+#!/usr/bin/env python3
+"""Nemotron Super3 pretraining data prep.
+
+Downloads nvidia/Nemotron-Pretraining-Specialized-v1.1 configs, tokenizes with
+our rustbpe/tiktoken tokenizer (trained by prepare.py), and writes
+shard_{NNNNN}.parquet files consumable by the existing training pipeline —
+identical layout to prepare.py: a single column named 'tokens' of dtype uint16,
+with rows of length equal to --tokens-per-row (default: all tokens in one row
+group, matching parquet convention used by training.py via _document_batches).
+
+Phase 1 (diversity blend): equal weight across all 5 configs.
+Phase 2 (quality blend):   weighted toward Multiple-Choice/Economics/Formal-Logic.
+
+Usage:
+    python prep_nemotron.py --phase phase1 --parts-per-config 8
+    python prep_nemotron.py --phase phase2 --parts-per-config 8 --shard-id-start 100
+
+The --shard-id-start flag lets phase 2 append shards without colliding with
+phase 1 output (phase 2 resumes from the checkpoint stored in HF Hub by
+entrypoint.py, so the shard ids just need to be unique on-disk).
+"""
+
+import argparse
+import os
+import pickle
+import shutil
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+from huggingface_hub import HfApi, hf_hub_download
+
+# ---------------------------------------------------------------------------
+# Import constants from prepare.py (tokenizer path, data dir, val shard id)
+# ---------------------------------------------------------------------------
+# prepare.py lives in the same directory; import at module level so
+# DATA_DIR / TOKENIZER_DIR are always available.
+import prepare as _p
+
+NEMOTRON_REPO = "nvidia/Nemotron-Pretraining-Specialized-v1.1"
+
+# The 5 configs per the Super3 recipe
+ALL_CONFIGS = [
+    "Nemotron-Pretraining-Code-Concepts",
+    "Nemotron-Pretraining-Unconditional-Algorithmic",
+    "Nemotron-Pretraining-Economics",
+    "Nemotron-Pretraining-Formal-Logic",
+    "Nemotron-Pretraining-Multiple-Choice",
+]
+
+CONFIGS_PHASE1: dict[str, float] = {
+    "Nemotron-Pretraining-Code-Concepts": 0.20,
+    "Nemotron-Pretraining-Unconditional-Algorithmic": 0.20,
+    "Nemotron-Pretraining-Economics": 0.20,
+    "Nemotron-Pretraining-Formal-Logic": 0.20,
+    "Nemotron-Pretraining-Multiple-Choice": 0.20,
+}
+
+CONFIGS_PHASE2: dict[str, float] = {
+    "Nemotron-Pretraining-Multiple-Choice": 0.45,   # MMLU-style: high quality
+    "Nemotron-Pretraining-Economics": 0.20,
+    "Nemotron-Pretraining-Formal-Logic": 0.15,
+    "Nemotron-Pretraining-Code-Concepts": 0.10,
+    "Nemotron-Pretraining-Unconditional-Algorithmic": 0.10,
+}
+
+# Parquet files in this repo follow: {config}/part_{NNNNNN}.parquet
+# Some configs also have plain 0.parquet, 1.parquet naming — handled by list_repo_files.
+_TEXT_COLUMN_CANDIDATES = ["text", "content", "prompt_completion", "body", "input"]
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _load_tokenizer() -> "_p.Tokenizer":
+    """Load the tiktoken tokenizer produced by prepare.py."""
+    tokenizer_pkl = os.path.join(_p.TOKENIZER_DIR, "tokenizer.pkl")
+    if not os.path.exists(tokenizer_pkl):
+        raise RuntimeError(
+            f"Tokenizer not found at {tokenizer_pkl}. "
+            "Run `python prepare.py --num-shards 1` first to train the BPE tokenizer."
+        )
+    with open(tokenizer_pkl, "rb") as f:
+        enc = pickle.load(f)
+    return _p.Tokenizer(enc)
+
+
+def download_nemotron_files(config: str, n_parts: int, token: str) -> list[str]:
+    """List parquet files for *config*, download up to *n_parts*. Return local paths."""
+    api = HfApi(token=token)
+    repo_files = list(api.list_repo_files(NEMOTRON_REPO, repo_type="dataset"))
+    prefix = f"{config}/"
+    config_files = sorted(
+        f for f in repo_files
+        if f.startswith(prefix) and f.endswith(".parquet")
+    )
+    if not config_files:
+        print(f"  [warn] no parquet files found under {prefix} in {NEMOTRON_REPO}", flush=True)
+        return []
+    config_files = config_files[:n_parts]
+    local_paths: list[str] = []
+    for remote_path in config_files:
+        local = hf_hub_download(
+            repo_id=NEMOTRON_REPO,
+            filename=remote_path,
+            repo_type="dataset",
+            token=token,
+        )
+        local_paths.append(local)
+        print(f"  [download] {remote_path} -> {local}", flush=True)
+    return local_paths
+
+
+def _detect_text_column(schema: pa.Schema) -> str:
+    """Return the name of the text column from a parquet schema."""
+    col_names = schema.names
+    for candidate in _TEXT_COLUMN_CANDIDATES:
+        if candidate in col_names:
+            return candidate
+    # Fallback: first string column
+    for i, field in enumerate(schema):
+        if pa.types.is_string(field.type) or pa.types.is_large_string(field.type):
+            return field.name
+    # Last resort: first column
+    return col_names[0]
+
+
+def tokenize_and_write_shards(
+    parquet_paths: list[str],
+    tokenizer: "_p.Tokenizer",
+    out_dir: str,
+    shard_id_start: int,
+    tokens_per_shard: int,
+) -> int:
+    """
+    Stream-tokenize all text from *parquet_paths*, write fixed-size token shards.
+
+    Shard format (identical to prepare.py):
+      - single column 'tokens', dtype uint16
+      - each row group contains *tokens_per_shard* tokens
+
+    Returns the next available shard_id (= shard_id_start + shards_written).
+    """
+    shard_id = shard_id_start
+    tokens_buf: list[int] = []
+
+    for path in parquet_paths:
+        pf = pq.ParquetFile(path)
+        text_col = _detect_text_column(pf.schema_arrow)
+        print(f"  [tokenize] {os.path.basename(path)} column='{text_col}'", flush=True)
+        for rg_idx in range(pf.num_row_groups):
+            rg = pf.read_row_group(rg_idx, columns=[text_col])
+            texts: list[str] = rg.column(text_col).to_pylist()
+            # encode_ordinary_batch is faster (no special-token handling needed)
+            # tokenizer.encode() wraps enc.encode_ordinary for str input
+            token_lists: list[list[int]] = tokenizer.encode(texts)
+            for ids in token_lists:
+                tokens_buf.extend(ids)
+                # Flush complete shards
+                while len(tokens_buf) >= tokens_per_shard:
+                    chunk = tokens_buf[:tokens_per_shard]
+                    tokens_buf = tokens_buf[tokens_per_shard:]
+                    _write_shard(out_dir, shard_id, chunk)
+                    shard_id += 1
+
+    # Flush final partial shard (if any meaningful data remains)
+    if len(tokens_buf) >= 1024:
+        _write_shard(out_dir, shard_id, tokens_buf)
+        shard_id += 1
+
+    return shard_id
+
+
+def _write_shard(out_dir: str, shard_id: int, tokens: list[int]) -> None:
+    filename = f"shard_{shard_id:05d}.parquet"
+    out_path = os.path.join(out_dir, filename)
+    tmp_path = out_path + ".tmp"
+    arr = pa.array(tokens, type=pa.uint16())
+    tbl = pa.table({"tokens": arr})
+    pq.write_table(tbl, tmp_path)
+    os.rename(tmp_path, out_path)
+    print(f"  [shard] wrote {filename} ({len(tokens):,} tokens)", flush=True)
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Nemotron Super3 data prep — tokenize and shard to prepare.py-compatible format"
+    )
+    parser.add_argument(
+        "--phase",
+        choices=["phase1", "phase2"],
+        required=True,
+        help="phase1 = equal blend; phase2 = quality-weighted blend",
+    )
+    parser.add_argument(
+        "--parts-per-config",
+        type=int,
+        default=4,
+        help="Base number of parquet parts to download per config (scaled by weight)",
+    )
+    parser.add_argument(
+        "--tokens-per-shard",
+        type=int,
+        default=10_000_000,
+        help="Tokens per output shard (default 10M, matching climbmix convention)",
+    )
+    parser.add_argument(
+        "--shard-id-start",
+        type=int,
+        default=0,
+        help="First shard index to write (use non-zero to append after phase1 shards)",
+    )
+    parser.add_argument(
+        "--hf-token",
+        default=os.environ.get("HF_TOKEN"),
+        help="HuggingFace token (also read from $HF_TOKEN)",
+    )
+    args = parser.parse_args()
+
+    if not args.hf_token:
+        # Try ~/.hf_token as fallback (per project convention)
+        hf_token_path = os.path.expanduser("~/.hf_token")
+        if os.path.exists(hf_token_path):
+            with open(hf_token_path) as f:
+                args.hf_token = f.read().strip()
+
+    configs = CONFIGS_PHASE1 if args.phase == "phase1" else CONFIGS_PHASE2
+
+    tokenizer = _load_tokenizer()
+    os.makedirs(_p.DATA_DIR, exist_ok=True)
+
+    shard_id = args.shard_id_start
+    for config, weight in configs.items():
+        # Scale parts proportionally to weight so heavier configs get more data
+        n_parts = max(1, round(args.parts_per_config * weight * len(configs)))
+        print(
+            f"\n[nemotron] {config} weight={weight:.2f} n_parts={n_parts}",
+            flush=True,
+        )
+        parquet_paths = download_nemotron_files(config, n_parts, args.hf_token)
+        if not parquet_paths:
+            print(f"  [skip] no files downloaded for {config}", flush=True)
+            continue
+        shard_id = tokenize_and_write_shards(
+            parquet_paths,
+            tokenizer,
+            _p.DATA_DIR,
+            shard_id,
+            args.tokens_per_shard,
+        )
+
+    # Write validation shard — use Multiple-Choice (highest quality) as val source.
+    # Reserve the same VAL_SHARD index as prepare.py (6542) so training.py picks it up.
+    print("\n[nemotron] writing validation shard ...", flush=True)
+    val_paths = download_nemotron_files(
+        "Nemotron-Pretraining-Multiple-Choice", 1, args.hf_token
+    )
+    if val_paths:
+        tokenize_and_write_shards(
+            val_paths,
+            tokenizer,
+            _p.DATA_DIR,
+            _p.VAL_SHARD,          # 6542 — matches prepare.py VAL_SHARD constant
+            args.tokens_per_shard,
+        )
+    else:
+        print("  [warn] could not download val shard; evaluation may fail", flush=True)
+
+    print(
+        f"\n[nemotron] done — wrote shards {args.shard_id_start}..{shard_id - 1}"
+        f" + val shard {_p.VAL_SHARD}",
+        flush=True,
+    )
+
+
+if __name__ == "__main__":
+    main()

overlay/subsystems/sdr_retina.py

@@ -54,10 +54,14 @@ RETINA_PATH = os.path.join(CACHE_DIR, "retina.npz")
 GRID_H = 128
 GRID_W = 128
 N_BITS = GRID_H * GRID_W          # 16384
-TARGET_SPARSITY = 0.02            # 2%
-# int(floor(N_BITS * TARGET_SPARSITY)) = 327, matches the Cortical.io-style
-# "exactly 2% (327 / 16384) active bits per SDR" specification.
-TARGET_ACTIVE = int(N_BITS * TARGET_SPARSITY)  # 327
+TARGET_SPARSITY = 0.02            # 2% (default, Cortical.io-style)
+# Default = int(floor(N_BITS * TARGET_SPARSITY)) = 327, matches Webber/Numenta.
+# Override via HYDRA_SDR_TARGET_ACTIVE env var. The cache key encodes
+# target_active, so changing this triggers automatic retina regeneration.
+TARGET_ACTIVE = int(os.environ.get(
+    "HYDRA_SDR_TARGET_ACTIVE",
+    str(int(N_BITS * TARGET_SPARSITY)),
+))
 
 CONTEXT_WINDOW = 8                # +/- 8 tokens
 TOP_K_FEATURES = 64               # top-K context features per token

overlay/subsystems/sdr_semantic.py

@@ -23,7 +23,10 @@ import torch.nn as nn
 
 
 DEFAULT_RETINA_PATH = os.path.expanduser("~/.cache/autoresearch/retina.npz")
-DEFAULT_TARGET_ACTIVE = 327  # 2% of 16384
+# Default 327 = 2% of 16384 (Webber/Numenta canonical).
+# Override via HYDRA_SDR_TARGET_ACTIVE env var (must match the value used when
+# the retina cache was built — sdr_retina.py TARGET_ACTIVE reads the same var).
+DEFAULT_TARGET_ACTIVE = int(os.environ.get("HYDRA_SDR_TARGET_ACTIVE", "327"))
 
 
 class _SDRSTE(torch.autograd.Function):

runtime_setup.sh

@@ -0,0 +1,72 @@
+#!/usr/bin/env bash
+# Runtime setup for the stock pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel image.
+# We avoid baking feather + mamba_ssm + htm_rust into a custom Docker image
+# because build-time baking on HF's cpu-basic builder reliably corrupts CUDA
+# state on h200 runtime ("Error 802: system not yet initialized" every time,
+# even in a fresh python -c subprocess). Installing at runtime, on the h200
+# itself, avoids that path and keeps CUDA healthy.
+#
+# Trade-off: ~5-8 min cold start per job vs ~1 min for a baked image. The
+# training run is 12h long, so the overhead is negligible.
+
+set -euo pipefail
+
+echo "[runtime] $(date -u +%H:%M:%S) starting feather runtime setup on $(hostname)"
+
+# 1. Confirm CUDA before we do anything else.
+python -c 'import torch; assert torch.cuda.is_available(), "cuda unavailable at runtime start"; print("[runtime] cuda OK —", torch.cuda.get_device_name(0))'
+
+# 2. Install system build deps (rustup/build-essential for htm_rust).
+apt-get update -qq
+apt-get install -y -qq --no-install-recommends git curl ca-certificates build-essential pkg-config libssl-dev
+# Rust toolchain for htm_rust
+curl -sSf https://sh.rustup.rs | bash -s -- -y --profile minimal --default-toolchain stable
+export PATH=/root/.cargo/bin:$PATH
+
+# 3. Install Python deps.
+pip install --quiet --upgrade pip setuptools wheel
+pip install --quiet \
+    maturin \
+    huggingface_hub \
+    requests \
+    pyarrow \
+    rustbpe \
+    pandas \
+    tiktoken \
+    pydantic \
+    ninja \
+    packaging \
+    einops
+
+# 4. Install mamba_ssm + causal_conv1d (prebuilt wheels, matching torch2.6/cu12).
+pip install --quiet \
+    'https://github.com/Dao-AILab/causal-conv1d/releases/download/v1.6.1.post4/causal_conv1d-1.6.1+cu12torch2.6cxx11abiFALSE-cp311-cp311-linux_x86_64.whl' \
+    'https://github.com/state-spaces/mamba/releases/download/v2.3.1/mamba_ssm-2.3.1+cu12torch2.6cxx11abiFALSE-cp311-cp311-linux_x86_64.whl'
+
+# 5. Graft Mamba3 from main (pure Triton, not in v2.3.1 release).
+SITE=/opt/conda/lib/python3.11/site-packages/mamba_ssm
+BASE=https://raw.githubusercontent.com/state-spaces/mamba/main
+curl -fsSL "$BASE/mamba_ssm/modules/mamba3.py" -o "$SITE/modules/mamba3.py"
+mkdir -p "$SITE/ops/triton/mamba3"
+for f in __init__.py angle_dt.py mamba3_mimo_rotary_step.py mamba3_mimo_utils.py \
+         mamba3_siso_bwd.py mamba3_siso_combined.py mamba3_siso_fwd.py \
+         mamba3_siso_step.py utils.py; do
+    curl -fsSL "$BASE/mamba_ssm/ops/triton/mamba3/$f" -o "$SITE/ops/triton/mamba3/$f"
+done
+# Replace the eager-init __init__.py with our minimal version.
+cp /workspace/feather/hf_jobs/feather_h200_image/mamba_ssm_init.py "$SITE/__init__.py"
+
+# 6. Confirm CUDA still works after all installs.
+python -c 'import torch; assert torch.cuda.is_available(), "cuda broken by installs"; print("[runtime] cuda OK after deps —", torch.cuda.get_device_name(0))'
+
+# 7. Build + install htm_rust with sm_90 PTX (h200 arch).
+cd /workspace/feather
+export HTM_CUDA_ARCH=sm_90
+export LD_LIBRARY_PATH=/usr/local/cuda/lib64:${LD_LIBRARY_PATH:-}
+maturin build --release --features gpu --manifest-path htm_rust/Cargo.toml 2>&1 | tail -5
+pip install --quiet htm_rust/target/wheels/htm_rust-*.whl
+
+# 8. Sanity: cuda still alive after htm_rust install.
+python -c 'import torch; assert torch.cuda.is_available(), "cuda broken by htm_rust"; import htm_rust; print("[runtime] htm_rust OK, cuda OK")'
+
+echo "[runtime] $(date -u +%H:%M:%S) runtime setup complete"