GEOPE relies on JAX's device parallelism (JIT, vmap) but does not expose any mechanism for running multiple independent optimisations in parallel across CPU cores or accelators
Proposed change:
We should investigate if we can enable sharding of the heavy subroutines. For example, when running on multiple cores the jacobian calculation could be sharded over devices, which would require sharding the individual pulses over all devices. This would also require padding
to ensure that parameters can always be sharded correctly.
Ideally we could rely on single axis sharding and access only `jax.devices()' on single nodes. Multi-node setups would be a long term goal.
Acceptance criteria:
- Geope tests runs correctly with
n_workers=1 (sequential fallback) and n_workers>1.
- Does not break JAX JIT caching within each worker process.
- At least one integration test.
GEOPE relies on JAX's device parallelism (JIT, vmap) but does not expose any mechanism for running multiple independent optimisations in parallel across CPU cores or accelators
Proposed change:
We should investigate if we can enable sharding of the heavy subroutines. For example, when running on multiple cores the jacobian calculation could be sharded over devices, which would require sharding the individual pulses over all devices. This would also require padding
to ensure that parameters can always be sharded correctly.
Ideally we could rely on single axis sharding and access only `jax.devices()' on single nodes. Multi-node setups would be a long term goal.
Acceptance criteria:
n_workers=1(sequential fallback) andn_workers>1.