[Tracing] Code AutoWrapper

[kylesayrs]

Purpose

• Reduce model support burden by automatically wrapping untraceable code
  • This is a programmatic implementation of all of the rules specified by the tracing guide
• Remove traceable definitions for Idefics3ForConditionalGeneration, LlavaForConditionalGeneration, MllamaForConditionalGeneration, LlavaForConditionalGeneration, Qwen2_5_VLForConditionalGeneration, and Qwen2VLForConditionalGeneration, Gemma3ForConditionalGeneration (all of them)

Fixes

• Tracing tests failing with latest transformers>4.52.0 #1457

Autowrap Patterns

These patterns match syntax which is untraceable and unlikely to call sequential targets (either directly or indirectly)

If statements whose conditions cannot be statically evaluated

This if statement can be statically evaluated, since its value can be evaluated in the context of {"self": LlamaModel(...)}

if self.config._attn_implementation != "eager":
    ...

If the statement cannot be statically evaluated, then it is wrapped

torch.fx.wrap
def wrapped(input_ids, inputs_embeds):
    if (input_ids is None) ^ (inputs_embeds is not None):
        raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

def forward(...):
    (,) = wrapped(input_ids, inputs_embeds)

Ignored functions (`_update_causal_mask`)

Any function or method names listed in the ignore list will be automatically wrapped

torch.fx.wrap
def wrapped(attention_mask, inputs_embeds, cache_position, ...):
    return self._update_causal_mask(attention_mask, inputs_embeds, cache_position, ...)

def forward(...):
    causal_mask = wrapped(attention_mask, inputs_embeds, cache_position, ...)

Starred tuple unpacking

Any use of iterated unpacking will be automatically wrapped

torch.fx.wrap
def wrapped(input_shape):
    return (*input_shape, -1, self.head_dim)

def forward(...):
    hidden_shape = wrapped(input_shape)

Starred argument unpacking

Any use of iterated unpacking into variadic args is automatically wrapped

torch.fx.wrap
def wrapped(attn_output, input_shape):
    return attn_output.reshape(*input_shape, -1)

def forward(...):
    attn_output = wrapped(attn_output, input_shape)

Autowrap Implementation Details

Wrapper arguments

Autowrapping a piece of code requires determining which variable names are used by that code and which variable names are produced by that code. This is done using the `NameAnalyzer`, which determines the unbound, assigned, and conditionally assigned names for a given piece of code.

# unbound := names which are read by node before being assigned
# assigned := names which are assigned by operations in node
# cond_assigned := names which may be assigned depending on execution
analyzer = NameAnalyzer(omit=self.namespace.keys())
unbound, assigned, conditionally_assigned = analyzer.analyze(node)

This information is then used to determine what the args, kwargs, and return names should be for the wrapping function.

# args := names which already existed and are needed for ops or wrapped return
# kwargs := names which are needed for return but did not already exist
# returns := names which are assigned or could be assigned
args = (unbound | conditionally_assigned) & self._local_names
kwargs = conditionally_assigned - self._local_names
returns = assigned | conditionally_assigned

Wrapping methods

Some untraceable code references `self` during execution. While normally, `self` would be an argument to the wrapped function, `self` is usually a `torch.nn.Module` which is not a handled type that can be passed around the graph. Instead, we treat `self` as a variable in the compiled python module namespace, and this namespace is automatically captured and executed by `torch.fx._symbolic_trace`

Unwrappable code

Some code cannot be wrapped because it contains control flow statements which must exist in a certain context. For example, we cannot wrap code that contains a `continue` without also wrapping the for loop that surrounds it.

for index, layer in enumerate(self.layers):
    # ---- cannot autowrap ----
    if index <= 10:
        continue
    # ---- cannot autowrap ----
    
    hidden_states = layer(hidden_states)

Future Extensions/ Improvements

Sequentially executing vision towers

Sequentially tracing vision towers is a lower priority, as the vision towers are typically fewer parameters and aren't quantization targets. However, in the event that they do become quantization targets, or memory(vision_tower + one target) > memory(one gpu), then the vision tower layers will need to be split up.

Some changes may be required to support this. Conditionally executing the vision tower is a very common pattern:

def forward(pixel_values, image_embeds, ...):
    if image_embeds is None:
        image_embeds = self.vision_tower(pixel_values)
    ...

Some approaches might be

1. Allowing names like image_embeds to be evaluated based on the sample input being passed
2. Pattern matching against self.{module_name}(), where module_name is determined to be a module through evaluation
3. Using type hinting analysis tools like jedi to track the types of all names, and to check if any names whose type is a module are called

Towards perfect autowrapping

As mentioned in in "Sequentially executing vision towers", it may be possible to use use type hinting analysis tools like jedi or pytype to infer whether any given code chunk calls a sequential target or target ancestor. If this can be done reliably (which will require extensions such as analysis of called functions), then all code that does not call sequential targets can be wrapped.

Towards removing tracing

If we can reliably determine if a given code chunk calls sequential any targets, it may be possible to define each autowrapped function as its own subgraph directly, without the need for tracing. However, this will require inference of model execution from the static code, which means unrolling loops, expanding any function/method calls, and resolving dynamic model structure (in the case of llava and idefics), which may be tricky.

For example, if you want to determine the execution of the LlamaDecoder layers of a llava model like pixtral, you'd need to evaluate self.language_model, then analyze the source of the caller's forward function, which is stored in a separate file.

def forward(...):
    self.language_model(...)  # the type of self.language_model is determined from the config

Another point of evaluation would be evaling any iteration of ModuleLists

def forward(...):
    for decoder_layer of self.layers:  # ModuleList isn't well typed and may contain different types
        decoder_layer(...)

The tracing system could be replaced with static code inference, both are different systems for solving the problem of determining model execution

Testing

• Able to trace all models in tests/llmcompressor/transformers/tracing/test_models.py without requiring traceable definitions
• Verified sequentially executed outputs are correct for LlamaForCausalLM
• Ran examples/quantization_w4a16/llama3_example.py to completion

[kylesayrs]

@brian-dellabetta The types of tracing failures are as follows:

1. The model fails to trace
2. The model traces but does not create the expected number of subgraphs
3. The model traces but the traced subgraphs do not produce the same output as the original model

(1) and (2) are covered by the tracing tests. (3) does not currently have a test (I'll add this as a follow up), but from my testing of the examples, it seems that this is not an issue

[brian-dellabetta]

discussed further over a call, impressive work!

[dsikka]

the tracing guide was so well written
im sad to see it go :(

[kylesayrs]

@dsikka I have no such attachment :)

[dsikka]

This LGTM but I was unfamiliar with the ast library so I used this source to try and understand the manipulation provided by the classe: https://github.com/xbeat/Machine-Learning/blob/main/Exploring%20Python's%20Abstract%20Syntax%20Tree%20Manipulation.md

We should potentially provide some sort of docs/resource in the source code as well before landing

[brian-dellabetta]

awesome stuff!

[kylesayrs]

Seems like there's a problem specifically with loading llama4 processors https://huggingface.co/meta-llama/Llama-4-Scout-17B-16E-Instruct/discussions/45 🫠