| |
| """Harness-independent compute claim: FLOPs per action for WorldDiT. |
| Their default = 20 flow steps (per the model card). Ours = 2 steps + text cached. |
| Separates the FLOP reduction (few-step + text-cache) from latency wins (CUDA graphs).""" |
| import os, sys, json |
| os.environ.setdefault("LIBERO_CONFIG_PATH", "/root/.libero_lat") |
| os.environ.setdefault("MUJOCO_GL", "osmesa"); os.environ.setdefault("PYOPENGL_PLATFORM", "osmesa") |
| sys.path.insert(0, "/root/LIBERO") |
| import torch |
| from pathlib import Path |
| from eval import load_model |
| from torch.utils.flop_counter import FlopCounterMode |
|
|
| m = load_model(Path("/workspace/worlddit"), "libero_spatial", torch.device("cuda")) |
| blob = torch.load("/workspace/worlddit/inputs.pt", map_location="cuda") |
| primary, wrist, state, text = (blob[k].cuda() for k in ("primary", "wrist", "state", "text")) |
| head = m.unified_action_world_head |
|
|
| def flops_at(steps): |
| _hf = head.forward |
| head.forward = (lambda ctx, sampling_steps=None, _hf=_hf, s=steps: _hf(ctx, sampling_steps=s)) |
| fc = FlopCounterMode(display=False) |
| with fc: |
| m(primary, wrist, state, text) |
| head.forward = _hf |
| return fc.get_total_flops() |
|
|
| f20 = flops_at(20) |
| f2 = flops_at(2) |
| v = (f20 - f2) / 18.0 |
| enc = f2 - 2 * v |
|
|
| fc = FlopCounterMode(display=False) |
| with fc: |
| m.clip_model.encode_text(text.flatten(0, 1)) |
| f_text = fc.get_total_flops() |
|
|
| opt = enc - f_text + 2 * v |
| G = 1e9 |
| res = { |
| "their_default_20step_GFLOPs": round(f20/G, 2), |
| "per_flow_step_GFLOPs": round(v/G, 3), |
| "encoders_run_once_GFLOPs": round(enc/G, 2), |
| "of_which_text_encoder_GFLOPs": round(f_text/G, 3), |
| "ours_2step_textcached_GFLOPs": round(opt/G, 2), |
| "FLOP_reduction_x": round(f20/opt, 2), |
| "note": "FLOP reduction < latency speedup; the rest of latency is CUDA-graphs (systems, no FLOP change)", |
| } |
| print("FLOPS " + json.dumps(res, indent=2), flush=True) |
| Path("/workspace/worlddit/results/flops.json").write_text(json.dumps(res, indent=2)) |
|
|
