Question about JAX data parallelism

[knightXun]

Hi Guys, I attempted to train a model using data parallelism with JAX. However, the speedup was a meager 1.14. Below is a portion of my code. Could you help me identify where I went wrong?

...
  @value_and_grad
  def train_forward(params, inputs, bboxes, labels, step): 
    outputs, params = model.run(params, inputs, train=True, rngs={'dropout': step})
    outs = [out.transpose((0,3,1,2)) for out in outputs]
    res = head_loss.run(head_loss_param, outs, bboxes, labels, method='loss_by_feat')
    loss_val, losses= parse_losses(res)
    return loss_val, params

  @jit
  def train_step(params, inputs, bboxes, labels, opt_state, clip_state, step):
    (loss_val, params), grads = train_forward(params, inputs, bboxes, labels, step) 
    new_grads, clip_state = clipper.update(grads['param'], clip_state)
    updates, opt_state = optimizer.update(new_grads, opt_state, params['param'])
    params['param'] = apply_updates(params['param'], updates)
    return params, opt_state, clip_state, loss_val

  devices = mesh_utils.create_device_mesh((2))
  mesh = Mesh(devices, axis_names=('x'))
  data_sharding = NamedSharding(mesh,  P('x', None, None, None))
  bboxes_sharding = NamedSharding(mesh,  P('x', None, None))
  labels_sharding = NamedSharding(mesh,  P('x', None))
  replicated_sharding = NamedSharding(mesh, P())

  start_time = time.time()
  for epoch in range(1):
      inputs, (gt_instances, gt_instances_ignore, batch_img_metas) = preprocess_data(data_batch)
      bboxes = jnp.stack([gt_instance.bboxes for gt_instance in gt_instances])
      labels = jnp.stack([gt_instance.labels for gt_instance in gt_instances])
      inputs = jax.device_put(inputs, data_sharding)
      bboxes = jax.device_put(bboxes, bboxes_sharding)
      labels = jax.device_put(labels, labels_sharding)      
      params = jax.device_put(params, replicated_sharding)
   
      params, opt_state, clip_state, loss_val = train_step(params, inputs, bboxes, labels, opt_state, clip_state, step)

...

[bmaxdk]

Your main issue is that you used jit instead of a parallelizing transform. To truly leverage multi-device parallelism, switch to pmap (for simple data parallelism) or pjit (for SPMD parallelism with sharding). This change will allow the computation to be divided and executed across multiple devices,

ex:

@partial(pmap, axis_name='devices')
def train_step(params, inputs, bboxes, labels, opt_state, clip_state, step):
    (loss_val, params), grads = train_forward(params, inputs, bboxes, labels, step)
    # .....

For pjit, you need to specifying how the data and parameters are sharded.

from jax.experimental.pjit import pjit
from jax.sharding import NamedSharding, Mesh, PartitionSpec as P

@pjit(
    in_shardings=(replicated_sharding, data_sharding, bboxes_sharding, labels_sharding, replicated_sharding, replicated_sharding, replicated_sharding),
    out_shardings=(replicated_sharding, replicated_sharding, replicated_sharding, replicated_sharding)
)
def train_step(params, inputs, bboxes, labels, opt_state, clip_state, step):
    (loss_val, params), grads = train_forward(params, inputs, bboxes, labels, step)

[knightXun]

I tried the pjit, but it has no effect, the speedup is same as the jit

[knightXun]

@bmaxdk After a thorough investigation, I found that the storage medium on my side is relatively slow, and the A100 is based on PCIe. Thanks very much.