Skip to content

oneshot entrypoint update #1445

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Conversation

ved1beta
Copy link
Contributor

@ved1beta ved1beta commented May 17, 2025

SUMMARY:
Updated the oneshot function to use explicit keyword arguments instead of relying solely on **kwargs.

TEST PLAN:
Changes were tested using mock tests that verify:
Correct handling of explicit arguments
Backward compatibility with kwargs
Proper behavior with mixed explicit args and kwargs
Verification that kwargs correctly override explicit parameters when duplicates exist
fixes
#1443

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.

@kylesayrs kylesayrs self-requested a review May 17, 2025 14:43
@kylesayrs kylesayrs added the ready When a PR is ready for review label May 17, 2025
kylesayrs
kylesayrs previously approved these changes May 17, 2025
Copy link
Collaborator

@kylesayrs kylesayrs left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks great! There are a few arguments missing from the signature like dvc_data_repository, but this is handled and by **kwargs and those arguments don't need to be explicit.

This will go a long way towards making LLM Compressor more user friendly :)

Returns:
PreTrainedModel: The calibrated model
"""
params = {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the future we could use something like this to avoid duplicating information

from typing import Tuple, List, Any, Dict
import inspect

def collect_function_arguments() -> Tuple[List[Any], Dict[str, Any]]:
    caller_frame = inspect.currentframe().f_back
    args_info = inspect.getargvalues(caller_frame)
    args = []
    kwargs = {}
    for key, value in args_info.locals.items():
        if key in args_info.args:
            args.append(value)
        elif key == args_info.varargs:
            args.extend(value)
        elif key == args_info.keywords:
            kwargs.update(value)
        else:
            assert False
        
    return args, kwargs
    
def oneshot(...):
    args, kwargs = collect_function_arguments()
    one_shot = Oneshot(*args, **kwargs)

@kylesayrs
Copy link
Collaborator

@ved1beta Could you please make style; make quality?

@dsikka dsikka changed the title onshort function update : ) oneshot entrypoint update May 17, 2025
Copy link
Collaborator

@brian-dellabetta brian-dellabetta left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks @ved1beta for taking a shot at this! I think some of the type hints need to be updated to account for actual classes being passed in other than string representations, but otherwise this is looking great!

one_shot = Oneshot(**kwargs)
def oneshot(
# Required model parameters
model: str,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Users can also pass in a pre-loaded model

Suggested change
model: str,
model: Union[str, PreTrainedModel],

Comment on lines 183 to 184
tokenizer: Optional[str] = None,
processor: Optional[str] = None,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These can also be the loaded up classes. I think these are the best base classes to use? cc @kylesayrs

Suggested change
tokenizer: Optional[str] = None,
processor: Optional[str] = None,
tokenizer: Optional[Union[str, transformers.PreTrainedTokenizerBase]] = None,
processor: Optional[Union[str, transformers.ProcessorMixin]] = None,

oneshot_device: str = "cuda:0",
model_revision: str = "main",
# Recipe parameters
recipe: Optional[str] = None,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This can also be a list of Modifiers

clear_sparse_session: bool = False,
stage: Optional[str] = None,
# Dataset parameters
dataset: Optional[str] = None,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this can be an actual loaded dataset

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.

3 participants