Repository: huggingface/transformers
Pull request: fix gptoss crash in tp


- quantization: @SunMarc @MekkCyber

 ```

[rank0]: During handling of the above exception, another exception occurred:

[rank0]: Traceback (most recent call last):
[rank0]:   File "/mnt/disk3/wangyi/transformers/test_tp.py", line 24, in <module>
[rank0]:     outputs = model.generate(**inputs.to(model.device), max_new_tokens=100, do_sample=False)
[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/disk0/wangyi/miniforge3/envs/transformers/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
[rank0]:     return func(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/disk3/wangyi/transformers/src/transformers/generation/utils.py", line 2638, in generate
[rank0]:     result = decoding_method(
[rank0]:              ^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/disk3/wangyi/transformers/src/transformers/generation/utils.py", line 2833, in _sample
[rank0]:     outputs = self._prefill(input_ids, generation_config, model_kwargs)
[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/disk3/wangyi/transformers/src/transformers/generation/utils.py", line 3822, in _prefill
[rank0]:     return self(**model_inputs, return_dict=True)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/disk0/wangyi/miniforge3/envs/transformers/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/disk0/wangyi/miniforge3/envs/transformers/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/disk3/wangyi/transformers/src/transformers/utils/generic.py", line 842, in wrapper
[rank0]:     output = func(self, *args, **kwargs)
[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/disk3/wangyi/transformers/src/transformers/models/gpt_oss/modeling_gpt_oss.py", line 660, in forward
[rank0]:     outputs: MoeModelOutputWithPast = self.model(
[rank0]:                                       ^^^^^^^^^^^
[rank0]:   File "/mnt/disk0/wangyi/miniforge3/envs/transformers/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/disk0/wangyi/miniforge3/envs/transformers/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/disk3/wangyi/transformers/src/transformers/utils/generic.py", line 1025, in wrapper
[rank0]:     raise original_exception
[rank0]:   File "/mnt/disk3/wangyi/transformers/src/transformers/utils/generic.py", line 1016, in wrapper
[rank0]:     outputs = func(self, *args, **kwargs)
[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/disk3/wangyi/transformers/src/transformers/models/gpt_oss/modeling_gpt_oss.py", line 499, in forward
[rank0]:     hidden_states = decoder_layer(
[rank0]:                     ^^^^^^^^^^^^^^
[rank0]:   File "/mnt/disk3/wangyi/transformers/src/transformers/modeling_layers.py", line 93, in __call__
[rank0]:     return super().__call__(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/disk0/wangyi/miniforge3/envs/transformers/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/disk0/wangyi/miniforge3/envs/transformers/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/disk3/wangyi/transformers/src/transformers/models/gpt_oss/modeling_gpt_oss.py", line 390, in forward
[rank0]:     hidden_states, _ = self.mlp(hidden_states)  # diff with llama: router scores
[rank0]:                        ^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/disk0/wangyi/miniforge3/envs/transformers/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/disk0/wangyi/miniforge3/envs/transformers/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/disk3/wangyi/transformers/src/transformers/integrations/mxfp4.py", line 446, in mlp_forward
[rank0]:     routed_out = self.experts(hidden_states, routing_data, gather_idx, scatter_idx)
[rank0]:                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/disk0/wangyi/miniforge3/envs/transformers/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/disk0/wangyi/miniforge3/envs/transformers/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1881, in _call_impl
[rank0]:     return inner()
[rank0]:            ^^^^^^^
[rank0]:   File "/mnt/disk0/wangyi/miniforge3/envs/transformers/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1829, in inner
[rank0]:     result = forward_call(*args, **kwargs)
[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: TypeError: Mxfp4GptOssExperts.forward() missing 1 required positional argument: 'scatter_idx'
```

Changed files:
filename: src/transformers/integrations/mxfp4.py
patch: @@ -443,7 +443,7 @@ def mlp_forward(self, hidden_states):
     with on_device(router_logits.device):
         routing_data, gather_idx, scatter_idx = routing(router_logits, self.router.top_k)
 
-    routed_out = self.experts(hidden_states, routing_data, gather_idx, scatter_idx)
+    routed_out = self.experts(hidden_states, routing_data, gather_idx, scatter_idx=scatter_idx)
     routed_out = routed_out.reshape(batch_size, -1, self.router.hidden_dim)
     return routed_out, router_logits