accelerate sim_matrix process in multi-GPU

[zsnoob]

I edit two main things:

1. Deleting the "loss.mean()" that do nothing. DDP provides automatically gradient synchronization.

2. Refer to this comment, Batch Sharding Details openai/CLIP#132 (comment) we will do every similarity calculation locally. This will use all negative samples in global batch and positive samples in local batch, so local sim_matrix will be shaped in (batch_size / n_gpu, batch_size).

• But this approach will cause another question that if the loss function always put the first local-batch-size columns as the diagonal elements, and local batch is not the first in global, the correct postive samples will locating at column range: local_rank * local_batch_size - (local_rank + 1) * local_batch_size. So i give the second parameter in torch.diag() which means the first positive sample's column.

By experiments, the model can converge as usual, and more efficient.

[zsnoob]

Also mentioned in this issue, #101 (comment).