diff --git a/extensions_built_in/captioner/Qwen3VLCaptioner.py b/extensions_built_in/captioner/Qwen3VLCaptioner.py index bc4792c..cf55dfa 100644 --- a/extensions_built_in/captioner/Qwen3VLCaptioner.py +++ b/extensions_built_in/captioner/Qwen3VLCaptioner.py @@ -5,6 +5,8 @@ from transformers import ( ) from collections import OrderedDict +import torch +import torch.nn.functional as F from optimum.quanto import freeze from toolkit.basic import flush from toolkit.util.quantize import quantize, get_qtype @@ -14,6 +16,29 @@ import transformers import logging import warnings + +def patch_qwen_vl_patch_embed(model): + """Qwen-VL's vision patch_embed is a Conv3d whose kernel == stride, i.e. a plain + linear projection of each flattened patch. bf16 Conv3d has no fast cuDNN kernel and + falls back to a slow, GPU-underutilizing path. Swap it for the equivalent F.linear + (a GEMM). The weight is read lazily so this survives later .to(device)/dtype moves. + Returns the number of patch_embed modules patched.""" + patched = 0 + for module in model.modules(): + proj = getattr(module, "proj", None) + if ( + isinstance(proj, torch.nn.Conv3d) + and tuple(proj.kernel_size) == tuple(proj.stride) + ): + def fast_forward(hidden_states, _proj=proj): + w = _proj.weight.reshape(_proj.weight.shape[0], -1) + x = hidden_states.view(-1, w.shape[1]).to(w.dtype) + return F.linear(x, w, _proj.bias) + + module.forward = fast_forward + patched += 1 + return patched + # transformers.logging.set_verbosity_error() warnings.filterwarnings("ignore") logging.disable(logging.WARNING) @@ -35,6 +60,8 @@ class Qwen3VLCaptioner(BaseCaptioner): dtype=self.torch_dtype, device_map="cpu", ) + # swap the slow bf16 Conv3d patch_embed for an equivalent fast linear + patch_qwen_vl_patch_embed(self.model) if not self.caption_config.low_vram: self.model.to(self.device_torch) if self.caption_config.quantize: