AI Engineer Exhaustive
HardPlusArcee Trinity

Full Forward Pass

Arcee Trinity

Hard · 🔒 Plus required

Complete Arcee Trinity forward pass

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

The Trinity Large forward pass is a decoder-only stack: embed, optional \sqrt{d} scale, sixty sandwich blocks, a final RMSNorm, then an untied LM head. Section 2 specifies the pieces; Table 2 fixes Large’s widths. The official AfmoeModel / AfmoeForCausalLM is the inference graph that matches those equations, minus training-only SMEBU updates.

What makes the pass different from a dense Llama stack is the per-layer branching: three sliding-window RoPE layers then one global NoPE layer, repeated; six dense SwiGLU layers then fifty-four coarse MoE layers; and dual RMSNorms around every sublayer. The token path has to carry two causal masks and apply RoPE only on local layers.

How it works

Tokens \mathrm{tok}_t index a 200{,}192-row embedding E\in\mathbb{R}^{V\times d} with d=3072. Equation (39) then scales the lookup when μP-style input scaling is on (checkpoint mup_enabled: true):

\mathbf{e}_{t}=\sqrt{d}\,E(\mathrm{tok}_{t}).

AfmoeModel.forward implements that as hidden_states * hidden_size**0.5 after embed_tokens. Position ids are the usual 0\ldots T-1 offset by cache length. One RoPE cache (\cos,\sin) is built for the whole batch from \theta=10^{4} and d_h=128; global layers will ignore it.

Two masks are built once: a full causal mask and a sliding-window causal mask of width w=4096. Layer \ell picks with layer_types[ℓ]. The default construction is local unless (\ell+1)\bmod 4=0, i.e. a 3:1 RoPE/NoPE schedule.

Each of the L=60 layers is equation (33) applied to attention, then again to the FFN/MoE. Layers \ell<6 use a dense SwiGLU of width 12288. Layers \ell\ge 6 use one shared expert plus 256 routed experts of width 3072, K_r=4, sigmoid scores renormalized over the chosen experts and scaled by 2.448. Attention is GQA (h_q=48, h_{kv}=8, d_h=128) with QK-norm, optional RoPE, optional window, then a sigmoid gate on the SDPA output before W^O.

After the last layer, equation (36):

\mathbf{z}=\mathrm{RMSNorm}_{\mathrm{LM}}(\mathbf{h}_{L}),

then logits W_{\mathrm{LM}}\mathbf{z} with W_{\mathrm{LM}}\in\mathbb{R}^{V\times d}, not tied to E (tie_word_embeddings: false). Linear layers have no biases.

A compact token path:

  1. \mathbf{h}\leftarrow\sqrt{d}\,E(\mathrm{tok})
  2. For \ell=0,\ldots,59: sandwich-attention (RoPE+window iff local), then sandwich-FFN (dense iff \ell<6 else MoE)
  3. \mathrm{logits}\leftarrow W_{\mathrm{LM}}\,\mathrm{RMSNorm}(\mathbf{h})

Training additionally maintains per-MoE-layer expert bias \mathbf{b} via SMEBU and a small sequence-wise balance loss. Those terms do not appear in the released generate path; expert_bias is a non-trainable buffer consumed by Top-$$ only.

Official code

arcee-ai/trinity-large-tech-report publishes the report, not a trainer. The forward is AfmoeModel.forward and AfmoeForCausalLM.forward in modeling_afmoe.py. Mask construction uses create_causal_mask / create_sliding_window_causal_mask. Layer schedule and widths are in Trinity-Large-Preview config.json (num_hidden_layers, num_dense_layers, layer_types, mup_enabled).

Watch-outs

Sources