LLM architecture (part 3): Assembling model
Bringing embeddings, transformer blocks and the output layer together—and tracing tensor dimensions and the model’s parameter count.
Updated
July 16, 2026
•
9
min read
•
View as Markdown
A
Aman Jain
Product manager learning to build AI from the inside out — from LLM architecture to product applications.
Part of series
Building a GPT Model From Scratch
At this point, we have almost everything we need to assemble a baseline GPT model. In this article, I will:
Bring the different components together and show how they form the complete model.
Do a quick walkthrough of the tensor dimensions at each stage to understand how the data changes as it moves through the network.
Calculate the total number of trainable parameters in the model. This parameter count is typically what people refer to when they describe the size of a model.
Assembling the components
Let us begin with a quick review of the components introduced in the previous articles:
Tokenization and embeddings
convert natural-language text into tokens and then represent those tokens as numerical vectors.
Attention
allows each token to incorporate context from the other relevant tokens in the sequence.
Activation layers
help the model identify and strengthen useful patterns within the data.
Normalization layers
keep the values flowing through the model within a stable range.
Dropout layers
randomly remove some information during training to reduce overfitting. I have left dropout out of the diagram below to keep it simple.
The following diagram represents how these components fit together in the baseline model:
The following piece of code does the heavy lifting:
class GPTModel(nn.Module):
    def __init__(self, cfg):
        super().__init__()

        self.tok_emb = nn.Embedding(
            cfg["vocab_size"],
            cfg["emb_dim"]
        )
        self.pos_emb = nn.Embedding(
            cfg["context_length"],
            cfg["emb_dim"]
        )
        self.drop = nn.Dropout(cfg["drop_rate"])

        # Transformer block repeated n_layers times
        self.trf_blocks = nn.Sequential(
            *[
                TransformerBlock(cfg)
                for _ in range(cfg["n_layers"])
            ]
        )

        # Final normalization layer
        self.final_norm = LayerNorm(cfg["emb_dim"])
        self.out_head = nn.Linear(
            cfg["emb_dim"],
            cfg["vocab_size"],
            bias=False
        )

    def forward(self, in_idx):
        batch_size, seq_len = in_idx.shape

        tok_embeds = self.tok_emb(in_idx)
        pos_embeds = self.pos_emb(
            torch.arange(seq_len, device=in_idx.device)
        )

        x = self.drop(tok_embeds + pos_embeds)
        x = self.trf_blocks(x)
        x = self.final_norm(x)
        logits = self.out_head(x)

        return logits
Tracking tensor dimensions
Recall the model configuration introduced in Part 1:
GPT_CONFIG_124M = {
    "vocab_size": 50_257,
    "context_length": 1_024,
    "emb_dim": 768,
    "n_heads": 12,
    "n_layers": 12,
    "drop_rate": 0.1,
    "qkv_bias": False,
}
In addition, let us assume that every input has been padded or truncated to
context_length
and that each input batch contains
batch_size
inputs. Intuitively, the final output should give us
batch_size
predicted tokens—one for each input.
To keep things simple, we do not need to follow every tensor transformation inside the model. At a high level, the input begins as a sequence of token IDs. Each token is then converted into an
emb_dim
-dimensional vector, and this overall shape is preserved as the data moves through the transformer blocks.
Inside each block, the attention and activation components temporarily reshape or expand the data, but the output returns to the same dimensions.
Finally, the output layer converts each token representation into a score for every token in the vocabulary.
Stage
What happens
Output dimensions
Tokenization
Each input is converted into a sequence of token IDs.
batch_size * context_length
Token embeddings
Each token ID is converted into a vector in the embedding space.
batch_size * context_length * emb_dim
Positional embeddings
Position information is added to each token. The overall dimensions remain unchanged.
batch_size * context_length * emb_dim
Input to transformer blocks
Token and positional embeddings are combined and passed into the transformer blocks.
batch_size * context_length * emb_dim
Attention
The model creates query, key and value representations to understand relationships between tokens. The attention output returns to the original dimensions.
batch_size * context_length * emb_dim
Activation expansion
The representation is temporarily expanded to four times the embedding dimension.
batch_size * context_length * (emb_dim * 4)
Activation projection
The expanded representation is reduced back to the original embedding dimension.
batch_size * context_length * emb_dim
Output of each transformer block
The original tensor dimensions are preserved before the data enters the next block.
batch_size * context_length * emb_dim
Output layer
Each token representation is converted into a score for every token in the vocabulary.
batch_size * context_length * vocab_size
Final-position logits
The vocabulary scores from the final token position are selected for next-token prediction.
batch_size * vocab_size
Token prediction
One next token is selected for each input in the batch.
batch_size * 1
Count of trainable parameters
There is one final piece to examine before we complete the model: the total number of trainable parameters.
These are the values the model adjusts during training as it learns patterns from the data. The overall parameter count is also commonly used to describe the size of a model. For example, GPT-2 Small is referred to as a 124-million-parameter model.
Let us create an instance of our GPT model and count its parameters:
torch.manual_seed(123)
model = GPTModel(GPT_CONFIG_124M)

#setting output layer to use same weights as embedding layer. 
# as it is same dictionary. 
model.out_head.weight = model.tok_emb.weight


tied_param_count = sum(p.numel() for p in model.parameters())
print(f"Total number of unique parameters: {tied_param_count}")
The model produces the following output:
Total number of parameters: 124402944
Let us examine this in more detail:
Model component
Trainable parameters
How it is calculated
Token embeddings
38,597,376
50,257 * 768
: one 768-dimensional embedding vector for each of the 50,257 tokens in the vocabulary
Positional embeddings
786,432
1,024 * 768
: one 768-dimensional embedding vector for each of the 1,024 positions in the context window
12 transformer blocks
85,017,600
7,084,800 * 12
: 12 transformer blocks, each containing attention, activation and normalization layers
Final normalization layer
1,536
768 + 768
: one trainable scale value and one trainable shift value for each of the 768 embedding dimensions
Output head
0
re-using token embedding weights.
Total in the current implementation
124402944
Approximately
124 million parameters
Each transformer block contains approximately 7.1 million parameters:
Component within one transformer block
Parameters
How it is calculated
Attention
2,359,296
Four projection matrices—query, key, value and output—each containing
768 * 768
weights
First activation layer
2,362,368
768 * 3,072
weights +
3,072
bias values; expands each token representation to four times the embedding dimension
Second activation layer
2,360,064
3,072 * 768
weights +
768
bias values; projects the expanded representation back to 768 dimensions
Two normalization layers
3,072
Each layer contains
768
scale values +
768
shift values, giving
1,536 * 2
parameters
Total per block
7,084,800
Attention + two activation layers + two normalization layers
This gives us approximately
124 million unique trainable parameters
, which is the convention used when describing GPT-2 Small as a 124M model.
Before we close, let us look at the structure of the model we have built. Most of these components should now look familiar: token and positional embeddings, multi-head attention, activation layers, normalization, and the final output layer.
model = GPTModel(GPT_CONFIG_124M)

model.load_state_dict(
    torch.load("model.pth", map_location=device)
)

model.eval()
print(model)
The output below has been shortened for readability. Each of the 12 transformer blocks follows the same structure.
GPTModel(
  (tok_emb): Embedding(50257, 768)
  (pos_emb): Embedding(256, 768)
  (drop): Dropout(p=0.1)

  (trf_blocks): Sequential(
    (0): TransformerBlock(
      (attn): MultiHeadAttention(
        (W_query): Linear(768 → 768, bias=False)
        (W_key): Linear(768 → 768, bias=False)
        (W_value): Linear(768 → 768, bias=False)
        (out_proj): Linear(768 → 768, bias=False)
        (dropout): Dropout(p=0.1)
      )

      (ff): FeedForward(
        (layers): Sequential(
          (0): Linear(768 → 3072)
          (1): GELU()
          (2): Linear(3072 → 768)
        )
      )

      (norm1): LayerNorm()
      (norm2): LayerNorm()
      (drop_shortcut): Dropout(p=0.1)
    )

    ...

    TransformerBlock structure repeated 12 times

    ...
  )

  (final_norm): LayerNorm()
  (out_head): Linear(768 → 50257, bias=False)
)
The model starts with a vocabulary of
50,257
tokens and represents each token using
768
dimensions. Every transformer block preserves this 768-dimensional representation, while the activation layers temporarily expand it to
3,072
dimensions. Finally, the output layer maps the representation back to
50,257
logits—one score for every token in the vocabulary.
Generating text from the untrained model
With that, we now have a complete GPT-style model assembled from the individual components covered in the previous articles.
We have also traced how the tensor dimensions change as data moves through the model and calculated the number of trainable parameters that define its size.
However, the model currently begins with randomly initialized parameters. It has the right architecture, but it has not yet learned anything meaningful about language.
Let us try generating some text:
start_text = "Polish my unrevealed"
encoded = tokenizer.encode(start_text)

# Add the batch dimension
encoded_tensor = torch.tensor(encoded).unsqueeze(0)

# Put the model into evaluation mode
model.eval()

out = generate_text_simple(
    model=model,
    idx=encoded_tensor,
    max_new_tokens=6,
    context_size=GPT_CONFIG_124M["context_length"]
)

decoded_text = tokenizer.decode(out.squeeze(0).tolist())
print(decoded_text)
The model generates:
Polish my unrevealeditters Aminreally pocket cumbers reprodu
As expected, a model with random parameters produces random and largely meaningless output.
What comes next?
In the next article, we will look at how a GPT model is trained, how its performance is evaluated during training and how text is generated from its predictions.
We will then download pretrained GPT-2 weights and load them into our custom-built architecture. This will allow us to compare a randomly initialized model with one that has already learned language patterns from a large body of text.
References
Sebastian Raschka — *Build a Large Language Model (From Scratch).
*
https://www.manning.com/books/build-a-large-language-model-from-scratch
Github repo -
https://github.com/amanjain454/llms-from-scratch-work
#
gpt
#
deeplearning
33
views