Fix OOM in WanAnimate BitsAndBytes Training Test#13777
Merged
Merged
Conversation
Signed-off-by: jiqing-feng <[email protected]>
Contributor
Author
|
Hi @sayakpaul . Would you please review this PR? Thanks! |
sayakpaul
approved these changes
May 21, 2026
sayakpaul
left a comment
Member
There was a problem hiding this comment.
Thanks! There seems to be some dequant test failures for this model. Would you like to fix them in a separate PR?
FAILED tests/models/transformers/test_models_transformer_wan_animate.py::TestWanAnimateTransformer3DTorchAo::test_torchao_dequantize[int4wo] - NotImplementedError: Int4Tensor dispatch: attempting to run unimplemented operator/function: func=<OpOverload...
FAILED tests/models/transformers/test_models_transformer_wan_animate.py::TestWanAnimateTransformer3DTorchAo::test_torchao_dequantize[int8dq] - NotImplementedError: LinearActivationQuantizedTensor dispatch: attempting to run unimplemented operator/funct...
Contributor
Author
I cannot see the tests: |
Member
|
You can find them in the |
Contributor
Author
Still cannot find: https://github.com/huggingface/diffusers/blob/main/tests/models/testing_utils/quantization.py#L936-L939. The link is from main of diffusers, only |
Member
|
Hmm you are right! |
DN6
pushed a commit
that referenced
this pull request
Jul 1, 2026
reduce input size for tests Signed-off-by: jiqing-feng <[email protected]> Co-authored-by: Sayak Paul <[email protected]>
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
2 participants
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.
Problem
TestWanAnimateTransformer3DBitsAndBytes::test_bnb_trainingfails withtorch.OutOfMemoryErroron CUDA/XPU (24 GB), attempting to allocate ~68.9 GiB.The root cause is oversized dummy inputs in
get_dummy_inputs():hidden_states(1, 36, 21, 64, 64)pose_hidden_states(1, 16, 20, 64, 64)face_pixel_values(1, 3, 77, 512, 512)The training test runs forward + backward with LoRA adapters under
autocast, so intermediate activations (especially from the motion encoder's multi-layer Conv2d chain on 512×512 faces) dominate memory.Fix
Reduce the BnB test dummy input dimensions to match the smaller shapes already used by
TestWanAnimateTransformer3DTorchAo:hidden_states(1, 36, 5, 16, 16)pose_hidden_states(1, 16, 4, 16, 16)face_pixel_values(1, 3, 13, 512, 512)face_pixel_valueskeeps 512×512 spatial resolution (required by the Hub model'smotion_encoder_size=512) but reduces frame count from 77 to 13. Latent spatial dims drop from 64 to 16 and frame counts are similarly reduced.