Pre-initialize torch._dynamo to prevent double-registration with peft torch.compile() call#1228
Merged
hychiang-git merged 7 commits intomainfrom Apr 11, 2026
Merged
Conversation
hychiang-git
requested review from
ChenhanYu,
Edwardf0t1,
kevalmorabia97 and
shengliangxu
and removed request for
ChenhanYu,
Edwardf0t1 and
shengliangxu
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThe Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes 🚥 Pre-merge checks | ✅ 4✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
Contributor
|
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1228 +/- ##
==========================================
+ Coverage 72.14% 77.46% +5.31%
==========================================
Files 350 350
Lines 40478 40480 +2
==========================================
+ Hits 29202 31356 +2154
+ Misses 11276 9124 -2152
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
shengliangxu
reviewed
Apr 9, 2026
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
modelopt/torch/__init__.py (1)
24-25: Consider deferringtorch._dynamoinitialization to avoid coupling all users to this private Torch submodule at import time.Line 25 imports
torch._dynamounconditionally for all users ofmodelopt.torch, even those not usingtorch.compile(). Since PEFT and compile features are optional integrations, and the codebase already uses lazy initialization elsewhere (e.g.,speculative/plugins/transformers.py:625), move this import behind a guard or defer it until needed:♻️ Proposed lazy approach
# Pre-initialize torch._dynamo to prevent double-registration with peft's torch.compile() call -importlib.import_module("torch._dynamo") +try: + importlib.import_module("torch._dynamo") +except ImportError: + _warnings.warn( + "torch._dynamo is unavailable; compile-related integrations may initialize it later.", + RuntimeWarning, + )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@modelopt/torch/__init__.py` around lines 24 - 25, The unconditional importlib.import_module("torch._dynamo") in modelopt.torch causes eager coupling to a private Torch submodule; change it to a lazy/guarded initialization so torch._dynamo is only imported when needed (e.g., when PEFT or torch.compile integrations are invoked). Add a small helper (e.g., ensure_dynamo_initialized or a guarded branch in the code paths that use torch.compile) that calls importlib.import_module("torch._dynamo") the first time it's required, and remove the top-level import from the module init so normal imports of modelopt.torch don't trigger torch._dynamo import.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@modelopt/torch/__init__.py`:
- Around line 24-25: The unconditional importlib.import_module("torch._dynamo")
in modelopt.torch causes eager coupling to a private Torch submodule; change it
to a lazy/guarded initialization so torch._dynamo is only imported when needed
(e.g., when PEFT or torch.compile integrations are invoked). Add a small helper
(e.g., ensure_dynamo_initialized or a guarded branch in the code paths that use
torch.compile) that calls importlib.import_module("torch._dynamo") the first
time it's required, and remove the top-level import from the module init so
normal imports of modelopt.torch don't trigger torch._dynamo import.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: d9336e4b-80aa-44c3-8dda-90ed3ad9371c
📒 Files selected for processing (1)
modelopt/torch/__init__.py
…'s torch.compile() call Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
hychiang-git
force-pushed
the
hungyuehc/fix_local_run
branch
from
e7f3bce to
8e1696e
Compare
kevalmorabia97
approved these changes
Apr 10, 2026
Collaborator
|
/ok to test 8e1696e |
Collaborator
|
/ok to test 89b860d |
kinjalpatel27
pushed a commit
that referenced
this pull request
Apr 13, 2026
…t` torch.compile() call (#1228) ### What does this PR do? Type of change: ? Bug fix This PR fixes the error `AssertionError: Artifact of type=precompile already registered in mega-cache artifact factory` when launching with local docker. The error message: ``` Qwen3-8B-0/0 Traceback (most recent call last): Qwen3-8B-0/0 File "/nemo_run/code/modules/Megatron-LM/examples/post_training/modelopt/quantize.py", line 21, in <module> Qwen3-8B-0/0 import modelopt.torch.quantization as mtq Qwen3-8B-0/0 File "/usr/local/lib/python3.12/dist-packages/modelopt/torch/__init__.py", line 23, in <module> Qwen3-8B-0/0 from . import distill, nas, opt, peft, prune, quantization, sparsity, speculative, utils Qwen3-8B-0/0 File "/usr/local/lib/python3.12/dist-packages/modelopt/torch/prune/__init__.py", line 24, in <module> Qwen3-8B-0/0 from . import fastnas, gradnas, plugins Qwen3-8B-0/0 File "/usr/local/lib/python3.12/dist-packages/modelopt/torch/prune/gradnas.py", line 71, in <module> Qwen3-8B-0/0 from transformers.models.bert.modeling_bert import BertAttention Qwen3-8B-0/0 File "/usr/local/lib/python3.12/dist-packages/transformers/models/bert/modeling_bert.py", line 30, in <module> Qwen3-8B-0/0 from ...generation import GenerationMixin Qwen3-8B-0/0 File "<frozen importlib._bootstrap>", line 1412, in _handle_fromlist Qwen3-8B-0/0 File "/usr/local/lib/python3.12/dist-packages/transformers/utils/import_utils.py", line 2317, in __getattr__ Qwen3-8B-0/0 module = self._get_module(self._class_to_module[name]) Qwen3-8B-0/0 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Qwen3-8B-0/0 File "/usr/local/lib/python3.12/dist-packages/transformers/utils/import_utils.py", line 2347, in _get_module Qwen3-8B-0/0 raise e Qwen3-8B-0/0 File "/usr/local/lib/python3.12/dist-packages/transformers/utils/import_utils.py", line 2345, in _get_module Qwen3-8B-0/0 return importlib.import_module("." + module_name, self.__name__) Qwen3-8B-0/0 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Qwen3-8B-0/0 File "/usr/lib/python3.12/importlib/__init__.py", line 90, in import_module Qwen3-8B-0/0 return _bootstrap._gcd_import(name[level:], package, level) Qwen3-8B-0/0 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Qwen3-8B-0/0 File "/usr/local/lib/python3.12/dist-packages/transformers/generation/utils.py", line 43, in <module> Qwen3-8B-0/0 from ..masking_utils import create_masks_for_generate Qwen3-8B-0/0 File "/usr/local/lib/python3.12/dist-packages/transformers/masking_utils.py", line 40, in <module> Qwen3-8B-0/0 from torch._dynamo._trace_wrapped_higher_order_op import TransformGetItemToIndex Qwen3-8B-0/0 File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/__init__.py", line 13, in <module> Qwen3-8B-0/0 from . import ( Qwen3-8B-0/0 File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/aot_compile.py", line 16, in <module> Qwen3-8B-0/0 from torch._dynamo.package import SystemInfo Qwen3-8B-0/0 File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/package.py", line 443, in <module> Qwen3-8B-0/0 @CacheArtifactFactory.register Qwen3-8B-0/0 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Qwen3-8B-0/0 File "/usr/local/lib/python3.12/dist-packages/torch/compiler/_cache.py", line 72, in register Qwen3-8B-0/0 assert artifact_cls.type() not in cls._artifact_types, ( Qwen3-8B-0/0 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Qwen3-8B-0/0 AssertionError: Artifact of type=precompile already registered in mega-cache artifact factory ``` ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - 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`: ❌ - Did you write any new necessary tests?: ❌ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ❌ <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Improved startup compatibility with certain Torch-based optimization tools by ensuring required runtime components are loaded during package initialization to prevent initialization errors. * No changes to public APIs or exported modules; behavior and interfaces remain unchanged. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Type of change: ? Bug fix
This PR fixes the error
AssertionError: Artifact of type=precompile already registered in mega-cache artifact factorywhen launching with local docker.The error message:
Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: ❌Summary by CodeRabbit