Support loading straight to device for JAX / FLAX

[adonath]

Feature request

It seems loading straight to a device is currently only supported for PyTorch. I would be great to have the support for JAX as well.

Motivation

Loading arrays into JAX currently goes via Numpy which forces the arrays into CPU memory and the creates a copy to load into the device memory. This introduces additional overhead.

Your contribution

As a temporary workaround I think one can to load the data using PyTorch and then donate the buffers to a JAX array using the dlpack protocol (which enables buffer donation on GPUs). This requires PyTorch to be installed.

[Narsil]

Indeed that would be great. However I cannot find any jax function loading directly from memory pointer (or assimilated) which would be necessary to load directly on device.

https://docs.jax.dev/en/latest/search.html?q=load

If there are alternatives API that could be used to skip CPU, I'd be glad to hear of it. (nothing really skip the CPUs really actually, but we could skip the kernel or an extra CPU copy).

[adonath]

Thanks @Narsil! I am not aware of any direct methods either. However my understanding is that PyTorch as well as JAX support the (DLPack exchange mechanisms)[https://data-apis.org/array-api/latest/design_topics/data_interchange.html#design-topics-data-interchange--page-root] defined by the Python array API. This offers support for zero copy data exchange in the case of CUDA devices and the following limitation:

Use of DLPack requires that the data can be represented by a strided, in-memory layout on a single device

So there would not be any support for distributed ("sharded") arrays and TPUs, but that's at least a starting point. On further investigation, it should actually work "out of the box" along the lines of:

import torch
from jax import numpy as jnp

array = torch.ones((1000, 1000, 1000), device="cuda:3")

array_jax = jnp.asarray(array)

print(array_jax.device)

# Modify the original PyTorch tensor
array[0, 0, 0] = 999.0

# Check if JAX array was affected
print(f"PyTorch value: {array[0, 0, 0].item()}")
print(f"JAX value: {array_jax[0, 0, 0].item()}")

Which prints:

cuda:3
PyTorch value: 999.0
JAX value: 999.0

So for now it is maybe just a question of documenting the following pattern:

from safetensors import safe_open
from jax import numpy as jnp

with safe_open(filename, framework="pt", device="cuda:0") as f:
    jax_array = jnp.asarray(f.get_tensor(...))

The jnp.asarray should default to non-copy on CUDA devices. But maybe it requires certain versions of PyTorch and JAX. I haven't checked.

[Narsil]

Well the problem is that depending on both torch and Jax is usually not a great pattern as far as I'm aware, whereas numpy is a dependency (or very close to being one given the doc I'm reading).

I think the example you've shown here is great and more people can find it.

I'm not sure we can put it into the library as adding extra dependencies "just" for that is imho an anti pattern.

[adonath]

Well the problem is that depending on both torch and Jax is usually not a great pattern as far as I'm aware, whereas numpy is a dependency (or very close to being one given the doc I'm reading).

I agree. Most users make the choice between either PyTorch and JAX.

The example I gave above could be turned into single paragraph in the docs. Not even showing the code, but just mentioning the device buffer donation as a possibility.