Sorry about the delay here.



 aclifton314:

Also, have I called mp.spawn correctly? I have args=(world_size,) but calc_cos_sims(rank, world_size) has the rank first followed by the world size and those go into dist.init_process_group('gloo', rank=rank, world_size=world_size).


Yes this is correct.

I read the link you posted about aggregating, but I’m not entirely sure how to implement it. How would that be done in this case?

It really depends on the type of aggregation you want to do across cosine_tensor. You basically have a total of 4 cosine_tensor (one on each rank) and you want to aggregate them. For example if you want to sum them all up and get the total sum on each rank you can do this:
dist.all_reduce(cosine_tensor)
# Now cosine_tensor would be overwritten to contain the sum of all the 4 `cosine_tensor`.

This tutorial will probably give a much better overview of how to use all of these collective APIs: Writing Distributed Applications with PyTorch — PyTorch Tutorials 1.9.1+cu102 documentation 3