Asynchronous CUDA Execution Issue with Hugging Face Transformers
Author: taehyunzzz
Pull request: no

### System Info

I am reaching out to you regarding an issue I've been experiencing with the Hugging Face Transformers library in a PyTorch environment. I'm encountering unexpected CUDA synchronizations while executing my code, which seems to be impairing the performance of my model. I am hopeful that you can provide some guidance on this matter.

As a bit of background, I am utilizing the Deepspeed integration for running inference of pre-trained Switch-Transformers with 8 experts model from the Hugging Face library for a project. This involves processing large amounts of data and hence necessitates efficient GPU utilization for timely results. To maximize the GPU's potential, I have been aiming to leverage CUDA's asynchronous execution feature.

In CUDA, as you know, the CPU queues kernels for execution on the GPU. Ideally, while the CPU is busy queueing up kernels, the GPU should be asynchronously running the kernels that have already been queued. This is the behavior I have previously observed when using the Megatron-LM library.

However, when using the Hugging Face Transformers library, I am finding that there seems to be a CUDA synchronize call after every kernel, which effectively serializes the CPU and GPU operations. This has led to a significant decrease in processing speed and efficiency, as the GPU is left idle while the CPU prepares the next kernel.

I am unsure whether this is due to an issue with my code or if it's an inherent characteristic of the Hugging Face Transformers library. I was wondering if you might have any insights into this issue or any suggestions for further troubleshooting steps I could take. Does the implementation of Switch Transformer Model have implicit synchronization that I might not be aware of?

### Who can help?

@ArthurZucker  @younesbelkada  @stas00 

### 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. Wrap the pretrained switch model with deepspeed initialize
2. Take the first Switch Layer
3. Initialize a random tensor to run forward on
4. Profile using Pytorch profiler
5. Behavior of profile seems as if CPU waits for the CUDA kernel to finish
![image](https://github.com/huggingface/transformers/assets/33646149/a26aad02-a0e3-4e67-91f3-c08390b6fd84)

Below is the code I ran :
```

from transformers import AutoTokenizer, SwitchTransformersForConditionalGeneration
from transformers.deepspeed import HfDeepSpeedConfig

# ds_config_file = "ds_zero_stage_0_config.json"
# ds_config_file = "ds_zero_stage_infinity-cpu.json"
ds_config_file = "ds_zero_stage_2_config.json"
with open(ds_config_file) as fin:
    ds_config = json.load(fin)

tokenizer = AutoTokenizer.from_pretrained("google/switch-base-8")
model = SwitchTransformersForConditionalGeneration.from_pretrained("google/switch-base-8", device_map="auto")

dschf = HfDeepSpeedConfig(ds_config)
model, optimizer, _, _ = deepspeed.initialize(
    config           = ds_config,
    model            = model,
    model_parameters = [{
                            "params": [p for n,p in list(model.named_parameters())],
                            "name" : "base",
                            "weight_decay" : 0.01
                        }]
)
model.eval()

# Get a switch layer
switch_layer = model.encoder.block[1].layer[1].mlp

batch_size = 1000
seq_len = 1000
d_model = model.module.config.d_model
d_ff = model.module.config.d_ff
din = torch.rand(size=(batch_size, seq_len, d_model),
                    dtype=torch.float32,
                    device="cuda:0"
                )

#hidden_states, (router_logits, expert_index)
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
        record_shapes=True, profile_memory=True) as prof:
    for i in range(10):
        dout = switch_layer(din)
prof.export_chrome_trace("trace_check.json")
```

### Expected behavior

I would like to know if there are implicit synchronization in the switch model implementation and why the profiled result show that CPU synchronizes with the CUDA runtime?

Comments:
Hey! Thanks for opening an issue. For other people who might wonder, this is because of the inner workings of accelerate, which are activated by the call to `device_map="auto"`, meaning that the model will be dispatched between both CPU and GPU.