Silent vision bypass on transformers ≥4.50: constant action for every image (fix inside)

#2
by jashshah999 - opened

On a current stack (transformers 4.56.2, torch 2.10), predict_action on this checkpoint returns a plausible but input-independent action — the same 7 tokens (31872 ×7 → [0.096, 0.1071, -0.0027, -0.0016, -0.015, -0.0193, 0.0]) for a real frame, a pure black image, a pure white image and random noise. vision_backbone and projector forward hooks fire 0 times. No error is raised, so it looks like a working policy.

Cause (in modeling_prismatic.py): prepare_inputs_for_generation slices the prompt whenever past_key_values is not None. transformers ≥4.50 passes an empty DynamicCache on step 0 instead of None, so the 35-token prompt collapses to 1 token, forward routes into its input_ids.shape[1] == 1 cached-generation branch, and pixel_values is never consumed.

Fix: gate the slice on actual cache length and forward an empty cache as None (the multimodal branch asserts past_key_values is None). Two smaller modern-stack blockers: _supports_sdpa declared as a property (evaluated before language_model exists) and GenerationMixin no longer auto-inherited.

Full write-up with the verification on the official AutoModelForVision2Seq load path, and the patched function: https://github.com/openvla/openvla/issues/346

A drop-in shim that applies the fixes without changing the weights, plus a one-command black-vs-white check (openvla-compat-check): https://github.com/jashshah999/vla-lite — also runs the model in nf4 at 4.6 GB VRAM.

Suggested regression test for the remote code: a black frame and a white frame must not yield the same action. Happy to open a PR here if that's preferred.

Sign up or log in to comment