[tests] fix anyflow tests#13855
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
| key = attn.norm_k(key) | ||
|
|
||
| # Make layerwise upcasting work. | ||
| value = value.to(query.dtype) |
There was a problem hiding this comment.
| value = value.to(query.dtype) | |
| query = query.to(target_dtype) | |
| key = key.to(target_dtype) |
After some debugging with Claude, I think this should fix all of the layerwise casting issues (where target_dtype = hidden_states.dtype is the effective compute dtype, cached before the to_q/to_k/to_v projections happen). I tested this out locally and all the tests pass after this change, including TestAnyFlowFARTransformer3DMemory::test_layerwise_casting_training and TestAnyFlowFARTransformer3DTraining::test_mixed_precision_training, which were still failing even with the PR.
My understanding of the issue is as follows:
- Because
norm_qandnorm_kare skipped due to_skip_layerwise_casting_patterns, and because they arediffusers.models.normalization.RMSNorms which upcast to FP32 during the variance calculation, they end up upcastingqueryandkeyto FP32. - Because
valueis not normed, it stays in the effective compute dtypetarget_dtype(for example, the autocast dtype inside atorch.amp.autocastregion like intest_layerwise_casting_training). - If we cast
valuetoquery's dtype, we will therefore end up castingvalueto FP32, regardless of the effective compute dtype. In this case, the Flex Attention operation will always be performed in FP32. - However, if we are training, during the backward pass the gradients from the first layer after Flex Attention, which is
to_out, will be intarget_dtype(for example becausetorch.amp.autocastwill automatically downcast it to the autocast dtype). - Whenever
target_dtype != torch.float32, we will then get a dtype mismatch error when running the Flex Attention backward pass, because the attention scores will be in FP32, but the incoming gradients will be intarget_dtype.
So casting query and key to target_dtype after the QK norm solves the underlying issue by putting query/key/value, and therefore the Flex Attention computation, in the effective compute dtype, which works for both inference and training.
| pytest.mark.skipif(torch_device == "cpu", "FlexAttention has no CPU backward kernel.") | ||
|
|
There was a problem hiding this comment.
| pytest.mark.skipif(torch_device == "cpu", "FlexAttention has no CPU backward kernel.") |
I think we can remove this because Flex Attention does have a CPU backward path, and the below test (test_training) actually passes after the fixes.
| pytest.mark.skipif(torch_device == "cpu", "FlexAttention has no CPU backward kernel.") | ||
|
|
There was a problem hiding this comment.
| pytest.mark.skipif(torch_device == "cpu", "FlexAttention has no CPU backward kernel.") |
Same reason as #13855 (comment).
| pytest.mark.skipif(torch_device == "cpu", "FlexAttention has no CPU backward kernel.") | ||
|
|
There was a problem hiding this comment.
| pytest.mark.skipif(torch_device == "cpu", "FlexAttention has no CPU backward kernel.") |
Same reason as #13855 (comment).
dg845
left a comment
There was a problem hiding this comment.
Thanks for the PR! Left some comments.
|
Can you please open a PR to this branch? I think it's just easier that way. |
|
Opened a PR to this branch with the suggested fix: #13863. |
Fix AnyFlow FAR causal transformer training layerwise / mixed precision tests
* fix anyflow tests * [tests] fix anyflow tests layerwise casting (#13863) Fix AnyFlow FAR causal transformer training layerwise / mixed precision tests --------- Co-authored-by: dg845 <[email protected]>
What does this PR do?
Fixes https://github.com/huggingface/diffusers/actions/runs/26819669538/job/79072404192?pr=13712. Note that the approach I have taken to fix the dtype mismatch error includes an explicit cast. We can disable the layerwise casting tests if that's preferred.
Additionally, removes the
unitteststuff and solely usespytest.