跳到内容

llmcompressor.pipelines.sequential.ast_helpers

函数

  • autowrap_forwards –

    将给定模块的 forward 方法替换为重新编译的版本，其中

autowrap_forward

autowrap_forward(module: Module, ignore: List[str])

将给定模块的 forward 方法替换为重新编译的版本，其中所有不可追踪的代码模式都被移除，并替换为 torch.fx 函数包装器。

有关不可追踪代码模式及其解释的列表，请参阅 https://github.com/vllm-project/llm-compressor/pull/1411

参数

  • module

    (Module) –

    应替换 forward 方法的模块

  • ignore

    (List[str]) –

    要包装的函数名称的显式列表

源代码位于 llmcompressor/pipelines/sequential/ast_helpers.py 
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
@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
    """
    # get source code of module forward
    source = inspect.getsource(module.forward)
    source = textwrap.dedent(source)
    tree = ast.parse(source)

    # construct namespace for our new code
    defining_module = sys.modules[module.__class__.__module__]
    namespace = defining_module.__dict__.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")
    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

autowrap_forwards

autowrap_forwards(modules: List[Module], ignore: List[str])

将给定模块的 forward 方法替换为重新编译的版本，其中所有不可追踪的代码模式都被移除，并替换为 torch.fx 函数包装器

参数

  • modules

    (List[Module]) –

    应替换 forward 方法的模块列表

  • ignore

    (List[str]) –

    要包装的函数名称的显式列表

源代码位于 llmcompressor/pipelines/sequential/ast_helpers.py 
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
@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