diff --git a/src/transformers/masking_utils.py b/src/transformers/masking_utils.py index 45e43fdaf3aa..5dede9865a46 100644 --- a/src/transformers/masking_utils.py +++ b/src/transformers/masking_utils.py @@ -610,6 +610,14 @@ def eager_mask( min_dtype = torch.finfo(dtype).min # we need 0s where the tokens should be taken into account, and -inf otherwise (mask is already of boolean type) mask = torch.where(mask, torch.tensor(0.0, device=mask.device, dtype=dtype), min_dtype) + # Prevent NaN from softmax([-inf, ..., -inf]) for rows where all keys are masked. This occurs when a + # padding query position's sliding window falls entirely within the padding region (e.g. a short sequence + # padded to match a much longer one in a batch). Unlike torch.nn.functional.scaled_dot_product_attention + # (which handles this natively since torch>=2.5), the manual softmax in eager attention forward does not. + # For such rows we allow all keys (bias=0), producing a valid uniform distribution. These rows always + # correspond to padding positions and are excluded from final pooling, so the result is unused. + all_masked_rows = (mask == min_dtype).all(dim=-1, keepdim=True) + mask = torch.where(all_masked_rows, torch.zeros_like(mask), mask) return mask