This is fixed in torch 2.2.0 I think, maybe I should check for it and skip the calls?

I think it is fine to leave. We should probably bump the requirement for SDPA to torch>=2.2 in the future.

This got me thinking, and I ran an additional set of benchmarking, given that FA2 is supported and the contiguous bug is fixed in 2.2.0: training and inference.

Both training and inference were ~5% faster with torch==2.2.0 (FA2 should be supported). I also tried out gating the .contiguous() requirement and saw an additional ~5-10% gain on top of that.

if version.parse(get_torch_version()) < version.parse("2.2.0")
    query_layer = query_layer.contiguous()
    key_layer = key_layer.contiguous()
    value_layer = value_layer.contiguous()

I'm leaning towards adding the if-statement to gate the call, so users who upgrade to torch=2.2.0 first can get the benefits right away (before we set the min torch version to 2.2.0). WDYT?

I added the if-statement for 2.2.0 in there. If you don't think it's a good idea, let me know and I'll remove it.

I would probably move the Copied from just to the __init__ and other methods, but not forward. For the forward, you can probably just add a comment that it is adapted from bert/roberta and once bridge_tower supports sdpa we can put back to copied from.

WDYT @ArthurZucker @amyeroberts

There seems to be 8 methods that copy-from BertMode#forward() exactly and has this section of change.

I won't mind adding SDPA to them as well once this goes in and reinstating the copy-from. It shouldn't be that difficult (famous last words)

I've removed the fix-copies from the instances, and so the logic for sdpa attention masks should only be in BertModel now.

@ArthurZucker there are create_extended_attention_mask_for_decoder, invert_attention_mask, get_extended_attention_mask methods in modeling_utils.py that should probably be deprecated / redirect to modeling_attn_mask_utils.py.

Yea, I agree.

It'd be great if we could mark those old methods as deprecated, and slowly update them once we verify that the old methods and the new methods are always returning the same results.

For the updated_attention_mask for sdpa, why can't we keep the previous logic and just do:

                # Attend to all tokens in masked rows from the causal_mask, for example the relevant first rows when
                # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
                # Details: https://github.com/pytorch/pytorch/issues/110213
                causal_mask = causal_mask.mul(~torch.all(causal_mask == torch.finfo(embedding_output.dtype).min, dim=-1, keepdim=True)).to(
                    dtype
                )

(from Llama)?
Not super fan of the complexity of _prepare_4d_causal_attention_mask_for_sdpa, and we should not add it in our new code IMO.

This code was changed to pass the fx tracing test (in common tests).

It would be good if you can help double check the logic here. I think the idea here is that we'll still have to use our own attention mask (rather than None) when tracing is active. The previous "pass" would cause the function to end without any return statements, which would have defaulted to None.

It looks OK to me, cc @fxmarty to confirm.

AFAICT, the difference here is coming from the additional isinstance(mask, torch.fx.Proxy) in the is_tracing_check. I don't believe the reworking to remove pass should affect anything - the new code is equivalent.

Yes it is fine, see

This fix was added due to a test failure that uncovered an existing bug.

The head was initialized but the weights weren't retied as necessary. This was causing self.decoder.bias to be different from self.bias. When loading the pretrained model with low_cpu_mem_usage=True, the self.decoder.bias had uninitiated params (with device=meta) whereas self.bias was set properly (with device=cpu)

I'm slightly concerned this will affect the output some users see when using this model. Please let me know what you think about this.

I pulled this out to its own PR here:
#28948

This issue is unrelated to SDPA, but was just uncovered by a SPDA test, so I just pulled it out to its own PR.

Addition looks OK to me - thanks for digging into this.

I'm slightly concerned this will affect the output some users see when using this model. Please let me know what you think about this.

Could you expand on what you think might be an issue?

I was initially concerned that users were loading and using the model with a wrong bias (ie. device=meta), and this fix to use the correct bias will cause the results to change between versions.

However, that seems unlikely after playing around with this a bit more - turns out it was quite difficult to run the model when the bias had device=meta, so I doubt anyone was actually running the model in this particular configuration before the fix.

The self._prepare_for_class is necessary to support the BertForMultipleChoice model.

It looks OK to me, cc @fxmarty to confirm.

AFAICT, the difference here is coming from the additional isinstance(mask, torch.fx.Proxy) in the is_tracing_check. I don't believe the reworking to remove pass should affect anything - the new code is equivalent.

Addition looks OK to me - thanks for digging into this.

I'm slightly concerned this will affect the output some users see when using this model. Please let me know what you think about this.

Could you expand on what you think might be an issue?