[Distributed] Use Tensor Parallel instead of Sequence Parallel

[kwen2501]

SP is one step further than TP in that it further distributes the layer norm computation.

There are a couple reason to turn the dial back to pure TP:
(1) LLM inference has a prefill phase and a decoding phase which have different seqlen. The decoding phase has a seqlen of 1, to which SP cannot be applied. We don't want to create two models and apply SP and TP separately.
(2) The major motivation of SP is to reduce activation envelope. While this is important in training (bc backward needs those intermediates), we are using torch.no_grad() in inference, in which case activations are not kept anyway.
(3) While it is true that AllReduce = AllGather + ReduceScatter, in small-size region latency dominates, so launching 1 collective may be better than launching two collectives.

• Removed SP feature from model
• Removed // sp_degree in PP activation size

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchchat/1160

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit 611de83 with merge base 4774eaf ():

NEW FAILURE - The following job has failed:

• pull / test-mps / macos-job (gh)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[lessw2020]

Thanks for the detailed justification for this. Looks good!