Make GraniteMoeHybridModel compatible with torch.export

[juanigp]

What does this PR do?

PR aims to make GraniteMoeHybridModel compatible with torch.export . See #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,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? 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, and
  here are tips on formatting docstrings.
• 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

[github-actions]

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

run-slow: granitemoehybrid

[IlyasMoutawwakil]

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

[juanigp]

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().