Understanding device placement behaviour of pickled DeviceArray

[JossWhittle]

If we pickle a python object that is a jnp.ndarray or contains one as a member variable, and then restore the pickled object, potentially on a different machine, what are the rules on the device placement of the loaded tensors?

If the original tensor is on GPU:0 of the original machine and the new machine also has a GPU is it guaranteed that the loaded tensor will be on the new machines GPU:0?

If the original tensor is on GPU:1 of the original multi-gpu machine, and the new machine has multiple GPUs, are there any guarantees on which GPU (if any) the loaded tensor will be placed on?

If the original tensor is on CPU:0 of the original machine (which may or may not also have one or more GPUs), and the new machine has one or more GPUs is there any guarantee about whether the loaded tensor will be on a CPU or GPU device?

This question also effects how JAX behaves w.r.t. using MPI4PY to communicate tensors over the network, as MPI4PY uses the
pickle protocol to serialize python objects to byte streams.

Thank you for any help. I can experiment empirically, but not for every device and hardware configuration so understanding if and what any guaranteed behaviour is would be really helpful.

[jakevdp]

Pickle is currently only partially supported in JAX. At the moment, JAX pickles arrays by forwarding the value to numpy's pickling mechanism: https://github.com/google/jax/blob/037cc2b8213dac63ed27c91c30139ad0caf0580a/jax/interpreters/xla.py#L1264-L1265
This means that unpickled objects will be numpy arrays rather than JAX arrays:

import jax.numpy as jnp
import pickle

serialized = pickle.dumps(jnp.arange(4))
arr = pickle.loads(serialized)
print(type(arr))
# <class 'numpy.ndarray'>

At this point, if used in a JAX function they will be pushed to the default device, or you can manually push to a device of your choosing using jax.device_put.

[JossWhittle]

Hi Jake, thank you for the quick response.

I have a large hierarchy of class data with JAX tensors dotted throughout the hierarchy and I had hoped to use Pickle for saving and restoring the full state, and for communicating/synchronizing some sub-trees of the hierarchy between multiple compute nodes.

Do you think it would be safe to implement the def __getstate__(self) and def __setstate__(self, state) functions of my classes? Get state fetches both the current device name and the numpy version of the tensor as a tuple and returns it as part of the state tuple instead of the DeviceArray itself. Set state unpacks it and attempts to jax.device_put the numpy array onto the desired device or default device if it does not exist.

Is there a plan for fully supporting pickle in the future?

For anyone else coming to this discussion later I also found the following section in the documentation about device placement rules which was very useful. https://jax.readthedocs.io/en/latest/faq.html#controlling-data-and-computation-placement-on-devices

[jakevdp]

Implementing __getstate__ and __setstate__ seems like a reasonable approach to me. It might be useful to put that into JAX to preserve device upon pickling, though we'd have to do that carefully to avoid the potential to break multiprocessing workflows that implicitly depend on the current pickle behavior.

[JossWhittle]

Initial experiments looked promising but there are some issues with this approach. For complex classes it can become cumbersome to overload __getstate__ and __setstate__ throughout the hierarchy. This also prevents pickling Python container types where the elements are jnp.ndarray directly.

Approaches using wrapper classes that hold a jnp.ndarray or inherit from it while providing __getstate__ and __setstate__ become clunky to use in practice as the unpickled class instance is now of the wrapper type not the underlying jnp.ndarray. Further the use of __slots__ in the jnp.ndarray implementation makes it difficult to dynamically change the type of the wrapper to be a jnp.ndarray during runtime.

I have found a better approach is to implement a JAX-aware Pickler and Unpickler classes which are used at the top-level and transparently handle encoding and decoding JAX tensors whenever they are found in a class hierarchy. This requires no changes to any user or JAX library code, and if Pickle is used as normal then there is no change to the current behaviour expected by JAX.

Suggested implementations of a JAX-aware Pickler and Unpickler.

import os
os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=8'

import numpy as np
import jax
import jax.numpy as jnp
import pickle

class JAXPickler(pickle.Pickler):

    def persistent_id(self, obj):
        
        if isinstance(obj, jnp.ndarray) and not isinstance(obj, np.ndarray):
            tensor = obj
            return ('jax.ndarray', str(tensor.device_buffer.device()), np.array(tensor))
        else:
            return None


class JAXUnpickler(pickle.Unpickler):
    
    def __init__(self, file, devices=None, raise_on_missing_device=True):
        super().__init__(file)
        self._devices = { str(device): device for device in jax.devices() }
        self._raise_on_missing_device = raise_on_missing_device
        
        # devices is a dict of str -> str which remaps the device names from the host where 
        # the object was Pickled to device names where the object is being Unpickled
        if devices is not None:
            
            # Remap device names
            devices_update = dict()
            for pickled_device_name, local_device_name in devices.items():
                if local_device_name not in self.devices:
                    if self.raise_on_missing_device:
                        raise pickle.UnpicklingError(f'device name {local_device_name} is not recognized in {list(self.devices.keys())}')
                    
                    # Use default device
                    devices_update[pickled_device_name] = None
                else:
                    devices_update[pickled_device_name] = self.devices[local_device_name]
                
            # Update device map at once
            self._devices.update(devices_update)
                
    @property
    def devices(self):
        return self._devices
    
    @property
    def raise_on_missing_device(self) -> bool:
        return self._raise_on_missing_device

    def persistent_load(self, obj):
        
        tag, device_name, values = obj
        
        if tag == 'jax.ndarray':
            if device_name not in self.devices:
                if self.raise_on_missing_device:
                    raise pickle.UnpicklingError(f'device name {device_name} is not recognized in {list(self.devices.keys())}')
                
                # Use default device
                device = None
            else:
                device = self.devices[device_name]

            # Recover the JAX tensor
            return jax.device_put(values, device=device)
        else:
            raise pickle.UnpicklingError('unsupported persistent object')

Pickle and Unpickle a jnp.ndarray directly.

tensor = jnp.eye(4)

print('tensor', type(tensor), str(tensor.device_buffer.device()), tensor.dtype)

with open('tensor.pkl', 'wb') as fp:
    JAXPickler(fp).dump(tensor)
    
with open('tensor.pkl', 'rb') as fp:
    tensor = JAXUnpickler(fp).load()
    
print('tensor', type(tensor), str(tensor.device_buffer.device()), tensor.dtype)

tensor <class 'jax.interpreters.xla._DeviceArray'> cpu:0 float32
tensor <class 'jax.interpreters.xla._DeviceArray'> cpu:0 float32

A complex class containing one or more jnp.ndarray.

class Test:
    def __init__(self):
        self.a = jax.device_put(jnp.eye(4), jax.devices()[0])
        self.b = jax.device_put(jnp.eye(4), jax.devices()[1])
        self.c = jax.device_put(jnp.eye(4), jax.devices()[2])
        
    def __repr__(self):
        return f'Test::\n' \
               f'\t{type(self.a)} {str(self.a.device_buffer.device())} {self.a.dtype}\n' \
               f'\t{type(self.b)} {str(self.b.device_buffer.device())} {self.b.dtype}\n' \
               f'\t{type(self.c)} {str(self.c.device_buffer.device())} {self.c.dtype}\n'

Pickle and Unpickle a complex class containing one or more jnp.ndarray.

test = Test()
print(test)

with open('tensor.pkl', 'wb') as fp:
    JAXPickler(fp).dump(test)
    
with open('tensor.pkl', 'rb') as fp:
    test = JAXUnpickler(fp).load()
    
print(test)

Test::
    <class 'jax.interpreters.xla._DeviceArray'> cpu:0 float32
    <class 'jax.interpreters.xla._DeviceArray'> cpu:1 float32
    <class 'jax.interpreters.xla._DeviceArray'> cpu:2 float32

Test::
    <class 'jax.interpreters.xla._DeviceArray'> cpu:0 float32
    <class 'jax.interpreters.xla._DeviceArray'> cpu:1 float32
    <class 'jax.interpreters.xla._DeviceArray'> cpu:2 float32

Pickle and Unpickle a complex class containing one or more jnp.ndarray where device placement is remapped during Unpickling.

test = Test()
print(test)

with open('tensor.pkl', 'wb') as fp:
    JAXPickler(fp).dump(test)
    
with open('tensor.pkl', 'rb') as fp:
    test = JAXUnpickler(fp, devices={ 'cpu:0': 'cpu:1', 'cpu:1': 'cpu:0' }).load()
    
print(test)

Test::
    <class 'jax.interpreters.xla._DeviceArray'> cpu:0 float32
    <class 'jax.interpreters.xla._DeviceArray'> cpu:1 float32
    <class 'jax.interpreters.xla._DeviceArray'> cpu:2 float32

Test::
    <class 'jax.interpreters.xla._DeviceArray'> cpu:1 float32
    <class 'jax.interpreters.xla._DeviceArray'> cpu:0 float32
    <class 'jax.interpreters.xla._DeviceArray'> cpu:2 float32

Default Pickle and Unpickle behaviour currently provided by JAX is maintained.

test = Test()
print(test)

with open('tensor.pkl', 'wb') as fp:
    pickle.dump(test, fp)
    
with open('tensor.pkl', 'rb') as fp:
    test = pickle.load(fp)

print(type(test))
print(type(test.a), test.a.dtype)
print(type(test.b), test.b.dtype)
print(type(test.c), test.c.dtype)

Test::
    <class 'jax.interpreters.xla._DeviceArray'> cpu:0 float32
    <class 'jax.interpreters.xla._DeviceArray'> cpu:1 float32
    <class 'jax.interpreters.xla._DeviceArray'> cpu:2 float32

<class '__main__.Test'>
<class 'numpy.ndarray'> float32
<class 'numpy.ndarray'> float32
<class 'numpy.ndarray'> float32

Thanks again for your help.