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

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

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

wants to merge 1 commit into from

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 from 5f12314 to 3e3939c Compare January 14, 2025 21:07
@mgoldfarb-nvidia
Consolidate the distributed fused attention tests to shared input gen…
d063b87 
…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 3e3939c to d063b87 Compare January 14, 2025 21:12
@mgoldfarb-nvidia
Copy link
Collaborator Author

/te-ci jax L1

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@zlsh80826 zlsh80826 zlsh80826 left review comments

At least 1 approving review is required to merge this pull request.

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