some fixs

Signed-off-by: Lehui Liu <[email protected]>

Author: liulehui

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

@@ -20,6 +20,7 @@
 @dataclass
 class JaxConfig(BackendConfig):
     use_tpu: bool = False
+    use_gpu: bool = False
 
     @property
     def backend_cls(self):
@@ -108,6 +109,7 @@ def on_start(self, worker_group: WorkerGroup, backend_config: JaxConfig):
                     master_addr_with_port=master_addr_with_port,
                     num_workers=len(worker_group),
                     index=i,
+                    resources_per_worker=worker_group.get_resources_per_worker()
                 )
             )
         ray.get(setup_futures)

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

@@ -132,6 +132,7 @@ def __init__(
         if not jax_config:
             jax_config = JaxConfig(
                 use_tpu=scaling_config.use_tpu,
+                use_gpu=scaling_config.use_gpu,
             )
         super(JaxTrainer, self).__init__(
             train_loop_per_worker=train_loop_per_worker,