fix(qwen-image): drop the txt_seq_lens kwarg removed by diffusers - #190
Merged
Conversation
6 tasks
Collaborator
|
LGTM :) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Stop passing
txt_seq_lensintoQwenImageTransformer2DModel.forward. It stays aper-sample input —
collate_cond_for_sample_batchstill reads it to derive the textmask and the padding width — it just no longer travels in the collated dict that the
trainer splats into the model.
Why
txt_seq_lenswas removed from bothQwenImageTransformer2DModel.forwardandQwenEmbedRope.forwardafter diffusers 0.39. Under the pin in #189(
f53d5520, post-0.39 main) the qwen-image training forward raisesTypeError: forward() got an unexpected keyword argument 'txt_seq_lens'.Numerically a no-op on 0.38
Removing the kwarg cannot move a single bit on the current pin:
deprecate()warning inQwenImageTransformer2DModel.forward— nothing read it.compute_text_seq_len_from_mask(identical in 0.38 andf53d5520), passed topos_embedasmax_txt_seq_len=text_seq_len. That path neverconsults
txt_seq_lens.QwenEmbedRope.forwardwould fall back tomax(txt_seq_lens)only whenmax_txt_seq_len is None, which the transformer never does.Verified bitwise under #189's diffusers pin
tests/e2e/short/test_qwenimage_pickscore_grpo_5xGPU.pyrun on this branch inside thepinned-diffusers image, compared against the standard recorded on 0.38:
radixark/miles_diffusion:test-pr-1890.40.0.dev0@f53d552036a0d1bd5570782a39cd40cfabf112bc7605529--num-rollout 2(the test's own args, unmodified)rollout/reward/raw_num_samplesrollout/reward/raw_meanrollout/reward/raw_medianrollout/reward/raw_stdtrain/log_prob_old_idx_0train/log_prob_new_idx_0train/log_prob_mean_abs_difftrain/model_output_mean_abs_difftrain/model_output_rel_maxtrain/grad_norm32/32 points bit-for-bit, so the 311 commits between
v0.38.0andf53d5520leave theqwen-image train and rollout numerics untouched. The e2e is the guard for this change,
hence no new unit test.
Not covered here
test_wan22_pickscore_grpo_17xGPU_single_node_4xGPU_proxystill OOMs under #189's pin.That failure is memory headroom, not numerics: at the OOM the trainer needed 1.31 GiB
with 1.25 GiB free while a rollout engine still held 25.3 GiB and the colocated
pickscore reward 4.7 GiB on the same device. The wan-side diffusers delta is a fp32
ropebuffer pin (~0.5 MiB) and one addedhidden_states.contiguous()(~32 MiB afterthe ulysses-4 split) — enough to tip a device already at 99% occupancy, and tracked
separately.