NVIDIA / NVIDIA/Megatron-LM

[QUESTION] Unexpected updates on unused token embeddings when using --untie-embeddings-and-output-weights

Open
#2,994 5 comments 0 reactions 1 assignee Assigned to @maanug-nv View on GitHub
community-request question waiting-on-customer
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

I have a question regarding embedding parameter updates during training, and I would like to confirm whether my understanding is correct or if there are other mechanisms involved.

**Setup**

- Vocabulary size: 128256
- During training, all tokens appearing in the dataset have token IDs strictly within [128000, 128256)
Tokens with IDs < 128000 never appear in the input data
- Training is done with the flag:
--untie-embeddings-and-output-weights

**My understanding**

Given the above setup, I would expect that:
For token IDs that never appear in any batch:

- Their corresponding rows in the embedding matrix should not receive gradients
- Therefore, their embedding parameters should remain unchanged throughout training

**Observed behavior**

However, I observe that:

- All rows of the embedding matrix change over the course of training, including those corresponding to token IDs that never appear in the data

- Moreover, these changes are not continuous per iteration
Instead, they seem to occur sporadically, e.g. once every hundreds or even thousands of iterations

I have double-checked the input data and can confirm that no token IDs outside [128000, 128256) are present

**Question**

I would like to understand what mechanisms could cause embedding weights of unused token IDs to change in this scenario. For example:

Could this be related to:

- Optimizer behavior (e.g. Adam moment updates, weight decay)?
- Is there any known behavior in Megatron-LM where embedding parameters may change even without direct gradient updates?

Any clarification or pointers to relevant code paths would be greatly appreciated. Thanks a lot for your help!

Best regards

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.