TransformerFeedForwardMoe

A sharded MoE Layer.

Abstract Signature:

TransformerFeedForwardMoe(input_dims: int, hidden_dims: int, num_experts: int, num_groups: int)

PyTorch

API: —
Strategy: Custom / Partial

PaxML / Praxis

API: praxis.layers.TransformerFeedForwardMoe
Strategy: Direct Mapping
Official Docs ↗