|
| 1 | +#!/usr/bin/env python3 |
| 2 | +# Copyright 2026 The LoongForge Authors. |
| 3 | +# SPDX-License-Identifier: Apache-2.0 |
| 4 | +"""Online-vs-online VAE-input parity check (DATA link). |
| 5 | +
|
| 6 | +Compares the LATENT_DUMP dumps of TWO separate online training runs (both |
| 7 | +launched with the SAME PARITY_DATA_SEED so the sampler visits identical anchors |
| 8 | +in identical order). For each matching rankR_stepN.pt it bit-compares: |
| 9 | +
|
| 10 | + - first_frame (decode + normalization + H2D) |
| 11 | + - video_frames (decode + normalization + H2D) |
| 12 | + - clean_full_latent (end-to-end: data + VAE compute) |
| 13 | +
|
| 14 | +A `torch.equal` pass on first_frame / video_frames proves the data pipeline |
| 15 | +(torchcodec decode + transform + collate) is bit-deterministic ACROSS PROCESSES |
| 16 | +-- i.e. re-fetching the same (episode, condition_frame) anchor in the offline |
| 17 | +precompute will get the exact same input pixels. Combined with |
| 18 | +verify_offline_latents.py (compute link), a pass on both scripts means the |
| 19 | +offline latent cache is end-to-end bit-identical to the online encode. |
| 20 | +
|
| 21 | +Usage |
| 22 | +----- |
| 23 | + python examples/embodied/motus/diff_online_dumps.py \ |
| 24 | + --dir-a /tmp/latent_dump_run1 --dir-b /tmp/latent_dump_run2 |
| 25 | +""" |
| 26 | +from __future__ import annotations |
| 27 | + |
| 28 | +import argparse |
| 29 | +import glob |
| 30 | +import os |
| 31 | + |
| 32 | +import torch |
| 33 | + |
| 34 | + |
| 35 | +def _cmp(a: torch.Tensor, b: torch.Tensor) -> str: |
| 36 | + if a.shape != b.shape: |
| 37 | + return f"SHAPE-MISMATCH {tuple(a.shape)} vs {tuple(b.shape)}" |
| 38 | + if torch.equal(a, b): |
| 39 | + return "equal" |
| 40 | + diff = (a.float() - b.float()).abs().max().item() |
| 41 | + return f"DIFFERS max|abs|={diff:.3e}" |
| 42 | + |
| 43 | + |
| 44 | +def main() -> int: |
| 45 | + ap = argparse.ArgumentParser() |
| 46 | + ap.add_argument("--dir-a", required=True, help="dump dir of online run #1") |
| 47 | + ap.add_argument("--dir-b", required=True, help="dump dir of online run #2") |
| 48 | + args = ap.parse_args() |
| 49 | + |
| 50 | + files_a = sorted(glob.glob(os.path.join(args.dir_a, "rank*_step*.pt"))) |
| 51 | + if not files_a: |
| 52 | + print(f"[FAIL] no dumps in {args.dir_a}") |
| 53 | + return 2 |
| 54 | + |
| 55 | + all_ok = True |
| 56 | + for fa in files_a: |
| 57 | + name = os.path.basename(fa) |
| 58 | + fb = os.path.join(args.dir_b, name) |
| 59 | + if not os.path.exists(fb): |
| 60 | + print(f"{name}: MISSING in dir-b") |
| 61 | + all_ok = False |
| 62 | + continue |
| 63 | + da = torch.load(fa, map_location="cpu") |
| 64 | + db = torch.load(fb, map_location="cpu") |
| 65 | + ff = _cmp(da["first_frame"], db["first_frame"]) |
| 66 | + vf = _cmp(da["video_frames"], db["video_frames"]) |
| 67 | + lat = _cmp(da["clean_full_latent"], db["clean_full_latent"]) |
| 68 | + ok = ff == "equal" and vf == "equal" |
| 69 | + all_ok = all_ok and ok |
| 70 | + print(f"{name}: first_frame={ff} | video_frames={vf} | latent={lat}") |
| 71 | + |
| 72 | + print() |
| 73 | + if all_ok: |
| 74 | + print("[PASS] data pipeline is bit-deterministic across processes.") |
| 75 | + print(" -> offline precompute will get identical VAE inputs per anchor.") |
| 76 | + return 0 |
| 77 | + print("[WARN] VAE inputs differ across runs -> decode/transform is not") |
| 78 | + print(" cross-process deterministic; offline cache keys must pin the") |
| 79 | + print(" exact decoded frames (store inputs, not just anchor ids).") |
| 80 | + return 1 |
| 81 | + |
| 82 | + |
| 83 | +if __name__ == "__main__": |
| 84 | + raise SystemExit(main()) |
0 commit comments