Description of the bug
When a model's attention layers use scaled_dot_product_attention (SDPA) instead of eager attention, qmodel_prep silently fails to attach QBmm to any attention matmuls — even with nbits_bmm1/nbits_bmm2 correctly configured. SDPA is traced as a single fused kernel call rather than decomposed bmm/matmul ops, so BMM detection finds nothing. The only indication is an easy-to-miss INFO-level log line:
Found 0 torch.bmm and 0 torch.matmul
No warning or error is raised, so a user can end up with bmm or matmul left entirely unquantized in an otherwise-quantized model. Bug was only found while looking at the quantized model output.
Platform
Python 3.12
fms-model-optimizer 0.8.5
transformers 5.12
Sample Code
For reproducibility, any HF transformers model whose attention defaults to SDPA (e.g. PatchTSTForPrediction on transformers>=4.53) run through qmodel_prep with BMM quantization enabled (nbits_bmm1/nbits_bmm2 set).
Expected behavior
torchscript tracer should have found 6 matmul operations. 0 was found as mentioned above.
Observed behavior
Found 0 torch.bmm and 0 torch.matmul
Additional context
A workaround is to set attn_implementation="eager" at model construction. A Warning can be issued in the interim, alerting users that if SDPA is used, no bmm and matmul ops may be found.
Description of the bug
When a model's attention layers use scaled_dot_product_attention (SDPA) instead of eager attention, qmodel_prep silently fails to attach QBmm to any attention matmuls — even with nbits_bmm1/nbits_bmm2 correctly configured. SDPA is traced as a single fused kernel call rather than decomposed bmm/matmul ops, so BMM detection finds nothing. The only indication is an easy-to-miss INFO-level log line:
No warning or error is raised, so a user can end up with bmm or matmul left entirely unquantized in an otherwise-quantized model. Bug was only found while looking at the quantized model output.
Platform
Python 3.12
fms-model-optimizer 0.8.5
transformers 5.12
Sample Code
For reproducibility, any HF transformers model whose attention defaults to SDPA (e.g.
PatchTSTForPredictionon transformers>=4.53) run throughqmodel_prepwith BMM quantization enabled (nbits_bmm1/nbits_bmm2 set).Expected behavior
torchscripttracer should have found 6 matmul operations. 0 was found as mentioned above.Observed behavior
Found 0 torch.bmm and 0 torch.matmulAdditional context
A workaround is to set
attn_implementation="eager"at model construction. AWarningcan be issued in the interim, alerting users that ifSDPAis used, no bmm and matmul ops may be found.