Changing RMS layer norm to accept DTensors.

[kolehma8]

Summary

Changing RMS layer norm to accept DTensors.

Details

RMS layer norm parameters are NOT sharded under typical tensor parallelism implementations but the inputs (and gradients) may become sharded (DTensors). This PR resolves the issue by gathering the input tensors and performing the full compute in each device.

Relates Issues:

• Tensor Parallel support #826
• Support PyTorch Tensor Parallelism #868

Testing Done

• Hardware Type:
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

[shimizust]

Can you run the example e2e training script just using FSDP2 (since TP won't work until all kernels are tp-compatible) to ensure the perf/loss is correct?

[shimizust]

For my understanding, pytorch native TP keeps activations as DTensors and lets subsequent ops decide what to do?

[kolehma8]

That is my understanding as well.

[shimizust]

Suggest renaming to infer_comm_backend

[kolehma8]

Changed as suggested.

[kolehma8]

Can you run the example e2e training script just using FSDP2 (since TP won't work until all kernels are tp-compatible) to ensure the perf/loss is correct?

I did not see any benchmarks using FSDP in the scripts folder. I think we can say to a very high degree of confidence that this PR will not impact FSDP since if it did, the op would have crashed earlier (e.g. if any inputs are DTensors). Vice versa, this PR only impacts input DTensors and behavior for regular tensors remains unchanged.