Open johannahaffner opened 2 months ago
PS: bit random thought, will try later: should I shard the residuals function instead?
Ah, I think this is something JAX still needs better docs for.
In response to the various points you raise:
filter_vmap
definitely shouldn't do anything to do with sharding. We use vmap to describe the logical computation, and sharding to describe the physical computation -- these two things are deliberately held separate. (Combining them into one is pretty much what {filter_}pmap
does, and there's a reason that JAX is moving away from that!)
The recommended approach to this these days is to write out your logical computation as if it was all running on one big device -- using vmaps and whatever else necessary to express your computation. Then, when you call it, pass in sharded arrays to describe where you would like the computation to be performed physically.
As such, best practice is pretty much what's currently shown in the Equinox parallelism tutorial. You say that you seem to be fitting the average time series -- this sounds like a bug in your logical computation, and as such would be unrelated to sharding (which is the physical computation). So I think tracking that down is probably the main thing you need to do.
In terms of your example code: I think you've forgotten to vmap your sharded estimator. I only see you doing a vmap in the unsharded case.
Personally I've always found JAX's multi-CPU parallelism to be fairly dodgy. This has never been very well supported in JAX's backend, and indeed sharding is mostly used to express multi-GPU or multi-TPU parallelism. Logic bugs aside, it's entirely possible that you'll see no speedup / that it only uses one core.
Hi Patrick,
thank you for your thoughtful response! I always learn so much from your explanations.
I'm trying to whittle it down to a smaller thing. So far I noticed that it does not raise an error if I do not shard the inputs, but only the vmapped estimator. Then the output is what I would expect. My toy model is too small to check if it is really running on all cores in that case, though.
I'm trying to get it to work on the non-toy example, but that adds quite a bit of complexity and I'm not there yet :)
Hi again! I have yet another question, this time to do with sharding.
Essentially, I'm trying to shard to improve performance, but I would like the output to be the same as that of
vmap
.Background: I currently vmap to parallelise across my data, and I want to get a single value for each input array. (I do parameter estimation for ODEs across individuals.) I noticed that this approach only ever uses about 20 % of my CPU, and confirmed that it essentially runs on a single CPU core. I'd like to change that to multicore CPU, and then scale to GPUs.
What I did so far:
eqx.filter_vmap
and I'm wondering if this could be broadened to allow for something like thePartitionSpecs
options JAX has, or if I should use decorators on top ofeqx.filter_shard
, similar to how it is done in the JAX shard map tutorial, or not use equinox at all.pmap-of-vmaps
, but I'd like to avoid that since I'm not sure whether pmap will be deprecated.Here is some code that replicates the behavior I am talking about. Apologies for length, I tried to condense it as much as I could.