Skip to content

feat: Add script to re-initialize near-zero HF embeddings#2671

Open
ashors1 wants to merge 4 commits into
mainfrom
ashors/reinit-hf-embeds
Open

feat: Add script to re-initialize near-zero HF embeddings#2671
ashors1 wants to merge 4 commits into
mainfrom
ashors/reinit-hf-embeds

Conversation

@ashors1
Copy link
Copy Markdown
Contributor

@ashors1 ashors1 commented Jun 2, 2026

What does this PR do ?

For certain models (e.g. nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-Base-BF16), the HF checkpoint has embedding weights which are zero or near-zero. These weights can cause inf/nan grad norms during post-training. This PR adds a script to reinitialize such embeddings. It also adds an inf grad norm check with an informative message.

Issues

List issues that this PR closes (syntax):

Usage

  • You can potentially add a usage example below
# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

  • Make sure you read and followed Contributor guidelines
  • Did you write any new necessary tests?
  • Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
  • Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

  • ...

ashors1 added 2 commits June 2, 2026 11:25
Signed-off-by: Anna Shors <ashors@nvidia.com>
Signed-off-by: Anna Shors <ashors@nvidia.com>
@copy-pr-bot
Copy link
Copy Markdown

copy-pr-bot Bot commented Jun 2, 2026

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@ashors1 ashors1 marked this pull request as ready for review June 3, 2026 19:32
@ashors1 ashors1 requested review from a team as code owners June 3, 2026 19:32
Signed-off-by: Anna Shors <ashors@nvidia.com>
@ashors1 ashors1 requested a review from a team as a code owner June 3, 2026 19:52
@github-actions github-actions Bot added the Documentation Improvements or additions to documentation label Jun 3, 2026
Signed-off-by: Anna Shors <ashors@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Documentation Improvements or additions to documentation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant