MoE概念

模型参数增加很多;计算量没有增加(gating+小FNN,比以前的大FNN计算量要小);收敛速度变快;

效果:PR-MoE > 普通MoE > DenseTransformer

MoE模型,可视为Sparse Model,因为每次参与计算的是一部分参数;

Expert并行,可以和其他并行方式,同时使用:

 ep_size指定了MoE进程组大小,一个模型replica的所有MoE都分布在mp_size这些GPU卡上。

下例,启动2个DP replica,每个DP replica包含2个MoE rank,每个MoE rank包含4个MoE FFN,每个GPU放置一个MoE rank;

import torch
import deepspeed
import deepspeed.utils.groups as groups
from deepspeed.moe.layer import MoE

WORLD_SIZE = 4
EP_WORLD_SIZE = 2
EXPERTS = 8

fc3 = torch.nn.Linear(84, 84)
fc3 = MoE(hidden_size=84, expert=self.fc3, num_experts=EXPERTS, ep_size=EP_WORLD_SIZE, k=1)
fc4 = torch.nn.Linear(84, 10)

节约显存的方式:

1. 配置好ZeRo

2. 配置好fp16

    "fp16": {
      "enabled": true,
      "fp16_master_weights_and_grads": true,
  }

支持PR-MoE

Pyramid: 金字塔 ;前面的层,Experts数量少些,后面的层Experts数量多些;

Residual:每层都过一个固定的MLP,并和选中的Expert输出结果,相加;

支持Random Token Selection

推理

import deepspeed
import torch.distributed as dist

# Set expert-parallel size
world_size = dist.get_world_size()
expert_parallel_size = min(world_size, args.num_experts)

# create the MoE model
moe_model = get_model(model, ep_size=expert_parallel_size)
...

# Initialize the DeepSpeed-Inference engine
ds_engine = deepspeed.init_inference(moe_model,
                                     mp_size=tensor_slicing_size,
                                     dtype=torch.half,
                                     moe_experts=args.num_experts,
                                     checkpoint=args.checkpoint_path,
                                     replace_with_kernel_inject=True,)
model = ds_engine.module
output = model('Input String')

如果Experts数目大于GPU数目, 均分在各个GPU上;如果小于,则Expert将被切片,均分到各个GPU上;

注意:ep_size, mp_size(模型并行),dtype为half, 优化版kernel的使用replace_with_kernel_inject=True;

generate_samples_gpt.py \
       --tensor-model-parallel-size 1 \
       --num-experts ${experts} \
       --num-layers 24 \
       --hidden-size 2048 \
       --num-attention-heads 32 \
       --max-position-embeddings 1024 \
       --tokenizer-type GPT2BPETokenizer \
       --load $checkpoint_path \
       --fp16 \
       --ds-inference \

PR-MoE, 前面那些层的experts个数少,后面的多:

experts="64 64 64 64 64 64 64 64 64 64 128 128"
generate_samples_gpt.py \
       --tensor-model-parallel-size 1 \
       --num-experts ${experts} \
       --mlp_type 'residual' \
       --num-layers 24 \
       --hidden-size 2048 \
       --num-attention-heads 16 \
       --max-position-embeddings 1024 \
       --tokenizer-type GPT2BPETokenizer \
       --load $checkpoint_path \
       --fp16 \
       --ds-inference \

--mlp_type指定使用PR-MoE,推理latency更块;

点赞(0) 打赏

评论列表 共有 0 条评论

暂无评论

微信公众账号

微信扫一扫加关注

发表
评论
返回
顶部