Repository navigation
[JAX] Compare the attention tutorial to its baseline on a shorter sequence - #3828
Conversation
…uence The native baseline in the attention tutorial materializes fp32 scores, so at the tutorial's 4K sequence its backward pass needs a ~49 GiB buffer. That does not fit when CI runs several test workers per GPU, and test_bshd_gqa_swa_matches_baseline fails with RESOURCE_EXHAUSTED. Compare TE to the baseline on a 1K prefix instead. Signed-off-by: Alex Y. Chan <alechan@nvidia.com>
 [Low impact] The PR appears safe to merge. SummaryThe PR compares attention outputs and gradients on a 1,024-token prefix to reduce GPU memory use.
Reviews (4) · Last reviewed commit: "[JAX] Generalize the attention tutorial ..." · Reviewed by Greptile |
|
/te-ci L1 jax |
KshitijLakhani
left a comment
KshitijLakhani
left a comment
There was a problem hiding this comment.
Thanks for this change, Alex !
This change also does not affect the rendered tutorial (4k) and the static numbers displayed for it (for 4k) as well. So the tutorial users will see a cohesive tutorial.
@aybchan test_single_gpu_benchmark also uses the same sizes but of course I do not think it might be retaining the native buffers for a comparison at the end and so you might not see the same issue that you are fixing in this PR for test_bshd_gqa_swa_matches_baseline intest_single_gpu_benchmark - but it would be good to confirm and if needed, add that change in this PR as well. However, I'd expect it be fine as is - in which case we should be good to go
| assert out.dtype == attention.dtype | ||
|
|
||
|
|
||
| @requires_hopper_or_newer |
There was a problem hiding this comment.
nit:
requires_hopper_or_newer = pytest.mark.skipif(
get_device_compute_capability(0) < 90,
reason="the 4K native JAX baseline requires more than 48 GiB of device memory",
)
Which would technically not be true for this case (1k instead of 4k) so maybe add a single line comment for this here.
| # The native baseline materializes fp32 attention scores, so at the tutorial's 4K | ||
| # sequence its backward pass needs a ~49 GiB buffer, which does not fit when CI runs | ||
| # several test workers per GPU. Compare on a 1K prefix instead (1/16 of the memory). |
There was a problem hiding this comment.
| # The native baseline materializes fp32 attention scores, so at the tutorial's 4K | |
| # sequence its backward pass needs a ~49 GiB buffer, which does not fit when CI runs | |
| # several test workers per GPU. Compare on a 1K prefix instead (1/16 of the memory). | |
| # The native baseline materializes fp32 attention scores, so at the tutorial's 4K | |
| # sequence its backward pass needs a much larger buffer, which does not fit when CI runs | |
| # several test workers per GPU. Compare on a smaller prefix instead |
nit: make this a bit generic with the wording as these numbers can change based on how XLA/JAX decides to build in the future so a general note is fine :)
|
I have not seen that failure on our side but I will bear it in mind if we do come across it. |
Describe the native baseline's backward buffer without a specific size, since the allocation XLA makes can change between versions, and note that the Hopper skip reason refers to the 4K baseline rather than the 1K comparison. Signed-off-by: Alex Y. Chan <alechan@nvidia.com>
The native baseline in the attention tutorial materializes fp32 scores, so at the tutorial's 4K sequence its backward pass needs a ~49 GiB buffer. That does not fit when CI runs several test workers per GPU, and test_bshd_gqa_swa_matches_baseline fails with RESOURCE_EXHAUSTED. Compare TE to the baseline on a 1K prefix instead.
https://github.com/NVIDIA/JAX-Toolbox/actions/runs/42178106901/job/127809077177#step:3:15983
Trace
Description
Please include a brief summary of the changes, relevant motivation and context.
Fixes # (issue)
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: