Restore from single slice

[findmyway]

1. Note that I'm still using an old version (v0.11.15) of slice_devices method. (This means that it will return the devices from single slice instead of single replica)
2. The basic idea is to keep the original implementation almost unchanged. And I created another dimension of replicas to broadcast data.
3. I also tried the original idea of simply restore from single replica. But I get the InvalidShardingError. The reason is obvious, the devices of a process are distributed across different replicas.

orbax/checkpoint/orbax/checkpoint/_src/serialization/type_handlers.py

Lines 1468 to 1474 in 9fc3716

	if primary_replica_ids != expected_primary_replica_ids:
	raise InvalidShardingError(
	'The provided sharding is not valid. The primary replica has the'
	f' following devices: {primary_replica_ids}, but process indices'
	' associated with primary replica devices are expected to be:'
	f' {primary_replica_pids}.'
	)

My questions:

1. Any obvious errors or potential improvements with my current implementation?
   • One datapoint from my latest test: ~30s on deserialization plus ~60s on broadcasting (only one broadcast in total).
2. Any idea on how to address the above InvalidShardingError? (My initial thought is that, the resharding should still work after the sum op even though here's a mismatch.)

[cpgaffney1]

I also tried the original idea of simply restore from single replica. But I get the InvalidShardingError. The reason is obvious, the devices of a process are distributed across different replicas.

Devices from one process belonging to different replicas does violate a fundamental assumption. I suppose the validation just needs to be modified to account for this possibility, if we are saying that it is indeed a possibility.

Ideally unit testing can be improved but it might be tricky to emulate this situation in a unit test. FWIW here are some test cases for SingleReplicaArrayHandler - it's private only because it runs with an internal TPU-based test harness. https://gist.github.com/cpgaffney1/35161a6e6f6e1bc7bf2ffd3df543efe5

[cpgaffney1]

I don't quite understand why this was incorrect. Isn't the contract that the replica_axis_index-th dimension of the mesh should be the replica dimension?

[cpgaffney1]

Ah sorry, realized the intention here. The idea is to always use a single slice to broadcast, even when n_replicas != n_slices?

[findmyway]

Yes

[cpgaffney1]

Unused now?

[findmyway]

Correct

[cpgaffney1]

Is there also a problem with this check?

[findmyway]

This is unrelated to the PR here.

Actually, it may take a relative long time (>500s) in some cases to finish the checking. So I disabled it temporarily.

Ideally we'd also have a config on it.

[findmyway]

Thanks!

Looks like you are using two processes for testing here.

Could you also add the (2, 4) mesh shape in the test below? Make sure the mesh is created from jax.experimental.mesh_utils.create_device_mesh to validate my assumption.

https://gist.github.com/cpgaffney1/35161a6e6f6e1bc7bf2ffd3df543efe5#file-type_handlers_test-L291-L295