[Feauture] MoE refactor

.github/workflows/build_on_pr.yml

@@ -90,7 +90,7 @@ jobs:
     runs-on: [self-hosted, gpu]
     container:
       image: hpcaitech/pytorch-cuda:2.1.0-12.1.0
-      options: --gpus all --rm -v /dev/shm -v /data/scratch/llama-tiny:/data/scratch/llama-tiny
+      options: --gpus all --rm -v /dev/shm -v /data/scratch:/data/scratch
     timeout-minutes: 90
     defaults:
       run:
@@ -165,6 +165,7 @@ jobs:
         env:
           LD_LIBRARY_PATH: /github/home/.tensornvme/lib:/usr/local/nvidia/lib:/usr/local/nvidia/lib64
           LLAMA_PATH: /data/scratch/llama-tiny
+          MOE_TENSOR_PATH: /data/scratch/moe_tensors
 
       - name: Collate artifact
         env:

.github/workflows/build_on_schedule.yml

@@ -69,6 +69,7 @@ jobs:
         env:
           LD_LIBRARY_PATH: /github/home/.tensornvme/lib:/usr/local/nvidia/lib:/usr/local/nvidia/lib64
           LLAMA_PATH: /data/scratch/llama-tiny
+          MOE_TENSOR_PATH: /data/scratch/moe_tensors
 
       - name: Notify Lark
         id: message-preparation

.github/workflows/compatiblity_test_on_dispatch.yml

@@ -92,3 +92,4 @@ jobs:
           DATA: /data/scratch/cifar-10
           LD_LIBRARY_PATH: /github/home/.tensornvme/lib:/usr/local/nvidia/lib:/usr/local/nvidia/lib64
           LLAMA_PATH: /data/scratch/llama-tiny
+          MOE_TENSOR_PATH: /data/scratch/moe_tensors

.github/workflows/compatiblity_test_on_pr.yml

@@ -87,3 +87,4 @@ jobs:
           DATA: /data/scratch/cifar-10
           LD_LIBRARY_PATH: /github/home/.tensornvme/lib:/usr/local/nvidia/lib:/usr/local/nvidia/lib64
           LLAMA_PATH: /data/scratch/llama-tiny
+          MOE_TENSOR_PATH: /data/scratch/moe_tensors

.github/workflows/compatiblity_test_on_schedule.yml

@@ -85,6 +85,7 @@ jobs:
           DATA: /data/scratch/cifar-10
           LD_LIBRARY_PATH: /github/home/.tensornvme/lib:/usr/local/nvidia/lib:/usr/local/nvidia/lib64
           LLAMA_PATH: /data/scratch/llama-tiny
+          MOE_TENSOR_PATH: /data/scratch/moe_tensors
 
       - name: Notify Lark
         id: message-preparation

applications/ColossalMoE/infer.py

@@ -10,7 +10,6 @@
 from colossalai.booster.plugin.moe_hybrid_parallel_plugin import MoeHybridParallelPlugin
 from colossalai.cluster import DistCoordinator
 from colossalai.moe.checkpoint import MoECheckpointIO
-from colossalai.shardformer.policies.mixtral import MixtralForCausalLMPolicy
 
 
 def parse_args():
@@ -70,7 +69,6 @@ def main():
             ep_size=ep_size,
             zero_stage=1,
             precision=args.precision,
-            custom_policy=MixtralForCausalLMPolicy(),
             checkpoint_io=MoECheckpointIO,
             enable_fused_normalization=args.use_layernorm_kernel,
             enable_jit_fused=args.use_kernel,

applications/ColossalMoE/train.py

@@ -15,7 +15,6 @@
 from colossalai.moe.checkpoint import MoECheckpointIO
 from colossalai.nn.lr_scheduler import CosineAnnealingWarmupLR
 from colossalai.nn.optimizer import HybridAdam
-from colossalai.shardformer.policies.mixtral import MixtralForCausalLMPolicy
 from colossalai.utils import get_current_device
 
 
@@ -155,7 +154,6 @@ def main():
             pp_size=args.pp_size,
             ep_size=args.ep_size,
             microbatch_size=args.microbatch_size,
-            custom_policy=MixtralForCausalLMPolicy(),
             enable_fused_normalization=args.use_layernorm_kernel,
             enable_jit_fused=args.use_kernel,
             precision=args.precision,

colossalai/booster/plugin/moe_hybrid_parallel_plugin.py

@@ -30,7 +30,7 @@
 from colossalai.shardformer.policies.base_policy import Policy
 from colossalai.zero.low_level import LowLevelZeroOptimizer
 
-PP_AXIS, DP_AXIS, EP_AXIS, TP_AXIS = 0, 1, 2, -1
+PP_AXIS, DP_AXIS, EP_AXIS, TP_AXIS = 0, 1, 2, 3
 
 
 class HybridParallelZeroOptimizer(LowLevelZeroOptimizer):
@@ -352,7 +352,9 @@ def seed_worker(worker_id):
 
     def get_checkpoint_io(self) -> MoECheckpointIO:
         if self.checkpoint_io is None:
-            self.checkpoint_io = MoECheckpointIO(self.global_dp_group, self.pp_group, self.tp_group, self.zero_stage)
+            self.checkpoint_io = MoECheckpointIO(
+                self.global_dp_group, self.pp_group, self.tp_group, self.ep_group, self.moe_dp_group, self.zero_stage
+            )
         else:
             self.checkpoint_io = self.checkpoint_io(
                 self.global_dp_group,

colossalai/moe/load_balance.py

@@ -7,8 +7,8 @@
 from torch.distributed import ProcessGroup
 
 from colossalai.cluster import ProcessGroupMesh
-from colossalai.moe.experts import MLPExperts
 from colossalai.moe.manager import MOE_MANAGER
+from colossalai.shardformer.layer.moe.layers import MLPExperts
 from colossalai.zero.low_level import LowLevelZeroOptimizer
 
 

colossalai/shardformer/layer/moe/experts.py

@@ -9,7 +9,7 @@
 from colossalai.moe.manager import MOE_MANAGER
 from colossalai.moe.utils import get_activation
 from colossalai.shardformer.layer.utils import Randomizer
-from colossalai.tensor.moe_tensor.api import get_ep_rank, get_ep_size, set_moe_tensor_info
+from colossalai.tensor.moe_tensor.api import get_ep_rank, get_ep_size
 
 if HAS_TRITON:
     from colossalai.kernel.triton.llama_act_combine_kernel import LlamaActCombine

colossalai/shardformer/layer/moe/layers.py

@@ -11,7 +11,6 @@
 from colossalai.moe.load_balance import LoadBalancer
 from colossalai.moe.utils import create_ep_hierarchical_group, get_noise_generator
 from colossalai.shardformer.layer.moe import MLPExperts
-from colossalai.shardformer.layer.moe.routers import MoeRouter, get_router_cls
 from colossalai.tensor.moe_tensor.api import get_dp_group, get_ep_group, get_ep_group_ranks, get_ep_size
 
 

colossalai/shardformer/layer/moe/routers.py

@@ -9,7 +9,7 @@
 from colossalai.moe.manager import MOE_MANAGER
 from colossalai.moe.utils import get_activation
 from colossalai.shardformer.layer.utils import Randomizer
-from colossalai.tensor.moe_tensor.api import get_ep_rank, get_ep_size, set_moe_tensor_info
+from colossalai.tensor.moe_tensor.api import get_ep_rank, get_ep_size
 
 if HAS_TRITON:
     from colossalai.kernel.triton.llama_act_combine_kernel import LlamaActCombine