Repository navigation
Keep attention_mask as None in OPT's causal mask creation - #49002
Conversation
`OPTDecoder` replaces a `None` `attention_mask` with a dense all-ones tensor so that the learned positional embeddings can be computed. That tensor was then also forwarded to `create_causal_mask`, which therefore had to inspect its values and could never fall back to sdpa's `is_causal` argument. Keep the dense tensor local to the positional embeddings and pass the original `attention_mask` to `create_causal_mask`.
|
Hi @vasqu . Would you please review this PR? Thanks! |
| config_and_inputs = self.model_tester.prepare_config_and_inputs() | ||
| self.model_tester.create_and_check_decoder_model_past_large_inputs(*config_and_inputs) | ||
|
|
||
| def test_attention_mask_is_not_overwritten_for_causal_mask(self): |
There was a problem hiding this comment.
let's make it mamba2 style like splitting into the tester but otherwise LGTM
There was a problem hiding this comment.
Done, the body moved into OPTModelTester.
| model(inputs_dict["input_ids"], attention_mask=None) | ||
|
|
||
| self.assertTrue(mocked_create_causal_mask.called) | ||
| self.assertIsNone(mocked_create_causal_mask.call_args.kwargs["attention_mask"]) |
There was a problem hiding this comment.
Tbh wouldnt it also be possible to force the FA backend and run one forward? Just as alternative maybe, feels a bit more appropriate
There was a problem hiding this comment.
You are right, and it works now that #48975 is in main. Added as test_compiled_forward_dispatches_to_flash_attention: a compiled forward under sdpa_kernel([SDPBackend.FLASH_ATTENTION]), which only runs if the mask creation was skipped, since flash attention rejects any attn_mask.
| eager | compiled | |
|---|---|---|
| main | passes | raises |
| this PR | passes | passes |
Eager cannot discriminate: _ignore_causal_mask_sdpa reads the all-ones tensor and returns None anyway. Only the compiled path sees the difference, because reading the values is declined while tracing.
I kept the mocked assertion as well, since the new test needs a GPU and is skipped on the CPU runners. Happy to drop it if you prefer just the one.
|
[For maintainers] Suggested jobs to run (before merge) run-slow: opt |
|
Nice thanks merging |
CI recapDashboard: View test results in Grafana |
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
Fixes #48924
Split out of #48975 on review request.
Context
The same change was proposed in #48929 and closed, because it does not move the SDPA path on its own. That is correct, and this PR does not claim otherwise: the skip only takes effect together with the
masking_utilsfix in #48975. Measured on one A100 80GB with torch 2.14.0,opt-125m, batch 8, sequence 1024, bf16, sdpa,torch.compile(dynamic=False), no padding mask:mainWhy it is still worth doing
Two things that may not be obvious from the earlier discussion.
The all-ones tensor is allocated on every forward, but the caller can already supply
position_ids. At this call siteposition_idsis always computed beforeself.embed_positions(...), soOPTLearnedPositionalEmbedding.forwardnever reads itsattention_maskargument. Replacing it with zeros leaves the output bit-identical. The tensor exists only to be passed to something that ignores it, and tocreate_causal_mask, which is exactly where it does harm.Separately,
create_causal_maskis the one consumer that can tell the difference between "no mask" and "a mask that happens to be all ones". Handing it a synthetic mask discards information the caller gave us. Today only the compiled path observes this, but the constraint is not specific totorch.compile.Change
The dense tensor stays local, in
position_attention_mask. The caller'sattention_maskreachescreate_causal_maskunchanged.This is a no-op numerically. A dump of 30 intermediate tensors across fp32/bf16, unpadded, left and right padded inputs, in eager and under
torch.compile, is bit-identical before and after.Tests
test_attention_mask_is_not_overwritten_for_causal_maskis new and fails onmain.tests/models/opt: 150 passed, 148 skipped.Note
Happy to close this if you would rather keep OPT as is until #48975 lands, or to fold it back in. It is also possible to drop the all-ones tensor entirely and derive the positions from
arange, but that is a larger change and I left it out here.