Fix DynamicLayer state leak after dynamo export

[Ti-Tai Wang (titaiwangms)]

Describe your changes

OnnxConversion with use_dynamo_exporter=True monkeypatches transformers.cache_utils.DynamicLayer.lazy_initialization at the class level (via _patch_dynamic_layer_for_export()) in order to make ONNX export work, but never restored the original method afterward. Because the patch is class-level (not instance-level), it silently affected every subsequent DynamicLayer/DynamicCache instance created later in the same process, for any model class.

The patched implementation also had its own bug: it always initialized self.values from key_states (ignoring value_states). This is harmless when key/value head dimensions match, but corrupts the KV cache shape for architectures where they differ (e.g. DeepSeek-V3's MLA, where qk_rope_head_dim + qk_nope_head_dim != v_head_dim), producing errors such as:

RuntimeError: Sizes of tensors must match except in dimension 2. Expected size 16 but got size 8

This was discovered as a cross-test pollution bug: running test/passes/onnx/test_common.py::test_resave_model (any model, dynamo export) before a DeepSeek-V3 test in the same pytest process (as CI does, since it runs the whole test/ directory in one process) leaves the patched lazy_initialization in place and corrupts the later test's cache handling — even though the two tests use entirely unrelated models.

Fix

• Converted _patch_dynamic_layer_for_export() into a contextlib.contextmanager that saves the original DynamicLayer.lazy_initialization, applies the patch, yields, and restores the original in a finally block, so the patch is undone even if export raises.
• Fixed the cache initialization bug to use value_states (falling back to key_states only if value_states is None) instead of always reusing key_states for both keys and values.
• Wrapped the torch.onnx.export(...) call site in with patch_context: (using contextlib.nullcontext() for the transformers <5.0 branch that doesn't need this patch).

Verification

• Added regression tests: one asserting DynamicLayer.lazy_initialization is restored after a normal patched export path and correctly preserves distinct key/value shapes, and one asserting it is restored even when the patched block raises.
• Reproduced the original cross-test pollution locally by running test/passes/onnx/test_common.py together with a DeepSeek-V3 GPTQ MoE test in the same pytest process; confirmed the failure occurs before this fix and is resolved after it.
• lintrunner clean.

Checklist before requesting a review

• Add unit tests for this change.
• Make sure all tests can pass.
• Update documents if necessary.
• Lint and apply fixes to your code by running lintrunner -a
• Is this a user-facing change? If yes, give a description of this change to be included in the release notes.

(Optional) Issue link

Discovered while investigating a CI-only deepseek_v3 failure reported on #2610 (unrelated to that PR's changes; caused by this pre-existing global-state leak, only reproducible when the whole test/ suite runs in one process, as CI does).

[Copilot]

Pull request overview

This PR fixes a global state leak in the ONNX dynamo-export path by ensuring a temporary transformers.cache_utils.DynamicLayer.lazy_initialization monkeypatch is reliably restored after export, and corrects the patch logic to preserve distinct key/value cache shapes.

Changes:

• Converted _patch_dynamic_layer_for_export() into a contextlib.contextmanager that restores the original DynamicLayer.lazy_initialization in a finally block.
• Fixed the patched cache initialization to build self.values from value_states (falling back to key_states only when value_states is None).
• Added regression tests to validate method restoration (normal and exception paths) and to ensure key/value shapes remain distinct.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated no comments.

File	Description
olive/passes/onnx/conversion.py	Makes the DynamicLayer patch scoped + exception-safe; fixes value cache initialization; wraps export in the patch context.
test/passes/onnx/test_conversion.py	Adds targeted unit tests for restoration and key/value shape preservation under the patch.
test/passes/onnx/test_common.py	Adds a regression assertion ensuring OnnxConversion(...use_dynamo_exporter=True) does not leave the patch applied.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.