Skip to content

[train][jax] Enable Jax trainer on GPU #58322

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
Open
wants to merge 32 commits into
base: master
Choose a base branch
from
Open

[train][jax] Enable Jax trainer on GPU #58322

wants to merge 32 commits into from
+128 −11

Conversation

@liulehui
Copy link
Contributor

@liulehui liulehui commented Oct 30, 2025 •
edited
Loading

Description

  1. this PR added multihost GPU support for Ray Train JaxTrainer
  2. Following Jax GPU distributed doc: if ScalingConfig.use_gpu == True, we add "cuda" as JAX_PLATFORMS.
  3. if cuda is the jax platform, add CUDA_VISIBLE_DEVICES and initialize jax distributed with https://docs.jax.dev/en/latest/_autosummary/jax.distributed.initialize.html#jax.distributed.initialize

Related issues

Additional information

Optional: Add implementation details, API changes, usage examples, screenshots, etc.

  1. Tested with script here: https://gist.github.com/liulehui/b0b25065d48b730f2898b712aa92e06e

@liulehui liulehui added the go add ONLY when ready to merge, run all tests label Oct 30, 2025
@liulehui
Copy link
Contributor Author

liulehui commented Nov 13, 2025

jax gpu image build on anyscale platform: https://gist.github.com/liulehui/bda2419e1b3245d40d8027053a8dd26c

@liulehui
test jax trainer
5d3a943 
Signed-off-by: Lehui Liu <[email protected]>
@liulehui
CUDA support
030e31a 
Signed-off-by: Lehui Liu <[email protected]>
@liulehui
add unit test
9681874 
Signed-off-by: Lehui Liu <[email protected]>
@liulehui
set JAX_PLATFORMS automatically
dd877f1 
Signed-off-by: Lehui Liu <[email protected]>
@liulehui
lint
db73a15 
Signed-off-by: Lehui Liu <[email protected]>
@liulehui
try to mock jax distributed
cf8a942 
Signed-off-by: Lehui Liu <[email protected]>
@liulehui
mock jax distributed
dcc757d 
Signed-off-by: Lehui Liu <[email protected]>
@liulehui
reset to cpu for vanilla tests
b651a8a 
Signed-off-by: Lehui Liu <[email protected]>
@liulehui
modify
041d9aa 
Signed-off-by: Lehui Liu <[email protected]>
@liulehui
remove
8e43fc1 
Signed-off-by: Lehui Liu <[email protected]>
@liulehui
experimental gpu jax
5f7eaf3 
Signed-off-by: Lehui Liu <[email protected]>
@liulehui
fix
f2d4f48 
Signed-off-by: Lehui Liu <[email protected]>
@liulehui
some fixs
97e4d18 
Signed-off-by: Lehui Liu <[email protected]>
@liulehui
shutdown gpu jax distributed
78ba180 
Signed-off-by: Lehui Liu <[email protected]>
@liulehui
gpu cuda env var
2b92238 
Signed-off-by: Lehui Liu <[email protected]>
@liulehui
unit tests
f728a2b 
Signed-off-by: Lehui Liu <[email protected]>
@liulehui
try to install jax[cuda] for gpu unit test
34c337d 
Signed-off-by: Lehui Liu <[email protected]>
@liulehui liulehui marked this pull request as ready for review November 22, 2025 01:58
@liulehui liulehui requested review from a team, matthewdeng and richardliaw as code owners November 22, 2025 01:58
cursor[bot]
cursor bot reviewed Nov 22, 2025
View reviewed changes
@liulehui
try to fix cuda jaxlib
6baa078 
Signed-off-by: Lehui Liu <[email protected]>
@liulehui
try jax 0.4.23
265a860 
Signed-off-by: Lehui Liu <[email protected]>
@ray-gardener ray-gardener bot added the train Ray Train Related Issue label Nov 22, 2025
@liulehui
pin to jax 0.4.23
81ae1c0 
Signed-off-by: Lehui Liu <[email protected]>
@liulehui
pin jax to 0.4.23
9c19af0 
Signed-off-by: Lehui Liu <[email protected]>
cursor[bot]
cursor bot reviewed Nov 22, 2025
View reviewed changes
cursor[bot]
cursor bot reviewed Nov 23, 2025
View reviewed changes
@liulehui
align with compiled
012e6ba 
Signed-off-by: Lehui Liu <[email protected]>
@liulehui
fix logging
ca15115 
Signed-off-by: Lehui Liu <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@cursor cursor[bot] cursor[bot] left review comments

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

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

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

Assignees

No one assigned

Labels

go add ONLY when ready to merge, run all tests train Ray Train Related Issue

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@liulehui