Implement lazy loading for traceable models

[kylesayrs]

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

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
  • Similar to lazy loading implementation by transformers

Testing

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

[github-actions]

👋 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.

[horheynm]

Doesn't the registry take care of lazy loading?

[kylesayrs]

Not if we want to maintain the api of from llmcompressor.transformers.tracer import TraceableX

[horheynm]

I see, can we do something like

# examples script
from transformers import LlavaForConditionalGeneration
model = LlavaForConditionalGeneration.from_pretrained(...)

oneshot(model=model, ...)

# in the backend map `LlavaForConditionalGeneration` to your `llmcompressor.transformers.tracer` using registry.
# So, ex.,`LlavaForConditionalGeneration` maps to `TraceableLlavaForConditionalGeneration`

This way the user can just use transformers instead of our llmcompressor Model.

[kylesayrs]

model = LlavaForConditionalGeneration.from_pretrained(...)

1. I'm not really sure what this API is pointing to, as from what you've written, there's no distinction between code which loads the normal vs traceable definitions.
2. As we spoke about before, we'd need to change upstream code. These changes almost certainly wouldn't be accepted by the transformers team, as it's outside of the responsibilities of the transformers library
3. What is traceable for LLM Compressor is not what is traceable for other users. This is because there is specialized code in LLM Compressor in order to make tracing easier

[kylesayrs]

Doesn't the registry take care of lazy loading?

We could implement lazy loading using a registry, but this would involve having to add registry code into the traceable definition.

@TracingRegistry.register(name="TraceableLlavaForConditionalGeneration")
class LlavaForConditionalGeneration:
    ...

And imho this makes for a clunkier top level interface

from llmcompressor.transformers.tracing import TracingRegistry

model = TracingRegistry.load_from_registry("LlavaForConditionalGeneration").from_pretrained("path")

[horheynm]

We don't need to push code to HF upstream, we just use what in HF right now.
In llm-comp, we just add code mapping to which multi-model do we point to your tracing model.

Very simply something like

{
"HFMODEL": "YOUR_TRACIABLE_MODEL"
"LlavaForConditionalGeneration": "TraceableLlavaForConditionalGeneration"
}

But using the registry.

Here we can just use HF code in the example UX, but in backend would use your traciable model

[kylesayrs]

In llm-comp, we just add code mapping to which multi-model do we point to your tracing model.
Here we can just use HF code in the example UX, but in backend would use your traciable model

I understand what a registry dictionary is. I don't understand how it is possible to implement a registry in LLM Compressor while only using the HF library at the top level

[kylesayrs]

If you're referring to dynamically replacing the model definition within oneshot, I consider that to be an antipattern which makes it unclear to the user what model they're really loading. This also opens up the unintended consequences of loading a model definition twice, such as if the user modifies the model config prior to oneshot

[kylesayrs]

Also, traceable definitions are not needed for recipes which do not use the sequential pipeline. Because the traceable definitions include things like processing error checks, they're useful to keep in for most cases to allow the user to better debug their dataloading