Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[JAX] Consolidate the distributed fused attention test code #1405

Merged
merged 2 commits into from
Jan 17, 2025
Merged

[JAX] Consolidate the distributed fused attention test code #1405

merged 2 commits into from
Jan 17, 2025
+323 −392

Conversation

mgoldfarb-nvidia
Copy link
Collaborator

Consolidate test code for distributed fused attention to used FusedAttnRunner so we can share the logic that generates the sequences and masking.

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refractor

Changes

Updates to fused attention tests.

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

@mgoldfarb-nvidia
Copy link
Collaborator Author

/te-ci jax L1

zlsh80826
zlsh80826 reviewed Jan 13, 2025
View reviewed changes
tests/jax/test_distributed_fused_attn.py Outdated Show resolved Hide resolved
tests/jax/test_fused_attn.py Outdated Show resolved Hide resolved
@mgoldfarb-nvidia mgoldfarb-nvidia force-pushed the mgoldfarb-nvidia/consolidate_fused_attn_tests_code branch from fe2b02b to 7eb6428 Compare January 13, 2025 23:01
@mgoldfarb-nvidia
Copy link
Collaborator Author

/te-ci jax L1

@zlsh80826
Copy link
Collaborator

LGTM, there are some failures on the CI. We can merge it once the pipeline clean.

@mgoldfarb-nvidia
Copy link
Collaborator Author

/te-ci jax L1

@mgoldfarb-nvidia mgoldfarb-nvidia force-pushed the mgoldfarb-nvidia/consolidate_fused_attn_tests_code branch 2 times, most recently from 3e3939c to d063b87 Compare January 14, 2025 21:12
@mgoldfarb-nvidia
Copy link
Collaborator Author

/te-ci jax L1

@mgoldfarb-nvidia mgoldfarb-nvidia force-pushed the mgoldfarb-nvidia/consolidate_fused_attn_tests_code branch from d063b87 to 561a323 Compare January 16, 2025 17:50
@mgoldfarb-nvidia
Consolidate the distributed fused attention tests to shared input gen…
b2ef2cb 
…eration and execition logic.

Signed-off-by: Michael Goldfarb <[email protected]>
@mgoldfarb-nvidia mgoldfarb-nvidia force-pushed the mgoldfarb-nvidia/consolidate_fused_attn_tests_code branch from 561a323 to b2ef2cb Compare January 16, 2025 17:51
@mgoldfarb-nvidia
Copy link
Collaborator Author

/te-ci jax L1

zlsh80826
zlsh80826 approved these changes Jan 17, 2025
View reviewed changes
Copy link
Collaborator

@zlsh80826 zlsh80826 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM!

@mgoldfarb-nvidia mgoldfarb-nvidia merged commit 6e84892 into NVIDIA:main Jan 17, 2025
14 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@zlsh80826 zlsh80826 zlsh80826 approved these changes

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@mgoldfarb-nvidia @zlsh80826