row wise model parallelism and table wise parallelism

[xiexbing]

Hello,

We are using torchrec and the two types of parallelism in our system. To have a better understanding on the details and seek the best communication primitives for our code, I want to know some details on these parallelisms. e.g., what's the specific communication primitives (operators) used on each step, (e.g., for the metadata transfer, the embedding lookup, and the pooling results, gradient transfer). Please point me to both the code and documents for these.

thanks!

[joshuadeng]

for input dist (sending input features to local rank): we use all-to-all communication
for TW:
https://github.com/pytorch/torchrec/blob/main/torchrec/distributed/sharding/tw_sharding.py#L224
for RW we bucketize the ids to match rows of shard on the local rank: https://github.com/pytorch/torchrec/blob/main/torchrec/distributed/sharding/rw_sharding.py#L245

for embedding output:

• TW we use all-to-all
  • https://github.com/pytorch/torchrec/blob/main/torchrec/distributed/sharding/tw_sharding.py#L262
• RW we use reduce scatter
  • https://github.com/pytorch/torchrec/blob/main/torchrec/distributed/sharding/rw_sharding.py#L331