[Distributed] hardware=gpu does not correctly configure process-per-node mode with Slurm

[olupton]

Bug report

The logic for distributed initialisation with hardware=gpu:

maxtext/src/maxtext/utils/max_utils.py

Lines 246 to 266 in 37ded59

	def initialize_jax_for_gpu(raw_keys):
	"""Jax distributed initialize for GPUs."""
	if os.environ.get("JAX_COORDINATOR_IP") is not None:
	coordinator_ip = str(os.getenv("JAX_COORDINATOR_IP"))
	coordinator_port = str(os.getenv("JAX_COORDINATOR_PORT"))
	devices = os.getenv("CUDA_VISIBLE_DEVICES")
	if devices is not None:
	try:
	devices = [int(x) for x in devices.split(",")]
	except (ValueError, TypeError) as e:
	max_logging.log(f"Error parsing CUDA_VISIBLE_DEVICES: {e}")
	devices = None
	jax.distributed.initialize(
	coordinator_address=f"{coordinator_ip}:{coordinator_port}",
	num_processes=int(os.getenv("NNODES")),
	process_id=int(os.getenv("NODE_RANK")),
	initialization_timeout=raw_keys["jax_distributed_initialization_timeout"],
	local_device_ids=devices,
	)
	max_logging.log(f"JAX global devices: {jax.devices()}")

Is targeted at running in process-per-node mode (i.e. 1 process driving all GPUs in the machine). However, if CUDA_VISIBLE_DEVICES is not set explicitly then this falls through to auto-detection, which assumes process-per-GPU mode on Slurm. The result is that only the 0th GPU on each node is used.

It would be more user-friendly if -- given that this MaxText code is quite explicitly targeted at process-per-node mode -- it did not require the user to explicitly set CUDA_VISIBLE_DEVICES.

There is also hardware=gpu_multiprocess that defers to JAX's default distributed initialisation and, on a Slurm cluster, correctly yields a process-per-GPU configuration.

Logs/Output

srun --container-image=...--container-remap-root sh -c 'CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 NNODES=2 NODE_RANK=${SLURM_PROCID} JAX_COORDINATOR_PORT=2222 JAX_COORDINATOR_IP=... python3 -m maxtext.trainers.pre_train.train ... hardware=gpu ici_fsdp_parallelism=8 ...'

will work, while

srun --container-image=...--container-remap-root sh -c 'NNODES=2 NODE_RANK=${SLURM_PROCID} JAX_COORDINATOR_PORT=2222 JAX_COORDINATOR_IP=... python3 -m maxtext.trainers.pre_train.train ... hardware=gpu ici_fsdp_parallelism=8 ...'

will yield

AssertionError: Number of devices per slice 1 does not match the product of the ICI parallelism 8

Environment Information

No response

Additional Context

No response

[shralex]

@gabeweisz since you contributed this code, could you please take a look on the request ?

[gabeweisz]

This change is desirable, but I don't think that this behavior is achievable without a change to Jax.

I do not think it is possible to automatically select all GPUs. The reason for this is that, as far as I know, we can't tell how many GPUs are present without calling jax.devices(). However, jax,distributed.initialize() must be called before the GPU backend is initialized, which means that we cannot call jax.devices() to determine how many GPUs are available.

It might be possible to check slurm environment variables like SLURM_JOB_GPUS to see what GPUs are available. Otherwise the behavior of jax.distributed.initialize() in Jax to achieve this. I'll investigate.

[gabeweisz]

Slurm sets the environment variable "SLURM_STEP_GPUS" to list GPUs assigned to the job. I will create a PR to check that if CUDA_VISIBLE_DEVICES is not available. You still will need to pass in SLURM_STEP_GPUs (from the local environment, set automatically by slurm) to your container, but that is because there's no way to detect what GPUs are available from Jax without running jax.devices() AFAIK. Without a container this will all work without setting any extra environment variables

[olupton]

JAX's jax_cuda_visible_devices has some logic for a special "all" value -- can that be used here?

I am not sure that SLURM_STEP_GPUS is reliable.

[gabeweisz]

SLURM_STEP_GPUS is part of the slurm documentation - so if you have evidence that it is not reliable you should take that up with them.

I've confirmed that setting jax_cuda_visible_devices to "all" does not work (at least under ROCm) - we only get one GPU in Slurm regardless of how many are allocated
I did this with:
jax._src.config.update("jax_cuda_visible_devices", "all")
jax._src.config.update("jax_rocm_visible_devices", "all")

I can add this code anyway, but don't have the ability to test it in CUDA. But I can file a bug in ROCm's jax plugin to support "all"

From what I can tell, the special handling for 'all' is to not set visible_devices if jax_cuda_visible_devices is 'all' (as seen here)

This is best effort without changing Jax and XLA - using "CUDA_VISIBLE_DEVICES" is better than not using anything (which would only get you one GPU), and using "SLURM_STEP_GPUS" if "CUDA_VISIBLE_DEVICES" is not set will allow some amount of auto-configuration in Slurm.

[gabeweisz]

Requested change made. Now this code will:

• use CUDA_VISIBLE_DEVICES if set
• use SLURM_VISIBLE_DEVICES if CUDA_VISIBLE_DEVICES is not set and SLURM_VISIBLE_DEVICES is set
• set "jax_cuda_visible_devices" and "jax_rocm_visible_devices" to "all" in jax.config" if neither are set

[gabeweisz]

None of this would be necessary if we could disable this line of code in Jax for GPUs