Fix degeneracy in attention kernel backward tests#258
Open
GMNGeoffrey wants to merge 2 commits into
Open
Conversation
The backward tests used loss = mean(out), which divided every input gradient by ~230k elements down to ~1e-6, far below the absolute tolerance they were compared against. So the tests passed regardless of whether the kernel backward was correct. To mitigate this, I reworked the comparisons in a few ways: - Use a sum instead of mean - Multiply the output by a random cotangent to get the loss, keeping gradients at full scale, but allowing some cancellation. - Use randn (normal mean 0, std 1) rather than rand (uniform [0, 1)) for inputs. I think this is more realistic and the mix of negative and positive values again avoids everything blowing up. - Scale down the pair bias so it doesn't dominate the actual inputs. - Use fp64 for the reference implementation. The reference implementation has its own floating point errors in lower precision dtypes, which forces very loose tolerances. I validated the tolerances used against the error of the lower-precision reference to ensure we're not using something too loose. The inputs are still generated in the target dtype and upcast, so they're exactly representable, avoiding introducing additional error just from casting the input. - Only check the input gradients (q, kv, bias). The weight gradients are a downstream autograd contraction that adds no kernel coverage while forcing a much looser tolerance. To avoid future issues, I added a custom assert `assert_close_nondegenerate` that tries to guard against this sort of degeneracy. In addition to checking tensors are close with torch's assert_close, it checks: - All values are finite (no infs or nans) - The expected output is not all close to zero (within the same atol and rtol used for the closeness check). This would've caught the issue here. - The expected and actual values are not suspiciously close to a very tight tolerance. This guards against accidentally testing the reference against itself, or something like that. I also seeded the RNG for all kernel tests via the seeded_rng fixture for determinism.
12144f6 to
f72733c
Compare
I ran cuequivariance on RTX 3090 and it need the same looser tolerances. I wasn't able to run Deepspeed, but we explicitly cast the inputs there to bf16, so I assume it will apply as well. So I just set the tolerances default to that for the half type and got rid of the more complicated stuff.
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.
Summary
The attention kernel backward tests were degenerate because the gradients they were comparing were smaller than the tolerances used for comparison (
loss = mean(out)-> gradients < 1e-6). Reworked these so they're actually testing something and made some other changes to avoid future issues. I made similar changes to the forward kernel tests, but didn't touch the e2e tests for now. I can update those to match as well, but wanted to get feedback first in case you dislike this approach.Changes
To avoid future issues, I added a custom assert
assert_close_nondegeneratethat tries to guard against this sort of degeneracy. In addition to checking tensors are close with torch's assert_close, it checks:I also seeded the RNG for all kernel tests via the seeded_rng fixture for determinism.
Testing
I ran the Triton kernel tests on Radeon 8060S and the cueq kernel tests on RTX 3090. I wasn't able to run the deepspeed kernels because they appear to not be supported on RTX 3090. Hopefully CI coverage there is sufficient.