Trainer do not move the model to GPU when doing evaluation with FSDP
Author: exhyy
Pull request: no
Labels: name: trainer

### System Info

- `transformers` version: 4.39.3
- Platform: Linux-5.4.0-164-generic-x86_64-with-glibc2.31
- Python version: 3.10.14
- Huggingface_hub version: 0.19.4
- Safetensors version: 0.4.2
- Accelerate version: 0.29.1
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.1+cu118 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: yes

### Who can help?

@muellerzr @pacman100 

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

1. evaluation scripts
my_evaluate.py:
```python
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
from datasets import load_dataset
from trl import SFTTrainer

def main():
    # args
    model_name = 'TinyLlama/TinyLlama-1.1B-Chat-v1.0'
    dataset_name = 'HuggingFaceH4/ultrachat_200k'
    output_dir = './tiny_llama_eval'
    
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        trust_remote_code=True,
    )
    
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    eval_dataset = load_dataset(dataset_name, split='test_sft')
    eval_dataset = eval_dataset.map(lambda x: {"formatted_chat": tokenizer.apply_chat_template(x["messages"], tokenize=False, add_generation_prompt=False)})
    
    training_args = TrainingArguments(
        output_dir=output_dir,
        per_device_eval_batch_size=8,
    )
    trainer = SFTTrainer(
        model=model,
        tokenizer=tokenizer,
        args=training_args,
        eval_dataset=eval_dataset,
        eval_packing=False,
        dataset_text_field='formatted_chat',
        max_seq_length=2048,
    )
    results = trainer.evaluate()
    print(results)

if __name__ == '__main__':
    main()
```
fsdp_config.yaml:
```yaml
compute_environment: LOCAL_MACHINE                                                                                                                                           
debug: true                                                                                                                                                                 
distributed_type: FSDP
downcast_bf16: 'no'
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch: BACKWARD_PRE
  fsdp_cpu_ram_efficient_loading: true
  fsdp_forward_prefetch: false
  fsdp_offload_params: false
  fsdp_sharding_strategy: FULL_SHARD
  fsdp_state_dict_type: SHARDED_STATE_DICT
  fsdp_sync_module_states: true
  fsdp_use_orig_params: false
machine_rank: 0
main_training_function: main
mixed_precision: 'bf16'
num_machines: 1
num_processes: 8
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false
```
2. run `accelerate launch --config_file fsdp_config.yaml my_evaluate.py`
3. get error: `RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:6!` in [this line](https://github.com/pytorch/pytorch/blob/9bb54c7f3c048fdf2e5294e9e49e3642f1de98d8/torch/nn/functional.py#L2266) because `weight` is on cpu while `input` is on gpu.

### Expected behavior

Expected behavior: the model should be on correct device when doing evaluation with FSDP.

What makes the problem:
`Trainer` uses its `accelerator` to prepare the model (see [this line](https://github.com/huggingface/transformers/blob/b109257f4fb8b1166e7c53cc5418632014ed53a5/src/transformers/trainer.py#L3589)). However `accelerate.Accelerator` will not wrap the model with FSDP when `evaluation_mode == True`. As a result, the model will stay on cpu. See [this](https://github.com/huggingface/accelerate/blob/210778370e269362afaff1806f5c9801a688a5fc/src/accelerate/accelerator.py#L1388) for details.

I can fix the problem by simply setting `evaluation_mode=False` in https://github.com/huggingface/transformers/blob/b109257f4fb8b1166e7c53cc5418632014ed53a5/src/transformers/trainer.py#L3585-L3589
But I am not sure if this would cause other problems during the above evaluation.

Comments:
Gentle ping @muellerzr @SunMarc 

Another ping @muellerzr @SunMarc 

It has been reported that evaluation could be slow and OOM when training does not in recent issues

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.