[train][jax] Enable Jax trainer on GPU by liulehui · Pull Request #58322 · ray-project/ray

liulehui · 2025-10-30T19:10:48Z

Description

this PR added multihost GPU support for Ray Train JaxTrainer
Following Jax GPU distributed doc: if ScalingConfig.use_gpu == True, we add "cuda" as JAX_PLATFORMS.
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.

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

liulehui · 2025-11-13T22:36:32Z

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

Signed-off-by: Lehui Liu <lehui@anyscale.com>

python/ray/train/v2/jax/jax_trainer.py

python/ray/train/v2/jax/config.py

Signed-off-by: Lehui Liu <lehui@anyscale.com>

1. Jax dependency is introduced in #58322 2. The current test environment is for CUDA 12.1, which limit jax version below 0.4.14. 3. jax <= 0.4.14 does not support py 3.12. 4. skip jax test if it runs against py3.12+. Signed-off-by: Lehui Liu <lehui@anyscale.com>

1. Jax dependency is introduced in ray-project#58322 2. The current test environment is for CUDA 12.1, which limit jax version below 0.4.14. 3. jax <= 0.4.14 does not support py 3.12. 4. skip jax test if it runs against py3.12+. Signed-off-by: Lehui Liu <lehui@anyscale.com>

Jax dependency is introduced in #58322 The current test environment is for CUDA 12.1, which limit jax version below 0.4.14. jax <= 0.4.14 does not support py 3.12. skip jax test if it runs against py3.12+. Signed-off-by: elliot-barn <elliot.barnwell@anyscale.com>

1. this PR added multihost GPU support for Ray Train JaxTrainer 2. Following Jax [GPU distributed doc](https://docs.jax.dev/en/latest/multi_process.html#gpu-example): 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 --------- Signed-off-by: Lehui Liu <lehui@anyscale.com> Signed-off-by: YK <1811651+ykdojo@users.noreply.github.com>

1. this PR added multihost GPU support for Ray Train JaxTrainer 2. Following Jax [GPU distributed doc](https://docs.jax.dev/en/latest/multi_process.html#gpu-example): 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 --------- Signed-off-by: Lehui Liu <lehui@anyscale.com>

1. Jax dependency is introduced in ray-project#58322 2. The current test environment is for CUDA 12.1, which limit jax version below 0.4.14. 3. jax <= 0.4.14 does not support py 3.12. 4. skip jax test if it runs against py3.12+. Signed-off-by: Lehui Liu <lehui@anyscale.com>

Jax dependency is introduced in ray-project#58322 The current test environment is for CUDA 12.1, which limit jax version below 0.4.14. jax <= 0.4.14 does not support py 3.12. skip jax test if it runs against py3.12+. Signed-off-by: elliot-barn <elliot.barnwell@anyscale.com>

## Description 1. Jax dependency is introduced in #58322 2. The current test environment is for CUDA 12.1, which limit jax version below 0.4.14. 3. jax <= 0.4.14 does not support py 3.12. 4. skip jax test if it runs against py3.12+. --------- Signed-off-by: Lehui Liu <lehui@anyscale.com>

## Description 1. Jax dependency is introduced in ray-project#58322 2. The current test environment is for CUDA 12.1, which limit jax version below 0.4.14. 3. jax <= 0.4.14 does not support py 3.12. 4. skip jax test if it runs against py3.12+. --------- Signed-off-by: Lehui Liu <lehui@anyscale.com>

## Description 1. Jax dependency is introduced in ray-project#58322 2. The current test environment is for CUDA 12.1, which limit jax version below 0.4.14. 3. jax <= 0.4.14 does not support py 3.12. 4. skip jax test if it runs against py3.12+. --------- Signed-off-by: Lehui Liu <lehui@anyscale.com> Signed-off-by: jeffery4011 <jefferyshen1015@gmail.com>

#60593) ## Description We added GPU (#58322) and multislice TPU (#58629) support for JaxTrainer, this PR is to update the corresponding docs. ## Additional information 1. tested with `make develop && make local` --------- Signed-off-by: Lehui Liu <lehui@anyscale.com>

## Description 1. Jax dependency is introduced in ray-project#58322 2. The current test environment is for CUDA 12.1, which limit jax version below 0.4.14. 3. jax <= 0.4.14 does not support py 3.12. 4. skip jax test if it runs against py3.12+. --------- Signed-off-by: Lehui Liu <lehui@anyscale.com>

ray-project#60593) ## Description We added GPU (ray-project#58322) and multislice TPU (ray-project#58629) support for JaxTrainer, this PR is to update the corresponding docs. ## Additional information 1. tested with `make develop && make local` --------- Signed-off-by: Lehui Liu <lehui@anyscale.com> Signed-off-by: Sirui Huang <ray.huang@anyscale.com>

#60593) ## Description We added GPU (#58322) and multislice TPU (#58629) support for JaxTrainer, this PR is to update the corresponding docs. ## Additional information 1. tested with `make develop && make local` --------- Signed-off-by: Lehui Liu <lehui@anyscale.com> Signed-off-by: elliot-barn <elliot.barnwell@anyscale.com>

#60593) ## Description We added GPU (#58322) and multislice TPU (#58629) support for JaxTrainer, this PR is to update the corresponding docs. ## Additional information 1. tested with `make develop && make local` --------- Signed-off-by: Lehui Liu <lehui@anyscale.com>

ray-project#60593) ## Description We added GPU (ray-project#58322) and multislice TPU (ray-project#58629) support for JaxTrainer, this PR is to update the corresponding docs. ## Additional information 1. tested with `make develop && make local` --------- Signed-off-by: Lehui Liu <lehui@anyscale.com> Signed-off-by: Adel Nour <ans9868@nyu.edu>

1. this PR added multihost GPU support for Ray Train JaxTrainer 2. Following Jax [GPU distributed doc](https://docs.jax.dev/en/latest/multi_process.html#gpu-example): 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 --------- Signed-off-by: Lehui Liu <lehui@anyscale.com> Signed-off-by: peterxcli <peterxcli@gmail.com>

1. Jax dependency is introduced in ray-project#58322 2. The current test environment is for CUDA 12.1, which limit jax version below 0.4.14. 3. jax <= 0.4.14 does not support py 3.12. 4. skip jax test if it runs against py3.12+. Signed-off-by: Lehui Liu <lehui@anyscale.com> Signed-off-by: peterxcli <peterxcli@gmail.com>

Jax dependency is introduced in ray-project#58322 The current test environment is for CUDA 12.1, which limit jax version below 0.4.14. jax <= 0.4.14 does not support py 3.12. skip jax test if it runs against py3.12+. Signed-off-by: elliot-barn <elliot.barnwell@anyscale.com> Signed-off-by: peterxcli <peterxcli@gmail.com>

## Description 1. Jax dependency is introduced in ray-project#58322 2. The current test environment is for CUDA 12.1, which limit jax version below 0.4.14. 3. jax <= 0.4.14 does not support py 3.12. 4. skip jax test if it runs against py3.12+. --------- Signed-off-by: Lehui Liu <lehui@anyscale.com> Signed-off-by: peterxcli <peterxcli@gmail.com>

ray-project#60593) ## Description We added GPU (ray-project#58322) and multislice TPU (ray-project#58629) support for JaxTrainer, this PR is to update the corresponding docs. ## Additional information 1. tested with `make develop && make local` --------- Signed-off-by: Lehui Liu <lehui@anyscale.com> Signed-off-by: peterxcli <peterxcli@gmail.com>

## Description 1. Jax dependency is introduced in ray-project#58322 2. The current test environment is for CUDA 12.1, which limit jax version below 0.4.14. 3. jax <= 0.4.14 does not support py 3.12. 4. skip jax test if it runs against py3.12+. --------- Signed-off-by: Lehui Liu <lehui@anyscale.com> Signed-off-by: peterxcli <peterxcli@gmail.com>

ray-project#60593) ## Description We added GPU (ray-project#58322) and multislice TPU (ray-project#58629) support for JaxTrainer, this PR is to update the corresponding docs. ## Additional information 1. tested with `make develop && make local` --------- Signed-off-by: Lehui Liu <lehui@anyscale.com> Signed-off-by: peterxcli <peterxcli@gmail.com>

liulehui added the go add ONLY when ready to merge, run all tests label Oct 30, 2025

liulehui added 14 commits November 21, 2025 10:56

test jax trainer

5d3a943

Signed-off-by: Lehui Liu <lehui@anyscale.com>

CUDA support

030e31a

Signed-off-by: Lehui Liu <lehui@anyscale.com>

add unit test

9681874

Signed-off-by: Lehui Liu <lehui@anyscale.com>

set JAX_PLATFORMS automatically

dd877f1

Signed-off-by: Lehui Liu <lehui@anyscale.com>

lint

db73a15

Signed-off-by: Lehui Liu <lehui@anyscale.com>

try to mock jax distributed

cf8a942

Signed-off-by: Lehui Liu <lehui@anyscale.com>

mock jax distributed

dcc757d

Signed-off-by: Lehui Liu <lehui@anyscale.com>

reset to cpu for vanilla tests

b651a8a

Signed-off-by: Lehui Liu <lehui@anyscale.com>

modify

041d9aa

Signed-off-by: Lehui Liu <lehui@anyscale.com>

remove

8e43fc1

Signed-off-by: Lehui Liu <lehui@anyscale.com>

experimental gpu jax

5f7eaf3

Signed-off-by: Lehui Liu <lehui@anyscale.com>

fix

f2d4f48

Signed-off-by: Lehui Liu <lehui@anyscale.com>

some fixs

97e4d18

Signed-off-by: Lehui Liu <lehui@anyscale.com>

shutdown gpu jax distributed

78ba180

Signed-off-by: Lehui Liu <lehui@anyscale.com>

liulehui force-pushed the jax-gpu branch from 54fdcfd to 78ba180 Compare November 21, 2025 19:16

liulehui added 3 commits November 21, 2025 11:19

gpu cuda env var

2b92238

Signed-off-by: Lehui Liu <lehui@anyscale.com>

unit tests

f728a2b

Signed-off-by: Lehui Liu <lehui@anyscale.com>

try to install jax[cuda] for gpu unit test

34c337d

Signed-off-by: Lehui Liu <lehui@anyscale.com>

liulehui marked this pull request as ready for review November 22, 2025 01:58

liulehui requested review from a team, matthewdeng and richardliaw as code owners November 22, 2025 01:58

cursor bot reviewed Nov 22, 2025

View reviewed changes

python/ray/train/v2/jax/jax_trainer.py Show resolved Hide resolved

python/ray/train/v2/jax/config.py Show resolved Hide resolved

python/ray/train/v2/jax/config.py Show resolved Hide resolved

python/ray/train/v2/jax/config.py Show resolved Hide resolved

liulehui added 2 commits November 21, 2025 18:24

try to fix cuda jaxlib

6baa078

Signed-off-by: Lehui Liu <lehui@anyscale.com>

try jax 0.4.23

265a860

Signed-off-by: Lehui Liu <lehui@anyscale.com>

ray-gardener bot added the train Ray Train Related Issue label Nov 22, 2025

liulehui added 2 commits November 22, 2025 12:54

pin to jax 0.4.23

81ae1c0

Signed-off-by: Lehui Liu <lehui@anyscale.com>

pin jax to 0.4.23

9c19af0

Signed-off-by: Lehui Liu <lehui@anyscale.com>

justinvyu enabled auto-merge (squash) November 24, 2025 18:56

richardliaw approved these changes Nov 24, 2025

View reviewed changes

justinvyu merged commit b88bcc1 into ray-project:master Nov 24, 2025
7 checks passed

liulehui mentioned this pull request Nov 25, 2025

[train][jax] Skip py3.12 for jax test #58979

Merged

elliot-barn mentioned this pull request Nov 26, 2025

[train][jax] skipping jax test in test_backend #58994

Merged

liulehui mentioned this pull request Jan 12, 2026

[train][jax] Skip multislice jax tests for py3.12+ #60054

Merged

liulehui mentioned this pull request Jan 29, 2026

[train][docs] update Jax doc to include GPU and multislice TPU support #60593

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

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

[train][jax] Enable Jax trainer on GPU#58322
justinvyu merged 32 commits intoray-project:masterfrom
liulehui:jax-gpu

liulehui commented Oct 30, 2025 •

edited

Loading

Uh oh!

liulehui commented Nov 13, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

liulehui commented Oct 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related issues

Additional information

Uh oh!

liulehui commented Nov 13, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

liulehui commented Oct 30, 2025 •

edited

Loading