AI Engineer Exhaustive
HardPlusArcee Trinity

Transformer Block

Arcee Trinity

Hard · 🔒 Plus required

Full Arcee Trinity block assembly

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

A Trinity decoder layer is a residual sandwich around two sublayers: gated grouped-query attention, then either a dense SwiGLU FFN or a sparse MoE. Section 2.4 writes the same residual form for every sublayer. The paper’s Figure 2 is the family diagram: RoPE only on local layers, sliding-window GQA on those layers, and an MoE trunk after a short dense prefix.

Trinity Large uses L=60 layers, hidden width d=3072, and k=6 initial dense FFNs before routed experts. The first six layers therefore share one block skeleton with the later fifty-four; only the feed-forward body changes. There are no biases on linear maps. Depth-scaled sandwich RMSNorm, QK-norm, and a 3:1 local/global attention schedule are what make the block distinct from a plain Llama residual.

How it works

Let \mathbf{x}_{\ell}\in\mathbb{R}^{B\times T\times d} enter layer \ell. Each sublayer \mathcal{M}_{\ell} (attention or FFN/MoE) is wrapped as

\mathbf{y}_{\ell}=\mathbf{x}_{\ell}+\mathrm{RMSNorm}^{(2)}_{\ell}\!\Big(\mathcal{M}_{\ell}\big(\mathrm{RMSNorm}^{(1)}_{\ell}(\mathbf{x}_{\ell})\big)\Big).

That is equation (33). The official decoder applies it twice. First \mathcal{M} is attention: project Q,K,V from the pre-normed residual, RMSNorm each query and key head (d_h=128), apply RoPE only if the layer is local, run causal (and windowed, if local) GQA, sigmoid-gate the concatenated heads, then W^O. Post-attention RMSNorm lands on the attention output before the residual add. Second \mathcal{M} is the FFN: a dense SwiGLU of intermediate width 12288 when \ell<6, otherwise one shared expert plus 256 routed experts of width 3072 with K_r=4. The same pre/post RMSNorm pair wraps that body.

Shapes that matter on Large: h_q=48, h_{kv}=8, so each query head maps to KV head j(i)=\lceil i\cdot h_{kv}/h_q\rceil. The attention gate W^G\in\mathbb{R}^{d\times d} is applied to the pre-normed token, then split into 48 vectors in \mathbb{R}^{128} and multiplied into the SDPA output (equations (11)–(14)). Local layers restrict sources to a window of $$; global layers use the full causal prefix and skip RoPE.

Equation (15) writes an extra residual inside the MoE sum. The released AfmoeSparseMoeBlock does not: it returns shared plus routed only. The sandwich residual after post_mlp_layernorm is the one that actually closes the layer. Treat (15)’s inner \mathbf{u}_t+ as the DeepSeek-style writeup; the implemented block follows (33).

Post-norm gains are initialized to 1/\sqrt{L} (equation (35)), about $$ at L=60. Pre-norm gains stay at 1. That scale lives in the checkpoint, not as a runtime multiply in the public modeling file.

Official code

The tech-report repo (arcee-ai/trinity-large-tech-report) is the PDF and README only. The runnable block is Hugging Face AfmoeDecoderLayer in src/transformers/models/afmoe/modeling_afmoe.py, generated from modular_afmoe.py. Attention is AfmoeAttention; the feed-forward is AfmoeMLP or AfmoeSparseMoeBlock according to layer_idx >= num_dense_layers. Checkpoint hyperparameters sit in configuration_afmoe.py / config.json on arcee-ai/Trinity-Large-Preview.

Watch-outs

Sources