Skip to content

Fix error for upgrading JAX to 0.7.2 #2368

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
Merged
merged 1 commit into from
Sep 20, 2025
Merged

Fix error for upgrading JAX to 0.7.2 #2368

merged 1 commit into from
Sep 20, 2025

Conversation

@NuojCheng
Copy link
Collaborator

@NuojCheng NuojCheng commented Sep 18, 2025 •
edited
Loading

Description

Manually set testing scope VMEM as 64Mb to disable incorrect VMEM OMM error when Jax>0.7.0.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

  • I have performed a self-review of my code.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run end-to-end tests tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed.

@NuojCheng NuojCheng marked this pull request as ready for review September 18, 2025 23:15
@NuojCheng NuojCheng force-pushed the chengnuojin-fix-vmem branch 8 times, most recently from 5e58bb7 to 33e6d1e Compare September 19, 2025 22:28
Rohan-Bierneni
Rohan-Bierneni approved these changes Sep 20, 2025
View reviewed changes
Copy link
Collaborator

@Rohan-Bierneni Rohan-Bierneni left a comment

Choose a reason for hiding this comment

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

LGTM.

gobbleturk
gobbleturk reviewed Sep 20, 2025
View reviewed changes
Copy link
Collaborator

@gobbleturk gobbleturk left a comment

Choose a reason for hiding this comment

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

LGTM, but you could update the description of this PR? Instead of disabling tests we are now setting global VMEM scoped limit on all of these tests

@copybara-service copybara-service bot merged commit 08d9f20 into main Sep 20, 2025
27 checks passed
@copybara-service copybara-service bot deleted the chengnuojin-fix-vmem branch September 20, 2025 06:56
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@gobbleturk gobbleturk gobbleturk left review comments

@RissyRan RissyRan RissyRan approved these changes

@Rohan-Bierneni Rohan-Bierneni Rohan-Bierneni approved these changes

@khatwanimohit khatwanimohit Awaiting requested review from khatwanimohit khatwanimohit is a code owner

@shralex shralex Awaiting requested review from shralex shralex is a code owner

@parambole parambole Awaiting requested review from parambole parambole is a code owner

@richjames0 richjames0 Awaiting requested review from richjames0 richjames0 is a code owner

@bvandermoon bvandermoon Awaiting requested review from bvandermoon bvandermoon is a code owner

@vipannalla vipannalla Awaiting requested review from vipannalla vipannalla is a code owner

@gagika gagika Awaiting requested review from gagika gagika is a code owner

@yangyuwei yangyuwei Awaiting requested review from yangyuwei

@SurbhiJainUSC SurbhiJainUSC Awaiting requested review from SurbhiJainUSC SurbhiJainUSC is a code owner

@hengtaoguo hengtaoguo Awaiting requested review from hengtaoguo hengtaoguo is a code owner

@A9isha A9isha Awaiting requested review from A9isha A9isha is a code owner

@aireenmei aireenmei Awaiting requested review from aireenmei aireenmei is a code owner

@jiangjy1982 jiangjy1982 Awaiting requested review from jiangjy1982 jiangjy1982 is a code owner

Assignees

No one assigned

Labels

pull ready

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@NuojCheng @RissyRan @gobbleturk @Rohan-Bierneni