[WIP][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)

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

* unbound, assigned, conditionally_assigned

Wrapping code that references `self`

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

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 any if statements and analyzing any submodule calls, 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 (todo programmatically check other model architectures)
• Ran examples/quantization_w4a16/llama3_example.py to completion