enter image description here
I try to use 4 GPUs to solve a Physic-informed nerual networks(PINN) problem. And I find that when I use one GPU, the utlization of GPU can reach 100% and the training speed is high (200it/s), but when I use 4 GPUs by implementing sharding and jax.jit strategies as shown in https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html#way-batch-data-parallelism
enter image description here
I shard all my data to different GPUs and replciate the parameters and state of function as shown in guidelines of auto parallelization. And I find my 4 GPUs utilization are quite low, as shown before 10%, and the training speed is quite low (20 it/s).