Controlling Eagerness of Data Transfer with jnp.Array

[ntenenz]

Hello,

I'm looking to process a moderately large amount of data (fits in host RAM but not GPU memory). As a result, a scan seemed like a promising approach, processing batchwise to limit the amount of data resident in GPU RAM at any given time. However, I noticed that with a scan operation, all data is immediately transferred to the device instead of on-demand.

A minimal example:

import jax
import jax.numpy as jnp

jax.config.update("jax_array", True)
jax.config.update("jax_transfer_guard", "log_explicit")

with jax.default_device(jax.devices("cpu")[0]):
    x = jnp.ones((8192, 1024, 1024))
    
def foo(_, x_):
    return None, x_.sum()

jax.lax.scan(foo, None, x)

2022-12-07 11:04:33.138949: W external/org_tensorflow/tensorflow/compiler/xla/python/transfer_guard_lib.cc:110] host-to-device transfer: type=<class 'numpy.ndarray'>, shape=(), dtype=float32, dst_device=TFRT_CPU_0
2022-12-07 11:04:36.686053: W external/org_tensorflow/tensorflow/compiler/xla/python/transfer_guard_lib.cc:125] device-to-device transfer: shape=(8192, 1024, 1024), dtype=float32, device=TFRT_CPU_0, dst_device=gpu:0
2022-12-07 11:05:01.030716: W external/org_tensorflow/tensorflow/compiler/xla/python/transfer_guard_lib.cc:140] device-to-host transfer: shape=(8192,), dtype=float32, device=gpu:0

(None,
Array([1048576., 1048576., 1048576., ..., 1048576., 1048576., 1048576.], dtype=float32))

Is possible currently, or are there plans in the future, to signal to XLA that it should be a bit less eager in the host->memory transfer?

[ntenenz]

Sorry for bumping, just wondering if this is something being considered. Leveraging the new sharding API, solving this would theoretically allow one to efficiently hoist iterative solutions to larger-than-memory problems into XLA.