Fix gather and reduce scatter ops on sequence dimension

[bclyang]

Fix issue where gather and reduce scatter operations are not correct when seq_dim != 0. This is because reduce_scatter_tensor and gather_tensor only work correctly across the primary dimension of the tensor. This PR switches to use reduce_scatter and gather directly. This PR also adds support for no_weight_tying=False when sequence_parallel=True

Tests

After this fix, the 410M test config converges correctly when sequence_parallel=True and sequence_parallel=False on EnWiki8 on 2 GPUs without weight tying:

Memory consumption is lower on WandB:

Works with weight tying:

Works with pipe_parallel_size=0 and pipe_parallel_size=1:

Works with use_qk_layernorm=True and use_qk_layernorm=False:

Works with MoE with moe_type="deepspeed"and moe_num_experts=2, but have not tested with any other settings including moe_type="megablocks", moe_expert_parallel_size > 1 or enable_expert_tensor_parallelism=True. There's a lot of settings here, so for now, recommend not supporting.

[haileyschoelkopf]

Great finds!