Qwen text encoder LoRA training crashes with shape mismatch during backward pass (gradient checkpointing)

[dxqb]

Training a Qwen LoRA with text encoder training enabled (text_encoder.train=true) crashes during loss.backward() with a tensor shape mismatch inside qwen2_5_vl self-attention.

Steps to reproduce:

Use the #qwen LoRA 16GB preset (base model: Qwen/Qwen-Image) with any LoRA output format and text_encoder.train=true.

Error:

Traceback (most recent call last):
  File "scripts/train.py", ...
    trainer.train()
  File "modules/trainer/GenericTrainer.py", line 759, in train
    loss.backward()
  ...
  File "modules/util/checkpointing_util.py", line 115, in __checkpointing_forward
    output = self.orig_forward(*args) if self.checkpoint is None else self.checkpoint(*args)
  File "transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py", line 810, in forward
    hidden_states, _ = self.self_attn(...)
  File "transformers/integrations/sdpa_attention.py", line 92, in sdpa_attention_forward
    attn_output = torch.nn.functional.scaled_dot_product_attention(...)
RuntimeError: The expanded size of the tensor (1092) must match the existing size (546)
at non-singleton dimension 3.  Target sizes: [2, 28, 546, 1092].  Tensor sizes: [2, 1, 546, 546]

Analysis:

The crash occurs during gradient checkpointing's recompute pass. The attention mask has shape [2, 1, 546, 546] (key dimension = 546), but during recompute the key sequence is 1092 tokens (2×546). This suggests the Qwen2.5-VL model concatenates visual and text tokens in a way that the mask cached from the first forward pass no longer matches the recomputed sequence length. OT's gradient checkpointing wrapper in checkpointing_util.py does not account for this.

Without text_encoder.train=false (the default), training works correctly.

Environment:

• PyTorch 2.12.0+cu130
• transformers 5.5.4

Drafted by Claude

[dxqb]

Most likely solution: remove text encoder training from the Qwen model. It isn't implemented anymore in newer models, because finetuning LLMs on an image MSE loss isn't useful.
It was implemented for Qwen before this decision was made. I guess nobody uses it.