Distributed Sharded Layernorm/RMSNorm

[kpaigwar]

Summary

This PR introduces sharded distributed LayerNorm for pre/post-all-gather stages, leveraging existing single device sharded LayerNorm host code. The implementation includes new kernels and comprehensive test coverage.

Key Changes

1. Implemented sharded distributed LayerNorm for pre/post-all-gather stages

   • Utilized existing single device sharded LayerNorm host code
   • Ensures compatibility with current implementation

2. Performance and Accuracy

   • No impact on performance or accuracy of existing single device sharded LayerNorm

3. New Kernel Development

   • Added new kernels specifically for pre/post-all-gather operations

4. Test Coverage

   • Implemented test coverage for pre/post-all-gather on single device
   • Added end-to-end (e2e) test coverage on galaxy (TG)

Performance Benchmarks

Full Activation Shape : [1, 1, 32, 8192]
Activation and Weights dtype : ttnn.bfloat16

Stages	Norm Type	Num Distributed Devices	Num Cores	FW Duration(ns)
pre_all_gather	RMSnorm	4	32	6855
post_all_gather	RMSnorm	4	32	6396
pre_all_gather	Layernorm	4	32	9441
post_all_gather	Layernorm	4	32	7944

Future Steps

1. Add normalization support for batch_size > 32
2. Create functions for CB creation to improve readability

Checklist

• Post commit CI passes https://github.com/tenstorrent/tt-metal/actions/runs/10886636636
• New/Existing tests provide coverage for changes

[TT-BrianLiu]

Change copyright to 2024 for all new files.