HuggingFaceTB

/

SmolLM-135M

like 158

The SmolLM models have tied embeddings (config.tie_word_embeddings = True), i.e. input and output embeddings are shared.

However, the model contains two identical embedding tensors lm_head.weight and model.embed_tokens.weight. This seems to defeat the purpose of having shared (aka tied) embeddings. Any ideas?

Clicking on the arrow on the right-hand side (see screenshot below) shows a summary of the parameters without the lm_head tensor, see the second screenshot below (model.norm.weight is the last tensor shown):

Here is code showing that the model contains two identical tensors for input and output embeddings:

from transformers import AutoModelForCausalLM
import torch

model = AutoModelForCausalLM.from_pretrained('HuggingFaceTB/SmolLM-135M')
print(torch.equal(model.lm_head.weight, model.model.embed_tokens.weight))

Above code returns True, because the two tensors are identical.

The issue might be in modeling_llama.py, which doesn't seem to fully support shared embeddings. Specifically, line 1210

            logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]).float()

Contrast this for example with OpenELM's code that seems to fully support tied and untied embeddings on line 878

        if self.lm_head is None:
            # shared
            logits = F.linear(hidden_states, weight=self.transformer.token_embeddings.weight)
        else:
            logits = self.lm_head(hidden_states)