Motivation

When trying to add support for SmolLM, I noticed that its config dict has "tie_word_embeddings": true which isn't supported by SGLang, and as a result, lm_head.weight is not loaded. This problem likely exists for some other models too (not just SmolLM).

Modifications

Within LlamaForCausalLM, at the end of the weight loading stage, I added the following code:

if hasattr(self.config, "tie_word_embeddings") and self.config.tie_word_embeddings:
    # Tie output embedding layer to input embedding layer, to solve issues where lm_head.weight is missing
    param = self.lm_head.weight
    weight_loader = getattr(param, "weight_loader", default_weight_loader)
    weight_loader(param, self.model.embed_tokens.weight)

This can probably be copy-pasted to other LM classes other than LlamaForCausalLM, but I'm less familiar with those, and will defer to maintainers to decide whether to do the copying.

The effect can be tested by running python3 -m sglang.bench_latency --correct --model HuggingFaceTB/SmolLM-135M-Instruct. This would output gibberish before the change, but output coherent answers after it.

Checklist

• Format your code according to the Contributor Guide.
• Add unit tests as outlined in the Contributor Guide.
• Update documentation as needed, including docstrings or example tutorials.