Does JAX reduce the total work time required to complete a job if multiple GPUs are used?

[sunvod]

If the calculation work is paralyzed between different GPUs, should I expect the total time to decrease depending on which GPUs are, because the work is spread over different processes or does this not happen? (for example with Trax, can I expect the calculation of a single epoch to scale with the number of GPUs used?). I ask because from some tests it seems to me not to happen.

[sudhakarsingh27]

should I expect the total time to decrease depending on which GPUs are, because the work is spread over different processes

TL;DR: Based on several factors, you may or may not see speedup in a parallel computation. And this has very less to do with JAX itself.

The factors:

1. If the work remains constant (see amdahl's law), the maximum speedup will depend on the serial part in the computation. In deep NNs, the serial part is forward->backward->optimizerUpdate. Within each of those sub-part we could achieve parallelism. But no matter how many devices you have in parallel, that overall sequence of serial computation doesn't change. We usually work around this by increasing the work as well (e.g. increase batch size, or increase features, etc.) while increasing the number of parallel devices.
2. Unless the computation is "embarrassingly parallel" (i.e. no dependency among the parallel pieces of computation), there will be communication costs involved. For example, in deep NNs on parallel devices, this takes the form of all_reduce for gradients on multiple devices (for "batch parallel"). More general case is "tensor parallel" that JAX provides which may involve several other communication primitives adding to overall communication cost.
3. We need to also factor in "preparation phase" (when all the parallel devices have the data they need to do computation) and the "merging phase" (when parallel devices are done with their computation). That means a single "leader" device would place the "slices" of data ("input tensors" in NNs) in all the "replica" devices for them to compute in parallel. Any framework including JAX would do this.