The GoBots to Hugging Face's Transformers
- Architectural Components
- LLM Architectures
- BaGL
In what follows, unless otherwise stated, we make make the assumption that inputs to the model take the form [batchsize, sequence_len, embedding_dimension]
with generic shape [n, m, d].
much of the insights provided bellow can be found here.
absolute positional encoding add a fixed vector to each element of the input sequence. A popular version is the sinusoidal encoding (introduced int he original transformer paper) which is applied to the inputs before being fed through the transformer blocks
where
pros:
- easy to implement.
- relatively fast, it is just one addition.
cons:
- does not seem to generalize well to unseen input sequence lengths
- does not encourage attention to any particular positions in the sequence.
- is not explicitly included in the attention mechanism.
Rotary positional encodings (RoPE), introduced in the RoFormer paper work by rotating the query and key vectors in the attention mechanism by angles proportional to the position in the sequence. The rotation matrix can be parameterized as
where
pros:
- widely adopted
- incorporates positional information directly into the attention mechanism
- does not introduce any parameters into the model
- combines some benefits of relative and absolute encodings
cons:
- more expensive then some encodings (multiplication slower than addition)
- does not scale as well to unseen sequence lengths
These positional encodings were introduced in the T5 paper. Here a learned bias is added to the attention scores. The parameters are shared across layers and attention heads. A fixed number of relative distance buckets are used (32 in the original implementation) up to a maximum distance where the buckets scale logarithmically with distance.
pros:
- incorporated in the attention mechanism
- each attention head gets its own set of embeddings which are shared across layers, making this a fast and parameter efficient encoding mechanism
- generalizes well to unseen sequence lengths
- biases attention has two peaks in relative distance for near and distant tokens.
cons:
- slower then NoPE, ALiBi, and RoPE
The term ALiBi comes from the phrase "attention with linear biases" and was introduced here. As the name suggests this method encodes relative positions in the attention mechanism by subtracting a fixed bias term to the attention score which scales linearly with distance. These biases are not learned and each attention head is given its own bias (shared by all layers).
pros:
- Parameter efficient, no learned parameters
- incorporated into the attention mechanism
- fast to compute
cons:
- biases attention toward nearby tokens.
As most LLMs are implemented as decoder only transformers explicit positional encoding can be completely forgone. The causal mask in the attention mechanism appears to be all one needs to include positional information. NoPE, which stands for no positional encoding is precisely this.
pros:
- easy to implement
- adds no overhead
- generalizes to any sequence length
cons:
- Test were performed on a simplified toy model
- may not be better overall in general
Below are a non-exhaustive list of activation functions. Some, like the sigmoid function, which are not necessarily used in LLMs are include to help explain other activations.
The sigmoid activation function is commonly defined as the logistic function
The rectified linear unit is widely popular in many architectures. It can lead to a problem of negative values being set to zero and thus passing forward little to no information.
The leaky ReLU activation attempts to fix the dying ReLU problem by allowing small negative values
Another attempt to solve the issues of ReLU is the exponential linear unit (ELU).
The Gaussian error linear unit which is approximated as
The sigmoid linear unit
The Swish activation is an extension to SiLU that adds a weighting to the sigmoid component
The Gated linear unit is an adaptation of the idea of the gate units from RNNs where the network learns what information to pass on
flowchart BT
input --> linear_layer1[Linear layer] --> sigma[Sigmoid activation] --> merge((⊗))
input --> linear_layer2[Linear layer] --> merge
merge --> linear_layer_3[Linear layer] --> Output
parameterized by the weights
SwiGLU is a combination of the swish and GLU activations where the sigmoid of GLU is replaced with swish
flowchart BT
input --> linear_layer1[Linear layer] --> sigma[Swish activation] --> merge((⊗))
input --> linear_layer2[Linear layer] --> merge
merge --> linear_layer_3[Linear layer] --> Output
The heart of any transformer model is the attention mechanism. In general an attention block takes three inputs, a query
Scaled-dot product attention can be thought of as a mechanism where the query and key attend to each other via the dot product. The attention score is calculated as softmax of the dot-product between
the query and key. A mask,
Multi-head attention (MHA) generalizes standard dot-product attention by combining several attenion calculations at every step.
for some weight matrices
Multi-Query attention (MQA) is similar to multi-head attention where there are multiple query heads and a single key and value head that are shared byt all query heads. This reduces the complexity of the model at the expense of performance.
flowchart TB
k --- v
v --> q1
v--> q2
v --> q3
v --> q4[...]
v --> qh
Grouped-query attention (GQA) can be seen as a trade-off beteen the accuracy of full multi-head attention and the efficiency of mult-query attention. In GQA there are
flowchart TB
k1 --- v1
v1 --> q1
v1--> q2
v1--> q3
v1 --> q4
k2 --- v2
v2 --> q5
v2--> q6
v2--> q7
v2 --> q8
ki[...] --- vi[...] --> qi[...]
vhk["v(kv)"] --- khk["k(kv)"] --> qhm3["q(h-3)"]
khk --> qhm2["q(h-2)"]
khk --> qhm1["q(h-1)"]
khk --> qh["q(h)"]
In multi-head latent attention (MLA) the querys, keys, and values are all compressed into lower rank spaces before applying the attention operation. The primary purpose of this operation is to lower the size of the KV cache at inference time. In order to incorporate positional embeddings while preserving the benefits of the low-rank projection the DeepSeek architects created a decoupled RoPE mechanism. Here we concatenate two vectors, one which carries the content information and another which bears the RoPE positional information
We then apply attention between the query, key and value represented by
The second common component of all transformer architectures is the feedforward layer. The feedforward layer is applied after the attention mechanism as a way to add extra information to the token embeddings and to prepare the output of the attention block for the next transformer block layer in the stack.
A dense feedforward layer is simply a dense neural network layer that is applied to all outputs of the attention mechanism. The layer is shared across the sequence. In many new LLMs the feedforward layer is a SwiGLU layer.
A mixture of experts (MoE) layer consists of two sets of parallel dense feedforward layers, known as experts. One set, the shared experts, are applied to all tokens containing
Several versions of MoE have been proposed, below we follow the one used in DeepSeek V3 when calculating the gating function
where
flowchart BT
input --> router["Router top-k"] --> routed1["RE 1"] --> merge@{shape: circle, label: " + "} --> output
router --> routed[...] --> merge
router --> routedn["RE n"] --> merge
input --> shared1["SE 1"] --> merge
input --> sharedd[...] --> merge
input --> shared2["SE m"] --> merge
input -- "skip" --> merge
A problem that might arise when training an MoE layer is that the model may learn to use one or a few of the experts exclusively. A popular solution to this problem is to add an auxiliary loss which enforces load balancing amongst experts (
Adding this loss may hamper model accuracy, an alternative that provides load balancing without an auxiliary loss can be found here). In this implementaton a bias is added when finding the top-k experts.
At training time the biases are initialized to zero and updated by adding or subtracting a small value depending on whether the given expert is over or under loaded.
Mixture of experts layers are notoriously difficult to train. Much of this stems from the exponentials used (softmax or sigmoid) when calculating the router weights. Several solutions have been proposed to alleviate the problem of numerical instability
- By forcing the model to do all of the routing calculations in full precision (32-bit float) overflow and underflow issues can be reduced.
- The addition of an extra loss term which is designed to keep the exponentials of router logits (which are used when calculating the routing weights) low. This loss is usually referred to as the z-loss and is given by
for a batch
- Careful weight initialization can help maintain stability. Typically, the weights are initialized by sampling from a normal distribution with mean
$\mu=0$ and standard deviation given my some scale factor. In the original switch transformer paper the authors suggest drawing weights from a truncate normal distribution with standard deviation$\sigma = \sqrt{s / d_{in}}$ where$s$ is a scale factor (the authors suggest$s=0.1$ ) and$d_{in}$ is the input dimension of the weight matrix being initialized.
A standard causal language model takes as input a sequence of tokens
Below is a schematic of the MTP head architecture used in DeepSeek V3:
flowchart BT
subgraph input_sequence[input sequence]
t1
t2
t3
dots1[...]
tn
end
input_sequence --> embed[Embedding layer]--> embeds1
subgraph embeds1[ebedded tokens]
e1
e2
e3
dots3[...]
en
end
embeds1 --> llm[Transformer stack] --> hidden_state1
subgraph hidden_state1[final hidden states]
h1
h2
h3
dots2[...]
hn
end
hn --> lm_head[LM head] -- argmax --> tnp1["t(n + 1)"] --> embed2[embeddingg layer] --> enp1
subgraph embeds2[ebedded tokens]
direction LR %%
ee2[e2] ~~~ ee3[e3] ~~~ ee4[e4] ~~~ dots4[...] ~~~ enp1["e(n+1)"]
end
hidden_state1 --> norm1
embeds2 --> norm2
subgraph mtp_head[MTP Head]
norm1[RMSNorm] --> concatenate
norm2[RMSNorm] --> concatenate
concatenate --> proj[linea layer] --> block[transformer block]
end
block --> lm_head2[LM head] --> out_tokens
subgraph out_tokens[output sequence]
tt3[t3]
tt4[t4]
tt5[t5]
dots5[...]
tnp2["t(n+2)"]
end
style mtp_head fill: lightblue
Several MPT heads can be chained to predict additional tokens. The input to the dth MTP is the hidden states of the previous layer and the embeddings of the dth through n+d tokens from the input sequence. New tokens are determined by taking the argmax of the result of applying the LLM output head to the final token in the current sequence.
The MTP loss is given by
where
Models in this collection share some common parameters:
attention_bias(booloptional defaults toFalse): whether to include a bias term in the attention blocks.attention_dropout(floatoptional defaults to0.0): dropout rate in the attention blocks.hidden_dim(int, optional default to2048): the embedding dimension for tokens.intermediate_dim(intoptional defaults to8192): dimension of the feedforward layers.max_position_embeddings(intoptional defaults to4096): maximum number of positional embeddings supported.mlp_bias(booloptional defaults toFalse): whether to include bias in the feedforward networks.num_attention_heads(intoptional defaults to32): number of attention heads per query.num_hidden_layers(intoptional ddefaults to16): number of attention blocks.num_key_value_heads(intoptional defaults to8): number of key/value heads in the attention mechanism if num_key_value_heads < num_attention_heads uses grouped query attention. If ``num_key_value_heads = 1this becomes MQA, fornum_key_value_heads = num_attention_heads` it is full MHA, and for any other value `0 < num_key_value_heads < num_attention_heads` it leads to GQA.rope_base(floatoptional defaults to500000.0): the base for the RoPE embedding.rms_norm_eps(floatoptional defaults to1E-05): regularizer for rms norm layersvocab_size(int, optional, defaults to 32000): vocabulary size for the model.
LLama 3 has a fairly straightforward architecture. Its stand out features are the use of GQA, SwiGLU feedforward layers, and RoPE encodings.
flowchart BT
text_input[Text input] --> Tokenizer["Tokenizer (vocab_size)"]
Tokenizer --> embedding["Token embedding layer (hidden_dim)"]
subgraph model [LLM Model]
embedding --- split1
subgraph block["Transformer Blocks (num_hidden_layers)"]
split1@{shape: f-circ} --> norm1[RMSNorm 1] --> attention["GQA (num_attention_heads, num_key_value_heads)"]
pos_emb[RoPE] --> attention --> merge1@{shape: circle, label: " + "}
split1 --> merge1 --- split2@{shape: f-circ} --> norm2[RMSNorm 2] --> ffn["SwiGLU (intermediate_dim)"] --> merge2@{shape: circle, label: " + "}
split2 --> merge2
end
merge2 --> norm_final[Final RMSNorm] --> output_layer["Linear output layer (vocab_size)"]
end
output_layer --> output[sequence decoder] --> Output
style model fill: lightblue
style block fill: pink
As an example a model with the architecture of Llama 3.2 1B can be created with the following configuration
from gobots.models.llama3_clone import Llama3Config, Llama3Model
config = Llama3Config(
vocab_size=128256,
hidden_dim=2048,
intermediate_dim=8192,
num_attention_heads=32,
num_hidden_layers=16,
num_key_value_heads=8,
attention_bias=False,
attention_dropout=0.0,
mlp_bias=False,
rms_norm_eps=1e-05,
max_position_embeddings=131072,
rope_base=500000.0,
)
model = Llama3Model(config)Qwen3 has a similar architecture to LLama 3 with GQA, SwiGLU feedforard layers, and RoPE encodings. The primary distinguishing feature is the inclusion of RMS norm layers applied to the query and key after the attention head projections but before the dot product.
flowchart BT
text_input[Text input] --> Tokenizer["Tokenizer (vocab_size)"]
Tokenizer --> embedding["Token embedding layer (hidden_dim)"]
subgraph model[LLM Model]
embedding --- split1
subgraph block["Transformer Blocks (num_hidden_layers)"]
split1@{shape: f-circ} --> norm1[RMSNorm 1] --> attention["GQA (num_attention_heads, num_key_value_heads)"]
pos_emb[RoPE] --> attention --> merge1@{shape: circle, label: " + "}
attn_norm[Q/K RMSNorm] --> attention
split1 --> merge1 --- split2@{shape: f-circ} --> norm2[RMSNorm 2] --> ffn["SwiGLU (intermediate_dim)"] --> merge2@{shape: circle, label: " + "}
split2 --> merge2
end
merge2 --> norm_final[Final RMSNorm] --> output_layer["Linear output layer (vocab_size)"]
end
output_layer --> output[sequence decoder] --> Output
style model fill: lightblue
style block fill: pink
As an example a model with the architecture of Qwen3 4B can be created with the following configuration
from gobots.models.qwen3_dense_clone import Qwen3DenseConfig, Qwen3DenseModel
config = Qwen3DenseConfig(
vocab_size=151936,
hidden_dim=2560,
intermediate_dim=9728,
num_attention_heads=32,
num_hidden_layers=36,
num_key_value_heads=8,
attention_bias=False,
attention_dropout=0.0,
mlp_bias=False,
rms_norm_eps=1e-06,
max_position_embeddings=40960,
rope_base=1000000,
)
model = Qwen3DenseModel(config)Similar to Llama 3 this model uses GQA and SwiGLU feedforward layers. The primary distinguishing feature is that it uses a mixture of RoPE and NoPE positional
encodings. Most layers use RoPE however, every no_rope_layer_interval layers the RoPE encoding is replaced with NoPE. The thought behind this choice is to
capture both the positional information of RoPE with the long context length of NoPE.
flowchart BT
text_input[Text input] --> Tokenizer["Tokenizer (vocab_size)"]
Tokenizer --> embedding["Token embedding layer (hidden_dim)"]
subgraph model[LLM Model]
embedding --- split1
subgraph block["Transformer Blocks (num_hidden_layers)"]
split1@{shape: f-circ} --> norm1[RMSNorm 1] --> attention["GQA (num_attention_heads, num_key_value_heads)"]
no_pos_emb[NoPE] --> split3@{shape: f-circ}
pos_emb[RoPE] --> split3 -- "no_rope_layer_interval" --> attention --> merge1@{shape: circle, label: " + "}
split1 --> merge1 --- split2@{shape: f-circ} --> norm2[RMSNorm 2] --> ffn["SwiGLU (intermediate_dim)"] --> merge2@{shape: circle, label: " + "}
split2 --> merge2
end
merge2 --> norm_final[Final RMSNorm] --> output_layer["Linear output layer (vocab_size)"]
end
output_layer --> output[sequence decoder] --> Output
style model fill: lightblue
style block fill: pink
As an example a model with the architecture of SmolLM3 3B can be created with the following configuration
from gobots.models.smollm3_clone import SmolLM3Config, SmolLM3Model
config = SmolLM3Config(
vocab_size=128256,
hidden_dim=2048,
intermediate_dim=11008,
num_attention_heads=16,
num_hidden_layers=36,
num_key_value_heads=4,
attention_bias=False,
attention_dropout=0.0,
mlp_bias=False,
rms_norm_eps=1e-06,
no_rope_layer_interval=4,
max_position_embeddings=65536,
rope_base=5000000.0,
)
model = SmolLM3Model(config)Llama 4 uses GQA and SwiGLU Feedforward layers. Every other layer uses an MoE layer as opposed to a dense layer. Each MoE has a single shared expert and uses one routed expert per token.
flowchart BT
text_input[Text input] --> Tokenizer["Tokenizer (vocab_size)"]
Tokenizer --> embedding["Token embedding layer (hidden_dim)"]
subgraph model[LLM Model]
embedding --- split1
subgraph block["Transformer Blocks (num_hidden_layers)"]
split1@{shape: f-circ} --> norm1[RMSNorm 1] --> attention["GQA (num_attention_heads, num_key_value_heads)"]
pos_emb[RoPE] --> attention --> merge1@{shape: circle, label: " + "}
split1 --> merge1 --- split2@{shape: f-circ} --> norm2[RMSNorm 2] --> ffn["SwiGLU (intermediate_dim) / MoE (num_experts_per_tok, num_local_experts)"] --> merge2@{shape: circle, label: " + "}
split2 --> merge2
end
merge2 --> norm_final[Final RMSNorm] --> output_layer["Linear output layer (vocab_size)"]
end
output_layer --> output[sequence decoder] --> Output
style model fill: lightblue
style block fill: pink
As an example a model with the architecture of Llama 4 17B with 16 experts per MoE layer:
from gobots.models.llama4_clone import Llama4Config, Llama4Model
config = Llama4Config(
vocab_size=202048,
hidden_dim=5120,
interleave_moe_layer_step=1,
intermediate_dim=8192,
intermediate_size_mlp=16384,
num_experts_per_tok=1,
num_local_experts=16,
num_attention_heads=40,
num_hidden_layers=48,
num_key_value_heads=8,
attention_bias=False,
attention_dropout=0.0,
mlp_bias=False,
use_qk_norm=True,
rms_norm_eps=1e-05,
max_position_embeddings=262144,
rope_base=500000.0,
)
model = Llama4Model(config)Another popular model which employs MoE feedforward layers is DeepSeek V3. The primary differences between this model and Llama 4 are the use of MLA instead of GQA, selects more than one routed expert per layer,and begins with first_k_dense_replace dense layers. Additionally this model adds
flowchart BT
text_input[Text input] --> Tokenizer["Tokenizer (vocab_size)"]
Tokenizer --> embedding["Token embedding layer (hidden_dim)"]
subgraph model[LLM Model]
embedding --- split1
subgraph block["Transformer Blocks (num_hidden_layers)"]
split1@{shape: f-circ} --> norm1[RMSNorm 1] --> attention["MLA (num_attention_heads, q_lora_rank, kv_lora_rank)"]
pos_emb[RoPE] --> attention --> merge1@{shape: circle, label: " + "}
split1 --> merge1 --- split2@{shape: f-circ} --> norm2[RMSNorm 2] --> ffn["SwiGLU (intermediate_dim) / MoE (num_experts_per_tok, num_local_experts)"] --> merge2@{shape: circle, label: " + "}
split2 --> merge2
end
merge2 --> norm_final[Final RMSNorm] --> output_layer["Linear output layer (vocab_size)"]
norm_final --> mtp["MTP heads (num_nextn_predict_layers)"] -- "extend sequence" --> output_layer
end
output_layer --> output[sequence decoder] --> Output
style model fill: lightblue
style block fill: pink
from gobots.models.deepseek_v3_clone import DeepSeekV3Config, DeepSeekV3Model
config = DeepSeekV3Config(
vocab_size=129280,
pad_token_id=2,
hidden_dim=7168,
intermediate_dim=18432,
moe_intermediate_size=2048,
num_experts_per_tok=8,
n_routed_experts=256,
num_attention_heads=128,
num_hidden_layers=61,
attention_bias=False,
attention_dropout=0.0,
mlp_bias=False,
rms_norm_eps=1e-06,
max_position_embeddings=163840,
rope_base=10000.0,
n_shared_experts=1,
first_k_dense_replace=3,
kv_lora_rank=512,
q_lora_rank=1536,
qk_nope_head_dim=128,
qk_rope_head_dim=64,
v_head_dim=128,
num_nextn_predict_layers=1,
initializer_range=0.02,
mtp_config={
"attention_type": "multi_head_latent_attention",
"d_model": 7168,
"num_heads": 128,
"v_head_dim": 128,
"q_lora_rank": 1536,
"kv_lora_rank": 512,
"qk_rope_head_dim": 128,
"qk_nope_head_dim": 64,
"dropout": 0.0,
"attention_bias": False,
"feedforward_type": "moe",
"n_shared_experts": 1,
"n_routed_experts": 256,
"intermediate_size": 2048,
"num_experts_per_token": 1,
},
)
model = DeepSeekV3Model(config)Ultimately, this repository can be used to train an LLM from scratch. The code can be used to train one of three versions of BaGL (BaGL is A Good Language model).
First modify trainer.Dockerfile to use the desired training config file. Configs for the dense, MOE, hybrid versions of the model are provided in ./configs. Next build the training image
docker compose -f docker-compose.yaml build trainer
Next run the training proccess
docker compose -f docker-compose.yaml up -d trainer
optionally the process can be monitored with trackio while running
docker compose -f docker-compose.yaml up -d trainer trackio
The evaluations use lighteval framework and can be run using the evaluator service
docker compose -f docker-compose.yaml build evaluator
docker compose -f docker-compose.yaml up -d evaluator
the results of training the three versions of BaGL are summarized in the following tables:
| model | truthfulqa mc1 | truthfulqa mc2 | winogrande (5-shot) | hellaswag (10-shot) | arc | ifeval (prompt strict) | ifeval (instruct strict) | ifeval (prompt loose) | ifeval (instruct loose) | gpqa | mmlu pro (5-shot) | bbh (3-shot) | musr |
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| dense | 0.2644 | 0.4520 | 0.4964 | 0.3174 | 0.2159 | ||||||||
| moe | 0.2521 | 0.4437 | 0.4957 | 0.2475 | 0.2218 | 0.1109 | 0.1966 | 0.1331 | 0.2530 | 0.2545 | 0.099 | 0.3075 | 0.4041 |
| hybrid | 0.2656 | 0.4575 | 0.5122 | 0.2536 | 0.2039 | 0.1146 | 0.2266 | 0.1201 | 0.2314 | 0.2500 | 0.1139 | 0.3087 | 0.3923 |
task-wise break down of bbh results:
| bbh subtask | moe | hybrid |
|---|---|---|
| causal_judgment | 0.5158 | 0.5158 |
| date_understanding | 0.0000 | 0.0000 |
| disambiguation_qa | 0.3101 | 0.3140 |
| geometric_shapes | 0.0917 | 0.1000 |
| logical_deduction_five_objects | 0.1980 | 0.2040 |
| logical_deduction_seven_objects | 0.1400 | 0.1429 |
| logical_deduction_three_objects | 0.3367 | 0.3333 |
| movie_recommendation | 0.2440 | 0.2600 |
| navigate | 0.5000 | 0.4970 |
| reasoning_about_colored_objects | 0.0795 | 0.0810 |
| ruin_names | 0.2478 | 0.2790 |
| salient_translation_error_detection | 0.2475 | 0.1784 |
| snarks | 0.4530 | 0.4641 |
| sports_understanding | 0.5010 | 0.5110 |
| temporal_sequences | 1.0000 | 1.0000 |
| tracking_shuffled_objects_five_objects | 0.2000 | 0.2008 |
| tracking_shuffled_objects_seven_objects | 0.1337 | 0.1429 |
| tracking_shuffled_objects_three_objects | 0.3367 | 0.3333 |
task-wise break down of musr results:
| musr sub task | moe | hybrid |
|---|---|---|
| murder_mysteries | 0.5000 | 0.4840 |
| object_placements | 0.3242 | 0.2969 |
| team_allocation | 0.3880 | 0.3960 |