Thanks for providing a code example. IIUC, you want to just compute a particular math function in parallel and not really train a model in data parallel fashion where you aggregate gradients across ranks in the backward pass? If so, you don’t really need to use DistributedDataParallel.



 aclifton314:

Once all the processes are finished, mp.spawn() will aggregate the results. Is my understanding correct?


No, each process would have its results local to that process. You can collect all the results in one process using a collective operation like all_gather. See: Distributed communication package - torch.distributed — PyTorch 1.9.0 documentation 1


Is the entire dataset read in 4 times with the tmp_df = pd.read_pickle('./embed_pairs_df_million.pkl') line? If so, is there a more efficient way to read it in for this setup?


Yes, it is read 4 times. I’m not very familiar with pandas, but if you can pre-split the data before hand and read only what is required on each process then you would avoid reading the data 4 times.


The forward() method returns the cosine similarity (or it will once I write it) between two embeddings. If calc_cos_sims() is copied to each process, would I need to replace the mp.spawn() line with all_cos_sims = mp.spawn() in order to store the results from all the GPUs?


As mentioned above, mp.spawn will not aggregate results for you. You will have to use something like all_gather to do that.