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


- model loading (from pretrained, etc): @CyrilVallez
- distributed: @3outeille @ArthurZucker

fix tp crash. crash stack is 

[rank0]: Traceback (most recent call last):
[rank0]:   File "/transformers/benchmark_v2/test_tp.py", line 29, in <module>
[rank0]:     outputs = model.generate(**inputs.to(model.device), max_new_tokens=100, do_sample=False)
[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/venv/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 124, in decorate_context
[rank0]:     return func(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/transformers/src/transformers/generation/utils.py", line 2671, in generate
[rank0]:     result = decoding_method(
[rank0]:              ^^^^^^^^^^^^^^^^
[rank0]:   File "/transformers/src/transformers/generation/utils.py", line 2866, in _sample
[rank0]:     outputs = self._prefill(input_ids, generation_config, model_kwargs)
[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/transformers/src/transformers/generation/utils.py", line 3855, in _prefill
[rank0]:     return self(**model_inputs, return_dict=True)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1776, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1787, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/transformers/src/transformers/utils/generic.py", line 835, in wrapper
[rank0]:     output = func(self, *args, **kwargs)
[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/transformers/src/transformers/models/gpt_oss/modeling_gpt_oss.py", line 659, in forward
[rank0]:     outputs: MoeModelOutputWithPast = self.model(
[rank0]:                                       ^^^^^^^^^^^
[rank0]:   File "/opt/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1776, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1787, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/transformers/src/transformers/utils/generic.py", line 1002, in wrapper
[rank0]:     outputs = func(self, *args, **kwargs)
[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/transformers/src/transformers/models/gpt_oss/modeling_gpt_oss.py", line 498, in forward
[rank0]:     hidden_states = decoder_layer(
[rank0]:                     ^^^^^^^^^^^^^^
[rank0]:   File "/transformers/src/transformers/modeling_layers.py", line 93, in __call__
[rank0]:     return super().__call__(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1776, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1787, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/transformers/src/transformers/models/gpt_oss/modeling_gpt_oss.py", line 374, in forward
[rank0]:     hidden_states, _ = self.self_attn(
[rank0]:                        ^^^^^^^^^^^^^^^
[rank0]:   File "/opt/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1776, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1787, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/transformers/src/transformers/models/gpt_oss/modeling_gpt_oss.py", line 332, in forward
[rank0]:     attn_output, attn_weights = attention_interface(
[rank0]:                                 ^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/transformers/src/transformers/models/gpt_oss/modeling_gpt_oss.py", line 262, in eager_attention_forward
[rank0]:     combined_logits = torch.cat([attn_weights, sinks], dim=-1)
[rank0]:                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: RuntimeError: Sizes of tensors must match except in dimension 3. Expected size 32 but got size 64 for tensor number 1 in the list.


Changed files:
filename: src/transformers/models/gpt_oss/configuration_gpt_oss.py
patch: @@ -36,7 +36,7 @@ class GptOssConfig(PreTrainedConfig):
         "layers.*.self_attn.k_proj": "colwise",
         "layers.*.self_attn.v_proj": "colwise",
         "layers.*.self_attn.o_proj": "rowwise",
-        "layers.*.self_attn.sinks": "rowwise",
+        "layers.*.self_attn.sinks": "colwise",
         "layers.*.mlp.router": "ep_router",
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
         "layers.*.mlp.experts.gate_up_proj_bias": "grouped_gemm",