Linear does not support TP comm overlap for Column Parallel mode

[parthmannan]

All output modes in Linear class only support Row parallel outputs (https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/module/linear.py#L358)

The Linear forward checks for ub_overlap_rs and changes the expected dim size to s/tp, h instead of also supporting column parallel mode which would have a different dim size for example, s, 3*h/tp - https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/module/linear.py#L268

Linear is used by TEColumnParallelLinear in Megatron-LM as well and supporting TP Comm overlap is highly important.
The alternative to use LayerNormLinear was considered but it does not allow skipping normalization either.