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 一样。

代码示例¶

>>> import paddle
>>> from paddle.incubate.nn.functional import fused_ec_moe

>>> paddle.set_device('gpu')
>>> x = paddle.randn([10, 128, 1024])
>>> gate = paddle.randn([10, 128, 8])
>>> bmm0_weight = paddle.randn([8, 1024, 4096])
>>> bmm0_bias = paddle.randn([8, 1024, 4096])
>>> bmm1_weight = paddle.randn([8, 1024, 4096])
>>> bmm1_bias = paddle.randn([8, 1024, 4096])
>>> out = fused_ec_moe(x, gate, bmm0_weight, bmm0_bias, bmm1_weight, bmm1_bias, act_type="gelu")
>>> print(out.shape)
[10, 128, 1024]