Mask2Former HungarianMatcherMaks2Former overflow on bad initial state when using 16-bit autocast
Author: pjh4993
Pull request: no

### System Info

- `transformers` version: 4.29.0
- Platform: Linux-5.15.0-79-generic-x86_64-with-glibc2.31
- Python version: 3.10.12
- Huggingface_hub version: 0.13.4
- Safetensors version: 0.3.2
- PyTorch version (GPU?): 2.0.0+cu117 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: yes

### Who can help?

@amy

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

1. import pair_wise_sigmoid_cross_entropy_loss from transformers.models.mask2former_modeling_mask2former
```
from transformers.models.mask2former.modeling_mask2former import pair_wise_sigmoid_cross_entropy_loss
```
2. initialize arbitrary input and labels. 
```
inputs = torch.randn([16,121445]).cuda() - 2 # To simulate bad initialization
labels = torch.randn([2, 121445]).cuda()
```
3. Compute loss with 16-bit autocast
```
with torch.autocast("cuda", dtype=torch.hafl):
     result = pair_wise_sigmoid_cross_entropy_loss(inputs, labels)
```
4. result got int everywhere
```
>>> result.isinf().all()
tensor(True, device='cuda:0')
```

I suggest normalizing `cross_entropy_loss_pos`, `cross_entropy_loss_neg` with `height_and_width` before computing matmul as follows:
```

def pair_wise_sigmoid_cross_entropy_loss(
    inputs: torch.Tensor, labels: torch.Tensor
) -> torch.Tensor:
    """
    A pair wise version of the cross entropy loss, see `sigmoid_cross_entropy_loss` for usage.

    Args:
        inputs (`torch.Tensor`):
            A tensor representing a mask.
        labels (`torch.Tensor`):
            A tensor with the same shape as inputs.
            Stores the binary classification labels for each element in inputs
            (0 for the negative class and 1 for the positive class).

    Returns:
        loss (`torch.Tensor`): The computed loss between each pairs.
    """

    height_and_width = inputs.shape[1]

    criterion = torch.nn.BCEWithLogitsLoss(reduction="none")
    # prevent overflow
    cross_entropy_loss_pos = criterion(inputs, torch.ones_like(inputs)) / height_and_width
    cross_entropy_loss_neg = criterion(inputs, torch.zeros_like(inputs)) / height_and_width

    loss_pos = torch.matmul(cross_entropy_loss_pos, labels.T)
    loss_neg = torch.matmul(cross_entropy_loss_neg, (1 - labels).T)
    loss = loss_pos + loss_neg
    return loss
```

This helps prevent overflow at matmul computation with 16-bit precision or lower.

### Expected behavior

non-overflow ce-loss result

Comments:
Regarding the script, it seems that you are forcing a bad initialisation, but the `Mask2FormerPreTrainedModel` 's `_init_weights` should be responsible of a correct initialisation, which is why I am not sure we should include such changes. Not familiar with the model / original loss might not use this, which can change expected results. Pinging @amyeroberts (which is the person you were tying to ping?) 


Hi @pjh4993, thanks for raising this issue! 

@ArthurZucker's right that casting should be handled through officially supported routes e.g. `Mask2FormerForUniversalSegmentation.from_pretrained(checkpoint, torch_dtype=torch.float16)`

That being said, your proposed change is a good one: it makes things more stable and equivalent to the previous functions logic. Would you like to open a PR with this change? This way you get the github contribution. 

@amyeroberts  Thx. I will open PR till UTC+9:00 2023-09-10 12:00:00

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.