Silent Overwrite of Custom Optimizer When Using DeepSpeed with Transformers Trainer
Author: VladPyzh
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.49.0
- Platform: Linux-5.4.0-216-generic-x86_64-with-glibc2.31
- Python version: 3.13.2
- Huggingface_hub version: 0.29.2
- Safetensors version: 0.5.3
- Accelerate version: 1.7.0
- Accelerate config:    not found
- DeepSpeed version: 0.16.8
- PyTorch version (GPU?): 2.6.0+cu124 (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?: No
- Using GPU in script?: Yes
- GPU type: Tesla V100-PCIE-32GB

### Who can help?

@zach-huggingface @SunMarc

### 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

To reproduce the problem:

1. Pass some DeepSpeed config file to the TrainingArguments
2. Pass custom optimizer to the Trainer

Example:

```
class CustomOptim(torch.optim.Adam):
    def step(self, *args, **kwargs):
        print("steped!")
        super().step(*args, **kwargs)

model_name = "facebook/opt-125m"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

dataset = load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1")
train_ds = dataset["train"]
eval_ds = dataset["validation"]

def tokenize(examples):
    return tokenizer(
        examples["text"], return_special_tokens_mask=True, truncation=True
    )

train_ds = train_ds.map(tokenize, batched=True, remove_columns=["text"])
eval_ds = eval_ds.map(tokenize, batched=True, remove_columns=["text"])

data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

training_args = TrainingArguments(
    output_dir="./opt-125m-finetuned",
    do_eval=False,
    overwrite_output_dir=True,
    num_train_epochs=1,
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    max_steps=1,
    push_to_hub=False,
    deepspeed="/data/vladislav/ds_config.json",
)

optimizer = CustomOptim(model.parameters())
scheduler = get_constant_schedule(optimizer)

trainer = Trainer(
    model=model,
    args=training_args,
    optimizers=(optimizer, scheduler),
    train_dataset=train_ds,
    eval_dataset=eval_ds,
    data_collator=data_collator,
)

trainer.train()
```
DeepSpeed config:
```
{
  "train_batch_size": 32,
  "gradient_accumulation_steps": 1,
  "gradient_clipping": 1.0,

  "zero_optimization": {
    "stage": 0,
    "contiguous_gradients": true,
    "overlap_comm": true,
    "reduce_bucket_size": 50000000,
    "allgather_bucket_size": 500000000
  }
  
}
```

In this example CustomOptimizer will be overwritten and no print would appear. 

### Expected behavior

Hi HuggingFace team,

While working with the `transformers.Trainer` to train a LLM, I encountered an unexpected behavior when combining a custom optimizer with a DeepSpeed config in TrainingArguments.

It seems that when a custom optimizer is passed to the Trainer, but a DeepSpeed config is also provided, the custom optimizer is silently ignored and overwritten by DeepSpeed. This was quite surprising, and from what I can tell, this behavior is not mentioned in the documentation for either Trainer or TrainingArguments.

To avoid confusion and help users catch this early, it would be helpful if an explicit warning or exception were raised in such scenarios.

Thanks for the great work on the library!

Comments:
Hey @VladPyzh, nice to see you here ! 

> It seems that when a custom optimizer is passed to the Trainer, but a DeepSpeed config is also provided, the custom optimizer is silently ignored and overwritten by DeepSpeed. This was quite surprising, and from what I can tell, this behavior is not mentioned in the documentation for either Trainer or TrainingArguments.

This shouldn't happen, if you check `deepspeed_init` and `deepspeed_optim_sched`, we create the DS optimizer if you added it in the ds_config, otherwise, your custom optimizer should be used. 

Good to see you too @SunMarc and thanks for super rapid response.

As I see `deepspeed_init` will update `trainer.optimizer` and `trainer.scheduler`. Inside of the `deepspeed_init` there is a ` trainer.optimizer = None` and new initialization of optimizer using `trainer.create_optimizer()` inside of `deepspeed_optim_sched` in case when there is no optimizer in ds_config. It seems like this function doesn't initialize custom optimizers and rely on `TrainingArguments` to get the type of optimizer. 

Maybe I miss something, does my example work for your correctly? 



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.