Feat/faster rms norm

[S1ro1]

Summary

Implements partial aggregation in rms_norm, similar to that in layer_norm, as described in #179 .

Testing Done

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

[S1ro1]

i have made gradients w.r.t weights work properly, except 1 singular test, which just calculates random number of elements wrong with no apparent cause (every other call the number and indices of elements are different), however it's always the same test (test/transformers/test_rms_norm.py::test_correctness[LlamaRMSNorm-0.0-llama-dtype1-0.2-0.2-16-1024-4096]), so help would be appreciated to fix this issue as im getting clueless. (setting the seed makes the elements wrong be deterministic, however still the same issue)

[lancerts]

@S1ro1 did a fix on the main branch rmsnorm, can you check if the same error persist after rebase? ty

[S1ro1]

E           Mismatch at index (150,): tensor1[(150,)] = 0.076171875, tensor2[(150,)] = -0.2099609375
E           Mismatch at index (806,): tensor1[(806,)] = -1.6953125, tensor2[(806,)] = -1.1875
E           Mismatch at index (2538,): tensor1[(2538,)] = -1.3125, tensor2[(2538,)] = -0.80078125
E           Mismatch at index (3000,): tensor1[(3000,)] = -0.7578125, tensor2[(3000,)] = -0.373046875
E           Mismatch at index (3927,): tensor1[(3927,)] = -0.1591796875, tensor2[(3927,)] = -0.71484375

@lancerts the error stays the same, just few elements mismatched in always 1 and the test same. I have no clue how to fix that to be fair, so I guess you can unassign me, I'm working on this on/off for 3/4 days but haven't got a sinlge step forward since the bug occurance. I guess I took a too big bite. The branch should have some base work done on the issue, but maybe going from scratch might be a better issue.

[lancerts]

#255 resolved by this PR