Model and Inference Engineering · Staff
The checkpoint loads and loss starts normally. Why do tied embeddings drift apart after training resumes?
The question
Interview question
An LM uses one parameter for its input token embeddings and output projection. A checkpoint conversion copies both arrays into a new distributed model. They have the same values at step zero, so a spot check passes. After several steps, the input and output matrices differ. Is the checkpoint corrupt?
Take a few minutes to form your approach. Then open a worked answer and compare the decisions.
Reveal a worked answer
Equal values are not the same as shared storage or a shared trainable parameter. With tied weights, the input lookup and output projection use the same underlying parameter, so gradient contributions from both paths are accumulated into its update. Two independent parameters initialized to identical numbers receive different gradients and separate after one step. Hugging Face's model documentation describes tying input and output embeddings when configured, and PyTorch's module documentation exposes parameter traversal with duplicate handling. Check the actual parameter identity in the framework and sharding setup, not only tensor equality in a serialized file.
I would inspect the model construction order, load operation, tie operation, distributed wrapping and optimizer creation. A converter might clone the output head. Or the training model ties weights before loading, but loading or wrapping materializes a second parameter afterwards. Some distributed and export paths represent shared parameters differently, so I would not assume a pointer comparison across wrappers tells the whole story. Test the documented alias invariant in the unwrapped model, then verify logical sharing through gradients and a single optimizer step under the actual training runtime. Enumerate optimizer parameter IDs and groups to ensure one logical tied weight is updated once, with intended weight decay and sharding.
The correction depends on whether the architecture really intends tied weights. If yes, reconstruct the alias at the right stage, load a canonical tensor, rebuild optimizer state to match the new parameter mapping and verify a round-trip save and resume. Do not average the two diverged arrays and declare recovery. Their optimizer moments may have diverged too. Decide whether a clean checkpoint is needed and compare token-level validation, especially rare tokens, before resuming a long run.
There is a reasonable interviewer pushback: maybe the model intentionally untied the head during fine-tuning. Fine, but that is an architecture change with its own optimizer and evaluation plan. Record it explicitly. This question is about an accidental break in a stated invariant. The new LoRA adapter has gradients. Why do its weights never change? has a LoRA weight missing from optimizer groups. Here both matrices can update, which is why the failure can hide under a plausible loss curve.
Continue reading
Related questions
Read beyond the question
Explore more model and inference engineering
Follow another question in this area, or search the complete Question Library.
Browse this area →Browse Question Library →