[Bugfix] Fix for sharding TP only

[zarzen]

Description

• When TP is enabled and PP is disabled, the last linear layer is sharded. Thus, we need ParallelCrossEntropy as the loss instead of the loss from HF.
• Using TP with ZeRO-1 optimizer in DS runtime requires mpu as input to the deespeed.initialize(...) function.
• When --disable_pipeline is used for disabling pipeline parallelism, we need to get the num_mp parameters from the args.

Known issue

• With the fix, the loss converge to about 7.x after 200 steps, which is still slower than ZeRO-1 or ZeRO-3 only.
• For now only tested with GPT model.

Checklist

• PR's title starts with a category (e.g. [Bugfix], [Model], [Tutorial], etc)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage
• Code is well-documented

[comaniac]

Otherwise LGTM.

[comaniac]

Thanks @zarzen

[comaniac]

Found that this PR breaks 3D scripts (which are not covered in CI yet). @zarzen please file a fix.