Get all shards, even copies of replicated shards, to ensure consistency (GDA or jax.Array)

[dlwh]

Hi,

Sorry if this is discussed elsewhere, but I couldn't immediately find it.

I'm debugging some sharded data loading code and want to see if different hosts that should be loading the same data are in fact loading the same data (and inverse: hosts that should be loading different data are in fact loading different data)

ATM, the code is still GDA-based (GDA.from_callback), but I'm happy to do the migration to jax.Array (which I've been putting off).

So the code is basically array = GlobalDeviceArray.from_callback(shape, mesh, pspec, local_data)

I've seen process_all_gather and copy_to_host_async, but they reasonably don't grab replicas

My current idea is to do something like: convert to jax.Array, loop through the global_shards, get the addressable buffer if it exists (doing something reasonable for the ones that don't exist), then do essentially a tiled process_all_gather on each shard index, and do my checks after that.

Is there anything better/preexisting?

[yashk2810]

So if I understand your question correctly, you have data that is replicated along hosts but sharded along the devices in the host?

If that's the case, you can just pull them to host via np.asarray(arr) and then use multihost_utils.assert_equal to check if each host has the same value?

Note that you will need jax and jaxlib 0.4.7 for this.

[dlwh]

i think i'm just being slow and maybe misunderstand what you mean by testing indices, and i appreciate your patience

Let's say I have:

shape = [32]
 jax.make_array_from_callback(shape, jax.sharding.MeshPspecSharding(mesh, P(None)), lambda index: jnp.full(shape, jax.process_index()))

Then I think I promised to jax that I created a fully replicated array with these buffers, but actually the data is different on all hosts (right?). If I just use index checks, I can't discover this bug?

(Of course, it's not actually full(process_index()) but something more or complex that depends on the process_index and the amount of data/tensor parallelism, but I want to be able to catch this case)

[dlwh]

(random: it'd be nice if broadcast_one_to_all supported a pytree prefix for is_source? would that be something y'all would be interested in my pr'ing?)

[yashk2810]

If I just use index checks, I can't discover this bug?

Ohh I see, yeah in that case checking the data is correct. I was hoping that there are not shenanigans like this happening 🤣

nice if broadcast_one_to_all supported a pytree prefix for is_source?

Yeah, sounds good. I am a bit unclear on what you mean here but I guess the PR will make it clear! Feel free to send it my way :)

[dlwh]

any issue would almost certainly be my fault!

Thanks for the help!

[dlwh]

Mostly as an fyi to posterity, here's my attempt. I did some light checking and it covers all my cases.

def check_sharded_consistency(tree: PyTree[ArrayImpl], check_disjoint_indices_are_different: bool = False):
    """Checks the following consistency conditions on an array:
    - all replicas have the same data
    - if check_disjoint_indices_are_different is True, then all shards with disjoint indices have different data
    """
    # index is a tuple[slice, ...], slices are obnoxiously not hashable so we have to convert to tuple

    def check_array(array: ArrayImpl):
        def _to_tuple(index: Tuple[slice, ...]) -> Tuple[Tuple[int, int], ...]:
            my_indices: Tuple[Tuple[int, int], ...] = tuple(
                s.indices(axis_size)[0:2] for axis_size, s in zip(array.shape, index)
            )

            return my_indices

        replicas_by_index = defaultdict(list)
        for shard in array.global_shards:
            replicas_by_index[_to_tuple(shard.index)].append(shard)

        # global shards is not necessarily sorted consistently, so we have to sort the indices
        sorted_indices = sorted(replicas_by_index.keys())

        # ok now get canonical versions of each index
        replica_0_arrays = {}

        for index in sorted_indices:
            shards = replicas_by_index[index]
            try:
                leader = next(s for s in shards if s.replica_id == 0)
            except StopIteration:
                raise ValueError("No replica 0 for index", index)

            data = leader.data
            if data is None:
                shard_shape = [s[1] - s[0] for s in index]
                data = jnp.zeros(shard_shape, dtype=array.dtype)

            replica_0_array = multihost_utils.broadcast_one_to_all(data, is_source=leader.data is not None)
            replica_0_arrays[index] = replica_0_array

        for shard in array.addressable_shards:
            replica_0_array = replica_0_arrays[_to_tuple(shard.index)]
            assert shard.data is not None

            if not jnp.array_equal(shard.data, replica_0_array, equal_nan=True):
                raise ValueError("Shard data does not match replica 0 data", shard, replica_0_array)

            if check_disjoint_indices_are_different:
                for other_index, other_array in replica_0_arrays.items():
                    if other_index == _to_tuple(shard.index):
                        continue

                    if shard.index != other_index:
                        if jnp.array_equal(shard.data, other_array, equal_nan=True):
                            raise ValueError(
                                "Shard data is the same as another shard with disjoint indices", shard, other_array
                            )

    for leaf in jax.tree_util.tree_leaves(tree):
        check_array(leaf)