Skip to content

Create maxtext gpu package and unit/integration tests #2463

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
Oct 15, 2025
Merged

Create maxtext gpu package and unit/integration tests #2463

merged 1 commit into from
Oct 15, 2025
+377 −5

Conversation

kanglant
Copy link
Collaborator

@kanglant kanglant commented Oct 7, 2025 •
edited
Loading

Description

This PR adds GPU extras to the maxtext package, allowing users to install the necessary dependencies for a GPU environment by running uv pip install -e .[cuda12] --resolution=lowest. It also integrates GPU package tests into the CI workflow.

FIXES: b/450562915

Tests

I've updated the package test and it is passed.

Checklist

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

  • I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
  • 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.

@kanglant kanglant force-pushed the kanglan/gpu_pyproj branch 5 times, most recently from 9db7c83 to f287808 Compare October 9, 2025 21:45
@kanglant kanglant marked this pull request as ready for review October 9, 2025 21:45
bvandermoon
bvandermoon approved these changes Oct 14, 2025
View reviewed changes
@kanglant
Create maxtext gpu package and unit/integration tests
8bf5cb1 
Use exclusion rules instead of direct pinning to disable unsupported JAX versions. This will allow newer JAX versions in future releases without requiring manual updates.
@copybara-service copybara-service bot merged commit 210e9d7 into main Oct 15, 2025
37 checks passed
@copybara-service copybara-service bot deleted the kanglan/gpu_pyproj branch October 15, 2025 07:01
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@bvandermoon bvandermoon bvandermoon approved these changes

@SurbhiJainUSC SurbhiJainUSC SurbhiJainUSC approved these changes

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

@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

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

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

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

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

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

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

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

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

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

@suexu1025 suexu1025 Awaiting requested review from suexu1025 suexu1025 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.

3 participants

@kanglant @SurbhiJainUSC @bvandermoon