Qwen Image prompt encoding is not padding to max seq len

[bghira]

Describe the bug

The pipeline method for QwenImagePipeline.encode_prompts is not padding correctly; it's padding by the longest sequence length in the batch, which leaves very very short embeds that are out of distribution for the models' training set.

The padding should remain at 1024 tokens even after the system prompt is dropped. The attention mask has to be expanded too.

Reproduction

Execute QwenImagePipeline.encode_prompts() and check resulting shape.

Logs

prompt_embeds.shape=torch.Size([1, 1024, 3584]), prompt_embeds_mask.shape=torch.Size([1, 1024])

^ after fixing.


prompt_embeds.shape=torch.Size([1, 5, 3584]) prompt_embeds_mask.shape=torch.Size([1, 5])

^ before.

prompt was simply minecraft

this leads to extremely high loss at training time unless very-long prompts are used.

at inference time, it causes patch embed artifacts because the RoPE is not accustomed to these positions.

System Info

Latest git main.

Who can help?

No response

[a-r-r-o-w]

A fix should be easy to add quickly if this is a problem, but I'll cc @naykun from Qwen team to take a look if this is intended or unexpected. I do see us not making use of max_sequence_length, so most likely unexpected

[bghira]

well, loss is at 3.6 without the fix, and normal range with it 🤔

the loss is due to my use of the VAE, not the text embed sequence length.

the text embed sequence length merely promotes visual artifacts and burnt learning.

[bghira]

the problem is that because the text positions and img positions share the same RoPE (tricky) i can't give the text embeds purely the max sequence length, because then there are no positions left for image tokens. it definitely needs some answers from Qwen team, maybe we just can't use short training captions on this model at all.

[bghira]

if we run the pipeline with a 1024x1024 image and a very long prompt, it simply crashes with the same kind of position errors. silently truncates.

[naykun]

Hi @bghira @a-r-r-o-w , apologies for the delayed response.

If this is solely for inference, it's a known limitation: the current implementation doesn't support variable-length text prompts in batched mode. As a result, padding isn't critical since the sequences will be truncated in the MMDiT anyway.

For training, the situation is more complex. Internally, we have a Megatron-based implementation that supports cuseq lengths with Flash Attention. However, integrating this may require further discussion and planning.

Let me know your thoughts or use case, and we can explore potential solutions.

[bghira]

can i allow users to decide how long their batch tokens can pad to? so they can do 512 tokens and it pads the batch? i am implementing training in simpletuner and want to provide the most in-distribution options for my users