llmcompressor.pipelines.sequential.ast_helpers

@contextlib.contextmanager
def autowrap_forward(module: torch.nn.Module, ignore: list[str]):
    """
    Replace the `forward` method of the given module with a recompiled version where
    all untraceble code patterns are removed and replaced with torch.fx function
    wrappers.

    For a list of untraceable code patterns and their explainations, see
    https://github.com/vllm-project/llm-compressor/pull/1411

    :param module: module whose forward method should be replaced
    :param ignore: explicit list of function names to wrap
    """
    # check forward method is implemented
    if module.forward.__name__ == "_forward_unimplemented":
        raise ValueError(
            "Cannot calibrate model which does not implement `forward` method. Please "
            "either implement a forward method on the model, or pass a submodule to "
            "`oneshot`. For example, `oneshot(model.thinker, ...)`"
        )

    # get source code of module forward
    forward_fn = get_unwrapped_forward(module)
    source = inspect.getsource(forward_fn)
    source = textwrap.dedent(source)
    tree = ast.parse(source)

    # construct namespace for our new code
    namespace = getattr(forward_fn, "__globals__", None)
    if namespace is None:
        defining_module = sys.modules[module.__class__.__module__]
        namespace = defining_module.__dict__
    namespace = namespace.copy()
    namespace.update({"torch.fx.wrap": torch.fx.wrap})
    namespace.update({"self": module})

    # autowrap untraceable code
    auto_wrapper = AutoWrapper(namespace, ignore)
    tree = auto_wrapper.auto_wrap(tree)
    source = ast.unparse(tree)

    # compile new forward function from autowrapped code
    filename = f"<Autowrapped {module.__class__.__name__} {id(module)}>"
    code = compile(source, filename=filename, mode="exec")
    with append_autowrap_source_on_fail():
        exec(code, namespace)  # ensure ns of functions is the same ns as torch.fx.wrap

    # enable better tracebacks if autowrapped code fails
    linecache.cache[filename] = (
        len(source),
        None,
        [line + "\n" for line in source.splitlines()],
        filename,
    )

    # patch forward with autowrapped forward
    new_forward = namespace["forward"].__get__(module)
    with patch_attr(module, "forward", new_forward):
        yield

@contextlib.contextmanager
def autowrap_forwards(modules: list[torch.nn.Module], ignore: list[str]):
    """
    Replace the `forward` method of the given modules with a recompiled version where
    all untraceble code patterns are removed and replaced with torch.fx function
    wrappers

    :param modules: list of modules whose forward methods should be replaced
    :param ignore: explicit list of function names to wrap
    """
    with contextlib.ExitStack() as stack:
        for module in modules:
            if not isinstance(module, (torch.nn.ModuleList, torch.nn.ModuleDict)):
                stack.enter_context(autowrap_forward(module, ignore))
        yield

def get_unwrapped_forward(module: torch.nn.Module) -> Callable:
    """
    Get the original function which implements the `forward` method of a module,
    stripping away any decorators which may have been applied to it.

    `inspect.unwrap` only follows the `__wrapped__` attribute, which is set by
    decorators which use `functools.wraps`. Decorators which do not use
    `functools.wraps`, such as transformers' `force_accelerate_hooks` (see
    `transformers/integrations/accelerate.py`), leave no `__wrapped__` attribute
    behind, meaning that `inspect.getsource` returns the source of the *wrapper*,
    which defines a function named `wrapped` rather than `forward`. In this case,
    fall back to searching the wrapper's closure cells for the original function.

    Note that the source of the original function includes its decorator lines,
    meaning that decorators are reapplied when the source is recompiled and their
    behavior is preserved (e.g. accelerate hook setup).

    :param module: module whose forward function should be retrieved
    :return: function which implements the module's forward method
    """
    forward = inspect.unwrap(module.forward)

    # `module.forward` is a bound method, whose underlying function holds the closure
    if inspect.ismethod(forward):
        forward = forward.__func__

    if getattr(forward, "__name__", None) == "forward":
        return forward

    for cell in getattr(forward, "__closure__", None) or ():
        try:
            contents = cell.cell_contents
        except ValueError:  # cell is empty
            continue

        if isinstance(contents, FunctionType) and contents.__name__ == "forward":
            return contents

    return forward