Error: GPT-OSS-20B QAT failed with "RecursionError: maximum recursion depth exceeded"

[tjdhg456]

GitHub Issue Template for NVIDIA Model-Optimizer

Describe the bug

From the "https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/gpt-oss" code, we try to utilize QAT after training with SFT.

When using QATSFTTrainer with a GPT-OSS MoE (Mixture of Experts) model, training fails with a RecursionError: maximum recursion depth exceeded during the evaluation phase.

The issue occurs in _quantized_bmm() function where torch._bmm is intercepted and redirected back to _quantized_bmm, causing infinite recursion (954+ repeated calls).

Impact: Blocker - Cannot perform QAT training on MoE architecture models.

Stack Trace (abbreviated):

RecursionError: maximum recursion depth exceeded

File "modelopt/torch/quantization/plugins/huggingface.py", line 626, in _quantized_bmm
    return torch._bmm(batch1, batch2)
  [Previous line repeated 954 more times]

File "modelopt/torch/quantization/plugins/huggingface.py", line 624, in _quantized_bmm
    batch1 = self.down_proj_input_quantizer(batch1) if self._down_proj_mul else batch1

File "modelopt/torch/quantization/nn/modules/tensor_quantizer.py", line 275, in pre_quant_scale
    if not hasattr(self, "_pre_quant_scale") or not self._enable_pre_quant_scale:

RecursionError: maximum recursion depth exceeded

Steps/Code to reproduce bug

1. Load a GPT-OSS MoE model using transformers
2. Apply QAT using QATSFTTrainer
3. Start training - error occurs during first evaluation

Minimal reproduction code:

from modelopt.torch.quantization.plugins import QATSFTTrainer, QuantizationArguments
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import SFTConfig, ScriptArguments, ModelConfig, TrlParser

# Load GPT-OSS MoE model
model = AutoModelForCausalLM.from_pretrained(
    "path/to/gpt-oss-moe-model",
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
)
tokenizer = AutoTokenizer.from_pretrained("path/to/gpt-oss-moe-model")

# Setup QAT trainer
trainer = QATSFTTrainer(
    model=model,
    args=training_args,  # SFTConfig with eval_strategy="steps"
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    processing_class=tokenizer,
    quant_args=QuantizationArguments(),  # default MXFP4 config
)

# This triggers the error during evaluation
trainer.train()

Model architecture details:

• The GPT-OSS model uses MoE with GptOssExperts layer
• The experts layer uses torch.bmm() for batched matrix multiplication:

  # transformers/models/gpt_oss/modeling_gpt_oss.py:130
  gate_up = torch.bmm(hidden_states, self.gate_up_proj) + self.gate_up_proj_bias[..., None, :]

Expected behavior

QAT training should complete successfully without recursion errors. The _quantized_bmm function should properly quantize torch.bmm operations without creating circular calls.

Who can help?

• @kevalmorabia97

---

System information

• Container used: Custom Ray runtime container with conda environment

• OS: Linux (RHEL-based, kernel 5.15.x)

• CPU architecture: x86_64

• GPU name: NVIDIA H100

• GPU memory size: 80GB

• Number of GPUs: 8 (multi-GPU training with DeepSpeed ZeRO-3)

• Library versions:

  • Python: 3.12
  • ModelOpt version: latest (installed via nvidia-modelopt[all])
  • CUDA: 12.8
  • PyTorch: 2.8.0
  • Transformers: 4.57.6
  • TRL: 0.27.2
  • Accelerate: 1.12.0
  • DeepSpeed: latest
  • TensorRT-LLM: N/A
  • ONNXRuntime: N/A
  • TensorRT: N/A

• Additional details:

  • Training framework: Ray + DeepSpeed + Accelerate
  • The error occurs at ~28% progress (step 8/29) during the first evaluation
  • All 8 ranks fail with the same RecursionError
  • Training phase before evaluation completes successfully

---

Questions

1. Does nvidia-modelopt officially support MoE (Mixture of Experts) architectures? If so, are there specific configuration requirements?

2. The _quantized_bmm function in huggingface.py creates a circular call when torch._bmm is intercepted. Is this a known issue with MoE models?

3. What is the recommended way to exclude specific layers (e.g., MoE experts using torch.bmm) from quantization while still applying QAT to the rest of the model?

4. Are there alternative quantization configurations that avoid quantizing torch.bmm operations?

---

Workarounds Attempted

1. Using manual mtq.quantize() instead of QATSFTTrainer - Not yet tested
2. Excluding experts layers from quantization - Not yet tested:

   from modelopt.torch.quantization import disable_quantization
   for name, module in model.named_modules():
       if "experts" in name:
           disable_quantization(module)

---

Related Files

• Error location: modelopt/torch/quantization/plugins/huggingface.py:624-626
• Model forward: transformers/models/gpt_oss/modeling_gpt_oss.py:130

[tjdhg456]

Bug Report: RecursionError in _QuantGptOssExperts when using torch.bmm

Describe the bug

When using QATSFTTrainer with a GPT-OSS MoE (Mixture of Experts) model, training/evaluation fails with RecursionError: maximum recursion depth exceeded.

The root cause is in _QuantGptOssExperts.functionals_to_replace property in modelopt/torch/quantization/plugins/huggingface.py. The _quantized_bmm function calls torch._bmm, but since replace_function dynamically sets torch._bmm during the forward pass, there's a race condition or re-entrant call that causes torch._bmm to redirect back to _quantized_bmm, creating infinite recursion.

Impact: Blocker - Cannot perform QAT training/evaluation on GPT-OSS MoE architecture models.

Error Message:

RecursionError: maximum recursion depth exceeded

File "modelopt/torch/quantization/plugins/huggingface.py", line 626, in _quantized_bmm
    return torch._bmm(batch1, batch2)
  [Previous line repeated 954 more times]

File "modelopt/torch/quantization/plugins/huggingface.py", line 624, in _quantized_bmm
    batch1 = self.down_proj_input_quantizer(batch1) if self._down_proj_mul else batch1

---

Root Cause Analysis

Problematic Code

# modelopt/torch/quantization/plugins/huggingface.py (lines 621-636)
@property
def functionals_to_replace(self):
    def _quantized_bmm(batch1, batch2):
        batch1 = self.down_proj_input_quantizer(batch1) if self._down_proj_mul else batch1
        self._down_proj_mul = not self._down_proj_mul
        return torch._bmm(batch1, batch2)  # <-- Problem: torch._bmm may redirect to _quantized_bmm

    def _tensor_matmul(self_t, other):
        self_t = self.down_proj_input_quantizer(self_t) if self._down_proj_mul else self_t
        self._down_proj_mul = not self._down_proj_mul
        return torch.matmul(self_t, other)  # <-- Same problem

    return [
        (torch, "bmm", _quantized_bmm),
        (torch.Tensor, "__matmul__", _tensor_matmul),
    ]

Why Recursion Occurs

The replace_function utility (in utils.py) does:

@contextmanager
def replace_function(package, name, new_func, og_func_cache_name=None):
    if og_func_cache_name is None:
        og_func_cache_name = "_" + name  # "_bmm"
    old_func = getattr(package, name)     # torch.bmm
    setattr(package, name, new_func)       # torch.bmm = _quantized_bmm
    setattr(package, og_func_cache_name, old_func)  # torch._bmm = original
    yield
    # restore...

The issue is that torch._bmm is set dynamically during forward pass. In MoE models with multiple expert calls or nested forward passes, the timing of when torch._bmm is accessed vs when it's set can cause it to point back to _quantized_bmm instead of the original function.

---

Proposed Fix

Cache the original torch.bmm and torch.matmul functions at module import time (before any replacement happens), and use these cached references instead of relying on torch._bmm.

Option 1: Fix in _QuantGptOssExperts class

# At module level (top of huggingface.py)
_ORIGINAL_BMM = torch.bmm
_ORIGINAL_MATMUL = torch.matmul

# In _QuantGptOssExperts class
@property
def functionals_to_replace(self):
    def _quantized_bmm(batch1, batch2):
        batch1 = self.down_proj_input_quantizer(batch1) if self._down_proj_mul else batch1
        self._down_proj_mul = not self._down_proj_mul
        return _ORIGINAL_BMM(batch1, batch2)  # Use cached original

    def _tensor_matmul(self_t, other):
        self_t = self.down_proj_input_quantizer(self_t) if self._down_proj_mul else self_t
        self._down_proj_mul = not self._down_proj_mul
        return _ORIGINAL_MATMUL(self_t, other)  # Use cached original

    return [
        (torch, "bmm", _quantized_bmm),
        (torch.Tensor, "__matmul__", _tensor_matmul),
    ]

Option 2: Fix in replace_function utility

Store the original function reference in a way that's guaranteed not to be overwritten:

@contextmanager
def replace_function(package, name, new_func, og_func_cache_name=None):
    # Use a unique name that won't conflict
    if og_func_cache_name is None:
        og_func_cache_name = f"_modelopt_original_{name}"
    # ... rest of implementation

---

Workaround (Current)

We've implemented a monkey-patch in our training script that caches the original functions before any ModelOpt code runs:

import torch

# Cache original functions before any replacement
_ORIGINAL_BMM = torch.bmm
_ORIGINAL_MATMUL = torch.matmul

def _patch_quant_gpt_oss_experts():
    """Patch _QuantGptOssExperts.functionals_to_replace to use cached original functions."""
    try:
        from modelopt.torch.quantization.plugins.huggingface import _QuantGptOssExperts
    except ImportError:
        return
    
    @property
    def patched_functionals_to_replace(self):
        def _quantized_bmm(batch1, batch2):
            batch1 = self.down_proj_input_quantizer(batch1) if self._down_proj_mul else batch1
            self._down_proj_mul = not self._down_proj_mul
            return _ORIGINAL_BMM(batch1, batch2)
        
        def _tensor_matmul(self_t, other):
            self_t = self.down_proj_input_quantizer(self_t) if self._down_proj_mul else self_t
            self._down_proj_mul = not self._down_proj_mul
            return _ORIGINAL_MATMUL(self_t, other)
        
        return [
            (torch, "bmm", _quantized_bmm),
            (torch.Tensor, "__matmul__", _tensor_matmul),
        ]
    
    _QuantGptOssExperts.functionals_to_replace = patched_functionals_to_replace

_patch_quant_gpt_oss_experts()

---

Steps to Reproduce

1. Load a GPT-OSS MoE model using transformers
2. Apply QAT using QATSFTTrainer with any quantization config
3. Run training or evaluation - error occurs during forward pass

Minimal reproduction:

from modelopt.torch.quantization.plugins import QATSFTTrainer, QuantizationArguments
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import SFTConfig

model = AutoModelForCausalLM.from_pretrained(
    "path/to/gpt-oss-moe-model",
    torch_dtype=torch.bfloat16,
)
tokenizer = AutoTokenizer.from_pretrained("path/to/gpt-oss-moe-model")

trainer = QATSFTTrainer(
    model=model,
    args=SFTConfig(output_dir="./output", eval_strategy="steps"),
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    processing_class=tokenizer,
    quant_args=QuantizationArguments(quant_cfg="MXFP4_MLP_WEIGHT_ONLY_CFG"),
)

trainer.train()  # RecursionError during evaluation

---

System Information

• OS: Linux (RHEL-based)
• CPU architecture: x86_64
• GPU: NVIDIA H100 (80GB) x 8
• Library versions:
  • Python: 3.12
  • nvidia-modelopt: 0.27.x (latest)
  • PyTorch: 2.8.0
  • Transformers: 4.57.6
  • TRL: 0.27.2
  • CUDA: 12.8
• Training setup: Ray + DeepSpeed ZeRO-3 + Accelerate

---

Expected Behavior

QAT training and evaluation should complete successfully without recursion errors. The _quantized_bmm function should correctly call the original torch.bmm implementation.

---

Additional Context

• The error occurs specifically with MoE (Mixture of Experts) models that use GptOssExperts layer
• The GptOssExperts.forward() uses torch.bmm() for batched matrix multiplication
• The error is reproducible 100% of the time during evaluation phase
• Training phase may complete some steps before hitting the error during eval

---

Related Files

• modelopt/torch/quantization/plugins/huggingface.py - _QuantGptOssExperts class (lines 580-665)
• modelopt/torch/quantization/plugins/custom.py - _QuantFunctionalMixin class
• modelopt/torch/quantization/utils.py - replace_function utility (lines 340-349)