[moe] support load balance by oahzxl · Pull Request #4914 · hpcaitech/ColossalAI

oahzxl · 2023-10-15T17:07:53Z

support load balance for moe
add a unit test

github-actions · 2023-10-15T18:15:43Z

The code coverage for the changed files is 72%.

Click me to view the complete report

Name                                                                 Stmts   Miss  Cover
----------------------------------------------------------------------------------------
colossalai/booster/plugin/hybrid_parallel_plugin.py                    214     13    94%
colossalai/booster/plugin/moe_hybrid_parallel_plugin.py                102     32    69%
colossalai/context/__init__.py                                           6      0   100%
colossalai/context/random/__init__.py                                    2      0   100%
colossalai/context/random/_helper.py                                    46      9    80%
colossalai/initialize.py                                               180    133    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                                           159     44    72%
colossalai/moe/checkpoint.py                                           133     31    77%
colossalai/moe/experts.py                                               97     16    84%
colossalai/moe/layers.py                                               136     34    75%
colossalai/moe/load_balance.py                                         201     13    94%
colossalai/moe/loss.py                                                  21     21     0%
colossalai/moe/manager.py                                               91      7    92%
colossalai/moe/routers.py                                              176     15    91%
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     32    93%
examples/language/openmoe/model/__init__.py                              0      0   100%
examples/language/openmoe/model/modeling_openmoe.py                    462    328    29%
examples/language/openmoe/model/openmoe_policy.py                      248    157    37%
tests/test_infer_ops/triton/test_llama_act_combine.py                   40     23    42%
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_checkpoint.py                                   69      2    97%
tests/test_moe/test_moe_ep_tp.py                                        48      1    98%
tests/test_moe/test_moe_group.py                                        54      2    96%
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      1    96%
tests/test_moe/test_moe_zero_fwd_bwd.py                                 77      3    96%
tests/test_moe/test_moe_zero_optim.py                                   66      9    86%
----------------------------------------------------------------------------------------
TOTAL                                                                 3799   1058    72%

github-actions · 2023-10-16T02:14:01Z

The code coverage for the changed files is 54%.

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                                           159     97    39%
colossalai/moe/checkpoint.py                                           133    106    20%
colossalai/moe/experts.py                                               97     20    79%
colossalai/moe/layers.py                                               136     44    68%
colossalai/moe/load_balance.py                                         211      7    97%
colossalai/moe/loss.py                                                  21     21     0%
colossalai/moe/manager.py                                               91     13    86%
colossalai/moe/routers.py                                              176     82    53%
colossalai/moe/utils.py                                                 79     41    48%
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      4    83%
colossalai/tensor/moe_tensor/moe_info.py                                13      0   100%
colossalai/zero/low_level/low_level_optim.py                           431    150    65%
tests/test_moe/moe_utils.py                                             97     53    45%
tests/test_moe/test_moe_load_balance.py                                124      9    93%
----------------------------------------------------------------------------------------
TOTAL                                                                 2489   1144    54%

github-actions · 2023-10-16T10:07:15Z

The code coverage for the changed files is 72%.

Click me to view the complete report

Name                                                                 Stmts   Miss  Cover
----------------------------------------------------------------------------------------
colossalai/booster/plugin/hybrid_parallel_plugin.py                    214     13    94%
colossalai/booster/plugin/moe_hybrid_parallel_plugin.py                102     32    69%
colossalai/context/__init__.py                                           6      0   100%
colossalai/context/random/__init__.py                                    2      0   100%
colossalai/context/random/_helper.py                                    46      9    80%
colossalai/initialize.py                                               180    133    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                                           159     44    72%
colossalai/moe/checkpoint.py                                           133     31    77%
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      7    92%
colossalai/moe/routers.py                                              176     15    91%
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     32    93%
examples/language/openmoe/model/__init__.py                              0      0   100%
examples/language/openmoe/model/modeling_openmoe.py                    462    328    29%
examples/language/openmoe/model/openmoe_policy.py                      248    157    37%
tests/test_infer_ops/triton/test_llama_act_combine.py                   40     23    42%
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_checkpoint.py                                   69      2    97%
tests/test_moe/test_moe_ep_tp.py                                        48      1    98%
tests/test_moe/test_moe_group.py                                        54      2    96%
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      1    96%
tests/test_moe/test_moe_zero_fwd_bwd.py                                 77      3    96%
tests/test_moe/test_moe_zero_optim.py                                   66      9    86%
----------------------------------------------------------------------------------------
TOTAL                                                                 3808   1062    72%

* add load balance * update test * update param exchange * pass test * update test * update test * update test * update test * fix ranks * update

oahzxl added 7 commits October 15, 2023 16:16

add load balance

0f24026

update test

3bb9a38

update param exchange

8be2fce

pass test

2ac2b82

update test

e863dc2

update test

5589539

update test

5429cac

oahzxl added 2 commits October 16, 2023 03:15

update test

063c279

fix ranks

c2ebb1f

ver217 reviewed Oct 16, 2023

View reviewed changes

Comment thread colossalai/moe/layers.py Outdated

ver217 reviewed Oct 16, 2023

View reviewed changes

Comment thread colossalai/moe/load_balance.py

update

d153572

ver217 approved these changes Oct 16, 2023

View reviewed changes

oahzxl merged commit 926ad74 into hpcaitech:feature/MoE Oct 16, 2023

oahzxl deleted the load branch October 18, 2023 02:21

oahzxl added a commit to oahzxl/ColossalAI that referenced this pull request Oct 26, 2023

[moe] support load balance (hpcaitech#4914)

5844f34

* add load balance * update test * update param exchange * pass test * update test * update test * update test * update test * fix ranks * update

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[moe] support load balance#4914

[moe] support load balance#4914
oahzxl merged 10 commits intohpcaitech:feature/MoEfrom
oahzxl:load

oahzxl commented Oct 15, 2023

Uh oh!

github-actions Bot commented Oct 15, 2023

Uh oh!

github-actions Bot commented Oct 16, 2023

Uh oh!

Uh oh!

Uh oh!

github-actions Bot commented Oct 16, 2023

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

oahzxl commented Oct 15, 2023

Uh oh!

github-actions Bot commented Oct 15, 2023

Uh oh!

github-actions Bot commented Oct 16, 2023

Uh oh!

Uh oh!

Uh oh!

github-actions Bot commented Oct 16, 2023

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants