Tensor Parallel examples for custom modules

distributed
1

Hi,

Tensor Parallel currently supports only

  • nn.Linear and nn.Embedding via RowwiseParallel and ColumnwiseParallel
  • nn.LayerNorm and nn.Dropout via SequenceParallel

I need to implement tensor parallelization for custom modules and other standard pytorch modules such as nn.Conv1d. However, this currently raises NotImplementedError.

Are there any examples or guidelines for how to implement ColumnwiseParallel and RowwiseParallel for other modules such as nn.Conv1d?

Thanks!

2

Good question. Do you mind sharing the full error log of NotImplementedError and the code to reproduce? This can be a feature request in Sign in to GitHub · GitHub
TP/SP are implemented for matmul / nn.Linear right now. The behavior for nn.Conv1d is not well defined

3

I am referring to this NotImplementedError for anything other than nn.Linear and nn.Embedding: pytorch/torch/distributed/tensor/parallel/style.py at main · pytorch/pytorch · GitHub

I have created an issue regarding common modules such as convolutions and groupnorm here: Tensor parallel for convolutions and groupnorm · Issue #133221 · pytorch/pytorch · GitHub

The implementation of tensor parallel covers really just a tiny fraction of common modules, which makes tensor parallel currently not applicable to most models other than llama. An alternative is to use the Megatron library, but I would prefer to directly use Pytorch distributed.

1 Like
4

I see all the related POCs are tagged in the issue. Thanks for explaining the context clearly