Skip to content

Do not disable weight tying on meta-init ranks - #856

Open
kylesayrs wants to merge 1 commit into
mainfrom
remove-tie-word-embeddings-meta-default
Open

Do not disable weight tying on meta-init ranks#856
kylesayrs wants to merge 1 commit into
mainfrom
remove-tie-word-embeddings-meta-default

Conversation

@kylesayrs

@kylesayrs kylesayrs commented Aug 26, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Purpose

  • Latest versions of transformers now throw warnings on non-source ranks for trying to load models with tied tensors
  • This skip is no longer necessary as of [Offloading] Support full disk offloading huggingface/transformers#46749 landing in 5.15
    • Do not pin transformers on CT. Many features are only compatible with latest transformers versions, but pinning can have downstream effects on vLLM

Changes

  • Remove tied_embeddings=False override when loading non-source ranks

Testing

  • Able to load gemma4 without issue
  • Added tests for distributed and non-distributed loading of all model types (tied, non-tied, tied but both in checkpoint, tied vlm)

@coderabbitai

coderabbitai Bot commented Aug 26, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Important

Review skipped

Auto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: da5e25ab-ef55-402e-9483-0252cc5c79fa

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • 🔍 Trigger review

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@kylesayrs kylesayrs mentioned this pull request Aug 26, 2026
Open
@mergify

mergify Bot commented Aug 26, 2026

Copy link
Copy Markdown

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews

Protection Waiting on
🔴 Require one maintainer review 👀 reviews

🔴 Require one maintainer review

Waiting for any of

  • approved-reviews-by=HDCharles
  • approved-reviews-by=brian-dellabetta
  • approved-reviews-by=dsikka
  • approved-reviews-by=kylesayrs
This rule is failing.

All PRs must have at least one approving review from a maintainer before merging.

  • any of:
    • approved-reviews-by=HDCharles
    • approved-reviews-by=brian-dellabetta
    • approved-reviews-by=dsikka
    • approved-reviews-by=kylesayrs
  • #changes-requested-reviews-by = 0

@kylesayrs kylesayrs added the ready When a PR is ready for full CI testing before merge label Aug 26, 2026
@kylesayrs @claude
Do not disable weight tying on meta-init ranks
31c7e79 
Non-source ranks load on the `meta` device and previously set
`tie_word_embeddings=False`. For models with tied word embeddings, that
turns the tied output projection (e.g. `lm_head.weight`) into a standalone
parameter absent from the checkpoint, so it is reported as a missing key
and never populated.

Remove the default so weights are tied consistently across ranks. Add
`test_load_no_missing_keys`, which loads a range of models (multimodal,
tied, and untied) under `load_offloaded_model` and asserts that
transformers reports no missing keys.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
@kylesayrs
kylesayrs force-pushed the remove-tie-word-embeddings-meta-default branch from 78874c7 to 31c7e79 Compare August 26, 2026 20:53
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

ready When a PR is ready for full CI testing before merge

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@kylesayrs