Fix wan2.2 VBench scorer cold-start failures: .mp4 staging, torch>=2.6 checkpoints, missing wget#403
Fix wan2.2 VBench scorer cold-start failures: .mp4 staging, torch>=2.6 checkpoints, missing wget#403wu6u3tw wants to merge 2 commits into
Conversation
|
MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅ |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request improves VBench evaluation robustness by ensuring staged videos always use the ".mp4" extension to avoid loading errors, adding pre-flight checks for system dependencies ("wget" and "unzip"), and monkeypatching "torch.load" to support older VBench checkpoints on PyTorch 2.6+. The review feedback suggests a cleaner approach to configuring PyTorch's loading behavior by setting the "TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD" environment variable instead of manually monkeypatching "torch.load".
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| def _default_torch_load_to_full_pickles() -> None: | ||
| """Make torch.load default to weights_only=False in this subprocess. | ||
|
|
||
| VBench's reference checkpoints (e.g. motion_smoothness AMT, RAFT) are full | ||
| pickled objects, written before torch 2.6 flipped torch.load's default to | ||
| weights_only=True; loading them with the new default fails with | ||
| UnpicklingError (e.g. "Unsupported global: typing.OrderedDict"). They are | ||
| the VBench-sanctioned reference weights, so loading them unrestricted here | ||
| matches upstream VBench behavior on the torch versions it was written for. | ||
| Callers that pass weights_only explicitly (e.g. torch.hub) are unaffected; | ||
| scoped to this subprocess only, the parent benchmark keeps stock semantics. | ||
| """ | ||
| orig_load = torch.load | ||
|
|
||
| @functools.wraps(orig_load) | ||
| def _load(*args, **kwargs): | ||
| kwargs.setdefault("weights_only", False) | ||
| return orig_load(*args, **kwargs) | ||
|
|
||
| torch.load = _load | ||
|
|
||
|
|
||
| _default_torch_load_to_full_pickles() |
There was a problem hiding this comment.
Instead of manually monkeypatching torch.load, you can use PyTorch's built-in environment variable TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD. Setting this environment variable to "1" achieves the exact same behavior (defaulting weights_only to False when not explicitly specified) in a much cleaner, safer, and more robust way without modifying core library functions. This also allows you to remove the unused import functools.
import os
# PyTorch >= 2.6 defaults to weights_only=True. VBench's reference checkpoints
# are full pickled objects and fail to load with this default. Instead of
# monkeypatching torch.load, we use PyTorch's built-in environment variable
# to default weights_only to False when not explicitly specified.
os.environ["TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD"] = "1"VBench's load_video dispatches purely on the file extension and raises NotImplementedError for anything but .mp4/.gif/frame dirs. trtllm-serve emits MJPEG .avi when ffmpeg is unavailable server-side, so a full accuracy run generated 248 videos and then failed at scoring. decord (libav) detects the real container by content, so an .avi symlinked under an .mp4 name decodes fine (verified on GB300: 81-frame 720x1280 MJPEG-AVI reads correctly). Always name the staged symlinks .mp4.
Two cold-start failures observed on aarch64 GB300 runs: - The accuracy venv resolves torch 2.12, where torch.load defaults to weights_only=True (since 2.6). VBench's reference checkpoints (motion_smoothness AMT, RAFT) are full pickles and fail with UnpicklingError. Default torch.load to weights_only=False in this subprocess only; explicit weights_only callers are unaffected. - vbench.utils.init_submodules downloads per-dimension weights via literal wget/unzip subprocesses; without them the run dies with an opaque FileNotFoundError after videos were already generated. Fail fast preflight with a structured error naming the missing tools.
wu6u3tw
force-pushed
the
wan22-vbench-deps
branch
from
6ab9beb to
38c44de
Compare
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #403 +/- ##
=======================================
Coverage ? 80.03%
=======================================
Files ? 129
Lines ? 17226
Branches ? 0
=======================================
Hits ? 13786
Misses ? 3440
Partials ? 0 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
2 participants
Fix wan2.2 VBench scorer cold-start failures: .mp4 staging, torch>=2.6 checkpoints, missing wget
Split from #401 (part 2 of 2; the drain-timeout config is in the companion PR).
Three independent failures in the VBench scoring path, each discovered only after a full generation pass (30-50 min on 18 GB300 nodes). With these fixes (initially applied as runtime monkeypatches), a full wan2.2 Offline accuracy run completed on GB300-NVL72 and scored 0.69998 over 248 samples.
Fix 1: stage videos as
.mp4regardless of source suffixVBenchScorer._stage_videoskept the source extension (src.suffix or '.mp4'). VBench'sload_videodispatches purely on the extension and raises bareNotImplementedErrorfor anything but.mp4/.gif/frame dirs. trtllm-serve emits MJPEG.aviwhen ffmpeg is not installed server-side, so a full accuracy pass generated all 248 videos and then died at scoring.decord (libav) detects the container by content, not extension; an MJPEG-AVI symlinked under an
.mp4name decodes correctly (verified on GB300: 81 frames, 720x1280). Staged symlinks are now always named{prompt}-{idx}.mp4.Fix 2:
vbench_runner.pyfails on torch >= 2.6 checkpoint loadingThe wan2.2 accuracy subproject resolves torch 2.12, where
torch.loaddefaults toweights_only=True(changed in 2.6). VBench's reference checkpoints (motion_smoothness AMT, RAFT) are full pickles and fail withUnpicklingError("Unsupported global: typing.OrderedDict"). The runner now defaultstorch.loadtoweights_only=Falsewithin the scorer subprocess only; callers passingweights_onlyexplicitly (e.g.torch.hub) are unaffected, and the parent benchmark process keeps stock semantics.Fix 3: opaque failure when
wget/unzipare missingvbench.utils.init_submodulesdownloads per-dimension weights via literalwget/unzipsubprocesses on a cold cache. In a minimal client container this dies mid-evaluation withFileNotFoundError: 'wget', after the videos were already generated. The runner now preflights both tools and exits early with a structured error naming what is missing and how to fix it.Testing
test_stage_videos_renames_non_mp4_sources_to_mp4.tests/unit/evaluation/test_scoring.pyand the full unit suite pass with zero new failures vs the unmodified base commit..avivideos end to end (score 0.69998), including cold-cache weight downloads.