昇腾社区首页
中文
注册

混合专家模型(MoE)

前期准备

  1. 参考软件安装模型迁移完成安装依赖、安装MindSpeed、获取Megatron-LM并指定commit id。
  2. MindSpeed目录下的tests_extend文件夹拷贝到Megatron-LM目录下。
    # {PATH_TO_MEGATRON_LM}: the root path of `Megatron-LM`
    cp -r MindSpeed/tests_extend {PATH_TO_MEGATRON_LM}

操作步骤

进入Megatron-LM目录,执行pretrain_llama_tp4_ep2_dp2.sh脚本。
cd {PATH_TO_MEGATRON_LM}
bash tests_extend/system_tests/llama2/pretrain_llama_tp4_ep2_dp2.sh

上述脚本pretrain_llama_tp4_ep2_dp2.sh,包括但不限于以下并行与优化特性

  • 并行配置
    --tensor-model-parallel-size    # 张量并行
    --pipeline-model-parallel-size  # 流水线并行
    --sequence-parallel             # 序列并行
    --expert-model-parallel-size    # 专家并行
  • 融合算子
    --use-flash-attn                # Flash Attention融合算子
    --normalization RMSNorm         # RMSNorm归一化
    --use-fused-rmsnorm             # RMSNorm融合算子
    --swiglu                        # SwiGLU激活函数
    --use-fused-swiglu              # SwiGLU融合算子
  • 计算优化
    --enable-token-rearrange-opt    # MoE Token重排性能优化