[SAM3] Compute masks once instead of per-layer, fix fa2 crash#42543
Conversation
|
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. |
molbap
left a comment
There was a problem hiding this comment.
Better already, thanks for the refacto!
| hidden_states = torch.cat([presence_token, hidden_states], dim=1) | ||
| query_pos = torch.cat([torch.zeros_like(presence_token), query_pos], dim=1) | ||
| # Prepend zeros to query_pos for presence token | ||
| query_pos = F.pad(query_pos, (0, 0, 1, 0), mode="constant", value=0) |
vasqu
left a comment
There was a problem hiding this comment.
LGTM overall, just some info on FA and a bit of fixes there.
Where are the relative biases introduced? Not important but moreso FMI. Sadly, most FA versions can't support this.
| if ( | ||
| self.config._attn_implementation == "flash_attention_2" | ||
| and attention_mask is not None | ||
| and attention_mask.dtype != torch.bool | ||
| ): | ||
| # Relative position bias tensors are represented as float masks and are incompatible with Flash Attention 2. | ||
| # Fallback to SDPA for this call only so the rest of the model can still benefit from FA2. | ||
| attention_interface = ALL_ATTENTION_FUNCTIONS["sdpa"] | ||
| logger.warning_once( | ||
| "Sam3Attention: falling back to SDPA for relative-position cross-attention because " | ||
| "Flash Attention 2 does not support additive bias masks." | ||
| ) |
There was a problem hiding this comment.
Is this common in the model? Maybe we should just deactivate FA support if that's the case --> it seldomly can work at all. But given the comment I doubt it?
| if ( | |
| self.config._attn_implementation == "flash_attention_2" | |
| and attention_mask is not None | |
| and attention_mask.dtype != torch.bool | |
| ): | |
| # Relative position bias tensors are represented as float masks and are incompatible with Flash Attention 2. | |
| # Fallback to SDPA for this call only so the rest of the model can still benefit from FA2. | |
| attention_interface = ALL_ATTENTION_FUNCTIONS["sdpa"] | |
| logger.warning_once( | |
| "Sam3Attention: falling back to SDPA for relative-position cross-attention because " | |
| "Flash Attention 2 does not support additive bias masks." | |
| ) | |
| if ( | |
| "flash" in self.config._attn_implementation | |
| and attention_mask is not None | |
| and attention_mask.dtype != torch.bool | |
| ): | |
| # Relative position bias tensors are represented as float masks and are incompatible with Flash Attention | |
| # Fallback to SDPA for this call only so the rest of the model can still benefit from FA | |
| attention_interface = ALL_ATTENTION_FUNCTIONS["sdpa"] | |
| logger.warning_once( | |
| "Sam3Attention: falling back to SDPA for relative-position cross-attention because " | |
| "Flash Attention does not support additive bias masks." | |
| ) |
For FA, you should only check for the "flash" keyword. The reason is that we also support different other attentions such as FA3 and kernel variations (+ FA4 soon I hope, working on that one).
There was a problem hiding this comment.
But yea, FA itself cannot support relative biases (except alibi). There was a different library which might support this #41882 (FA with biases).
There was a problem hiding this comment.
Thanks for the info! I'll put "flash" to handle fa3/4.
Is this common in the model? Maybe we should just deactivate FA support if that's the case --> it seldomly can work at all. But given the comment I doubt it?
It's only used in DETR decoder layer attention, so other attentions (there are quite a few) can still benefit from fa.
| _checkpoint_conversion_mapping = { | ||
| r"detector_model.(.+)": r"\1" # the regex allows to remove the prefix, and add it back in revert mode | ||
| } | ||
| _checkpoint_conversion_mapping = {"detector_model.": ""} |
There was a problem hiding this comment.
I know Cyril worked a lot on this to work, just want to double check if this works :D
There was a problem hiding this comment.
This must have been changed back by mistake during a merge, thanks for catching that!
|
[For maintainers] Suggested jobs to run (before merge) run-slow: sam3 |
…gface#42543) * Compute masks once instead of per-layer, fix fa2 crash * nit * Change after review
…gface#42543) * Compute masks once instead of per-layer, fix fa2 crash * nit * Change after review
Some small improvements for SAM3. There will be more to come as we discussed @molbap ;), including using modular.
This PR adds the following improvements: