fix(export): save models with legacy list-style _tied_weights_keys (NVBug 6518665)

[Edwardf0t1]

What does this PR do?

Type of change: Bug fix

Fixes NVBug 6518665 / OMNIML-5583: PTQ of Step-3.7-Flash (--recipe general/ptq/nvfp4_mlp_only-kv_fp8, transformers 5.5.4) calibrates fine and then dies while writing the checkpoint:

File "modelopt/torch/export/unified_export_hf.py", line 1510, in export_hf_checkpoint
    model.save_pretrained(
File "modelopt/torch/opt/plugins/transformers.py", line 155, in _save_pretrained_with_checks
File "transformers/modeling_utils.py", line 3352, in save_pretrained
    state_dict = remove_tied_weights_from_state_dict(state_dict, model_to_save)
File "transformers/modeling_utils.py", line 338, in _get_tied_weight_keys
    tied_weight_keys.extend([f"{name}.{k}" if name else k for k in tied.keys()])
AttributeError: 'list' object has no attribute 'keys'

Root cause is version drift in transformers, not in quantization. transformers 5.0 changed _tied_weights_keys from list[str] to a {target: source} dict, and save_pretrained → _get_tied_weight_keys calls .keys() on every submodule's declaration without a type check. stepfun-ai/Step-3.7-Flash's remote code still uses the 4.x list format — _tied_weights_keys = ["lm_head.weight"] on Step3p7TextModel, Step3p7Model and Step3p7ForConditionalGeneration.

The load-time consumer (get_expanded_tied_weights_keys) returns early when config.tie_word_embeddings is false, which it is for this checkpoint, so the model loads and calibrates normally and only fails at the end of the run, after the expensive part. The same crash reproduces with a plain AutoModel.from_pretrained(..., trust_remote_code=True).save_pretrained(...), and transformers 5.12 tolerates None there but still not a list, so it is not fixed upstream either.

Fix

_save_pretrained_with_checks — the entry point every ModelOpt HF save routes through (unified export, ModelOptHFTrainer, user-called save_pretrained) — now normalizes a list-style declaration to the equivalent dict for the duration of the save and restores the original attribute afterwards. Mapping each entry to itself preserves the legacy semantics: those list entries were exactly the dedup patterns _get_tied_weight_keys is expected to return. No-op on transformers<5, where the list format is native, and no-op for models that already declare a dict.

The restore tracks whether the instance owned the attribute, so for the usual class-level declaration nothing is left shadowing it.

Usage

python hf_ptq.py --model /local/Step-3.7-Flash --recipe general/ptq/nvfp4_mlp_only-kv_fp8 \
    --dataset /local/cnn_dailymail --calib_size 32 --export_path /local/Step-3.7-Flash-nvfp4 --trust_remote_code

Testing

Two new tests in tests/unit/torch/opt/plugins/test_transformers_save_load.py:

• test_save_pretrained_with_legacy_tied_weights_keys (parametrized over tie_word_embeddings) — a quantized tiny-Llama declaring list-style keys saves and round-trips; without the fix it fails with the exact reported AttributeError at modeling_utils.py:338. It asserts the saved model.safetensors keys directly, so no weight can be silently dropped by the save-time dedup: model.embed_tokens.weight is always present, and lm_head.weight is present iff the weights are not actually tied (when they are, transformers drops the alias and re-ties it on load — the transformers 4.x behavior).
• test_legacy_tied_weights_keys_as_dict_restores_class_attribute — the shim leaves no instance attribute shadowing a class-level declaration (skipped on transformers<5).

Ran locally against both transformers majors:

• transformers 5.5.4 (the version in the bug report) + torch 2.11: tests/unit/torch/opt/plugins/ 27 passed, tests/unit/torch/export/ 124 passed.
• transformers 4.57.6: tests/unit/torch/opt/plugins/ 26 passed (the shim no-ops, new dict-restore test skips).

Not yet re-run end-to-end on the real Step-3.7-Flash checkpoint (1.4 TB / 8×B200); QA can re-run the reported command against this branch.

Known limitation

The shim covers the save side only. A legacy-list model that also sets tie_word_embeddings=True still crashes on the load side, in get_expanded_tied_weights_keys during PreTrainedModel.post_init — before ModelOpt has a model object to patch. Such a checkpoint cannot be loaded by transformers 5 at all, with or without ModelOpt, so it is out of scope here. Step-3.7-Flash is unaffected (tie_word_embeddings is unset, i.e. false).

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: ✅
• Did you get Claude approval on this PR?: ❌

Additional Information

The upstream-correct fix is for StepFun to migrate their modeling code to the transformers 5 dict format; this shim unblocks every not-yet-migrated trust_remote_code checkpoint in the meantime.

🤖 Generated with Claude Code

Summary by CodeRabbit

• Bug Fixes

  • Improved compatibility with Transformers 5 when exporting models with legacy tied-weight declarations.
  • Preserved tied-weight settings after saving, including class-level declarations.
  • Ensured saved models reload with equivalent state and outputs.

• Tests

  • Added coverage for quantized models, tied and untied embeddings, and legacy tied-weight formats.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 11b485a0-7a4e-4fbc-9030-36128771b823

📥 Commits

Reviewing files that changed from the base of the PR and between 439fc5d and 87ee221.

📒 Files selected for processing (1)

• CHANGELOG.rst

🚧 Files skipped from review as they are similar to previous changes (1)

• CHANGELOG.rst

---

📝 Walkthrough

Walkthrough

The Transformers plugin now supports legacy collection-form _tied_weights_keys declarations during checkpoint export. It restores the original declarations after saving. Tests cover tied and untied embeddings, class-level declarations, and round-trip equivalence.

Changes

Tied-weight export compatibility

Layer / File(s)	Summary
Tied-weight compatibility context modelopt/torch/opt/plugins/transformers.py	The plugin detects Transformers 5 and temporarily converts legacy collection-form _tied_weights_keys attributes to dictionaries while preserving their original state.
Save integration and regression coverage modelopt/torch/opt/plugins/transformers.py, tests/unit/torch/opt/plugins/test_transformers_save_load.py, CHANGELOG.rst	Model saving uses the compatibility context. Tests verify restoration, class-level declarations, state/output equivalence, and safetensor keys. The changelog records the fix.

Estimated code review effort: 2 (Simple) | ~10 minutes

Suggested reviewers: ajrasane, kevalmorabia97

🚥 Pre-merge checks | ✅ 6

✅ Passed checks (6 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main fix for saving models with legacy list-style _tied_weights_keys.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns	✅ Passed	The PR adds no prohibited load, pickle, hardcoded trust_remote_code, eval/exec, or # nosec pattern, and changes no dependency manifest.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/legacy-tied-weights-keys-save

---

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

[github-actions]

PR Preview Action v1.8.1
Preview removed because the pull request was closed.
2026-08-07 20:20 UTC

[coderabbitai]

Warning

CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.

Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.

👉 Steps to fix this

Actionable comments posted: 1

🤖 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.

Inline comments:
In `@modelopt/torch/opt/plugins/transformers.py`:
- Around line 173-188: Serialize the temporary _tied_weights_keys mutation and
the protected save operation so concurrent save_pretrained calls on the same
model cannot overlap. Update the surrounding patch/restore context, using a
shared reentrant lock or equivalent per-model ownership tracking, and ensure the
lock remains held through the Transformers save that consumes the temporary
dictionaries before restoration occurs.

🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 188323ca-232c-4ca1-92c6-1a129acbd1a1

📥 Commits

Reviewing files that changed from the base of the PR and between fed1980 and a4f862c.

📒 Files selected for processing (3)

• CHANGELOG.rst
• modelopt/torch/opt/plugins/transformers.py
• tests/unit/torch/opt/plugins/test_transformers_save_load.py

[codecov]

Codecov Report

❌ Patch coverage is 90.00000% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 78.08%. Comparing base (089109d) to head (87ee221).

Files with missing lines	Patch %	Lines
modelopt/torch/opt/plugins/transformers.py	90.00%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2071      +/-   ##
==========================================
- Coverage   78.73%   78.08%   -0.65%     
==========================================
  Files         522      522              
  Lines       60157    60175      +18     
==========================================
- Hits        47363    46988     -375     
- Misses      12794    13187     +393     

Flag	Coverage Δ
examples	42.91% <65.00%> (-0.18%)	⬇️
gpu	58.66% <65.00%> (-0.62%)	⬇️
regression	14.94% <65.00%> (+0.08%)	⬆️
unit	55.39% <90.00%> (+0.02%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[cjluo-nv]

Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.

Small, well-scoped shim (+91/-5) that normalizes legacy list-style _tied_weights_keys to {key: key} for the duration of _save_pretrained_with_checks, with a finally restore that correctly distinguishes instance-owned vs. class-level declarations. Two new unit tests, CHANGELOG entry present, no licensing surface, no design-review trigger (single context manager, no new subsystem). No prompt-injection content in the PR metadata. Mechanics look right; the reason to have a human look is the semantics of the normalization rather than the plumbing:

• {key: key} is a self-referential tie, and transformers 5 treats dict keys as alias patterns and values as the canonical weight. If remove_tied_weights_from_state_dict drops alias-side keys because "the canonical copy is in the state dict", a self-map means the alias is the canonical, so the tensor could be dropped entirely — turning a loud AttributeError into a silently incomplete checkpoint. That's precisely the reported case (Step-3.7-Flash has tie_word_embeddings=False, so lm_head.weight is a real, untied weight). The author's local run on transformers 5.5.4 suggests no removal happens, but the PR explicitly notes it was not re-run end-to-end on the real checkpoint, and this behavior is transformers-version-dependent (the body already documents 5.12 diverging from 5.5 on None handling). Worth a maintainer who knows the transformers 5 save path confirming, and worth considering the more conservative normalization ({} — declare nothing tied — falling back to HF's existing data_ptr dedup for genuinely shared storage) if removal turns out to be unconditional.
• The new test only pins this indirectly. test_save_pretrained_with_legacy_tied_weights_keys relies on tf_modelopt_state_and_output_tester logits equality to catch a dropped weight. An explicit assertion that the saved model.safetensors/index still contains lm_head.weight and model.embed_tokens.weight would pin the actual invariant (no weight silently dropped) and make a future transformers version change fail loudly with an obvious message.
• Minor/no action needed: the patch loop runs before the try, so an exception mid-loop would leave earlier modules patched (unlikely in practice); the shim covers save only — a legacy-list model with tie_word_embeddings=True would still hit the same crash on the load side (get_expanded_tied_weights_keys), which is arguably out of scope but not stated in the PR body; the repo's existing _collect_canonical_tied_patterns in modelopt/torch/export/model_utils.py deliberately skips list-style declarations, so consider whether that comment/behavior should reference this shim for consistency.

[Edwardf0t1]

Thanks — I checked the normalization semantics against the transformers source and empirically on 5.5.4 (the version in the bug report). Summary: {key: key} does not risk dropping a weight, and the suggested {} alternative would actively break genuinely tied models. Pushed 439fc5d with the test hardening you asked for.

1. {key: key} cannot silently drop a weight. In remove_tied_weights_from_state_dict (modeling_utils.py, 5.5.4):

• Deletion candidates come only from shared_ptrs = {ptr: names for ptr, names in ptrs.items() if len(names) > 1} — tensors that actually share storage. An untied lm_head.weight (the reported Step-3.7-Flash case, tie_word_embeddings=False) forms a group of one and is never eligible for deletion, regardless of what the declaration says.
• Even inside a shared group, deletion is guarded by if found < len(names), so at least one name in every group always survives — a group can never be emptied.
• The dict values are never read on the save path: _get_tied_weight_keys returns tied.keys() only, used as regex patterns. The alias→canonical semantics you're describing live in get_expanded_tied_weights_keys, which is a load/tie_weights path and is not called during save_pretrained.

Measured on transformers 5.5.4, with _tied_weights_keys = ["lm_head.weight"] declared list-style:

config	lm_head.weight in file	model.embed_tokens.weight in file	reload logits
tie_word_embeddings=False (reported case)	kept	kept	identical
tie_word_embeddings=True (real tie, shared storage)	dropped (alias)	kept (canonical)	identical

The second row is the intended behavior and matches transformers 4.x: the alias is meant to be dropped and is re-tied on load.

2. {} would be worse, not more conservative. With no patterns, to_delete_names stays empty, so for a genuinely tied model the shared pair falls through to the error branch:

RuntimeError: The weights trying to be saved contained shared tensors
[{'lm_head.weight', 'model.embed_tokens.weight'}] which are not properly defined.
We found all the potential target tied weights keys to be: set().

So {} trades a crash for a different crash on exactly the models the declaration exists to describe.

3. Test hardening — done. test_save_pretrained_with_legacy_tied_weights_keys is now parametrized over tie_word_embeddings and asserts the saved model.safetensors keys directly (model.embed_tokens.weight always present; lm_head.weight present iff the weights are not tied) instead of relying on logits equality. A future transformers change to the dedup fails with an obvious message.

4. Patch loop before try — fixed in the same commit; the walk now runs inside the try, so a failure partway through still restores the modules already patched.

5. Load side (tie_word_embeddings=True + legacy list) — correct that it is not covered, and it isn't fixable here: that crash happens inside PreTrainedModel.post_init while the model is being constructed, before ModelOpt has an object to patch. Such a checkpoint cannot be loaded by transformers 5 at all, with or without ModelOpt. Added to the PR body as a stated limitation.

6. _collect_canonical_tied_patterns — that skip is deliberate and orthogonal: it needs the alias/canonical distinction to decide state-dict ordering for DiffusionGemma, and a legacy list genuinely does not carry it. This shim only supplies dedup patterns, which the list does carry, so no consistency issue between them.

[cjluo-nv]

Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.

Re-review of PR #2071 (3 files, +106/-5). All previously-raised findings have been resolved in the current diff; the remaining items are judgment calls for a maintainer familiar with the transformers 5 save path.

Addressed since last round:

• 💬 Author replied with a detailed read of remove_tied_weights_from_state_dict on 5.5.4 (deletion candidates come only from shared_ptrs with len(names) > 1, guarded by if found < len(names), and dict values are never read on the save path) — the {key: key} self-tie cannot empty a group, and the suggested {} alternative would instead trip the "shared tensors not properly defined" RuntimeError for genuinely tied models. Confirmed in code: the shim only holds for the duration of the save, so the alias→canonical semantics of get_expanded_tied_weights_keys (load path) are not exercised.
• Test hardening landed: test_save_pretrained_with_legacy_tied_weights_keys is now parametrized over tie_word_embeddings and asserts the saved model.safetensors keys directly (model.embed_tokens.weight always present, lm_head.weight present iff untied), so a future transformers dedup change fails loudly instead of silently dropping a weight. The class-attribute restore test pins the del vs. reassign branch.
• Patch loop moved inside the try, so a failure partway through the model.modules() walk still restores already-patched modules.
• Load-side limitation (legacy list + tie_word_embeddings=True crashing in post_init) is now stated in the PR body.
• CodeRabbit's lock/race finding was explicitly withdrawn after the author's reply (adjacent _patch_revert_weight_conversion mutates transformers module globals unguarded; concurrent save_pretrained on one instance is unsupported).

Why a human should still glance:

• 💬 Author states the fix was not re-run end-to-end on the real Step-3.7-Flash checkpoint (1.4 TB / 8×B200) and defers to QA — the reported failure mode is only reproduced indirectly via a tiny-Llama unit test, and the save-time dedup behavior the reasoning depends on is transformers-version-dependent (the body itself notes 5.12 diverging from 5.5 on None handling). The new test is a good tripwire, but it only exercises whichever transformers version CI installs.
• Minor, no action strictly needed: the shim silently coerces tuple/set declarations too (fine, but untested); assert ("lm_head.weight" in saved_keys) is not tie_word_embeddings reads a little cryptically vs. !=; and the CHANGELOG entry lands under 0.46 Bug Fixes rather than the empty 0.47 section — worth confirming that matches the intended release/cherry-pick target.

No licensing surface, no new subsystem (single context manager), and no prompt-injection content in the PR metadata or comments.