#!/usr/bin/env python3 """throughput_bench.py — PORTABLE inference-throughput benchmark for WorldDiT. Measures the number that matters for robot deployment: sustained per-replan latency and throughput of one model() call, under the chosen optimization config. It captures one real input set once (saved to inputs.pt) so it can then run on ANY machine WITHOUT the LIBERO sim — e.g. a Jetson Orin. Robot-control framing: WorldDiT chunks (predicts 7, executes 3), so effective action rate = replan_rate x 3. # on box2 (capture + measure): python throughput_bench.py --steps 4 --bf16 --compile --iters 200 # on the Orin (no sim; uses the shipped inputs.pt): python throughput_bench.py --steps 4 --bf16 --compile --inputs inputs.pt --iters 200 """ import os, sys, time, json, argparse, contextlib, copy, statistics from pathlib import Path import numpy as np, torch def main(): ap = argparse.ArgumentParser() ap.add_argument("--suite", default="libero_spatial") ap.add_argument("--steps", type=int, default=4) ap.add_argument("--bf16", action="store_true") ap.add_argument("--compile", action="store_true") ap.add_argument("--iters", type=int, default=200) ap.add_argument("--warmup", type=int, default=30) ap.add_argument("--chunk", type=int, default=3, help="actions executed per replan") ap.add_argument("--control-hz", type=float, default=20.0, help="target robot control rate for the verdict") ap.add_argument("--model-root", default="/workspace/worlddit") ap.add_argument("--libero-path", default="/root/LIBERO") ap.add_argument("--inputs", default="inputs.pt", help="cached model inputs; captured from sim if missing") ap.add_argument("--cache-text", action="store_true", help="memoize CLIP text encoding (constant per episode)") ap.add_argument("--profile", action="store_true", help="per-component latency breakdown (eager)") ap.add_argument("--quant", default="none", choices=["none","int8","int8dyn","fp8","fp8dyn"], help="quantize vision/perceiver/head") ap.add_argument("--out", default="") a = ap.parse_args() dev = torch.device("cuda" if torch.cuda.is_available() else "cpu") sys.path.insert(0, a.libero_path) from eval import PolicyRunner, load_model model = load_model(Path(a.model_root), a.suite, dev) head = model.unified_action_world_head if a.bf16: head.bfloat16() try: model.vision_encoder.bfloat16() except Exception as e: print("bf16 vision skip:", e) if a.quant != "none": from torchao.quantization import (quantize_, int8_weight_only, int8_dynamic_activation_int8_weight, float8_weight_only, float8_dynamic_activation_float8_weight) q = {"int8": int8_weight_only, "int8dyn": int8_dynamic_activation_int8_weight, "fp8": float8_weight_only, "fp8dyn": float8_dynamic_activation_float8_weight}[a.quant]() for name in ("vision_encoder", "perceiver_resampler"): mod = getattr(model, name, None) if mod is not None: try: quantize_(mod, q) except Exception as e: print(f"quant {name} skip:", e) try: quantize_(head, q) except Exception as e: print("quant head skip:", e) if a.compile: head._velocity = torch.compile(head._velocity, mode="reduce-overhead", dynamic=False) for name in ("vision_encoder", "perceiver_resampler"): mod = getattr(model, name, None) if mod is not None: setattr(model, name, torch.compile(mod, dynamic=False)) _hf = head.forward head.forward = (lambda context, sampling_steps=None, _hf=_hf, s=a.steps: _hf(context, sampling_steps=s)) # text-embedding caching: instruction is constant per episode, so memoize CLIP encode_text if a.cache_text: _enc = model.clip_model.encode_text _tc = {} def _cached_encode(tok, _enc=_enc): key = hash(tok.detach().to("cpu").numpy().tobytes()) v = _tc.get(key) if v is None: v = _enc(tok); _tc[key] = v return v model.clip_model.encode_text = _cached_encode # per-component profiler (eager only; adds cuda-syncs that would break graphs) PROF = {} def _wrap(obj, attr, name): orig = getattr(obj, attr); PROF[name] = [0.0, 0] def w(*ar, __orig=orig, __n=name, **kw): torch.cuda.synchronize(); t0 = time.time() r = __orig(*ar, **kw) torch.cuda.synchronize(); PROF[__n][0] += (time.time()-t0)*1000; PROF[__n][1] += 1 return r setattr(obj, attr, w) if a.profile: _wrap(model.vision_encoder, "forward", "vision") _wrap(model.clip_model, "encode_text", "text") _wrap(model.perceiver_resampler, "forward", "perceiver") _wrap(head, "_velocity", "head_step") # --- inputs: load cached (portable) or capture one real set from the sim --- if Path(a.inputs).is_file(): blob = torch.load(a.inputs, map_location=dev) primary, wrist, state, text = (blob[k].to(dev) for k in ("primary","wrist","state","text")) print(f"loaded cached inputs from {a.inputs}") else: primary, wrist, state, text = _capture(a, model, dev) torch.save({"primary":primary.cpu(),"wrist":wrist.cpu(),"state":state.cpu(),"text":text.cpu()}, a.inputs) print(f"captured + saved inputs to {a.inputs} shapes: {tuple(primary.shape)} {tuple(text.shape)}") # --- timed loop over the full model() call --- with torch.no_grad(): for _ in range(a.warmup): model(primary, wrist, state, text) if dev.type == "cuda": torch.cuda.synchronize() times = [] for _ in range(a.iters): t0 = time.time() model(primary, wrist, state, text) if dev.type == "cuda": torch.cuda.synchronize() times.append((time.time()-t0)*1000.0) times.sort() p50 = statistics.median(times); p99 = times[int(0.99*len(times))-1] replan_hz = 1000.0/p50; action_hz = replan_hz*a.chunk gpu = torch.cuda.get_device_name(0) if dev.type=="cuda" else "cpu" res = {"device":gpu, "steps":a.steps, "bf16":a.bf16, "compile":a.compile, "p50_ms":round(p50,2), "p99_ms":round(p99,2), "replan_hz":round(replan_hz,1), "effective_action_hz":round(action_hz,1), "meets_control_rate": action_hz >= a.control_hz, "target_control_hz":a.control_hz} if a.profile: print("=== per model() call component breakdown (ms) ===") for k, (tot, n) in PROF.items(): print(f" {k:12s}: {tot/max(a.iters,1):7.2f} ms ({n/max(a.iters,1):.2f} sub-calls/fwd)") print("THROUGHPUT " + json.dumps(res), flush=True) if a.out: Path(a.out).write_text(json.dumps(res)+"\n") def _capture(a, model, dev): import json as _j pkg = Path(a.libero_path)/"libero"/"libero" cfg = Path("/root/.libero_lat"); cfg.mkdir(parents=True, exist_ok=True) (cfg/"config.yaml").write_text(_j.dumps({"assets":str(pkg/"assets"),"bddl_files":str(pkg/"bddl_files"), "benchmark_root":str(pkg),"datasets":str(pkg.parent/"datasets"),"init_states":str(pkg/"init_files")})) os.environ["LIBERO_CONFIG_PATH"]=str(cfg); os.environ.setdefault("MUJOCO_GL","osmesa") from eval import PolicyRunner from libero.libero import benchmark from libero.libero.envs import OffScreenRenderEnv grabbed = {} _orig = model.forward def grab(primary, wrist, state, text, *ar, **kw): grabbed.setdefault("v", (primary.detach(), wrist.detach(), state.detach(), text.detach())) return _orig(primary, wrist, state, text, *ar, **kw) model.forward = grab with open(os.devnull,"w") as q, contextlib.redirect_stdout(q), contextlib.redirect_stderr(q): suite = benchmark.get_benchmark_dict()[a.suite]() task = suite.get_task(0) bddl = pkg/"bddl_files"/task.problem_folder/task.bddl_file inits = torch.load(pkg/"init_files"/task.problem_folder/task.init_states_file, weights_only=False) env = OffScreenRenderEnv(bddl_file_name=str(bddl), camera_heights=128, camera_widths=128, render_gpu_device_id=0) env.reset(); env.seed(100); env.reset(); env.set_init_state(inits[0]) warm=[] for _ in range(5): o,*_=env.step(np.zeros(7)); warm.append(copy.deepcopy(o)) runner = PolicyRunner(model, a.temperature if hasattr(a,"temperature") else 0.01, 3, 600) runner.reset(); runner.prefill(warm) obs = warm[-1] while "v" not in grabbed: act = runner.act(obs, task.language, 0) obs,_,done,_ = env.step(act) env.close() model.forward = _orig return grabbed["v"] if __name__ == "__main__": main()