| |
| """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)) |
|
|
| |
| 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 |
|
|
| |
| 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") |
|
|
| |
| 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)}") |
|
|
| |
| 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() |
|
|
