mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

[jax.distributed] Allow setting local device ids via env var

Browse Source
This commit is contained in:
Georg Stefan Schmid 2024-07-10 10:15:35 +00:00
parent a8b425cac5
commit f9bc4c643b
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
jax/_src/distributed.py
View File

@ -45,10 +45,12 @@ class State:
initialization_timeout: int = 300,
coordinator_bind_address: str | None = None):
coordinator_address = (coordinator_address or
os.environ.get('JAX_COORDINATOR_ADDRESS', None))
os.environ.get('JAX_COORDINATOR_ADDRESS'))
if isinstance(local_device_ids, int):
local_device_ids = [local_device_ids]
if local_device_ids is None and (env_ids := os.environ.get('JAX_LOCAL_DEVICE_IDS')):
local_device_ids = list(map(int, env_ids.split(",")))
(coordinator_address, num_processes, process_id, local_device_ids) = (
clusters.ClusterEnv.auto_detect_unset_distributed_params(