File size: 3,243 Bytes
1e05592 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | """Run make_cot_belief_cache with patched Qwen3VLVisionPatchEmbed.
Replaces the Conv3d patch projection with an equivalent Linear layer (math
identical, but ~64× faster because of a cuDNN slow-path bug for tiny Conv3d
on bf16). Saves whole-cache time from ~6 days to ~2 hours.
Usage: identical to make_cot_belief_cache, just call this instead.
python tools/run_qwen3_cache_fast.py \\
--ckpt_dir checkpoints/VLA/qwen3vl4b_cot_belief_perframe/best \\
--base_model models/Qwen3-VL-4B-Instruct \\
--split val \\
--out data/belief_cache_perframe_qwen3vl4b/multisrc_val.pt \\
--n_frames 8 --sampling last_biased --source_filter all \\
--batch_size 8 --num_workers 4 --chunk_size 2000
"""
import sys
sys.path.insert(0, ".")
import torch
import torch.nn as nn
from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLVisionPatchEmbed
# ─── Lazy-replacement: first forward call replaces Conv3d with Linear ─────
_PATCH_APPLIED = {}
def _fast_patch_embed_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
"""Mathematically equivalent to original Conv3d-based forward, but
routes through nn.Linear (which avoids the cuDNN slow-path bug on tiny
Conv3d inputs)."""
target_dtype = self.proj.weight.dtype
# First call on this instance: convert Conv3d → Linear in place.
if isinstance(self.proj, nn.Conv3d):
conv = self.proj
out_dim = conv.out_channels
in_dim = (conv.in_channels * conv.kernel_size[0]
* conv.kernel_size[1] * conv.kernel_size[2])
# Conv3d weight: (out, in, k_t, k_h, k_w) → flatten last 4 dims
w_flat = conv.weight.detach().reshape(out_dim, in_dim).contiguous()
bias = conv.bias.detach().clone() if conv.bias is not None else None
new_proj = nn.Linear(in_dim, out_dim, bias=bias is not None)
new_proj.weight.data.copy_(w_flat)
if bias is not None:
new_proj.bias.data.copy_(bias)
new_proj.to(device=conv.weight.device, dtype=conv.weight.dtype)
self.proj = new_proj
if id(self) not in _PATCH_APPLIED:
_PATCH_APPLIED[id(self)] = True
print(f"[fast_patch] patched Qwen3VLVisionPatchEmbed @ id={id(self)}: "
f"Conv3d({in_dim}→{out_dim}) → Linear({in_dim}→{out_dim})",
flush=True)
# Now self.proj is nn.Linear. Input may be (N, 1536) flat or (N, 3, 2, 16, 16).
if hidden_states.dim() > 2 or hidden_states.shape[-1] != self.proj.in_features:
hidden_states = hidden_states.reshape(-1, self.proj.in_features)
hidden_states = hidden_states.to(dtype=target_dtype)
return self.proj(hidden_states)
# Apply class-level patch BEFORE any model is instantiated
Qwen3VLVisionPatchEmbed.forward = _fast_patch_embed_forward
print("[fast_patch] Qwen3VLVisionPatchEmbed.forward replaced "
"(lazy Conv3d → Linear conversion)", flush=True)
# Hand off to the original cache builder ───────────────────────────────────
from training.Policy import make_cot_belief_cache # noqa: E402
if __name__ == "__main__":
sys.exit(make_cot_belief_cache.main())
|