Skip to content

Commit 0357cb9

Browse files
authored
Pre-initialize torch._dynamo to prevent double-registration with peft 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>
1 parent 915262a commit 0357cb9

1 file changed

Lines changed: 14 additions & 1 deletion

File tree

modelopt/torch/__init__.py

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,25 @@
1515

1616
"""Model optimization and deployment subpackage for torch."""
1717

18+
import importlib
1819
import warnings as _warnings
1920

2021
from packaging.version import Version as _Version
2122
from torch import __version__ as _torch_version
2223

23-
from . import distill, nas, opt, peft, prune, quantization, sparsity, speculative, utils
24+
# Pre-initialize torch._dynamo to prevent double-registration with peft's torch.compile() call
25+
importlib.import_module("torch._dynamo")
26+
from . import ( # noqa: E402
27+
distill,
28+
nas,
29+
opt,
30+
peft,
31+
prune,
32+
quantization,
33+
sparsity,
34+
speculative,
35+
utils,
36+
)
2437

2538
if _Version(_torch_version) < _Version("2.9"):
2639
_warnings.warn(

0 commit comments

Comments
 (0)