Commit 0357cb9
authored
Pre-initialize torch._dynamo to prevent double-registration with
### 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>peft torch.compile() call (#1228)1 parent 915262a commit 0357cb9
1 file changed
Lines changed: 14 additions & 1 deletion
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
15 | 15 | | |
16 | 16 | | |
17 | 17 | | |
| 18 | + | |
18 | 19 | | |
19 | 20 | | |
20 | 21 | | |
21 | 22 | | |
22 | 23 | | |
23 | | - | |
| 24 | + | |
| 25 | + | |
| 26 | + | |
| 27 | + | |
| 28 | + | |
| 29 | + | |
| 30 | + | |
| 31 | + | |
| 32 | + | |
| 33 | + | |
| 34 | + | |
| 35 | + | |
| 36 | + | |
24 | 37 | | |
25 | 38 | | |
26 | 39 | | |
| |||
0 commit comments