Saving model in safetensors format through Trainer fails for Gemma 2 due to shared tensors
Author: oranshayer
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.44.2
- Platform: Linux-5.10.220-209.869.amzn2.x86_64-x86_64-with-glibc2.26
- Python version: 3.10.14
- Huggingface_hub version: 0.25.1
- Safetensors version: 0.4.5
- Accelerate version: 0.34.2
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.4.1+cu121 (True)
- 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?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA A10G

### Who can help?

@muellerz @SunMarc 

### Information

- [X] 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

I am finetuning `google/gemma-2-2b` and these are the arguments and trainer call:

```
text_model = AutoModelForCausalLM.from_pretrained("google/gemma-2-2b", token=token, attn_implementation='eager')

training_args = TrainingArguments(
    output_dir=args.log_dir,
    num_train_epochs=args.epochs,
    per_device_train_batch_size=args.train_batch_size,
    per_device_eval_batch_size=args.eval_batch_size,
    warmup_steps=args.warmup_steps,
    learning_rate=args.learning_rate,
    evaluation_strategy="no",
    logging_dir=args.log_dir,
    logging_steps=50,
    save_strategy="steps",
    save_steps=2000,
    report_to="mlflow",
    run_name=args.run_name,
)

trainer = Trainer(
    model=text_model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    compute_metrics=compute_metrics,
)
```

I am getting the following error when trainer tries to save the model:

```
RuntimeError: 
            Some tensors share memory, this will lead to duplicate memory on disk and potential differences when loading them again: [{'text_model.model.embed_tokens.weight', 'text_model.lm_head.weight'}].
            A potential way to correctly save your model is to use `save_model`.
```

I have currently disabled saving as safetensors through the training arguments:
`save_safetensors=False,`

### Expected behavior

Should save in safetensors without raising an error.

Comments:
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.

cc @muellerzr @SunMarc !

Hey @oranshayer, could you share a minimal reproducer ? This shouldn't happens as we make sure to remove the shared tensors prior to saving. Having the entire traceback will also help us figuring out here the problem is. Thanks ! 

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.

Any solution for this issue, I'm facing the same issue while saving a fine tuned Gemma3 4B model with safe_serialization=True, if I make serialization = False, then the inference is returning weird result: crashing crashing crashing.

I am facing exactly the same issues while saving a fine tuned Gemma3 4B. Interestingly it works with Gemma 1B and 270M. Not sure what is the root cause