fused_ec_moe¶

参数¶

• x (Tensor) - 输入 Tensor，形状是 [bsz, seq_len, d_model] 。

• gate (Tensor) - 用于选择专家的 gate Tensor，形状是 [bsz, seq_len, e] 。

• bmm0_weight (Tensor) - 第一个 batch matmul 的权重数据，形状是 [e, d_model, d_feed_forward]。

• bmm0_bias (Tensor) - 第一个 batch matmul 的偏置数据，形状是 [e, 1, d_feed_forward]。

• bmm1_weight (Tensor) - 第二个 batch matmul 的权重数据，形状是 [e, d_model, d_feed_forward]。

• bmm1_bias (Tensor) - 第二个 batch matmul 的偏置数据，形状是 [e, 1, d_feed_forward]。

• act_type (string) - 激活函数类型，目前仅支持 gelu , relu。

返回¶

• Tensor，输出 Tensor，数据类型与 x 一样。

代码示例¶

# required: gpu
import paddle
from paddle.incubate.nn.functional import fused_ec_moe

batch = 10
seq_len = 128
d_model = 1024
d_feed_forward = d_model * 4
num_expert = 8

x = paddle.randn([batch, seq_len, d_model])
gate = paddle.randn([batch, seq_len, num_expert])
bmm0_weight = paddle.randn([num_expert, d_model, d_feed_forward])
bmm0_bias = paddle.randn([num_expert, d_model, d_feed_forward])
bmm1_weight = paddle.randn([num_expert, d_model, d_feed_forward])
bmm1_bias = paddle.randn([num_expert, d_model, d_feed_forward])
out = fused_ec_moe(x, gate, bmm0_weight, bmm0_bias, bmm1_weight, bmm1_bias, act_type="gelu")

print(out.shape) # [batch, seq_len, num_expert]