Skip to content

llmcompressor.pipelines.sequential.ast_utils

Modules:

Classes:

  • AutoWrapper –

    Automatically wraps untracable code according to the following patterns:

  • ControlFlowAnalyzer –

    Used to determine if a piece of code can be wrapped into a function. Includes any

  • NameAnalyzer –

    Determines the unbound, assigned, and conditionally assigned names associated with

AutoWrapper 

AutoWrapper(namespace: dict[str, Any], ignore: list[str])

Bases: NodeTransformer

Automatically wraps untracable code according to the following patterns:

The following patterns are automatically wrapped 1. If statements whose conditions cannot be statically evaluated 2. Ignored functions (_update_causal_mask) 3. Starred tuple unpacking 4. Starred argument unpacking

See also: https://github.com/vllm-project/llm-compressor/pull/1411

Methods:

  • auto_wrap –

    Modify ast by automatically wrapping any untraceable code segments. Segments to

  • visit_Call –

    Wrap any functions which use (4) variadic arguments or (2) match the ignore list

  • visit_Delete –

    Remove any deleted names from self._local_names,

  • visit_FunctionDef –

    Remove decorators which prevent forward function recompilation

  • visit_If –

    Attempt to statically evaluate the condition of the if statement. If the

  • visit_Name –

    Add any new names in self._local_names,

  • visit_Tuple –

    (3) Wrap any tuples which use starred unpacking

Source code in llmcompressor/pipelines/sequential/ast_utils/auto_wrapper.py 
24
25
26
27
28
29
def __init__(self, namespace: dict[str, Any], ignore: list[str]):
    self.namespace = namespace
    self.ignore = ignore
    self._wrapper_fn_defs: list[ast.FunctionDef] = list()
    self._local_names = set()
    self._wrapped_counter = 0

auto_wrap 

auto_wrap(tree: Module) -> ast.Module

Modify ast by automatically wrapping any untraceable code segments. Segments to wrap are determined through analysis of the code and basic pattern matching

Parameters:

  • tree

    (Module) –

    module containing a definition to an original forward function

Returns:

  • Module –

    module with added wrapper function definitions and function calls

Source code in llmcompressor/pipelines/sequential/ast_utils/auto_wrapper.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
def auto_wrap(self, tree: ast.Module) -> ast.Module:
    """
    Modify ast by automatically wrapping any untraceable code segments. Segments to
    wrap are determined through analysis of the code and basic pattern matching

    :param tree: module containing a definition to an original forward function
    :return: module with added wrapper function definitions and function calls
    """
    tree = self.visit(tree)
    for fn_def in self._wrapper_fn_defs:
        tree.body.insert(0, fn_def)

    return ast.fix_missing_locations(tree)

visit_Call 

visit_Call(node: Call) -> ast.Call

Wrap any functions which use (4) variadic arguments or (2) match the ignore list

Source code in llmcompressor/pipelines/sequential/ast_utils/auto_wrapper.py 
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
def visit_Call(self, node: ast.Call) -> ast.Call:
    """
    Wrap any functions which use (4) variadic arguments or (2) match the ignore list
    """
    # check for variadic starred
    if any(isinstance(elem, ast.Starred) for elem in node.args):
        return self._wrap_if_possible(node)

    # attempt to evaluate caller and check against ignore list
    try:
        caller = self._eval_expr(node.func)

    except Exception:
        caller = None

    finally:
        if (
            isinstance(caller, (FunctionType, MethodType))
            and caller.__name__ in self.ignore
        ):
            return self._wrap_if_possible(node)

    return super().generic_visit(node)

visit_Delete 

visit_Delete(node: Delete)

Remove any deleted names from self._local_names, which are used to determine function wrapper arguments

Source code in llmcompressor/pipelines/sequential/ast_utils/auto_wrapper.py 
86
87
88
89
90
91
92
93
94
95
96
97
def visit_Delete(self, node: ast.Delete):
    """
    Remove any deleted names from `self._local_names`,
    which are used to determine function wrapper arguments
    """
    ret = super().generic_visit(node)

    for target in node.targets:
        if isinstance(target, ast.Name):
            self._local_names.remove(target.id)

    return ret

visit_FunctionDef 

visit_FunctionDef(node: FunctionDef) -> ast.FunctionDef

Remove decorators which prevent forward function recompilation For example, add_start_docstrings_to_model_forward

Because _wrapper_fn_defs are appended after visit finishes, this function will not affect wrapper functions

Parameters:

  • node

    (FunctionDef) –

    function definition whose decorators will be stripped

Returns:

  • FunctionDef –

    function definition without decorators

Source code in llmcompressor/pipelines/sequential/ast_utils/auto_wrapper.py 
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
def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef:
    """
    Remove decorators which prevent forward function recompilation
    For example, add_start_docstrings_to_model_forward

    Because `_wrapper_fn_defs` are appended after `visit` finishes, this function
    will not affect wrapper functions

    :param node: function definition whose decorators will be stripped
    :return: function definition without decorators
    """
    node.decorator_list = [
        decorator_name
        for decorator_name in node.decorator_list
        if isinstance(decorator_name, ast.Name)
        and decorator_name.id in ("can_return_tuple",)  # modifies func signature
    ]

    if node.name == "forward":
        for arg in node.args.args:
            self._local_names.add(arg.arg)
        for arg in node.args.posonlyargs:
            self._local_names.add(arg.arg)
        for arg in node.args.kwonlyargs:
            self._local_names.add(arg.arg)
        if node.args.vararg:
            self._local_names.add(node.args.vararg.arg)
        if node.args.kwarg:
            self._local_names.add(node.args.kwarg.arg)
    return super().generic_visit(node)

visit_If 

visit_If(node: If) -> ast.If | ast.Assign

Attempt to statically evaluate the condition of the if statement. If the condition can not be statically evaluated (1), then attmept to wrap the if statement

Parameters:

  • node

    (If) –

    if statement which may be wrapped

Returns:

  • If | Assign –

    if the if statement cannot be statically evaluated, return the if statement with the condition replaced by True or False. Otherwise, return a wrapper function call

Source code in llmcompressor/pipelines/sequential/ast_utils/auto_wrapper.py 
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
def visit_If(self, node: ast.If) -> ast.If | ast.Assign:
    """
    Attempt to statically evaluate the condition of the `if` statement. If the
    condition can not be statically evaluated (1), then attmept to wrap the `if`
    statement

    :param node: `if` statement which may be wrapped
    :return: if the `if` statement cannot be statically evaluated, return the
        `if` statement with the condition replaced by `True` or `False`.
        Otherwise, return a wrapper function call
    """
    try:
        value = bool(self._eval_expr(node.test))

        # force a wrap if any assignments occur within the if statement
        for expr in ast.walk(node):
            if isinstance(expr, ast.NamedExpr):
                raise Exception("If statement contains assignment")

    except Exception:
        return self._wrap_if_possible(node)

    else:
        node.test = ast.Constant(value=value)
        return super().generic_visit(node)

visit_Name 

visit_Name(node: Name)

Add any new names in self._local_names, which are used to determine function wrapper arguments

Source code in llmcompressor/pipelines/sequential/ast_utils/auto_wrapper.py 
76
77
78
79
80
81
82
83
84
def visit_Name(self, node: ast.Name):
    """
    Add any new names in `self._local_names`,
    which are used to determine function wrapper arguments
    """
    if isinstance(node.ctx, ast.Store):
        self._local_names.add(node.id)

    return super().generic_visit(node)

visit_Tuple 

visit_Tuple(node: Tuple) -> ast.Tuple | ast.Call

(3) Wrap any tuples which use starred unpacking

Source code in llmcompressor/pipelines/sequential/ast_utils/auto_wrapper.py 
125
126
127
128
129
130
131
132
def visit_Tuple(self, node: ast.Tuple) -> ast.Tuple | ast.Call:
    """
    (3) Wrap any tuples which use starred unpacking
    """
    if any(isinstance(elem, ast.Starred) for elem in node.elts):
        return self._wrap_if_possible(node)

    return super().generic_visit(node)

ControlFlowAnalyzer

Bases: NodeVisitor

Used to determine if a piece of code can be wrapped into a function. Includes any code which includes return, continue, break, await, or yield without the proper context.

For example, this code can be wrapped: 

while True:
    if some_condition:
        break

Whereas the inner code cannot be wrapped without the while context 

def wrapped():
    if some_condition:
        break  # this control statement is now invalid

while True:
    wrapped()

Methods:

  • is_valid –

    Returns False if a node contains control statements that are not in their

is_valid 

is_valid(node: AST) -> bool

Returns False if a node contains control statements that are not in their proper control flow context

Parameters:

  • node

    (AST) –

    code to analyze

Returns:

  • bool –

    True iff the code does not contain invalid control statements

Source code in llmcompressor/pipelines/sequential/ast_utils/control_flow_analyzer.py 
32
33
34
35
36
37
38
39
40
41
42
43
def is_valid(self, node: ast.AST) -> bool:
    """
    Returns False if a node contains control statements that are not in their
    proper control flow context

    :param node: code to analyze
    :return: True iff the code does not contain invalid control statements
    """
    self._contexts = []
    self._is_valid = True
    self.visit(node)
    return self._is_valid

NameAnalyzer 

NameAnalyzer(omit: set[str])

Bases: NodeVisitor

Determines the unbound, assigned, and conditionally assigned names associated with a piece of code. This information is used to determine the arguments and return values of the wrapper function

For example, for the following piece of code 

b = a + 1
if some_condition:
    c = 5

a is unbound, meaning that it must be an input of wrapper function b is assigned, meaning that it must be an output of the wrapper function c is conditionally assigned, meaning that it must be an output of the wrapper function, and might be an input iff c already exists in the namespace

Note that names which are assigned to before being read are not considered unbound 

a = 2  # no longer unbound
b = a + 1

Methods:

  • analyze –

    Analyzes the use of names in the given piece of code

Source code in llmcompressor/pipelines/sequential/ast_utils/name_analyzer.py 
37
38
def __init__(self, omit: set[str]):
    self._omit = builtins.__dict__.keys() | omit

analyze 

analyze(node: AST) -> tuple[set[str], set[str], set[str]]

Analyzes the use of names in the given piece of code

Parameters:

  • node

    (AST) –

    code to analyze return: tuple of unbound names, assigned names, and conditionally assigned names

Source code in llmcompressor/pipelines/sequential/ast_utils/name_analyzer.py 
40
41
42
43
44
45
46
47
48
49
50
51
52
def analyze(self, node: ast.AST) -> tuple[set[str], set[str], set[str]]:
    """
    Analyzes the use of names in the given piece of code

    :param node: code to analyze
    return: tuple of unbound names, assigned names, and conditionally assigned names
    """
    self._unbound = set()
    self._assigned = set()
    self._conditionally_assigned = set()
    self.visit(node)

    return self._unbound, self._assigned, self._conditionally_assigned