[Tracing] Raise _is_compiling_flag while tracing

[kylesayrs]

Purpose

• Reduce model support burden by skipping untraceable code which has been marked by implementers using is_torchdynamo_compiling and other checks

transformers/models/llava/modeling_llava.py

if not is_torchdynamo_compiling() and inputs_embeds[special_image_mask].numel() != image_features.numel():
    n_image_tokens = (input_ids == self.config.image_token_id).sum()
    n_image_features = image_features.shape[0] * image_features.shape[1]
    raise ValueError(
        f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
    )

torch/compiler/init.py

def is_compiling() -> bool:
    """
    Indicates whether a graph is executed/traced as part of torch.compile() or torch.export().

    Example::

        >>> def forward(self, x):
        >>>     if not torch.compiler.is_compiling():
        >>>        pass # ...logic that is not needed in a compiled/traced graph...
        >>>
        >>>     # ...rest of the function...
    """

Changes

• Set _is_compiling_flag to True while tracing

Testing

• As a part of [WIP][Tracing] Mistral3ForConditionalGeneration #1387, was able to trace Mistral3ForConditionalGeneration by skipping untraceable sections

[github-actions]

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.

Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed.

[brian-dellabetta]

👍