New release 4.53.0 breaks HF trainer/model
Author: liujch1998
Pull request: no
Labels: name: bug

### System Info

transformers==4.53.0
linux ubuntu
python==3.12.9
pytorch==2.6.0
accelerate==1.8.1

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

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

### Reproduction

With 4.53.0, Trainer throws an error during model forward pass. I'm training with FSDP (haven't ablated if this is related). When I pin to transformers==4.52.4 the issue is gone.
Model: allenai/OLMo-2-1124-7B

Code snippet:
```
    tokenizer = AutoTokenizer.from_pretrained(base_model_name, add_bos_token=False, add_eos_token=False)
    model = AutoModelForCausalLM.from_pretrained(base_model_name, torch_dtype=torch.float32) # let FSDP handle device placement

    training_args = TrainingArguments(
        output_dir=save_dir,
        num_train_epochs=1,
        per_device_train_batch_size=per_device_train_batch_size,
        gradient_accumulation_steps=gradient_accumulation_steps,
        save_strategy="no",
        # save_steps=1000,
        # save_total_limit=1,
        logging_steps=100,
        learning_rate=5e-5,
        lr_scheduler_type="linear",
        warmup_ratio=0.1,
        weight_decay=0.01,
        bf16=True,  # Enable bf16 mixed precision training
        optim="adamw_torch",  # Use PyTorch's AdamW implementation
        adam_beta1=0.9,
        adam_beta2=0.95,
        adam_epsilon=1e-8,
        max_grad_norm=1.0,  # Gradient clipping
        seed=42,  # Set random seed for reproducibility

        # FSDP configuration
        fsdp="full_shard auto_wrap",  # Enable full sharding with auto wrapping
        fsdp_config={
            "transformer_layer_cls_to_wrap": ["Olmo2DecoderLayer"],
            "backward_prefetch": "backward_pre",
        },
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=hf_dataset,
    )

    trainer.train()
```

Error log:
```
2025-06-29T19:18:32.802Z 
  0%|          | 0/24 [00:00<?, ?it/s][rank3]: Traceback (most recent call last):
2025-06-29T19:18:32.802Z [rank3]:   File "/gantry-runtime/cont-pretrain.py", line 217, in <module>
2025-06-29T19:18:32.802Z [rank3]:     main()
2025-06-29T19:18:32.802Z [rank3]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/click/core.py", line 1442, in __call__
2025-06-29T19:18:32.802Z [rank3]:     return self.main(*args, **kwargs)
2025-06-29T19:18:32.802Z [rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.802Z [rank3]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/click/core.py", line 1363, in main
2025-06-29T19:18:32.802Z [rank3]:     rv = self.invoke(ctx)
2025-06-29T19:18:32.802Z [rank3]:          ^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.802Z [rank3]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/click/core.py", line 1226, in invoke
2025-06-29T19:18:32.802Z [rank3]:     return ctx.invoke(self.callback, **ctx.params)
2025-06-29T19:18:32.802Z [rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.802Z [rank3]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/click/core.py", line 794, in invoke
2025-06-29T19:18:32.802Z [rank3]:     return callback(*args, **kwargs)
2025-06-29T19:18:32.802Z [rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.802Z [rank3]:   File "/gantry-runtime/cont-pretrain.py", line 159, in main
2025-06-29T19:18:32.802Z [rank3]:     trainer.train()
2025-06-29T19:18:32.802Z [rank3]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/transformers/trainer.py", line 2207, in train
2025-06-29T19:18:32.802Z [rank3]:     return inner_training_loop(
2025-06-29T19:18:32.802Z [rank3]:            ^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.802Z [rank3]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/transformers/trainer.py", line 2549, in _inner_training_loop
2025-06-29T19:18:32.802Z [rank3]:     tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
2025-06-29T19:18:32.802Z [rank3]:                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.802Z [rank3]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/transformers/trainer.py", line 3750, in training_step
2025-06-29T19:18:32.802Z [rank3]:     loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
2025-06-29T19:18:32.802Z [rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.802Z [rank3]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/transformers/trainer.py", line 3837, in compute_loss
2025-06-29T19:18:32.802Z [rank3]:     outputs = model(**inputs)
2025-06-29T19:18:32.802Z [rank3]:               ^^^^^^^^^^^^^^^
2025-06-29T19:18:32.802Z [rank3]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
2025-06-29T19:18:32.802Z [rank3]:     return self._call_impl(*args, **kwargs)
2025-06-29T19:18:32.802Z [rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.802Z [rank3]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
2025-06-29T19:18:32.802Z [rank3]:     return forward_call(*args, **kwargs)
2025-06-29T19:18:32.802Z [rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.802Z [rank3]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/accelerate/utils/operations.py", line 818, in forward
2025-06-29T19:18:32.802Z [rank3]:     return model_forward(*args, **kwargs)
2025-06-29T19:18:32.802Z [rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.802Z [rank3]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/accelerate/utils/operations.py", line 806, in __call__
2025-06-29T19:18:32.802Z [rank3]:     return convert_to_fp32(self.model_forward(*args, **kwargs))
2025-06-29T19:18:32.802Z [rank3]:                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.802Z [rank3]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/torch/amp/autocast_mode.py", line 44, in decorate_autocast
2025-06-29T19:18:32.802Z [rank3]:     return func(*args, **kwargs)
2025-06-29T19:18:32.802Z [rank3]:            ^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.802Z [rank3]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/transformers/utils/generic.py", line 943, in wrapper
2025-06-29T19:18:32.802Z [rank3]:     output = func(self, *args, **kwargs)
2025-06-29T19:18:32.802Z [rank3]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.802Z [rank3]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/transformers/models/olmo2/modeling_olmo2.py", line 537, in forward
2025-06-29T19:18:32.802Z [rank3]:     outputs: BaseModelOutputWithPast = self.model(
2025-06-29T19:18:32.802Z [rank3]:                                        ^^^^^^^^^^^
2025-06-29T19:18:32.802Z [rank3]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
2025-06-29T19:18:32.802Z [rank3]:     return self._call_impl(*args, **kwargs)
2025-06-29T19:18:32.802Z [rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.802Z [rank3]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
2025-06-29T19:18:32.802Z [rank3]:     return forward_call(*args, **kwargs)
2025-06-29T19:18:32.802Z [rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.802Z [rank3]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/transformers/utils/generic.py", line 943, in wrapper
2025-06-29T19:18:32.802Z [rank3]:     output = func(self, *args, **kwargs)
2025-06-29T19:18:32.802Z [rank3]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.802Z [rank3]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/transformers/models/olmo2/modeling_olmo2.py", line 390, in forward
2025-06-29T19:18:32.802Z [rank3]:     inputs_embeds = self.embed_tokens(input_ids)
2025-06-29T19:18:32.802Z [rank3]:                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.802Z [rank3]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
2025-06-29T19:18:32.802Z [rank3]:     return self._call_impl(*args, **kwargs)
2025-06-29T19:18:32.802Z [rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.802Z [rank3]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
2025-06-29T19:18:32.802Z [rank3]:     return forward_call(*args, **kwargs)
2025-06-29T19:18:32.802Z [rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.802Z [rank3]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/torch/nn/modules/sparse.py", line 190, in forward
2025-06-29T19:18:32.802Z [rank3]:     return F.embedding(
2025-06-29T19:18:32.802Z [rank3]:            ^^^^^^^^^^^^
2025-06-29T19:18:32.802Z [rank3]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/torch/nn/functional.py", line 2551, in embedding
2025-06-29T19:18:32.802Z [rank3]:     return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
2025-06-29T19:18:32.802Z [rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.802Z [rank3]: RuntimeError: 'weight' must be 2-D
2025-06-29T19:18:32.806Z [rank5]: Traceback (most recent call last):
2025-06-29T19:18:32.806Z [rank5]:   File "/gantry-runtime/cont-pretrain.py", line 217, in <module>
2025-06-29T19:18:32.806Z [rank5]:     main()
2025-06-29T19:18:32.806Z [rank5]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/click/core.py", line 1442, in __call__
2025-06-29T19:18:32.806Z [rank5]:     return self.main(*args, **kwargs)
2025-06-29T19:18:32.806Z [rank5]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.806Z [rank5]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/click/core.py", line 1363, in main
2025-06-29T19:18:32.806Z [rank5]:     rv = self.invoke(ctx)
2025-06-29T19:18:32.806Z [rank5]:          ^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.806Z [rank5]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/click/core.py", line 1226, in invoke
2025-06-29T19:18:32.806Z [rank5]:     return ctx.invoke(self.callback, **ctx.params)
2025-06-29T19:18:32.806Z [rank5]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.806Z [rank5]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/click/core.py", line 794, in invoke
2025-06-29T19:18:32.806Z [rank5]:     return callback(*args, **kwargs)
2025-06-29T19:18:32.806Z [rank5]:            ^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.806Z [rank5]:   File "/gantry-runtime/cont-pretrain.py", line 159, in main
2025-06-29T19:18:32.806Z [rank5]:     trainer.train()
2025-06-29T19:18:32.806Z [rank5]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/transformers/trainer.py", line 2207, in train
2025-06-29T19:18:32.806Z [rank5]:     return inner_training_loop(
2025-06-29T19:18:32.806Z [rank5]:            ^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.806Z [rank5]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/transformers/trainer.py", line 2549, in _inner_training_loop
2025-06-29T19:18:32.806Z [rank5]:     tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
2025-06-29T19:18:32.806Z [rank5]:                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.806Z [rank5]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/transformers/trainer.py", line 3750, in training_step
2025-06-29T19:18:32.806Z [rank5]:     loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
2025-06-29T19:18:32.806Z [rank5]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.806Z [rank5]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/transformers/trainer.py", line 3837, in compute_loss
2025-06-29T19:18:32.806Z [rank5]:     outputs = model(**inputs)
2025-06-29T19:18:32.806Z [rank5]:               ^^^^^^^^^^^^^^^
2025-06-29T19:18:32.806Z [rank5]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
2025-06-29T19:18:32.806Z [rank5]:     return self._call_impl(*args, **kwargs)
2025-06-29T19:18:32.806Z [rank5]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.806Z [rank5]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
2025-06-29T19:18:32.806Z [rank5]:     return forward_call(*args, **kwargs)
2025-06-29T19:18:32.806Z [rank5]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.806Z [rank5]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/accelerate/utils/operations.py", line 818, in forward
2025-06-29T19:18:32.806Z [rank5]:     return model_forward(*args, **kwargs)
2025-06-29T19:18:32.806Z [rank5]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.806Z [rank5]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/accelerate/utils/operations.py", line 806, in __call__
2025-06-29T19:18:32.806Z [rank5]:     return convert_to_fp32(self.model_forward(*args, **kwargs))
2025-06-29T19:18:32.806Z [rank5]:                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.806Z [rank5]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/torch/amp/autocast_mode.py", line 44, in decorate_autocast
2025-06-29T19:18:32.806Z [rank5]:     return func(*args, **kwargs)
2025-06-29T19:18:32.806Z [rank5]:            ^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.806Z [rank5]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/transformers/utils/generic.py", line 943, in wrapper
2025-06-29T19:18:32.806Z [rank5]:     output = func(self, *args, **kwargs)
2025-06-29T19:18:32.806Z [rank5]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.806Z [rank5]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/transformers/models/olmo2/modeling_olmo2.py", line 537, in forward
2025-06-29T19:18:32.806Z [rank5]:     outputs: BaseModelOutputWithPast = self.model(
2025-06-29T19:18:32.806Z [rank5]:                                        ^^^^^^^^^^^
2025-06-29T19:18:32.806Z [rank5]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
2025-06-29T19:18:32.806Z [rank5]:     return self._call_impl(*args, **kwargs)
2025-06-29T19:18:32.806Z [rank5]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.806Z [rank5]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
2025-06-29T19:18:32.806Z [rank5]:     return forward_call(*args, **kwargs)
2025-06-29T19:18:32.806Z [rank5]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.806Z [rank5]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/transformers/utils/generic.py", line 943, in wrapper
2025-06-29T19:18:32.806Z [rank5]:     output = func(self, *args, **kwargs)
2025-06-29T19:18:32.806Z [rank5]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.806Z [rank5]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/transformers/models/olmo2/modeling_olmo2.py", line 390, in forward
2025-06-29T19:18:32.806Z [rank5]:     inputs_embeds = self.embed_tokens(input_ids)
2025-06-29T19:18:32.806Z [rank5]:                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.806Z [rank5]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
2025-06-29T19:18:32.806Z [rank5]:     return self._call_impl(*args, **kwargs)
2025-06-29T19:18:32.806Z [rank5]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.806Z [rank5]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
2025-06-29T19:18:32.806Z [rank5]:     return forward_call(*args, **kwargs)
2025-06-29T19:18:32.806Z [rank5]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.806Z [rank5]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/torch/nn/modules/sparse.py", line 190, in forward
2025-06-29T19:18:32.806Z [rank5]:     return F.embedding(
2025-06-29T19:18:32.806Z [rank5]:            ^^^^^^^^^^^^
2025-06-29T19:18:32.806Z [rank5]:   File "/opt/conda/envs/gantry/lib/python3.12/site-packages/torch/nn/functional.py", line 2529, in embedding
2025-06-29T19:18:32.806Z [rank5]:     assert padding_idx < weight.size(
2025-06-29T19:18:32.806Z [rank5]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-06-29T19:18:32.806Z [rank5]: AssertionError: Padding_idx must be within num_embeddings
```

### Expected behavior

Trainer should run w/o error.

Comments:
Arf, that is annoying 😢 cc @matej-svejda and @SunMarc if you can have a look! 

Can you try with the patch that was release 3 days ago ? 

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](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.