The embedding layer is a 2D matrix stored as the tensor model.embed_tokens.weight. Its shape is [vocab_size, hidden_size] — for Llama 3.1 8B, that's [128,256 × 4,096]. Each of the 128,256 rows is a learned vector of 4,096 floating-point numbers representing one token. The entire matrix is the model's "dictionary" — but instead of definitions, each word gets a coordinate in 4,096-dimensional space.

Unlike most layers in the transformer, the embedding layer does no matrix multiplication. It's a pure table lookup: given a token ID, it returns the corresponding row of the matrix. If the tokenizer converts "Hello" to ID 9906, the embedding layer simply reads row 9906 from the matrix — a 4,096-dimensional vector. This vector is the model's internal representation of "Hello" and gets passed into the first transformer layer.

During training, the embedding matrix learns to place tokens with similar meaning or usage near each other in the 4,096-dimensional space. "dog" and "cat" end up close together. "Python" and "JavaScript" cluster near each other. "happy" and "joyful" are neighbors. This isn't programmed — it emerges from the training data as the model learns to predict next tokens. The geometry captures semantic relationships: king - man + woman ≈ queen.

Vocabulary size directly controls the embedding matrix's first dimension. Llama 2 had 32,000 tokens; Llama 3 expanded to 128,256. This 4× increase in vocab size meant a 4× increase in embedding parameters — from ~131M to ~525M parameters just for the embedding layer. Larger vocabularies encode text more efficiently (fewer tokens per sentence) but increase memory for the embedding and output matrices.

At the output end, the model needs to convert its internal 4,096-dimensional representation back into a probability distribution over the vocabulary. The lm_head.weight tensor does this: it's a matrix of shape [vocab_size, hidden_size] — the same shape as the embedding. It multiplies the final hidden state by this matrix to produce a "logit" score for each of the 128,256 tokens. Softmax converts these logits into probabilities.

Weight tying means the embedding matrix and the output head share the same physical parameters — they point to the same memory. This was introduced by Press & Wolf (2017) and used in GPT-2 and BERT. The benefit: you save vocab_size × hidden_size parameters. For Llama 3.1 8B, that would save ~525M parameters (~1 GB in BF16). The config field tie_word_embeddings controls this.

Modern open-source LLMs like Llama and Mistral do NOT tie embeddings — they use tie_word_embeddings: false. The embedding and lm_head are separate tensors. This gives the model more capacity: the input embedding can specialize in encoding meaning, while the output head can specialize in predicting the next token.

For Llama 3.1 8B with untied embeddings, the embedding + lm_head together contain about 1.05 billion parameters out of 8.03 billion total — roughly 13% of the model. In BF16, that's ~2.1 GB out of ~16 GB. For models with smaller vocabularies (32K), this drops to about 3%. For models with larger hidden sizes (like Llama 3.1 70B with hidden_size=8192), the absolute size doubles but the percentage drops because the per-layer tensors grow even more.

Inspect token representations: Extract a row from the embedding matrix to see how the model "thinks about" a word. Compare cosine similarity between rows to find semantic neighbors.

Detect tokenizer mismatches: If the embedding matrix has 128,256 rows but your tokenizer has 32,000 tokens, you have a mismatch — loading will fail or produce garbage.

Estimate overhead: Switching from 32K to 128K vocabulary adds ~1.5 GB to the model (with untied embeddings in BF16). Know this before choosing a model.