ignore _update_mamba_mask for AWQ sequential tracing

[toncao]

SUMMARY:

In models with mamba-2 layers e.g., nvidia/NVIDIA-Nemotron-Nano-12B-v2, Qwen/Qwen3-Next-80B-A3B-Instruct, tracing _update_mamba_masks would lead to

  File "NemotronHModel_8045287568680_autowrapped", line 57, in forward
  File "/mnt/LinuxDrive/huggingface/modules/transformers_modules/NVIDIA_hyphen_Nemotron_hyphen_Nano_hyphen_12B_hyphen_v2/modeling_nemotron_h.py", line 1461, in _update_mamba_mask
    if cache_position[0] > 0 or (attention_mask is not None and torch.all(attention_mask == 1)):
       ^^^^^^^^^^^^^^^^^^^^^
  File "/home/toncao/anaconda3/envs/llm-compressor_v1/lib/python3.12/site-packages/transformers/utils/fx.py", line 674, in __bool__
    return super().__bool__()
           ^^^^^^^^^^^^^^^^^^
  File "/home/toncao/anaconda3/envs/llm-compressor_v1/lib/python3.12/site-packages/torch/fx/proxy.py", line 577, in __bool__
    return self.tracer.to_bool(self)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/toncao/anaconda3/envs/llm-compressor_v1/lib/python3.12/site-packages/torch/fx/proxy.py", line 388, in to_bool
    raise TraceError(
torch.fx.proxy.TraceError: symbolically traced variables cannot be used as inputs to control flow

from the function:

def _update_mamba_mask(self, attention_mask, cache_position):
--
"""
No need for zeroing states when
1. Cached forward
2. Attending to all inputs
"""
mamba_mask = attention_mask
if cache_position[0] > 0 or (attention_mask is not None and torch.all(attention_mask == 1)):
mamba_mask = None
return mamba_mask

And thus, adding _update_mamba_masks to the ignore tracing list makes AWQ sequential tracing works.

TEST PLAN:

local make test results:

===================================================== short test summary info =====================================================
FAILED tests/llmcompressor/modeling/test_calib_deepseek_v3.py::test_calib_deepseekv3_module - torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 56.00 MiB. GPU 0 has a total capacity of 23.57 GiB of which 14.1...
FAILED tests/llmcompressor/utils/test_helpers.py::test_disable_cache[MllamaForConditionalGeneration-meta-llama/Llama-3.2-11B-Vision-Instruct] - huggingface_hub.errors.GatedRepoError: 403 Client Error. (Request ID: Root=1-68ee275c-378c35b1649b823602164fc0;24ebe331-9031-4...
FAILED tests/lmeval/test_lmeval.py::TestLMEval::test_lm_eval[None] - TypeError: argument should be a str or an os.PathLike object where __fspath__ returns a str, not 'NoneType'
====================================== 3 failed, 242 passed, 4 skipped in 129.47s (0:02:09) =======================================

[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.

[gemini-code-assist]

Summary of Changes

Hello @toncao, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request resolves a critical "TraceError" that occurs when performing AWQ sequential tracing on models that incorporate Mamba-2 layers. The error stemmed from symbolic tracing variables being used in control flow within the "_update_mamba_mask" function. By explicitly adding this function to the tracing ignore list, the PR ensures that these models can now be successfully traced without encountering the previously reported issues.

Highlights

• Mamba-2 Layer Tracing Fix: Added "_update_mamba_mask" to the list of functions ignored during AWQ sequential tracing to resolve a "TraceError" encountered in models utilizing Mamba-2 layers.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request correctly addresses a tracing error with Mamba-2 models by adding _update_mamba_mask to the tracing_ignore list. This change is necessary and well-justified. I have one suggestion to improve code maintainability.

[brian-dellabetta]

Thanks @toncao , makes sense to me