IndexError: tuple index out of range when using Tensor Parallelism with FSDP2 on GPT-OSS 20B (tensor_parallel.py, line 510)

[JdRion]

System Info

• transformers version: 4.57.1
• Platform: Linux-5.15.92-2.el8.navix.ncc.x86_64-x86_64-with-glibc2.39
• Python version: 3.12.3
• Huggingface_hub version: 0.35.3
• Safetensors version: 0.5.2
• Accelerate version: 1.11.0
• Accelerate config: not found
• DeepSpeed version: not installed
• PyTorch version (accelerator?): 2.8.0+cu128 (CUDA)
• Tensorflow version (GPU?): not installed (NA)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Using distributed or parallel set-up in script?:
• Using GPU in script?:
• GPU type: NVIDIA H100 80GB HBM3

Who can help?

@3outeille @ArthurZucker

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

I’m using the official example script from the Accelerate repository:
examples/torch_native_parallelism/nd_parallel_trainer.py.

I only made two small changes to the example:
1. Changed the model name to GPT-OSS-20B
2. Added a quantization configuration

quantization_config = Mxfp4Config(dequantize=True)
model = AutoModelForCausalLM.from_pretrained(args.model_name, quantization_config=quantization_config, use_cache=False, **model_kwargs)

distributed_type: FSDP
mixed_precision: bf16

fsdp_config:
  fsdp_activation_checkpointing: false
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_cpu_ram_efficient_loading: false
  fsdp_offload_params: false
  fsdp_reshard_after_forward: true
  fsdp_state_dict_type: SHARDED_STATE_DICT
  fsdp_version: 2

parallelism_config:
  parallelism_config_cp_size: 1
  parallelism_config_dp_replicate_size: 1
  parallelism_config_dp_shard_size: 16
  parallelism_config_tp_size: 8

Expected behavior

When training the GPT-OSS 20B model with Tensor Parallel (TP) and FSDP2, a runtime error occurs inside the tensor parallel integration code of transformers.
The error trace shows an IndexError: tuple index out of range raised from transformers/integrations/tensor_parallel.py during forward pass (specifically in _prepare_input_fn).

[rank2]:   File "/usr/local/lib/python3.12/dist-packages/transformers/trainer.py", line 2328, in train
[rank2]:     return inner_training_loop(
[rank2]:            ^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/transformers/trainer.py", line 2672, in _inner_training_loop
[rank2]:     tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
[rank2]:                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/trl/trainer/sft_trainer.py", line 1161, in training_step
[rank2]:     return super().training_step(*args, **kwargs)
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/transformers/trainer.py", line 4009, in training_step
[rank2]:     loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/trl/trainer/sft_trainer.py", line 1079, in compute_loss
[rank2]:     (loss, outputs) = super().compute_loss(
[rank2]:                       ^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/transformers/trainer.py", line 4099, in compute_loss
[rank2]:     outputs = model(**inputs)
[rank2]:               ^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
[rank2]:     return self._call_impl(*args, **kwargs)
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1879, in _call_impl
[rank2]:     return inner()
[rank2]:            ^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1827, in inner
[rank2]:     result = forward_call(*args, **kwargs)
[rank2]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/transformers/utils/generic.py", line 940, in wrapper
[rank2]:     output = func(self, *args, **kwargs)
[rank2]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/transformers/models/gpt_oss/modeling_gpt_oss.py", line 663, in forward
[rank2]:     outputs: MoeModelOutputWithPast = self.model(
[rank2]:                                       ^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
[rank2]:     return self._call_impl(*args, **kwargs)
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1784, in _call_impl
[rank2]:     return forward_call(*args, **kwargs)
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/transformers/utils/generic.py", line 1064, in wrapper
[rank2]:     outputs = func(self, *args, **kwargs)
[rank2]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/transformers/models/gpt_oss/modeling_gpt_oss.py", line 502, in forward
[rank2]:     hidden_states = decoder_layer(
[rank2]:                     ^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/transformers/modeling_layers.py", line 94, in __call__
[rank2]:     return super().__call__(*args, **kwargs)
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
[rank2]:     return self._call_impl(*args, **kwargs)
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1879, in _call_impl
[rank2]:     return inner()
[rank2]:            ^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1827, in inner
[rank2]:     result = forward_call(*args, **kwargs)
[rank2]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/transformers/utils/deprecation.py", line 172, in wrapped_func
[rank2]:     return func(*args, **kwargs)
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/transformers/models/gpt_oss/modeling_gpt_oss.py", line 366, in forward
[rank2]:     hidden_states, _ = self.self_attn(
[rank2]:                        ^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
[rank2]:     return self._call_impl(*args, **kwargs)
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1784, in _call_impl
[rank2]:     return forward_call(*args, **kwargs)
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/torch/distributed/algorithms/_checkpoint/checkpoint_wrapper.py", line 171, in forward
[rank2]:     return self.checkpoint_fn(  # type: ignore[misc]
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/torch/_compile.py", line 53, in inner
[rank2]:     return disable_fn(*args, **kwargs)
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py", line 929, in _fn
[rank2]:     return fn(*args, **kwargs)
[rank2]:            ^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/torch/utils/checkpoint.py", line 495, in checkpoint
[rank2]:     ret = function(*args, **kwargs)
[rank2]:           ^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
[rank2]:     return self._call_impl(*args, **kwargs)
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1879, in _call_impl
[rank2]:     return inner()
[rank2]:            ^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1816, in inner
[rank2]:     args_result = hook(self, args)
[rank2]:                   ^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/transformers/integrations/tensor_parallel.py", line 381, in <lambda>
[rank2]:     module.register_forward_pre_hook(lambda mod, inputs: input_fn(mod, inputs, device_mesh))
[rank2]:                                                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/transformers/integrations/tensor_parallel.py", line 510, in _prepare_input_fn
[rank2]:     input_tensor = inputs[0]
[rank2]:                    ~~~~~~^^^
[rank2]: IndexError: tuple index out of range

[3outeille]

i'm taking a look

[3outeille]

took some time to check, can you reproduce the error with fsdp 2 + tp 2 + Mxfp4Config(dequantize=True) on a tinier version of GPT-oss ?

[JdRion]

I currently don’t have access to GPUs, so I should be able to share more details or logs by tomorrow.

However, I can confirm that the same setup worked correctly with Qwen3-14B-Base and NousResearch/Hermes-3-Llama-3.1-8B models.
On the other hand, when I tested Qwen3-30B-A3B-Thinking, I encountered a different type of error compared to GPT-OSS-20B.

Based on this, I am assuming that the issue is likely related to MoE models specifically, rather than the environment or configuration.

[JdRion]

@3outeille Hi, I took some time to reproduce the issue.

Environment

• 4 nodes, each with 8× H100 80GB (total 32 GPUs)
• TP, FSDP2 enabled (fully_shard)
• quantization_config = Mxfp4Config(dequantize=True)
• model: GPT-OSS-20B (GPT-OSS 20B is the smallest available version, so I wasn’t able to test with a smaller model.)

Accelerate Config

distributed_type: FSDP
fsdp_config:
  fsdp_activation_checkpointing: false
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_cpu_ram_efficient_loading: false
  fsdp_offload_params: false
  fsdp_reshard_after_forward: true
  fsdp_state_dict_type: SHARDED_STATE_DICT
  fsdp_version: 2
mixed_precision: bf16
parallelism_config:
  parallelism_config_cp_size: 1
  parallelism_config_dp_replicate_size: 4
  parallelism_config_dp_shard_size: 1
  parallelism_config_tp_size: 8
rdzv_backend: static

# TP=2, dp_shard_size=16
During from_pretrained(...) loading phase, the model tries to dequantize parameters before FSDP sharding happens.
This leads to a CUDA OOM during hf_quantizer.create_quantized_param():

Loading checkpoint shards:  67%|██████████████       | 2/3 [00:08<00:04,  4.00s/it]
[rank5]: Traceback (most recent call last):
[rank5]:   File "test.py", line 83, in <module>
[rank5]:     main()
[rank5]:   File "test.py", line 52, in main
[rank5]:     model = AutoModelForCausalLM.from_pretrained(args.model_name, quantization_config=quantization_config, use_cache=False, **model_kwargs)
[rank5]:             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]:   File "/usr/local/lib/python3.12/dist-packages/transformers/models/auto/auto_factory.py", line 604, in from_pretrained
[rank5]:     return model_class.from_pretrained(
[rank5]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]:   File "/usr/local/lib/python3.12/dist-packages/transformers/modeling_utils.py", line 277, in _wrapper
[rank5]:     return func(*args, **kwargs)
[rank5]:            ^^^^^^^^^^^^^^^^^^^^^
[rank5]:   File "/usr/local/lib/python3.12/dist-packages/transformers/modeling_utils.py", line 5048, in from_pretrained
[rank5]:     ) = cls._load_pretrained_model(
[rank5]:         ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]:   File "/usr/local/lib/python3.12/dist-packages/transformers/modeling_utils.py", line 5468, in _load_pretrained_model
[rank5]:     _error_msgs, disk_offload_index = load_shard_file(args)
[rank5]:                                       ^^^^^^^^^^^^^^^^^^^^^
[rank5]:   File "/usr/local/lib/python3.12/dist-packages/transformers/modeling_utils.py", line 843, in load_shard_file
[rank5]:     disk_offload_index = _load_state_dict_into_meta_model(
[rank5]:                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]:   File "/usr/local/lib/python3.12/dist-packages/torch/utils/_contextlib.py", line 120, in decorate_context
[rank5]:     return func(*args, **kwargs)
[rank5]:            ^^^^^^^^^^^^^^^^^^^^^
[rank5]:   File "/usr/local/lib/python3.12/dist-packages/transformers/modeling_utils.py", line 740, in _load_state_dict_into_meta_model
[rank5]:     hf_quantizer.create_quantized_param(
[rank5]:   File "/usr/local/lib/python3.12/dist-packages/transformers/quantizers/quantizer_mxfp4.py", line 246, in create_quantized_param
[rank5]:     dequantize(module, param_name, param_value, target_device, dq_param_name, **shard_kwargs)
[rank5]:   File "/usr/local/lib/python3.12/dist-packages/transformers/integrations/mxfp4.py", line 354, in dequantize
[rank5]:     setattr(module, proj, torch.nn.Parameter(dequantized.to(target_device)))
[rank5]:                                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 508.00 MiB. GPU 1 has a total capacity of 79.11 GiB of which 97.94 MiB is free. Process 20971 has 17.07 GiB memory in use. Process 20973 has 15.35 GiB memory in use. Process 20967 has 31.21 GiB memory in use. Process 20969 has 15.35 GiB memory in use. Of the allocated memory 16.56 GiB is allocated by PyTorch, and 640.00 KiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

# TP=4, dp_shard_size=8
loading passes, but fails later at FSDP2 setup:

[rank5]:   File "/usr/local/lib/python3.12/dist-packages/transformers/trainer.py", line 2325, in train
[rank5]:     return inner_training_loop(
[rank5]:            ^^^^^^^^^^^^^^^^^^^^
[rank5]:   File "/usr/local/lib/python3.12/dist-packages/transformers/trainer.py", line 2480, in _inner_training_loop
[rank5]:     model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer)
[rank5]:                             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]:   File "/usr/local/lib/python3.12/dist-packages/accelerate/accelerator.py", line 1555, in prepare
[rank5]:     result = self._prepare_fsdp2(*args)
[rank5]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]:   File "/usr/local/lib/python3.12/dist-packages/accelerate/accelerator.py", line 1687, in _prepare_fsdp2
[rank5]:     model = fsdp2_prepare_model(self, model)
[rank5]:             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]:   File "/usr/local/lib/python3.12/dist-packages/accelerate/utils/fsdp_utils.py", line 670, in fsdp2_prepare_model
[rank5]:     fully_shard(module, **fsdp2_kwargs)
[rank5]:   File "/usr/local/lib/python3.12/dist-packages/torch/distributed/_composable/contract.py", line 150, in wrapper
[rank5]:     updated = func(inp_module, *args, **kwargs)
[rank5]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]:   File "/usr/local/lib/python3.12/dist-packages/torch/distributed/fsdp/_fully_shard/_fully_shard.py", line 223, in fully_shard
[rank5]:     _move_states_to_device(params, buffers, device)
[rank5]:   File "/usr/local/lib/python3.12/dist-packages/torch/distributed/fsdp/_fully_shard/_fsdp_init.py", line 233, in _move_states_to_device
[rank5]:     raise AssertionError(
[rank5]: AssertionError: Expects DTensor to be moved to cuda but got cuda:1

TP=8, dp_shard_size=4 -> Same behavior

[3outeille]

i'll check if I can reproduce this on 1 node, otherwise it will be a pain to debug

[3outeille]

Right now, this is not possible to create a tiny GPT-OSS with dequantized=True (due to weight loading issues) which makes things very annoying to debug. However, once this PR that refactor weight loading will land (cf #41580), I will come back to it so stay tuned

[3outeille]

Btw, I tried a tiny GPT-oss training TP + FSDP with quantization on and it was working fine. I can take a look at the other bug you had with Qwen3-30B-A3B-Thinking (since it is doesnt have quantization logic tied to it) ?

[JdRion]

I currently don’t have access to GPUs, so it might take around 3–4 days before I can test it.
Also, could you share the GPT-OSS training setup you used?

[3outeille]

Oh wait, I am able to reproduce the error with TP + FSDP with quantization on. (I might have hallucinated earlier). Taking a look

• examples/torch_native_parallelism/nd_parallel_trainer.py && torchrun --nproc_per_node 4 nd_parallel_trainer.py --sequence-length=512

- MODEL_ID = "NousResearch/Hermes-3-Llama-3.1-8B"
+ MODEL_ID = "openai/gpt-oss-20b"

....
-    pc = ParallelismConfig()
+    pc = ParallelismConfig(dp_shard_size=2, tp_size=2) # FSDP=2 + TP=2 

...
-     if pc.tp_enabled:
-        model_kwargs["tp_plan"] = "auto"
-       model_kwargs["device_mesh"] = pc.build_device_mesh(args.device_type)
+    if pc.tp_enabled:
+       model_kwargs["tp_plan"] = "auto"
+        model_kwargs["device_mesh"] = pc.build_device_mesh(args.device_type)
+       model_kwargs["num_hidden_layers"] = 2

0]:            ^^^^^^^
[rank0]:   File "/fsx/ferdinandmom/ferdinand-hf/transformers_pr/env_v4.57.1/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1818, in inner
[rank0]:     args_result = hook(self, args)
[rank0]:                   ^^^^^^^^^^^^^^^^
[rank0]:   File "/fsx/ferdinandmom/ferdinand-hf/transformers_pr/work/v4.57.1-fix-gptoss-tp-fsdp/src/transformers/integrations/tensor_parallel.py", line 381, in <lambda>
[rank0]:     module.register_forward_pre_hook(lambda mod, inputs: input_fn(mod, inputs, device_mesh))
[rank0]:                                                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/fsx/ferdinandmom/ferdinand-hf/transformers_pr/work/v4.57.1-fix-gptoss-tp-fsdp/src/transformers/integrations/tensor_parallel.py", line 510, in _prepare_input_fn
[rank0]:     input_tensor = inputs[0]
[rank0]:                    ~~~~~~^^^
[rank0]: IndexError: tuple index out of range

[3outeille]

Just to let you know as well that we are soon releasing HF transformer backend compatible with torchtitan (huggingface/torchtitan#1)

[JdRion]

Thanks for letting me know I’m really looking forward to it!

[3outeille]

I quickly checked and there is 2 issues. First issue is in gpt-oss at the attention layer, all the function arguments are passed as kwargs but tensor parallel api expects args. This will trigger : IndexError: tuple index out of range but it is easy to fix by just passing hidden_states as args and not kwargs at self.atnlayer.

The second issue is at the attention_sinks that is a Dtensor despite us specifiying not to be a Dtensor (through local_rowise). I am not sure why this happens but will find the time to fix

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[akshan-main]

Hi @3outeille, I have implemented the fixes you suggested (passing hidden_states as a positional argument and handling the DTensor sinks in eager attention).

I am opening a PR now to resolve this.