[moe] support bf16 and update benchmark

[oahzxl]

• support bf16 to avoid overflow
• update benchmark
  • support mock dataset
  • use tokenizor in dataset
  • add load balance option

[github-actions]

The code coverage for the changed files is 66%.

Click me to view the complete report

Name                                                                 Stmts   Miss  Cover
----------------------------------------------------------------------------------------
colossalai/booster/plugin/hybrid_parallel_plugin.py                    214    157    27%
colossalai/booster/plugin/moe_hybrid_parallel_plugin.py                102    102     0%
colossalai/context/__init__.py                                           6      0   100%
colossalai/context/random/__init__.py                                    2      0   100%
colossalai/context/random/_helper.py                                    46     14    70%
colossalai/initialize.py                                               180    134    26%
colossalai/kernel/triton/__init__.py                                    14      3    79%
colossalai/kernel/triton/llama_act_combine_kernel.py                    89     67    25%
colossalai/legacy/engine/gradient_handler/__init__.py                    6      0   100%
colossalai/legacy/engine/gradient_handler/_moe_gradient_handler.py      20     20     0%
colossalai/moe/__init__.py                                               6      0   100%
colossalai/moe/_operation.py                                           177     49    72%
colossalai/moe/checkpoint.py                                           133    106    20%
colossalai/moe/experts.py                                               97     16    84%
colossalai/moe/layers.py                                               135     34    75%
colossalai/moe/load_balance.py                                         211     17    92%
colossalai/moe/loss.py                                                  21     21     0%
colossalai/moe/manager.py                                               91     12    87%
colossalai/moe/routers.py                                              176     33    81%
colossalai/moe/utils.py                                                 79     26    67%
colossalai/nn/layer/moe/__init__.py                                     12      0   100%
colossalai/nn/loss/__init__.py                                           0      0   100%
colossalai/tensor/moe_tensor/__init__.py                                 0      0   100%
colossalai/tensor/moe_tensor/api.py                                     24      2    92%
colossalai/tensor/moe_tensor/moe_info.py                                13      0   100%
colossalai/zero/low_level/low_level_optim.py                           431    146    66%
tests/test_moe/moe_utils.py                                             97      0   100%
tests/test_moe/test_grad_handler.py                                     59      1    98%
tests/test_moe/test_kernel.py                                           57      1    98%
tests/test_moe/test_moe_ep_tp.py                                        48      1    98%
tests/test_moe/test_moe_hybrid_zero.py                                  66      2    97%
tests/test_moe/test_moe_load_balance.py                                124      9    93%
tests/test_moe/test_moe_local.py                                        49      1    98%
tests/test_moe/test_moe_router.py                                       25      4    84%
tests/test_moe/test_moe_zero_fwd_bwd.py                                 77      3    96%
tests/test_moe/test_moe_zero_optim.py                                   66      9    86%
----------------------------------------------------------------------------------------
TOTAL                                                                 2953    990    66%