chore: publish rSkill OpenRAL/rskill-robometer_4b-any-general-nf4 v0.1.0

_vendor/build_experiment.py

@@ -0,0 +1,340 @@
+"""Experiment: can we produce a PRE-quantized Robometer checkpoint that reloads
+directly as 4-bit (no bf16 base+ckpt read, no requantize)? (Robometer NF4 load fix)
+
+Measures the current load path's cost, then tries save_pretrained + reload.
+Run: /tmp/robometer-env/bin/python rskills/robometer-4b/_vendor/build_experiment.py
+"""
+
+from __future__ import annotations
+
+import os
+import pathlib
+import resource
+import time
+
+# Deterministic cuBLAS so the reward ramp is byte-stable across process launches:
+# without it, cuBLAS heuristic algo selection depends on process warmup history
+# (a warmed process and a cold one differ ~0.006). Must precede CUDA init.
+os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
+
+import torch
+
+torch.backends.cudnn.allow_tf32 = False
+torch.use_deterministic_algorithms(True, warn_only=True)
+# Force the math SDP kernel: flash/mem-efficient selection is process-state
+# dependent (a warmed vs cold process picks different kernels -> ~0.006 drift).
+torch.backends.cuda.enable_flash_sdp(False)
+torch.backends.cuda.enable_mem_efficient_sdp(False)
+torch.backends.cuda.enable_math_sdp(True)
+
+OUT = pathlib.Path("/tmp/robometer-nf4-ckpt")
+MIN_PARAMS = 4_000_000
+
+
+def _rss_gb() -> float:
+    return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1e6  # KB->GB on linux
+
+
+def _quantize_nf4_in_place(root: torch.nn.Module, compute_dtype: torch.dtype) -> int:
+    import bitsandbytes as bnb
+
+    n = 0
+
+    def _replace(m: torch.nn.Module) -> None:
+        nonlocal n
+        for name, child in list(m.named_children()):
+            if isinstance(child, torch.nn.Linear) and child.weight.numel() >= MIN_PARAMS:
+                new = bnb.nn.Linear4bit(
+                    child.in_features,
+                    child.out_features,
+                    bias=child.bias is not None,
+                    compute_dtype=compute_dtype,
+                    quant_type="nf4",
+                )
+                new.weight = bnb.nn.Params4bit(
+                    child.weight.data.clone(), requires_grad=False, quant_type="nf4"
+                )
+                if child.bias is not None:
+                    new.bias = torch.nn.Parameter(
+                        child.bias.data.clone().to(compute_dtype), requires_grad=False
+                    )
+                setattr(m, name, new)
+                n += 1
+            else:
+                _replace(child)
+
+    _replace(root)
+    return n
+
+
+def main() -> int:
+    from dataclasses import fields
+
+    import yaml
+    from huggingface_hub import hf_hub_download
+    from robometer.configs.experiment_configs import ExperimentConfig
+    from robometer.utils.save import resolve_checkpoint_path
+    from robometer.utils.setup_utils import setup_model_and_processor
+
+    t0 = time.monotonic()
+    print("[exp] load bf16 on CPU via VANILLA path (use_unsloth=False) ...", flush=True)
+    # Replicate load_model_from_hf's config assembly but force use_unsloth=False so
+    # the model is built with vanilla Qwen3VLModel naming (matches the meta reload).
+    resolved = resolve_checkpoint_path("robometer/Robometer-4B")
+    cfg_yaml = hf_hub_download("robometer/Robometer-4B", "config.yaml")
+    raw = yaml.safe_load(open(cfg_yaml))
+    valid = {f.name for f in fields(ExperimentConfig)}
+    exp_config = ExperimentConfig(**{k: v for k, v in raw.items() if k in valid})
+    exp_config.model.use_unsloth = False
+    tokenizer, processor, model = setup_model_and_processor(
+        exp_config.model, str(resolved), peft_config=None
+    )
+    model = model.to("cpu").eval()
+    t_load = time.monotonic() - t0
+    print(f"[exp] LOAD took {t_load:.1f}s; peak RSS {_rss_gb():.1f} GB", flush=True)
+
+    t1 = time.monotonic()
+    n = _quantize_nf4_in_place(model, compute_dtype=torch.bfloat16)
+    model.to("cuda")
+    torch.cuda.synchronize()
+    t_quant = time.monotonic() - t1
+    print(
+        f"[exp] QUANTIZE+to(cuda) took {t_quant:.1f}s; {n} modules; "
+        f"{torch.cuda.memory_allocated() / 1e9:.2f} GB VRAM",
+        flush=True,
+    )
+
+    # Save the packed state_dict explicitly (keys match a vanilla meta reload).
+    from safetensors.torch import save_file
+
+    OUT.mkdir(parents=True, exist_ok=True)
+    sd = {
+        k: (v.detach().contiguous() if hasattr(v, "detach") else v)
+        for k, v in model.state_dict().items()
+    }
+    # Fold in the NON-persistent buffers (rotary inv_freq etc.) that state_dict()
+    # omits, so the meta reload loads them bit-identically instead of recomputing
+    # them (recompute drifts ~0.004 on the progress series). assign=True on reload
+    # restores them exactly.
+    n_extra = 0
+    for name, buf in model.named_buffers():
+        if name not in sd and buf is not None:
+            sd[name] = buf.detach().contiguous()
+            n_extra += 1
+    print(f"[exp] folded {n_extra} non-persistent buffers into the checkpoint")
+    save_file(sd, str(OUT / "model.safetensors"))
+    # Save the SELF-CONTAINED processor/tokenizer/config (resized vocab + added
+    # progress token) so the meta reload never touches the base model.
+    processor.save_pretrained(str(OUT))
+    tokenizer.save_pretrained(str(OUT))
+    model.config.save_pretrained(str(OUT))
+    size_gb = (OUT / "model.safetensors").stat().st_size / 1e9
+    vis_keys = [k for k in sd if "visual.blocks.0.mlp.linear_fc1" in k]
+    print(f"[exp] SAVE OK: {size_gb:.2f} GB, {len(sd)} tensors")
+    print(f"[exp] vision mlp keys present (sample): {vis_keys}")
+
+    # Reference forward on the SAME 10 frames the reload test uses → series_A.
+    series = _forward_series(model, tokenizer, processor, exp_config)
+    print(
+        f"[exp] REFERENCE progress series (bf16+quantize, LIVE processor): "
+        f"{[round(float(x), 4) for x in series]}",
+        flush=True,
+    )
+
+    # Same weights, but recompute with the processor/tokenizer RELOADED from the
+    # checkpoint dir (what the meta reload uses). Isolates any processor
+    # save_pretrained round-trip drift from any weight/construction drift.
+    from transformers import AutoProcessor, AutoTokenizer
+
+    ck_proc = AutoProcessor.from_pretrained(str(OUT))
+    ck_tok = AutoTokenizer.from_pretrained(str(OUT))
+    series_ck = _forward_series(model, ck_tok, ck_proc, exp_config)
+    print(
+        f"[exp] REFERENCE progress series (bf16+quantize, CKPT processor): "
+        f"{[round(float(x), 4) for x in series_ck]}",
+        flush=True,
+    )
+
+    # --- In-process round-trip diagnostic: does from_prequantized reconstruct the
+    # SAME 4-bit weights as the in-place quantization the build forward used? ---
+    import bitsandbytes as bnb
+    from safetensors.torch import load_file as _load_file
+
+    probe = "model.language_model.layers.0.mlp.gate_proj"
+    live_mod = dict(model.named_modules())[probe]
+    saved = _load_file(str(OUT / "model.safetensors"), device="cuda")
+    stats = {
+        s.lstrip("."): saved[f"{probe}.weight{s}"]
+        for s in (
+            ".absmax",
+            ".quant_map",
+            ".nested_absmax",
+            ".nested_quant_map",
+            ".quant_state.bitsandbytes__nf4",
+        )
+        if f"{probe}.weight{s}" in saved
+    }
+    reloaded_w = bnb.nn.Params4bit.from_prequantized(
+        data=saved[f"{probe}.weight"], quantized_stats=stats, requires_grad=False, device="cuda"
+    )
+    dq_live = bnb.functional.dequantize_4bit(
+        live_mod.weight.data, live_mod.weight.quant_state
+    ).float()
+    dq_reload = bnb.functional.dequantize_4bit(reloaded_w.data, reloaded_w.quant_state).float()
+    max_abs = (dq_live - dq_reload).abs().max().item()
+    print(
+        f"[exp] DEQUANT max|live - from_prequantized| for {probe}: {max_abs:.3e} "
+        f"(0.0 => bit-identical 4-bit round-trip)",
+        flush=True,
+    )
+
+    # Report the global numerics flags this (build) process is running under, so we
+    # can compare them against the meta-reload process.
+    print(
+        f"[exp] tf32 matmul={torch.backends.cuda.matmul.allow_tf32} "
+        f"cudnn.tf32={torch.backends.cudnn.allow_tf32} "
+        f"fp32_precision(matmul)={torch.get_float32_matmul_precision()}",
+        flush=True,
+    )
+
+    # --- DEFINITIVE same-process A/B: construct the model the meta-reload way from
+    # the file we just saved, run the SAME forward, compare to the build forward.
+    # Identical process => identical env => isolates construction-path drift. ---
+    del live_mod, dq_live, dq_reload, reloaded_w, saved
+    torch.cuda.empty_cache()
+    series_meta = _meta_reload_series(exp_config, base_id="Qwen/Qwen3-VL-4B-Instruct")
+    print(
+        f"[exp] META-RELOAD progress series (same process): "
+        f"{[round(float(x), 4) for x in series_meta]}",
+        flush=True,
+    )
+    import numpy as np
+
+    dmax = float(np.abs(np.asarray(series) - np.asarray(series_meta)).max())
+    print(f"[exp] max|build - meta_reload| (same process) = {dmax:.3e}", flush=True)
+    return 0
+
+
+def _meta_reload_series(exp_config, base_id):
+    """Build via meta + install_prequantized from OUT, return the progress series."""
+    import bitsandbytes as bnb
+    from robometer.models.rbm import RBM
+    from safetensors.torch import load_file
+    from transformers import AutoConfig, AutoProcessor, AutoTokenizer
+
+    config = AutoConfig.from_pretrained(str(OUT))
+    processor = AutoProcessor.from_pretrained(str(OUT))
+    tokenizer = AutoTokenizer.from_pretrained(str(OUT))
+    for c in (config, getattr(config, "text_config", None), getattr(config, "vision_config", None)):
+        if c is not None:
+            c._attn_implementation = "sdpa"
+    with torch.device("meta"):
+        model = RBM(
+            config,
+            processor,
+            tokenizer,
+            base_model=None,
+            base_model_id=base_id,
+            model_config=exp_config.model,
+        )
+
+    # Linear4bit shells for numel>=MIN_PARAMS (same rule as the build quantizer).
+    def _shells(m):
+        for name, child in list(m.named_children()):
+            if isinstance(child, torch.nn.Linear) and child.weight.numel() >= MIN_PARAMS:
+                setattr(
+                    m,
+                    name,
+                    bnb.nn.Linear4bit(
+                        child.in_features,
+                        child.out_features,
+                        bias=child.bias is not None,
+                        compute_dtype=torch.bfloat16,
+                        quant_type="nf4",
+                    ),
+                )
+            else:
+                _shells(child)
+
+    _shells(model)
+
+    state = load_file(str(OUT / "model.safetensors"), device="cuda")
+    sufs = (
+        ".absmax",
+        ".quant_map",
+        ".nested_absmax",
+        ".nested_quant_map",
+        ".quant_state.bitsandbytes__nf4",
+        ".quant_state.bitsandbytes__fp4",
+    )
+    consumed: set[str] = set()
+    for prefix, module in model.named_modules():
+        if not isinstance(module, bnb.nn.Linear4bit):
+            continue
+        wkey = f"{prefix}.weight"
+        if wkey not in state:
+            continue
+        stats = {s.lstrip("."): state[f"{wkey}{s}"] for s in sufs if f"{wkey}{s}" in state}
+        for s in sufs:
+            consumed.add(f"{wkey}{s}")
+        consumed.add(wkey)
+        module.weight = bnb.nn.Params4bit.from_prequantized(
+            data=state[wkey], quantized_stats=stats, requires_grad=False, device="cuda"
+        )
+        bkey = f"{prefix}.bias"
+        if module.bias is not None and bkey in state:
+            module.bias = torch.nn.Parameter(state[bkey].to("cuda"), requires_grad=False)
+            consumed.add(bkey)
+    leftover = {k: v for k, v in state.items() if k not in consumed}
+    model.load_state_dict(leftover, strict=False, assign=True)
+    for bname, buf in list(model.named_buffers()):
+        if not buf.is_meta:
+            continue
+        parent = model.get_submodule(bname.rsplit(".", 1)[0]) if "." in bname else model
+        if bname in state:
+            parent.register_buffer(
+                bname.rsplit(".", 1)[-1], state[bname].to("cuda"), persistent=False
+            )
+    for _nm, mod in model.named_modules():
+        ifb = getattr(mod, "inv_freq", None)
+        if ifb is not None and hasattr(mod, "original_inv_freq"):
+            mod.original_inv_freq = ifb
+    return _forward_series(model, tokenizer, processor, exp_config)
+
+
+def _forward_series(model, tokenizer, processor, exp_config):
+    """Run the discrete-mode progress forward on a fixed 10-frame clip."""
+    import decord
+    import numpy as np
+    from robometer.data.dataset_types import ProgressSample, Trajectory
+    from robometer.evals.eval_server import compute_batch_outputs
+    from robometer.utils.setup_utils import setup_batch_collator
+
+    model.eval()
+    vr = decord.VideoReader("/tmp/robometer_example.mp4")
+    step = max(1, int(round(vr.get_avg_fps() / 3.0)))
+    idx = list(range(0, len(vr), step))[:10]
+    frames = vr.get_batch(idx).asnumpy().astype(np.uint8)
+    collator = setup_batch_collator(processor, tokenizer, exp_config, is_eval=True)
+    traj = Trajectory(
+        frames=frames,
+        frames_shape=tuple(frames.shape),
+        task="Pick up the object and place it in the container",
+        id="0",
+        metadata={"subsequence_length": int(frames.shape[0])},
+        video_embeddings=None,
+    )
+    batch = collator([ProgressSample(trajectory=traj, sample_type="progress")])
+    inp = batch["progress_inputs"]
+    for k, v in inp.items():
+        if hasattr(v, "to"):
+            inp[k] = v.to("cuda")
+    with torch.no_grad():
+        res = compute_batch_outputs(
+            model, tokenizer, inp, sample_type="progress", is_discrete_mode=True, num_bins=100
+        )
+    return np.asarray(res["progress_pred"][0], dtype=np.float32)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

_vendor/keydiff.py

@@ -0,0 +1,60 @@
+"""Diagnose meta-skeleton ↔ checkpoint key divergence (Robometer NF4 load fix)."""
+
+from __future__ import annotations
+
+import pathlib
+from dataclasses import fields
+
+import torch
+import yaml
+
+CKPT = pathlib.Path("/tmp/robometer-nf4-ckpt/model.safetensors")
+
+
+def main() -> int:
+    from huggingface_hub import hf_hub_download
+    from robometer.configs.experiment_configs import ExperimentConfig
+    from robometer.models.rbm import RBM
+    from safetensors import safe_open
+    from transformers import AutoConfig
+
+    base_id = "Qwen/Qwen3-VL-4B-Instruct"
+    raw = yaml.safe_load(open(hf_hub_download("robometer/Robometer-4B", "config.yaml")))
+    valid = {f.name for f in fields(ExperimentConfig)}
+    exp_config = ExperimentConfig(**{k: v for k, v in raw.items() if k in valid})
+
+    for cfg_src in ("robometer/Robometer-4B", base_id):
+        config = AutoConfig.from_pretrained(cfg_src)
+        with torch.device("meta"):
+            model = RBM(
+                config,
+                None,
+                None,
+                base_model=None,
+                base_model_id=base_id,
+                model_config=exp_config.model,
+            )
+        mkeys = set(model.state_dict().keys())
+        vis_fc1 = [k for k in mkeys if "visual.blocks.0.mlp" in k]
+        vt = type(dict(model.named_modules()).get("model.visual.blocks.0.mlp.linear_fc1"))
+        lin = dict(model.named_modules()).get("model.visual.blocks.0.mlp.linear_fc1")
+        numel = lin.weight.numel() if lin is not None else None
+        print(f"\n=== config from {cfg_src} ===")
+        print(f"  model.visual.blocks.0.mlp keys: {vis_fc1}")
+        print(f"  linear_fc1 type={vt} numel={numel}")
+
+    with safe_open(str(CKPT), framework="pt") as f:
+        skeys = set(f.keys())
+    print("\n=== checkpoint file ===")
+    print(f"  total keys: {len(skeys)}")
+    print(f"  visual.blocks.0.mlp keys: {sorted(k for k in skeys if 'visual.blocks.0.mlp' in k)}")
+    # model keys (last config = base) vs file
+    only_model = sorted(mkeys - skeys)[:8]
+    only_file = sorted(skeys - mkeys)[:8]
+    print(f"\n  in model not file (sample): {only_model}")
+    print(f"  in file not model (sample): {only_file}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

_vendor/probe.py

@@ -0,0 +1,116 @@
+"""Phase 0 load+forward probe for robometer/Robometer-4B (reward rSkill gating spike).
+
+Mirrors robometer's scripts/example_inference_local.py but builds a tiny synthetic
+clip programmatically (no example_videos needed, since we pip-installed the package
+rather than cloning). Goal: confirm the model loads with inference-only deps and
+print the exact output field names / shapes / value ranges for progress + success.
+
+Run inside the isolated env:
+    /tmp/robometer-env/bin/python rskills/robometer-4b/_vendor/probe.py
+Optional: --weights <dir> to probe a quantized checkpoint, --device cuda.
+"""
+
+from __future__ import annotations
+
+import argparse
+
+import numpy as np
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model-path", default="robometer/Robometer-4B")
+    ap.add_argument("--device", default="cpu")
+    ap.add_argument("--frames", type=int, default=8)
+    ap.add_argument("--hw", type=int, default=224)
+    ap.add_argument("--task", default="pick up the cube and place it in the bowl")
+    ap.add_argument("--discrete", action="store_true")
+    ap.add_argument("--num-bins", type=int, default=100)
+    args = ap.parse_args()
+
+    from robometer.data.dataset_types import ProgressSample, Trajectory
+    from robometer.evals.eval_server import compute_batch_outputs
+    from robometer.utils.save import load_model_from_hf
+    from robometer.utils.setup_utils import setup_batch_collator
+
+    print(f"[probe] loading {args.model_path} on {args.device} ...", flush=True)
+    exp_config, tokenizer, processor, reward_model = load_model_from_hf(
+        model_path=args.model_path,
+        device=args.device,
+    )
+    print(f"[probe] loaded model class = {type(reward_model).__name__}", flush=True)
+    print(f"[probe] exp_config type = {type(exp_config).__name__}", flush=True)
+
+    batch_collator = setup_batch_collator(processor, tokenizer, exp_config, is_eval=True)
+
+    T, H, W, C = args.frames, args.hw, args.hw, 3
+    video_frames = np.random.randint(0, 255, (T, H, W, C), dtype=np.uint8)
+    traj = Trajectory(
+        frames=video_frames,
+        frames_shape=tuple(video_frames.shape),
+        task=args.task,
+        id="0",
+        metadata={"subsequence_length": T},
+        video_embeddings=None,
+    )
+    progress_sample = ProgressSample(trajectory=traj, sample_type="progress")
+    batch = batch_collator([progress_sample])
+    # The collator returns a wrapper dict; the model inputs live under "progress_inputs".
+    progress_inputs = batch["progress_inputs"]
+    for key, value in progress_inputs.items():
+        if hasattr(value, "to"):
+            progress_inputs[key] = value.to(args.device)
+    reward_model.eval()
+
+    print(f"[probe] progress_inputs keys = {sorted(progress_inputs.keys())}", flush=True)
+    print("[probe] running compute_batch_outputs ...", flush=True)
+    results = compute_batch_outputs(
+        reward_model,
+        tokenizer,
+        progress_inputs,
+        sample_type="progress",
+        is_discrete_mode=args.discrete,
+        num_bins=args.num_bins,
+    )
+
+    print("\n=== RESULT KEYS ===")
+    print(sorted(results.keys()))
+
+    # success lives nested under outputs_success
+    outputs_success = results.get("outputs_success")
+    if isinstance(outputs_success, dict):
+        print(f"[probe] outputs_success keys = {sorted(outputs_success.keys())}")
+        for k, v in outputs_success.items():
+            try:
+                arr = np.asarray(v[0] if (isinstance(v, list) and v) else v, dtype=np.float32)
+                print(
+                    f"  outputs_success[{k}]: shape={arr.shape} "
+                    f"min={arr.min():.4f} max={arr.max():.4f}"
+                    if arr.size
+                    else f"  outputs_success[{k}]: empty"
+                )
+            except Exception as exc:
+                print(f"  outputs_success[{k}]: <{type(v).__name__}> ({exc})")
+
+    def describe(name: str) -> None:
+        val = results.get(name)
+        if val is None:
+            print(f"  {name}: <absent>")
+            return
+        arr = np.asarray(val[0] if (isinstance(val, list) and val) else val, dtype=np.float32)
+        if arr.size:
+            print(
+                f"  {name}: shape={arr.shape} min={arr.min():.4f} max={arr.max():.4f} "
+                f"first5={np.round(arr.flatten()[:5], 4).tolist()}"
+            )
+        else:
+            print(f"  {name}: empty")
+
+    print("\n=== PROGRESS / SUCCESS ===")
+    for key in ("progress_pred", "success_probs", "success_pred", "preference", "pref_logits"):
+        describe(key)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

_vendor/quant_probe.py

@@ -0,0 +1,140 @@
+"""Phase 2 NF4 quantization probe for Robometer-4B.
+
+Replicates the repo's NF4 rewrite rule from openral_sim._quantization
+(quantize_nf4_in_place: nn.Linear with weight.numel() >= 4M -> bnb.nn.Linear4bit,
+quant_type="nf4", compute_dtype=bf16; pack happens on .to(cuda)). The rule is
+inlined here because the isolated robometer venv cannot import openral_sim.
+
+Loads RBM bf16 on CPU, quantizes, moves to CUDA, runs one forward, and reports
+peak VRAM — the empirical answer to "how easy to quantize" + "does it leave 8 GB
+headroom for a parallel VLA?".
+
+    /tmp/robometer-env/bin/python rskills/robometer-4b/_vendor/quant_probe.py
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+
+MIN_PARAMS = 4_000_000  # DEFAULT_MIN_PARAMS_TO_QUANTIZE (openral_sim._quantization)
+
+
+def quantize_nf4_in_place(root: torch.nn.Module, compute_dtype: torch.dtype) -> int:
+    import bitsandbytes as bnb
+
+    n = 0
+
+    def _replace(module: torch.nn.Module, prefix: str = "") -> None:
+        nonlocal n
+        for name, child in list(module.named_children()):
+            if isinstance(child, torch.nn.Linear) and child.weight.numel() >= MIN_PARAMS:
+                new = bnb.nn.Linear4bit(
+                    child.in_features,
+                    child.out_features,
+                    bias=child.bias is not None,
+                    compute_dtype=compute_dtype,
+                    quant_type="nf4",
+                )
+                new.weight = bnb.nn.Params4bit(
+                    child.weight.data.clone(),
+                    requires_grad=False,
+                    quant_type="nf4",
+                )
+                if child.bias is not None:
+                    new.bias = torch.nn.Parameter(
+                        child.bias.data.clone().to(compute_dtype),
+                        requires_grad=False,
+                    )
+                setattr(module, name, new)
+                n += 1
+            else:
+                _replace(child, f"{prefix}.{name}" if prefix else name)
+
+    _replace(root)
+    return n
+
+
+def main() -> int:
+    from robometer.data.dataset_types import ProgressSample, Trajectory
+    from robometer.evals.eval_server import compute_batch_outputs
+    from robometer.utils.save import load_model_from_hf
+    from robometer.utils.setup_utils import setup_batch_collator
+
+    assert torch.cuda.is_available(), "need CUDA for the VRAM measurement"
+
+    print("[quant] loading bf16 on CPU ...", flush=True)
+    exp_config, tokenizer, processor, reward_model = load_model_from_hf(
+        model_path="robometer/Robometer-4B",
+        device="cpu",
+    )
+    reward_model.eval()
+
+    print("[quant] rewriting large Linears -> NF4 ...", flush=True)
+    n = quantize_nf4_in_place(reward_model, compute_dtype=torch.bfloat16)
+    print(f"[quant] rewrote {n} Linear modules to NF4", flush=True)
+
+    torch.cuda.reset_peak_memory_stats()
+    print("[quant] moving to CUDA (packs nf4) ...", flush=True)
+    reward_model.to("cuda")
+    torch.cuda.synchronize()
+    resident = torch.cuda.memory_allocated() / 1e9
+    print(f"[quant] NF4 weights resident on CUDA: {resident:.2f} GB", flush=True)
+
+    # one forward to confirm correctness post-quant
+    batch_collator = setup_batch_collator(processor, tokenizer, exp_config, is_eval=True)
+    T = 8
+    frames = np.random.randint(0, 255, (T, 224, 224, 3), dtype=np.uint8)
+    traj = Trajectory(
+        frames=frames,
+        frames_shape=tuple(frames.shape),
+        task="pick up the cube",
+        id="0",
+        metadata={"subsequence_length": T},
+        video_embeddings=None,
+    )
+    batch = batch_collator([ProgressSample(trajectory=traj, sample_type="progress")])
+    progress_inputs = batch["progress_inputs"]
+    for k, v in progress_inputs.items():
+        if hasattr(v, "to"):
+            progress_inputs[k] = v.to("cuda")
+
+    print("[quant] running forward (discrete mode) ...", flush=True)
+    with torch.no_grad():
+        results = compute_batch_outputs(
+            reward_model,
+            tokenizer,
+            progress_inputs,
+            sample_type="progress",
+            is_discrete_mode=True,
+            num_bins=100,
+        )
+    torch.cuda.synchronize()
+    peak = torch.cuda.max_memory_allocated() / 1e9
+
+    prog = np.asarray(
+        results["progress_pred"][0]
+        if isinstance(results["progress_pred"], list)
+        else results["progress_pred"],
+        dtype=np.float32,
+    )
+    succ = results.get("outputs_success", {}).get("success_probs")
+    succ = (
+        np.asarray(succ[0] if isinstance(succ, list) and succ else succ, dtype=np.float32)
+        if succ is not None
+        else np.array([])
+    )
+
+    print("\n=== NF4 RESULT ===")
+    print(f"  modules quantized: {n}")
+    print(f"  NF4 resident VRAM: {resident:.2f} GB")
+    print(f"  peak VRAM (incl. 8-frame forward activations): {peak:.2f} GB")
+    print(f"  progress_pred: shape={prog.shape} range=[{prog.min():.4f},{prog.max():.4f}]")
+    if succ.size:
+        print(f"  success_probs: shape={succ.shape} range=[{succ.min():.4f},{succ.max():.4f}]")
+    print(f"  GPU total 8.0 GB -> headroom after peak: {8.0 - peak:.2f} GB")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

_vendor/reload_experiment.py

@@ -0,0 +1,275 @@
+"""Experiment: reload the pre-quantized 3.32 GB checkpoint WITHOUT the bf16 base
+load — build the RBM skeleton from config (random init), install Linear4bit, then
+load_state_dict the packed NF4 weights. Verify forward reproduces the ramp.
+
+Prereq: run build_experiment.py first (writes /tmp/robometer-nf4-ckpt/model.safetensors).
+Run: /tmp/robometer-env/bin/python rskills/robometer-4b/_vendor/reload_experiment.py
+"""
+
+from __future__ import annotations
+
+import os
+import pathlib
+import resource
+import time
+from dataclasses import fields
+
+import numpy as np
+import yaml
+
+# Match build_experiment: deterministic cuBLAS for byte-stable cross-process output.
+os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
+
+import torch
+
+torch.backends.cudnn.allow_tf32 = False
+torch.use_deterministic_algorithms(True, warn_only=True)
+torch.backends.cuda.enable_flash_sdp(False)
+torch.backends.cuda.enable_mem_efficient_sdp(False)
+torch.backends.cuda.enable_math_sdp(True)
+
+CKPT = pathlib.Path("/tmp/robometer-nf4-ckpt")
+MIN_PARAMS = 4_000_000
+
+
+def _rss_gb() -> float:
+    return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1e6
+
+
+_BNB_META_SUFFIXES = (
+    ".absmax",
+    ".quant_map",
+    ".nested_absmax",
+    ".nested_quant_map",
+    ".quant_state.bitsandbytes__nf4",
+    ".quant_state.bitsandbytes__fp4",
+)
+
+
+def _install_prequantized(policy, state, device):
+    """Inlined openral_sim._quantization.install_prequantized_linears."""
+    import bitsandbytes as bnb
+
+    consumed: set[str] = set()
+    count = 0
+    for prefix, module in policy.named_modules():
+        if not isinstance(module, bnb.nn.Linear4bit):
+            continue
+        wkey = f"{prefix}.weight"
+        if wkey not in state:
+            continue
+        stats = {}
+        for suf in _BNB_META_SUFFIXES:
+            full = f"{wkey}{suf}"
+            if full in state:
+                stats[suf.lstrip(".")] = state[full]
+                consumed.add(full)
+        consumed.add(wkey)
+        module.weight = bnb.nn.Params4bit.from_prequantized(
+            data=state[wkey], quantized_stats=stats, requires_grad=False, device=device
+        )
+        bkey = f"{prefix}.bias"
+        if module.bias is not None and bkey in state:
+            module.bias = torch.nn.Parameter(state[bkey].to(device), requires_grad=False)
+            consumed.add(bkey)
+        count += 1
+    return count, consumed
+
+
+def _quantize_structure(root, compute_dtype):
+    """Replace large Linears with Linear4bit (empty packed params) — structure only."""
+    import bitsandbytes as bnb
+
+    n = 0
+
+    def _replace(m):
+        nonlocal n
+        for name, child in list(m.named_children()):
+            if isinstance(child, torch.nn.Linear) and child.weight.numel() >= MIN_PARAMS:
+                new = bnb.nn.Linear4bit(
+                    child.in_features,
+                    child.out_features,
+                    bias=child.bias is not None,
+                    compute_dtype=compute_dtype,
+                    quant_type="nf4",
+                )
+                setattr(m, name, new)
+                n += 1
+            else:
+                _replace(child)
+
+    _replace(root)
+    return n
+
+
+def main() -> int:
+    from huggingface_hub import hf_hub_download
+    from robometer.configs.experiment_configs import ExperimentConfig
+    from robometer.models.rbm import RBM
+    from safetensors.torch import load_file
+    from transformers import AutoConfig, AutoProcessor, AutoTokenizer
+
+    base_id = "Qwen/Qwen3-VL-4B-Instruct"
+    # cheap pieces — no big model weights
+    cfg_yaml = hf_hub_download("robometer/Robometer-4B", "config.yaml")
+    raw = yaml.safe_load(open(cfg_yaml))
+    valid = {f.name for f in fields(ExperimentConfig)}
+    exp_config = ExperimentConfig(**{k: v for k, v in raw.items() if k in valid})
+
+    # Load config + processor + tokenizer from the SELF-CONTAINED checkpoint dir
+    # (resized vocab 151674 + robometer's added progress token), NOT the base.
+    config = AutoConfig.from_pretrained(str(CKPT))
+    processor = AutoProcessor.from_pretrained(str(CKPT))
+    tokenizer = AutoTokenizer.from_pretrained(str(CKPT))
+    # Direct construction (not from_pretrained) skips transformers' attn auto-select
+    # and defaults to "eager"; production's setup_model_and_processor uses "sdpa"
+    # (flash-attn absent). Force sdpa on every (sub)config so the meta path is
+    # numerically identical to the bf16+quantize reference.
+    for c in (config, getattr(config, "text_config", None), getattr(config, "vision_config", None)):
+        if c is not None:
+            c._attn_implementation = "sdpa"
+
+    print(
+        f"[reload] tf32 matmul={torch.backends.cuda.matmul.allow_tf32} "
+        f"cudnn.tf32={torch.backends.cudnn.allow_tf32} "
+        f"fp32_precision(matmul)={torch.get_float32_matmul_precision()}",
+        flush=True,
+    )
+
+    t0 = time.monotonic()
+    print("[reload] building RBM skeleton on META (instant, no weights) ...", flush=True)
+    with torch.device("meta"):
+        model = RBM(
+            config,
+            processor,
+            tokenizer,
+            base_model=None,
+            base_model_id=base_id,
+            model_config=exp_config.model,
+        )
+    n = _quantize_structure(model, compute_dtype=torch.bfloat16)
+    print(
+        f"[reload] meta skeleton + {n} Linear4bit shells in "
+        f"{time.monotonic() - t0:.1f}s; peak RSS {_rss_gb():.1f} GB",
+        flush=True,
+    )
+
+    t1 = time.monotonic()
+    state = load_file(str(CKPT / "model.safetensors"), device="cuda")
+    # 4-bit modules: rebuild packed weights directly on CUDA (no bf16 alloc).
+    n_q, consumed = _install_prequantized(model, state, device="cuda")
+    # everything else (embeddings, norms, heads): assign cuda tensors to meta params.
+    leftover = {k: v for k, v in state.items() if k not in consumed}
+    missing, unexpected = model.load_state_dict(leftover, strict=False, assign=True)
+    # load_state_dict SKIPS non-persistent buffers (rotary inv_freq) — it reports
+    # them as `unexpected` and never assigns them. So assign them by hand from the
+    # checkpoint, by dotted name, bit-identically (no recompute). Only truly-absent
+    # buffers fall back to recompute.
+    loaded_bufs, recomputed_meta = 0, 0
+    for bname, buf in list(model.named_buffers()):
+        if not buf.is_meta:
+            continue
+        parent = model.get_submodule(bname.rsplit(".", 1)[0]) if "." in bname else model
+        leaf = bname.rsplit(".", 1)[-1]
+        if bname in state:  # persisted in the checkpoint -> exact restore
+            parent.register_buffer(leaf, state[bname].to("cuda"), persistent=False)
+            loaded_bufs += 1
+        elif hasattr(parent, "rope_init_fn") and hasattr(parent, "config"):
+            inv_freq, scaling = parent.rope_init_fn(parent.config, "cuda")
+            parent.register_buffer(leaf, inv_freq, persistent=False)
+            if hasattr(parent, "attention_scaling"):
+                parent.attention_scaling = scaling
+            recomputed_meta += 1
+        else:  # Qwen3VLVisionRotaryEmbedding closed form
+            dim = 2 * buf.shape[0]
+            inv_freq = 1.0 / (
+                10000.0 ** (torch.arange(0, dim, 2, dtype=torch.float, device="cuda") / dim)
+            )
+            parent.register_buffer(leaf, inv_freq, persistent=False)
+            recomputed_meta += 1
+    # Rebind the dangling non-buffer rope attributes off the (now-real) inv_freq.
+    for _nm, mod in model.named_modules():
+        ifb = getattr(mod, "inv_freq", None)
+        if ifb is None:
+            continue
+        if hasattr(mod, "original_inv_freq"):
+            mod.original_inv_freq = ifb
+        if (
+            hasattr(mod, "rope_init_fn")
+            and hasattr(mod, "config")
+            and getattr(mod, "attention_scaling", None) is None
+        ):
+            _, scaling = mod.rope_init_fn(mod.config, "cuda")
+            mod.attention_scaling = scaling
+    print(
+        f"[reload] rotary buffers loaded-from-ckpt={loaded_bufs} "
+        f"recomputed-from-meta(expect 0)={recomputed_meta}"
+    )
+
+    # any params/buffers still on meta (not in the checkpoint)?
+    still_meta = [n for n, p in model.named_parameters() if p.is_meta]
+    still_meta_buf = [n for n, b in model.named_buffers() if b.is_meta]
+    print(f"[reload] meta params left: {still_meta[:6]}")
+    print(f"[reload] meta buffers left after rotary fix: {still_meta_buf}")
+    # sanity: is a vision weight real (install handled it)?
+    vw = dict(model.named_parameters()).get("model.visual.blocks.0.mlp.linear_fc1.weight")
+    print(
+        f"[reload] vision fc1 weight is_meta={vw.is_meta if vw is not None else 'absent'} "
+        f"dtype={vw.dtype if vw is not None else '-'}"
+    )
+    torch.cuda.synchronize()
+    print(
+        f"[reload] install_prequantized({n_q}) + load_state_dict in {time.monotonic() - t1:.1f}s; "
+        f"{torch.cuda.memory_allocated() / 1e9:.2f} GB VRAM; "
+        f"missing={len(missing)} unexpected={len(unexpected)} "
+        f"meta_params_left={len(still_meta)} meta_bufs_left={len(still_meta_buf)}",
+        flush=True,
+    )
+    if missing:
+        print(f"[reload] sample missing keys: {missing[:5]}")
+    if unexpected:
+        print(f"[reload] sample unexpected keys: {unexpected[:5]}")
+
+    # forward on the real video → expect the ramp
+    model.eval()
+    import decord
+    from robometer.data.dataset_types import ProgressSample, Trajectory
+    from robometer.evals.eval_server import compute_batch_outputs
+    from robometer.utils.setup_utils import setup_batch_collator
+
+    vr = decord.VideoReader("/tmp/robometer_example.mp4")
+    step = max(1, int(round(vr.get_avg_fps() / 3.0)))
+    idx = list(range(0, len(vr), step))[:10]
+    frames = vr.get_batch(idx).asnumpy().astype(np.uint8)
+    collator = setup_batch_collator(processor, tokenizer, exp_config, is_eval=True)
+    traj = Trajectory(
+        frames=frames,
+        frames_shape=tuple(frames.shape),
+        task="Pick up the object and place it in the container",
+        id="0",
+        metadata={"subsequence_length": int(frames.shape[0])},
+        video_embeddings=None,
+    )
+    batch = collator([ProgressSample(trajectory=traj, sample_type="progress")])
+    inp = batch["progress_inputs"]
+    for k, v in inp.items():
+        if hasattr(v, "to"):
+            inp[k] = v.to("cuda")
+    with torch.no_grad():
+        res = compute_batch_outputs(
+            model, tokenizer, inp, sample_type="progress", is_discrete_mode=True, num_bins=100
+        )
+    prog = np.asarray(res["progress_pred"][0], dtype=np.float32)
+    print(
+        f"[reload] progress series (meta+prequantized path): {[round(float(x), 4) for x in prog]}"
+    )
+    print(
+        "[reload] COMPARE to the build_experiment REFERENCE series — they must "
+        "match element-wise to ~1e-3 for the prequantized load to be trusted.",
+        flush=True,
+    )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())