Fix MaskFormerSwin attention mask dtype to follow hidden states - #49032
kaixuanliu wants to merge 4 commits into
Conversation
Signed-off-by: kaixuanliu <kaixuan.liu@intel.com>
guarin
left a comment
guarin
left a comment
There was a problem hiding this comment.
LGTM, thanks for the pr! Seems we already have the same dtype conversion for the other models based on swin.
The only caveat I see is that we now have a different argument order because of BC:
def get_attn_mask(self, input_resolution, device=None, dtype=None):
the others have dtype before device:
def get_attn_mask(self, height: int, width: int, dtype: torch.dtype, device: torch.device) -> torch.Tensor | None:
cc @vasqu
|
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. |
|
agreed that arg order is best kept as it was before! |
The arg order is already kept as before so all is good. The only issue is that the order is different from other models but this is required because of backwards compatibility. |
Signed-off-by: kaixuanliu <kaixuan.liu@intel.com>
vasqu
left a comment
vasqu
left a comment
There was a problem hiding this comment.
Small nit but otherwise LGTM
| w_region = (w_idx >= width - self.window_size).long() + (w_idx >= width - self.shift_size).long() | ||
| img_mask = (h_region[None, :, None, None] * 3 + w_region[None, None, :, None]).float() | ||
| img_mask = (h_region[None, :, None, None] * 3 + w_region[None, None, :, None]).to( | ||
| dtype=dtype if dtype is not None else torch.float32 |
There was a problem hiding this comment.
why not default the dtype arg to float instead then? then we avoid the ternary no?
There was a problem hiding this comment.
Well good advice. Looks better now.
|
Re the order: rather keep BC than breaking the order. its a bit awkward but not worth to break, especiallly considering that the model is quite old |
|
[For maintainers] Suggested jobs to run (before merge) run-slow: maskformer |
CI recapDashboard: View test results in Grafana |
This branch has not been deployed
5 participants
MaskFormerSwinLayer.get_attn_maskhardcodes the shifted-window attention mask to float32 via.float(), ignoring the model dtype. When the model runs in bfloat16/float16, adding this fp32 mask promotes attention_scores to fp32, so the subsequent torch.matmul(attention_probs, value_layer) mixes fp32 with a bf16 value_layer.This silently works in eager on some backends, but fails under torch.compile, where the fake-tensor meta kernel strictly validates dtypes:
@molbap @guarin pls help review, thx!