AI Engineer Exhaustive
HardPlusArcee Trinity

Coarse-Grained MoE FFN

Arcee Trinity

Hard · 🔒 Plus required

MoE with grouped intermediate dimensions

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

Trinity’s MoE follows DeepSeekMoE: many routed SwiGLU experts plus one always-on shared expert, with sigmoid token-choice routing. Nano and Mini keep that recipe fine-grained (expert widths 256 and 1024). Large deliberately coarsens it. Section 2.3 and the Large paragraph under Table 2 say they activate only four routed experts and make each expert larger “due to throughput requirements,” while raising routed-expert count to $$ so total capacity stays high at a low active-parameter budget.

“Grouped intermediate dimensions” is how the official tensors are stored, not a second routing algorithm. Each routed expert’s gate and up maps share one intermediate axis of length d_{\mathrm{ff}}^{\mathrm{moe}}, packed as a single [N_r,\,2d_{\mathrm{ff}}^{\mathrm{moe}},\,d] matrix. Tokens that hit an expert are gathered and run as a grouped GEMM against that slice. The paper specifies the widths; the grouping is the released AfmoeExperts layout (and the checkpoint flag use_grouped_mm).

How it works

Let \mathbf{u}_t\in\mathbb{R}^{d} be the MoE input after the pre-MLP sandwich norm. Paper equation (15) is

\mathbf{h}'_{t}=\mathbf{u}_{t}+\sum_{i=1}^{N_{s}}\mathrm{FFN}^{(s)}_{i}(\mathbf{u}_{t})+\sum_{i=1}^{N_{r}}g_{i,t}\,\mathrm{FFN}^{(r)}_{i}(\mathbf{u}_{t}).

On Large, N_s=1, N_r=256, K_r=4, d=3072, and each expert’s SwiGLU width (“expert size”) is 3072. The dense prefix uses a different intermediate width, 12288. Shared and routed experts share the MoE width: the shared expert is one AfmoeMLP with intermediate_size = moe_intermediate_size * num_shared_experts.

Each routed expert is SwiGLU without bias:

\mathrm{FFN}(\mathbf{u})=\mathrm{W}_{\mathrm{down}}\big(\mathrm{SiLU}(\mathrm{W}_{\mathrm{gate}}\mathbf{u})\odot\mathrm{W}_{\mathrm{up}}\mathbf{u}\big).

In AfmoeExperts the two up-projections are concatenated on the intermediate axis, so one matmul yields [\mathrm{gate}\,\|\,\mathrm{up}]\in\mathbb{R}^{2d_{\mathrm{ff}}^{\mathrm{moe}}} which is then chunked. The down map is stored as [N_r,\,d,\,d_{\mathrm{ff}}^{\mathrm{moe}}]. That is the grouped intermediate layout: expert index is the leading batch dimension of a 3D weight, not $$ separate Linear modules.

Gating g_{i,t} is not the Top-K score that includes expert bias. Bias b_i participates only in selection (equation (17)); the sigmoid scores of the chosen experts are renormalized to a simplex (equation (18)) and, in the official router, multiplied by a constant route scale $$. Those g_{i,t} scale the expert outputs, not the inputs.

The inner \mathbf{u}_t+ in (15) is not present in AfmoeSparseMoeBlock.forward, which returns shared_output + routed_output. The layer residual is the outer sandwich add after post_mlp_layernorm. If you implement (15) literally and equation (33), you double the skip.

Granularity contrast from Table 2: Nano activates 8 of 128 experts of size 256; Mini activates 8 of 128 of size 1024; Large activates 4 of 256 of size 3072. Same hidden width as the model (d=d_{\mathrm{ff}}^{\mathrm{moe}}=3072) makes each expert GEMM square-ish and much fatter than Nano’s 256-wide experts — that is the coarse-grain throughput bet.

Official code

No MoE kernels live in arcee-ai/trinity-large-tech-report. Look at AfmoeSparseMoeBlock, AfmoeExperts, and AfmoeMLP in modeling_afmoe.py (modular_afmoe.py is the source; AfmoeExperts subclasses the Qwen2-MoE 3D expert pack). Widths are moe_intermediate_size, num_experts, num_experts_per_tok, num_shared_experts in Trinity-Large-Preview config.json.

Watch-outs

Sources