Skip to content

Implement lazy loading for traceable models#1105

Closed
kylesayrs wants to merge 9 commits intomainfrom
kylesayrs/lazy-tracing-import
Closed

Implement lazy loading for traceable models#1105
kylesayrs wants to merge 9 commits intomainfrom
kylesayrs/lazy-tracing-import

Conversation

@kylesayrs
Copy link
Collaborator

@kylesayrs kylesayrs commented Jan 28, 2025

Purpose

  • Some models may contain imports to libraries that are not part of the base installation. We do not want users who want to use one traceable model to be forced to install libraries used by another traceable model
  • Model definitions are large, so it's better to load them only when needed
  • Adding new model definitions which rely on newer transformers versions affects other model definitions. For example, attempting to use transformers<4.47.0 raises an error because idefics had not been implemented at that point. While we encourage users to update to the latest transformers versions, this PR allows users to continue using the same environment without unnecessarily forcing an upgrade.

For an example of this lazy module pattern, see transformers/models/llama/__init__.py.

from typing import TYPE_CHECKING
from ...utils import _LazyModule
from ...utils.import_utils import define_import_structure

if TYPE_CHECKING:
    from .configuration_llama import *
    from .modeling_flax_llama import *
    from .modeling_llama import *
    from .tokenization_llama import *
    from .tokenization_llama_fast import *
else:
    import sys

    _file = globals()["__file__"]
    sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)

Changes

  • Implemented _AliasableLazyModule which extends _LazyModule to allow aliases
  • Dynamically replace llmcompressor.transformers.tracing with an instance of _AliasableLazyModule which lazily loads submodules as they are needed

Testing

  • Added passing tests in tests/llmcompressor/transformers/tracing/test_init.py

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
@github-actions
Copy link

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.

Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed.

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
@kylesayrs kylesayrs self-assigned this Jan 28, 2025
@kylesayrs kylesayrs added the ready When a PR is ready for review label Jan 28, 2025
@kylesayrs kylesayrs marked this pull request as ready for review January 28, 2025 17:37
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
@kylesayrs
Copy link
Collaborator Author

Traceable model definitions have been removed as of #1411

@kylesayrs kylesayrs closed this May 29, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ready When a PR is ready for review

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants