Bloom 176B with deepspeed-inference: Cuda illegal memory access
### System Info

- `transformers` version: 4.21.1
- Platform: Linux-3.10.0-1160.49.1.el7.x86_64-x86_64-with-glibc2.17
- Python version: 3.8.13
- Huggingface_hub version: 0.9.0
- PyTorch version (GPU?): 1.10.2+cu113 (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, deepspeed


### Who can help?

@stas00 @patrickvonplaten  @patil-suraj 

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

I am running `bigscience/bloom` on a distributed cluster of 4 nodes of 4x40GB A100 each using deepspeed-inference. I am using a modified version of this script here https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/main/scripts/inference/bloom-ds-inference.py

At batch size 1, everything is fine for 18 batches (more or less reproducible, once it cleared 19), then I get the error below.
Not sure whether to submit this here, to deepspeed, or to https://github.com/bigscience-workshop/Megatron-DeepSpeed

Here has the best visibility, but I might submit to all three and refer back to this. I'm also running on Determined AI and kubernetes, which could be relevant.

Here some key parts of the code:

```python
   with deepspeed.OnDevice(dtype=torch.float16, device="meta"):
    model = AutoModelForCausalLM.from_pretrained(
        model_args.model_name_or_path,
        torch_dtype=torch.bfloat16
    )

model = deepspeed.init_inference(
    model,
    mp_size=training_args.world_size,
    dtype=torch.float16,
    checkpoint=checkpoint_filename,
    replace_with_kernel_inject=True
)

data_loader = DataLoader(
    dataset,
    sampler=SequentialSampler(dataset),
    batch_size=training_args.per_device_eval_batch_size,
    collate_fn=collate,
    drop_last=False,
    num_workers=0,
    pin_memory=False,
)

for step, batch in enumerate(data_loader):
    outputs = model.generate(
        batch["input_ids"].to(torch.cuda.current_device()),
        max_new_tokens=150,
        do_sample=True
    )
```

And here is the error, which occurs after 18-19 examples:
```
!!!! kernel execution error. (m: 3584, n: 609, k: 14336, error: 13)
!!!! kernel execution error. (m: 14336, n: 609, k: 3584, error: 13)
Traceback (most recent call last):
  File "./ds_inference_core.py", line 186, in <module>
    main(core_context, trial) 
  File "./ds_inference_core.py", line 132, in main 
    outputs = model.generate( 
  File "/opt/conda/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 28, in decorate_context 
    return func(*args, **kwargs) 
  File "/opt/conda/lib/python3.8/site-packages/transformers/generation_utils.py", line 1326, in generate 
    return self.sample( 
  File "/opt/conda/lib/python3.8/site-packages/transformers/generation_utils.py", line 1944, in sample 
    outputs = self( 
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl 
    return forward_call(*input, **kwargs) 
  File "/opt/conda/lib/python3.8/site-packages/deepspeed/inference/engine.py", line 521, in forward 
    outputs = self.model_orig_fwd(*inputs, **kwargs) 
  File "/opt/conda/lib/python3.8/site-packages/transformers/models/bloom/modeling_bloom.py", line 821, in forward 
    transformer_outputs = self.transformer( 
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl 
    return forward_call(*input, **kwargs) 
  File "/opt/conda/lib/python3.8/site-packages/transformers/models/bloom/modeling_bloom.py", line 709, in forward 
    outputs = block( 
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl 
    return forward_call(*input, **kwargs) 
  File "/opt/conda/lib/python3.8/site-packages/deepspeed/ops/transformer/inference/transformer_inference.py", line 842, in forward 
    output = self.mlp(attention_output, input, inp_norm, self.attention.attn_ob) 
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl 
    return forward_call(*input, **kwargs) 
  File "/opt/conda/lib/python3.8/site-packages/deepspeed/ops/transformer/inference/transformer_inference.py", line 712, in forward 
    return DeepSpeedMLPFunction.apply(input, 
  File "/opt/conda/lib/python3.8/site-packages/deepspeed/ops/transformer/inference/transformer_inference.py", line 646, in forward 
    dist.all_reduce(output, group=mp_group) 
  File "/opt/conda/lib/python3.8/site-packages/deepspeed/comm/comm.py", line 126, in log_wrapper 
    return func(*args, **kwargs) 
  File "/opt/conda/lib/python3.8/site-packages/deepspeed/comm/comm.py", line 513, in all_reduce 
    return cdb.all_reduce(tensor, op, group, async_op) 
  File "/opt/conda/lib/python3.8/site-packages/deepspeed/comm/torch.py", line 45, in all_reduce 
    return torch.distributed.all_reduce(tensor=tensor, 
  File "/opt/conda/lib/python3.8/site-packages/torch/distributed/distributed_c10d.py", line 1287, in all_reduce 
    work = group.allreduce([tensor], opts) 
RuntimeError: CUDA error: an illegal memory access was encountered 
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect. 
For debugging consider passing CUDA_LAUNCH_BLOCKING=1. 
Fatal Python error: Aborted 
Thread 0x00007fbdd7fff700 (most recent call first):
  File "/opt/conda/lib/python3.8/socket.py", line 669 in readinto
  File "/opt/conda/lib/python3.8/http/client.py", line 277 in _read_status
  File "/opt/conda/lib/python3.8/http/client.py", line 316 in begin
  File "/opt/conda/lib/python3.8/http/client.py", line 1348 in getresponse
  File "/opt/conda/lib/python3.8/site-packages/urllib3/connectionpool.py", line 444 in _make_request
  File "/opt/conda/lib/python3.8/site-packages/urllib3/connectionpool.py", line 703 in urlopen
  File "/opt/conda/lib/python3.8/site-packages/requests/adapters.py", line 489 in send
  File "/opt/conda/lib/python3.8/site-packages/requests/sessions.py", line 701 in send
  File "/opt/conda/lib/python3.8/site-packages/requests/sessions.py", line 587 in request
  File "/run/determined/pythonuserbase/lib/python3.8/site-packages/determined/common/requests.py", line 41 in request
  File "/run/determined/pythonuserbase/lib/python3.8/site-packages/determined/common/api/request.py", line 129 in do_request
  File "/run/determined/pythonuserbase/lib/python3.8/site-packages/determined/common/experimental/session.py", line 65 in _do_request
  File "/run/determined/pythonuserbase/lib/python3.8/site-packages/determined/common/experimental/session.py", line 86 in get
  File "/run/determined/pythonuserbase/lib/python3.8/site-packages/determined/core/_preempt.py", line 57 in _get_preemption
  File "/run/determined/pythonuserbase/lib/python3.8/site-packages/determined/core/_preempt.py", line 89 in run
  File "/opt/conda/lib/python3.8/threading.py", line 932 in _bootstrap_inner
  File "/opt/conda/lib/python3.8/threading.py", line 890 in _bootstrap
[2022-08-24 10:39:18] [b57bcf88] [rank=0]
Thread 0x00007fbde3fff700 (most recent call first):
  File "/opt/conda/lib/python3.8/threading.py", line 306 in wait
  File "/opt/conda/lib/python3.8/queue.py", line 179 in get
  File "/opt/conda/lib/python3.8/site-packages/tensorboard/summary/writer/event_file_writer.py", line 227 in run
  File "/opt/conda/lib/python3.8/threading.py", line 932 in _bootstrap_inner
  File "/opt/conda/lib/python3.8/threading.py", line 890 in _bootstrap
[2022-08-24 10:39:18] [b57bcf88] [rank=0]
Current thread 0x00007fc2a776e6c0 (most recent call first):
<no Python frame> 
```

### Expected behavior

I expect the inference to complete successfully :)
Maybe related: bigscience-workshop/Megatron-Deepspeed#324<|||||>You probably want to open Deepspeed-Inference-related issues at https://github.com/microsoft/DeepSpeed since these aren't related to `transformers` other than that `transformers` are used from inside Deepspeed-Inference.

That project contains custom cuda kernels so it's very difficult to debug and when you file the issue you want to tag @RezaYazdaniAminabadi who wrote them. 

Please help him out by providing a minimal reproducible example and he will surely fix the problem.

But I'm closing this Issue since there is nothing we can do here.