Skip to content

Fix eager attention selection during model loading - #1375

Merged
fangyangci merged 3 commits into
mainfrom
fangyangci/fix-eager-attention-export
Sep 2, 2026
Merged

Fix eager attention selection during model loading#1375
fangyangci merged 3 commits into
mainfrom
fangyangci/fix-eager-attention-export

Conversation

@fangyangci

Copy link
Copy Markdown
Contributor

Summary

  • select the configured Transformers attention implementation before from_pretrained constructs attention modules
  • thread export compatibility through both build and direct export paths
  • cover loader and build propagation while keeping export test doubles compatible

Motivation

DistilBERT exports could instantiate DistilBertSdpaAttention and only switch the config to eager immediately before ONNX export. This left the module implementation and config inconsistent, causing integer attention masks to reach SDPA and fail with:

Expected attn_mask dtype to be bool or float or to match query dtype

Validation

  • uv run pytest tests/unit/loader/test_load_hf_model.py tests/unit/build/test_hf.py tests/unit/export/test_htp_exporter_attention_compat.py tests/unit/commands/test_config_value_priority.py tests/unit/commands/test_export.py -q --tb=short
  • 175 tests passed
  • retried all 6 affected QNN GPU models: the attention-mask dtype error was eliminated in all 6; 5 completed perf and accuracy successfully, while one progressed to a separate duplicate metadata_props ONNX validation failure

@fangyangci
fangyangci requested a review from a team as a code owner August 31, 2026 09:10
Comment thread src/winml/modelkit/build/hf.py Outdated
Comment thread src/winml/modelkit/commands/export.py
@fangyangci
fangyangci merged commit 3de717f into main Sep 2, 2026
9 checks passed
@fangyangci
fangyangci deleted the fangyangci/fix-eager-attention-export branch September 2, 2026 02:55
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants