pmap to single GPU
#22961
Replies: 1 comment
-
pmap is in maintenance mode. We recommend using
Nope. |
Beta Was this translation helpful? Give feedback.
0 replies
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
-
Hi,
I know that
pmap
is used to parallelize a function over multiple instances of GPU.Is it possible to do so with single instance too?
I mean - to consider the single instance of the GPU as multiple instances (similar to what
xla_force_host_platform_device_count
does)?Assuming the memory footprint of that function is small and multiple instances of it can fit the GPU - the performance is going to be better?
EDIT:
I am reading about the Jax sharding api (which is amazing), to try figure out myself an answer.
Meanwhile, I want to add another question which is related to my original one-
Is sharding capabilities relevant for me in any way if I am using only one gpu? (again, if we can consider one gpu as multiple instances, maybe it is?)
Beta Was this translation helpful? Give feedback.
All reactions