[shardformer] refactored layernorm

[FrankLeeeee]

📌 Checklist before creating the PR

• I have created an issue for this PR for traceability
• The title follows the standard format: [doc/gemini/tensor/...]: A concise description
• I have added relevant tags if possible for us to better distinguish different PRs

🚨 Issue number

Link this PR to your issue with words like fixed to automatically close the linked issue upon merge

e.g. fixed #1234, closed #1234, resolved #1234

📝 What does this PR do?

Summarize your work here.
if you have any plots/diagrams/screenshots/tables, please attach them here.

💥 Checklist before requesting a review

• I have linked my PR to an issue (instruction)
• My issue clearly describes the problem/feature/proposal, with diagrams/charts/table/code if possible
• I have performed a self-review of my code
• I have added thorough tests.
• I have added docstrings for all the functions/methods I implemented

⭐️ Do you enjoy contributing to Colossal-AI?

• 🌝 Yes, I do.
• 🌚 No, I don't.

Tell us more if you don't enjoy contributing to Colossal-AI.

[github-actions]

The code coverage for the changed files is 84%.

Click me to view the complete report

Name                                                                                       Stmts   Miss  Cover
--------------------------------------------------------------------------------------------------------------
colossalai/auto_parallel/tensor_shard/node_handler/node_handler.py                           164     82    50%
colossalai/auto_parallel/tensor_shard/node_handler/strategy/matmul_strategy_generator.py     388     76    80%
colossalai/auto_parallel/tensor_shard/utils/misc.py                                           45      9    80%
colossalai/checkpoint_io/utils.py                                                            243     43    82%
colossalai/device/device_mesh.py                                                             178     14    92%
colossalai/lazy/lazy_init.py                                                                 299     40    87%
colossalai/nn/layer/base_layer.py                                                             36     15    58%
colossalai/nn/layer/parallel_1d/_operation.py                                                 53     26    51%
colossalai/shardformer/__init__.py                                                             1      0   100%
colossalai/shardformer/_utils.py                                                              42     15    64%
colossalai/shardformer/layer/__init__.py                                                       7      0   100%
colossalai/shardformer/layer/_operation.py                                                   152     51    66%
colossalai/shardformer/layer/dropout.py                                                       21      0   100%
colossalai/shardformer/layer/embedding.py                                                    118      5    96%
colossalai/shardformer/layer/layernorm.py                                                     30      7    77%
colossalai/shardformer/layer/linear.py                                                       154     22    86%
colossalai/shardformer/layer/linear_conv.py                                                  186     24    87%
colossalai/shardformer/layer/loss.py                                                          49      8    84%
colossalai/shardformer/layer/parallel_module.py                                               76     21    72%
colossalai/shardformer/layer/utils.py                                                         52      0   100%
colossalai/shardformer/model/__init__.py                                                       0      0   100%
colossalai/shardformer/model/modeling_bert.py                                                 20     20     0%
colossalai/shardformer/policies/__init__.py                                                    0      0   100%
colossalai/shardformer/policies/autopolicy.py                                                 26      2    92%
colossalai/shardformer/policies/basepolicy.py                                                 40      5    88%
colossalai/shardformer/policies/bert.py                                                      108      2    98%
colossalai/shardformer/policies/gpt2.py                                                       60      2    97%
colossalai/shardformer/policies/llama.py                                                      32      2    94%
colossalai/shardformer/policies/t5.py                                                         27      2    93%
colossalai/shardformer/shard/__init__.py                                                       4      0   100%
colossalai/shardformer/shard/shard_config.py                                                  12      0   100%
colossalai/shardformer/shard/sharder.py                                                       74      9    88%
colossalai/shardformer/shard/shardformer.py                                                   22      1    95%
colossalai/tensor/comm_spec.py                                                               253     93    63%
colossalai/tensor/d_tensor/__init__.py                                                         4      0   100%
colossalai/tensor/d_tensor/api.py                                                            136     18    87%
colossalai/tensor/d_tensor/comm_spec.py                                                      151     35    77%
colossalai/tensor/d_tensor/layout.py                                                          38      1    97%
colossalai/tensor/d_tensor/layout_converter.py                                               195     12    94%
colossalai/tensor/d_tensor/utils.py                                                           38      7    82%
colossalai/tensor/shape_consistency.py                                                       294    120    59%
colossalai/tensor/sharding_spec.py                                                           139     13    91%
colossalai/testing/__init__.py                                                                 4      0   100%
colossalai/testing/comparison.py                                                              54      9    83%
tests/kit/model_zoo/registry.py                                                               17      0   100%
tests/kit/model_zoo/transformers/__init__.py                                                   6      0   100%
tests/kit/model_zoo/transformers/bert.py                                                      42      0   100%
tests/kit/model_zoo/transformers/gpt.py                                                       30      0   100%
tests/kit/model_zoo/transformers/llama.py                                                     26      2    92%
tests/kit/model_zoo/transformers/t5.py                                                        24      0   100%
tests/test_autochunk/test_autochunk_diffuser/test_autochunk_unet.py                           36     10    72%
tests/test_booster/test_mixed_precision/test_fp16_torch.py                                    30      1    97%
tests/test_booster/test_plugin/test_gemini_plugin.py                                          74     10    86%
tests/test_booster/test_plugin/test_low_level_zero_plugin.py                                  60      6    90%
tests/test_booster/test_plugin/test_torch_ddp_plugin.py                                       78      0   100%
tests/test_booster/test_plugin/test_torch_fsdp_plugin.py                                      43      0   100%
tests/test_checkpoint_io/test_gemini_checkpoint_io.py                                         80      0   100%
tests/test_device/test_device_mesh.py                                                         58     36    38%
tests/test_device/test_init_logical_pg.py                                                     27      1    96%
tests/test_fx/test_tracer/test_hf_model/hf_tracer_utils.py                                    21      2    90%
tests/test_fx/test_tracer/test_hf_model/test_hf_albert.py                                     17      1    94%
tests/test_fx/test_tracer/test_hf_model/test_hf_bert.py                                       15      1    93%
tests/test_fx/test_tracer/test_hf_model/test_hf_diffuser.py                                   50     28    44%
tests/test_fx/test_tracer/test_hf_model/test_hf_gpt.py                                        17      1    94%
tests/test_fx/test_tracer/test_hf_model/test_hf_opt.py                                        15      1    93%
tests/test_fx/test_tracer/test_hf_model/test_hf_t5.py                                         17      1    94%
tests/test_fx/test_tracer/test_timm_model/test_timm_model.py                                  36     24    33%
tests/test_fx/test_tracer/test_torchaudio_model/test_torchaudio_model.py                      14      5    64%
tests/test_fx/test_tracer/test_torchrec_model/test_deepfm_model.py                            39      3    92%
tests/test_fx/test_tracer/test_torchrec_model/test_dlrm_model.py                              41      4    90%
tests/test_fx/test_tracer/test_torchvision_model/test_torchvision_model.py                    31      1    97%
tests/test_lazy/lazy_init_utils.py                                                            72     14    81%
tests/test_lazy/test_distribute.py                                                            73      3    96%
tests/test_lazy/test_models.py                                                                13      1    92%
tests/test_shardformer/__init__.py                                                             0      0   100%
tests/test_shardformer/test_layer/test_dist_crossentropy.py                                   27      1    96%
tests/test_shardformer/test_layer/test_dropout.py                                             31      1    97%
tests/test_shardformer/test_layer/test_embedding.py                                           30      1    97%
tests/test_shardformer/test_layer/test_layernorm.py                                           27      1    96%
tests/test_shardformer/test_layer/test_linear_1d.py                                           85      1    99%
tests/test_shardformer/test_layer/test_linearconv_1d.py                                       73      1    99%
tests/test_shardformer/test_layer/test_vocab_parallel_embedding_1d.py                         32      1    97%
tests/test_shardformer/test_model/__init__.py                                                  0      0   100%
tests/test_shardformer/test_model/_utils.py                                                   22      0   100%
tests/test_shardformer/test_model/test_shard_bert.py                                          37      1    97%
tests/test_shardformer/test_model/test_shard_gpt2.py                                          37      1    97%
tests/test_shardformer/test_model/test_shard_llama.py                                         41      1    98%
tests/test_shardformer/test_model/test_shard_t5.py                                            35      1    97%
tests/test_tensor/test_dtensor/test_comm_spec.py                                              78      1    99%
tests/test_tensor/test_dtensor/test_dtensor.py                                                65      5    92%
tests/test_tensor/test_dtensor/test_layout_converter.py                                       91      1    99%
tests/test_tensor/test_shape_consistency.py                                                   50      2    96%
tests/test_tensor/test_sharded_linear.py                                                     130      1    99%
tests/test_tensor/test_sharding_spec.py                                                       13      1    92%
--------------------------------------------------------------------------------------------------------------
TOTAL                                                                                       6031    989    84%