Make GraniteMoeHybridModel compatible with torch.export
Author: juanigp
Pull request: yes

# What does this PR do?

PR aims to make GraniteMoeHybridModel compatible with `torch.export` . See [#42317](https://github.com/huggingface/transformers/pull/42317) for related export and fixes.

Initial change is to fix `_update_mamba_mask`, replacing a mask of None with a tensor of ones.


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@IlyasMoutawwakil 


Comments:
Hi ! thanks for the proposition, unfortunately torch.cond is still very experimental as  mentioned in torch docs, it also doesn't  support training https://docs.pytorch.org/docs/stable/generated/torch.cond.html
```
torch.cond is a prototype feature in PyTorch. It has limited support for input and output types and doesn’t support training currently. Please look forward to a more stable implementation in a future version of PyTorch. Read more about feature classification at: https://pytorch.org/blog/pytorch-feature-classification-changes/#prototype
```

Hi @IlyasMoutawwakil thanks for the reply! Do you think of any possible alternative? The changes aim to keep the original implementation (sorta) if not is_torchdynamo_compiling().

Last commit aims to achieve the goal without having torch.cond. @IlyasMoutawwakil would appreciate feedback (:

Hi @juanigp i'm not sure i follow.. the code is kinda overcomplicating the issue here, when compiling you simply wanna return the attention mask as is, without turning it into None (which is an optimization that's data dependent), i.e. _update_mamba_mask can simply become:
```
        mamba_mask = attention_mask
        if not is_torchcompiling() and (cache_position[0] > 0 or (attention_mask is not None and torch.all(attention_mask == 1))):
            mamba_mask = None
        return mamba_mask
```

> Hi @juanigp i'm not sure i follow.. the code is kinda overcomplicating the issue here, when compiling you simply wanna return the attention mask as is, without turning it into None (which is an optimization that's data dependent), i.e. _update_mamba_mask can simply become:
> 
> ```
>         mamba_mask = attention_mask
>         if not is_torchcompiling() and (cache_position[0] > 0 or (attention_mask is not None and torch.all(attention_mask == 1))):
>             mamba_mask = None
>         return mamba_mask
> ```

Thanks, this is way cleaner. When I naively replaced the original conditional with torch.cond it was complaining about the None value, but none of that is needed ... 

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: granitemoehybrid