GptOss experts implementation
Author: IlyasMoutawwakil
Pull request: yes

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

[wip] Fixes #43193 @vasqu @mattteochen
<img width="6300" height="4500" alt="generate" src="https://github.com/user-attachments/assets/60ee7037-de6d-4101-bae8-302d829481c2" />


## 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.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker @Cyrilvallez
- vision models: @yonigozlan @molbap
- audio models: @eustlb @ebezzam @vasqu
- multimodal models: @zucchini-nlp
- graph models: @clefourrier

Library:

- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- continuous batching: @remi-or @ArthurZucker @McPatate
- pipelines: @Rocketknight1
- tokenizers: @ArthurZucker and @itazap
- trainer: @SunMarc
- attention: @vasqu @ArthurZucker @CyrilVallez
- model loading (from pretrained, etc): @CyrilVallez
- distributed: @3outeille @ArthurZucker
- CIs: @ydshieh

Integrations:

- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization: @SunMarc @MekkCyber
- kernels: @MekkCyber @drbh
- peft: @BenjaminBossan @githubnemo

Devices/Backends:

- AMD ROCm: @ivarflakstad
- Intel XPU: @IlyasMoutawwakil
- Ascend NPU: @ivarflakstad 

Documentation: @stevhliu

Research projects are not maintained and should be taken as is.

 -->


Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_43227). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

run-slow: gpt_oss

> Not against it but we will need to solve the loading/reverse issue either way, not necessarily here

yes and the transposition can be achieved and works fine (all tests passà with the temporary weight renaming trick, however the biggest problem here is that megablocks expect the transposed format of expert weights so it might make more sense to transpose the other MoEs to be compatible with megablocks (if it's worth the speedup)

> the biggest problem here is that megablocks expect the transposed format of expert weights

Argh, yea this will be hard to solve otherwise (would need to rewrite that kernel). Maybe we transpose the other way around for the other models at some point - depends on how much improvement that kernel provides 

On another note, is the transpose expensive on runtime? I would expect slight perf regression having to call transpose here and there (in `grouped_mm` implementation, I see a few). 

> On another note, is the transpose expensive on runtime?

no, the grouped_mm work on strided tensors so the transpose is just a view (no op).

> LGTM, I guess the outputs slightly change now with the interface's grouped_mm/batched_mm or even the eager rewrite?

the eager for training is the same, I kept everything as is when i realized that the slow tests are brokens (see https://github.com/huggingface/transformers/pull/43246), I removed the batched implementation from there as it is redundant, so users will have to set experts impl to batched to get the same eval output as before.

Probably merging next week, currently don't have the permission to merge with red CI :smile: 

Should be good to merge 

#43353 leads me to believe that a simple weight converter won't be enough. Somehow under the concurrent workers we hit the same 16byte alignment issue on our CI with (qwen 3 omni moe).

I'll take a proper look tomorrow but it might make sense to add the fallback regardless sadly. Just wanted to put it out before we hastly merge 

@vasqu I think what you are dealing with might be the cpu+safetensors memory misalignement issue i have seen with gptoss ; there are two kinds of misalignment issues i have encountered so far:
- on cuda: this only happens if the strided tensor has shapes that are not divisible by 8 (for fp16/bf16) or 4 (for fp32), otherwise the cuda allocator will allocate an aligned tensor. That's the issue I was solving in the first experts PR, because I was testing on the dgx so tests were automatically running on cuda.
- on cpu: Safetensors adds an 8 bytes-aligned metadata on top of the file ([see this issue from which the 8 bytes were decided](https://github.com/huggingface/safetensors/issues/231)),  so when reloaded on cpu the metadata can "misalign" the tensors in the memmaped file. This is what the weight converter defined here solves.

I would suggest you try the weight converter defined here because the ci runner uses cpu (I had to set TRANSFORMERS_TEST_DEVICE=cpu to reproduce the gpt_oss issue locally)

To sum up the conversation I had with @vasqu internally:

> All moes that have a converter that creates new tensors, for example all the ones that merge their experts using cat/stack, will create 16-bytes aligned tensors on cpu. This is because they will use the allocator to create *new* merged tensors, and [the allocator will return a 16-bytes aligned memory](https://stackoverflow.com/questions/70692795/why-is-malloc-16-byte-aligned). 
If any MoE fails because of 16 bytes misalignment, it's gonna be because it's not allocating its tensors ; for example using safetensors memmaped tensors (which are only 8 bytes aligned). Those are the ones that will need the Force16BytesAlignement converter.

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

run-slow: gpt_oss

View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=43227&sha=711a65