fix: re-tie word embeddings before save to avoid duplicate lm_head - #2798
Conversation
… offloaded models Signed-off-by: EdalatiAli <aliedalati@cohere.com>
|
👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review. Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed. |
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
WalkthroughThe PR fixes a bug where accelerate offloading conversion splits tied weights (such as ChangesOffloaded weights re-tying
Estimated code review effort🎯 1 (Trivial) | ⏱️ ~3 minutes Suggested labels
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Warning Review ran into problems🔥 ProblemsLinked repositories: Your configuration references 1 linked repositories, but your current plan allows 0. Analyzed ``, skipped 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. Comment |
Merge ProtectionsYour pull request matches the following merge protections and will not be merged until they are valid. 🟢 Require two reviewsWonderful, this rule succeeded.PRs labelled "two-reviews" must have at least two approving reviews before merging.
|
coderabbitai
Bot
added
bug
EdalatiAli
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
src/llmcompressor/transformers/compression/compressed_tensors_utils.py (2)
108-118: ⚖️ Poor tradeoffConsider adding error handling for robustness during
tie_weights()invocation.The implementation safely checks for the existence and callability of
tie_weights()but doesn't handle potential exceptions raised during its invocation. While propagating exceptions is generally appropriate (as a failingtie_weights()indicates a model issue), adding defensive error handling with logging could aid debugging without silently swallowing errors.♻️ Optional: Add defensive logging for tie_weights failures
def _retie_offloaded_weights(model: PreTrainedModel): """Re-tie weights split by offload conversion so the tied head isn't saved twice. Offloading gives the input embeddings and a tied head (e.g. ``lm_head``) separate parameters, defeating transformers' save-time de-duplication. ``tie_weights`` restores the shared parameter and is a no-op for untied models. """ tie_weights = getattr(model, "tie_weights", None) if callable(tie_weights): - tie_weights() + try: + tie_weights() + except Exception as e: + logger.warning( + f"Failed to re-tie weights after offload conversion: {e}. " + "Model may save duplicate tied weights." + ) + # Re-raise to fail fast if this is critical + raise🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/llmcompressor/transformers/compression/compressed_tensors_utils.py` around lines 108 - 118, The _retie_offloaded_weights function currently calls model.tie_weights() without handling exceptions; wrap the callable check and invocation of tie_weights (on model) in a try/except Exception block, obtain a module logger (e.g., logging.getLogger(__name__)), log the failure with logger.exception including context that tie_weights failed for this model, and then re-raise the exception so failures are visible rather than silently swallowed.
77-79: _retie_offloaded_weights should preserve explicit untying; optionally add a defensive config guard.The save flow calls
_retie_offloaded_weights(model)afterto_accelerate(model), and_retie_offloaded_weightscurrently just unconditionally callsmodel.tie_weights(). In Hugging Face Transformers,PreTrainedModel.tie_weights()is gated byconfig.tie_word_embeddings, so ifuntie_word_embeddingssetconfig.tie_word_embeddings = False,tie_weights()should be a no-op and explicit untying should be preserved.Optional robustness: make that dependency explicit in
_retie_offloaded_weights(also aligning with the helper docstring).🛡️ Defensive check to preserve explicit untyling
def _retie_offloaded_weights(model: PreTrainedModel): """Re-tie weights split by offload conversion so the tied head isn't saved twice. Offloading gives the input embeddings and a tied head (e.g. ``lm_head``) separate parameters, defeating transformers' save-time de-duplication. ``tie_weights`` restores the shared parameter and is a no-op for untied models. """ + # Only retie if embeddings were originally tied + if hasattr(model.config, "tie_word_embeddings") and not model.config.tie_word_embeddings: + return + tie_weights = getattr(model, "tie_weights", None) if callable(tie_weights): tie_weights()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/llmcompressor/transformers/compression/compressed_tensors_utils.py` around lines 77 - 79, The _retie_offloaded_weights call currently always invokes model.tie_weights(), which will override an explicit untying; update _retie_offloaded_weights to check the model.config.tie_word_embeddings flag before calling model.tie_weights() so it becomes a no-op when untied (preserving explicit untying), and optionally add a defensive guard/logging in _retie_offloaded_weights's docstring or runtime to clarify this behavior; reference the _retie_offloaded_weights function and the model.tie_weights() call and the model.config.tie_word_embeddings flag when making the change.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@src/llmcompressor/transformers/compression/compressed_tensors_utils.py`:
- Around line 108-118: The _retie_offloaded_weights function currently calls
model.tie_weights() without handling exceptions; wrap the callable check and
invocation of tie_weights (on model) in a try/except Exception block, obtain a
module logger (e.g., logging.getLogger(__name__)), log the failure with
logger.exception including context that tie_weights failed for this model, and
then re-raise the exception so failures are visible rather than silently
swallowed.
- Around line 77-79: The _retie_offloaded_weights call currently always invokes
model.tie_weights(), which will override an explicit untying; update
_retie_offloaded_weights to check the model.config.tie_word_embeddings flag
before calling model.tie_weights() so it becomes a no-op when untied (preserving
explicit untying), and optionally add a defensive guard/logging in
_retie_offloaded_weights's docstring or runtime to clarify this behavior;
reference the _retie_offloaded_weights function and the model.tie_weights() call
and the model.config.tie_word_embeddings flag when making the change.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 550f9d57-b1f0-467e-8bee-54fb72d20b84
📒 Files selected for processing (1)
src/llmcompressor/transformers/compression/compressed_tensors_utils.py
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a helper function _retie_offloaded_weights to re-tie weights that were split during offload conversion, preventing duplicate parameters (such as lm_head) from being saved. This helper is integrated into the save_pretrained_wrapper flow. I have no feedback to provide.
brian-dellabetta
left a comment
brian-dellabetta
left a comment
There was a problem hiding this comment.
thanks! Will confirm with @kylesayrs
There was a problem hiding this comment.
Will write up more information in a bit, but temporarily blocking this for now for the following changes:
- This logic needs a guard on
if model.config.tie_embeddings:. Sometimesoneshotwill untie embeddings, as is the case for spinquant, so checking for the presence oftie_weightsis not alone. - Some models implement
tie_weights, but the implementation raises an error. We should try to catch this case and warn if possible
|
Outside of the required changes mentioned above, I'll also mention that this usage creates a data configuration that
However, typically accelerate expects all tied tensors to not only share the same meta tensor value, but also that you specify them in an explicit This is nbd, but it might break in the future if accelerate changes how it implements I think a more stable pattern might be to perform the tying before converting to accelerate if model.config.tie_embeddings: # might need a more robust check
_retie_offloaded_weights(model) # add warning/logger if weights are retied
to_accelerate(model) |
Signed-off-by: EdalatiAli <aliedalati@cohere.com>
|
Thanks for the detailed review @kylesayrs @brian-dellabetta ! Addressed the three points:
I also added a regression test ( |
brian-dellabetta
left a comment
brian-dellabetta
left a comment
There was a problem hiding this comment.
Thanks @EdalatiAli !
|
The quality checks have failed. Please run |
|
The quality checks have failed. Please run |
|
@EdalatiAli can you run |
Signed-off-by: EdalatiAli <aliedalati@cohere.com>
|
Apparently our smoke model is incorrectly configured, its config has I will confirm with original creator of smoke model that changing won't break anything or get overwritten by some other process |
|
Updating the model config to |
Signed-off-by: Brian Dellabetta <bdellabe@redhat.com>
|
Thank you @brian-dellabetta for fixing the issue with the failed test! Would be great to merge the PR if there are no other concerns |
Thanks @EdalatiAli , I'll cover this with Kyle when we meet this afternoon, hope to get it in today |
|
I think this solution may be incorrect, as it messes with cases where only the input embedding is quantized |
Yeah, this PR and the PR that introduced embedding quantization ( #2830 and vllm-project/compressed-tensors#718) were considered independently, so we're hitting an issue at the intersection of them. But we can resolve in #2835 |
…cate lm_head (#2835) ## Summary Quantizing a **tied** model's embedding is currently counterproductive. When embeddings are targeted, `QuantizationMixin.start_calibration` unties them (existing upstream behavior) so the input and output embeddings can be quantized independently — but for a tied model the one shared matrix then becomes a packed table **plus** a new full-precision `lm_head`, so the checkpoint grows. This matters most for small SOTA models (`Qwen3-0.6B`, `LFM2.5-350M`) where that shared matrix dominates file size. Per review (thanks @kylesayrs), this adds a **save-time** step that re-ties the input and output embeddings when they were quantized to identical values — the same matrix quantized twice — so they serialize as a single shared table: - **Re-tie identically-quantized embeddings before saving** (`compressed_tensors_utils.py`). If the (compressed) input and output embeddings have identical tensors, the output embedding's storage is aliased to the input's so transformers' save-time de-duplication writes a **single shared table**, and `tie_word_embeddings` is restored so the tie is reconstructed at load. If they differ — quantized differently, or only one was quantized — they are left untied and both are kept, preserving the model's integrity. The decision is driven purely by tensor equality, so nothing needs to track the original tie state. The re-tie uses the existing compressed-tensors offload API — `setattr`/`getattr` under `OffloadCache.disable_onloading()` — so the two modules share one tensor with no copy and no new compressed-tensors API. (The `disable_onloading` context can be removed once compressed-tensors#709 makes `__setitem__` a non-copying replacement.) ### Relationship to #2798 Builds on (does not duplicate) the merged #2798. #2798 re-ties **uncompressed** tied weights that offloading splits apart, via `model.tie_weights()`. That path can't handle a **compressed** embedding, which has no plain `.weight` to tie to (`tie_weights()` raises and logs a warning), and it only runs while `tie_word_embeddings` is still set. This PR covers exactly that gap. The two run back-to-back at save and are complementary: `_retie_offloaded_weights` handles the dense case, `_retie_quantized_embeddings` the compressed one. ### Why this is useful The single packed table is reused for **both** the input lookup and the logits matmul by runtimes that support it — vLLM consumer support: vllm-project/vllm#45535. Accuracy (lm-eval): tied **W8** embedding is near-lossless (Qwen3-0.6B wikitext bits-per-byte −0.0%, LFM2.5-350M +0.45%); **W4** is heavier on these compression-sensitive models (+1.2% / +4.4%). ## Test Plan ``` pytest tests/llmcompressor/transformers/compression/test_compress_tensor_utils.py -k tied ``` `test_tied_quantized_embedding_no_duplicate_head` quantizes `Qwen3-0.6B`'s embedding **and** `lm_head` (data-free) and asserts the saved checkpoint stays tied (`tie_word_embeddings=True`) with the packed embedding present and **no** `lm_head` weight tensors. ## Test Result - New test passes. - Existing `test_no_duplicate_tied_lm_head_on_save` (tied/untied × offload) still passes — an explicitly-untied model keeps its separate head (the re-tie is scoped to the compressed case). - End-to-end: the produced `Qwen3-0.6B` checkpoint has a single shared packed table (no fp16 head), config declares both embeddings quantized, `tie_word_embeddings=True`, 930 MB vs 1.5 GB fp16 — loads and generates coherently in vLLM (with #45535). - `ruff check` / `ruff format --check` clean. --- This change was developed with AI assistance (Claude). All changed lines were reviewed by the submitter. --------- Signed-off-by: Karthik Kothuri <karthikkothuri2009@gmail.com> Co-authored-by: Kyle Sayers <kylesayrs@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
…2798) ## Summary When a model with tied word embeddings is saved while offloaded, the offload conversion (`to_accelerate`) splits `embed_tokens` and `lm_head` into separate parameters. This defeats transformers' save-time de-duplication, so a redundant, bit-identical `lm_head.weight` gets written. This re-ties the weights right after `to_accelerate` and before saving, so the tied head is dropped as intended. It's a cheap attribute reassignment (no onloading) and a no-op for models with untied embeddings. ## Test plan - Exported an AWQ (sequential-pipeline) checkpoint and confirmed `lm_head.weight` no longer appears in `model.safetensors.index.json`. - Verified datafree/untied exports are unaffected. - passed the added test: `pytest tests/llmcompressor/transformers/compression/test_compress_tensor_utils.py -k no_duplicate_tied_lm_head -v` --------- Signed-off-by: EdalatiAli <aliedalati@cohere.com> Signed-off-by: Brian Dellabetta <bdellabe@redhat.com> Co-authored-by: Brian Dellabetta <brian-dellabetta@users.noreply.github.com> Co-authored-by: Brian Dellabetta <bdellabe@redhat.com>
…cate lm_head (#2835) ## Summary Quantizing a **tied** model's embedding is currently counterproductive. When embeddings are targeted, `QuantizationMixin.start_calibration` unties them (existing upstream behavior) so the input and output embeddings can be quantized independently — but for a tied model the one shared matrix then becomes a packed table **plus** a new full-precision `lm_head`, so the checkpoint grows. This matters most for small SOTA models (`Qwen3-0.6B`, `LFM2.5-350M`) where that shared matrix dominates file size. Per review (thanks @kylesayrs), this adds a **save-time** step that re-ties the input and output embeddings when they were quantized to identical values — the same matrix quantized twice — so they serialize as a single shared table: - **Re-tie identically-quantized embeddings before saving** (`compressed_tensors_utils.py`). If the (compressed) input and output embeddings have identical tensors, the output embedding's storage is aliased to the input's so transformers' save-time de-duplication writes a **single shared table**, and `tie_word_embeddings` is restored so the tie is reconstructed at load. If they differ — quantized differently, or only one was quantized — they are left untied and both are kept, preserving the model's integrity. The decision is driven purely by tensor equality, so nothing needs to track the original tie state. The re-tie uses the existing compressed-tensors offload API — `setattr`/`getattr` under `OffloadCache.disable_onloading()` — so the two modules share one tensor with no copy and no new compressed-tensors API. (The `disable_onloading` context can be removed once compressed-tensors#709 makes `__setitem__` a non-copying replacement.) ### Relationship to #2798 Builds on (does not duplicate) the merged #2798. #2798 re-ties **uncompressed** tied weights that offloading splits apart, via `model.tie_weights()`. That path can't handle a **compressed** embedding, which has no plain `.weight` to tie to (`tie_weights()` raises and logs a warning), and it only runs while `tie_word_embeddings` is still set. This PR covers exactly that gap. The two run back-to-back at save and are complementary: `_retie_offloaded_weights` handles the dense case, `_retie_quantized_embeddings` the compressed one. ### Why this is useful The single packed table is reused for **both** the input lookup and the logits matmul by runtimes that support it — vLLM consumer support: vllm-project/vllm#45535. Accuracy (lm-eval): tied **W8** embedding is near-lossless (Qwen3-0.6B wikitext bits-per-byte −0.0%, LFM2.5-350M +0.45%); **W4** is heavier on these compression-sensitive models (+1.2% / +4.4%). ## Test Plan ``` pytest tests/llmcompressor/transformers/compression/test_compress_tensor_utils.py -k tied ``` `test_tied_quantized_embedding_no_duplicate_head` quantizes `Qwen3-0.6B`'s embedding **and** `lm_head` (data-free) and asserts the saved checkpoint stays tied (`tie_word_embeddings=True`) with the packed embedding present and **no** `lm_head` weight tensors. ## Test Result - New test passes. - Existing `test_no_duplicate_tied_lm_head_on_save` (tied/untied × offload) still passes — an explicitly-untied model keeps its separate head (the re-tie is scoped to the compressed case). - End-to-end: the produced `Qwen3-0.6B` checkpoint has a single shared packed table (no fp16 head), config declares both embeddings quantized, `tie_word_embeddings=True`, 930 MB vs 1.5 GB fp16 — loads and generates coherently in vLLM (with #45535). - `ruff check` / `ruff format --check` clean. --- This change was developed with AI assistance (Claude). All changed lines were reviewed by the submitter. --------- Signed-off-by: Karthik Kothuri <karthikkothuri2009@gmail.com> Co-authored-by: Kyle Sayers <kylesayrs@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
3 participants
Summary
When a model with tied word embeddings is saved while offloaded, the offload conversion (
to_accelerate) splitsembed_tokensandlm_headinto separate parameters. This defeats transformers' save-time de-duplication, so a redundant, bit-identicallm_head.weightgets written.This re-ties the weights right after
to_accelerateand before saving, so the tied head is dropped as intended. It's a cheap attribute reassignment (no onloading) and a no-op for models with untied embeddings.Test plan
lm_head.weightno longer appears inmodel.safetensors.index.json.pytest tests/llmcompressor/transformers/compression/test_compress_tensor_utils.py -k no_duplicate_tied_lm_head -v