Repository: pytorch
Issue: CUBLAS_STATUS_EXECUTION_FAILED when calling cublasGemmEx
Labels: module: cuda,triaged,module: cublas

### 🐛 Describe the bug

I met a problem similar to #94294 when using torch.multiprocessing
```python
RuntimeError: CUDA error: CUBLAS_STATUS_EXECUTION_FAILED when calling `cublasGemmEx( handle, opa, opb, m, n, k, &falpha, a, CUDA_R_16BF, lda, b, CUDA_R_16BF, ldb, &fbeta, c, CUDA_R_16BF, ldc, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)
```
The following code should be quite easy to reproduce. All you need to do is  make sure CUDA VISIBLE DEVICES >1 .
```python

import torch
from torch import bfloat16
import torch.distributed as dist
import torch.multiprocessing as mp

from torch.utils.data import Dataset,DataLoader
import functools
from transformers import AutoTokenizer,DefaultDataCollator,GenerationConfig,PreTrainedModel,AutoModelForSeq2SeqLM,AutoModelForCausalLM,AutoConfig,DataCollatorWithPadding
import logging
from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
from tqdm import tqdm
# from accelerate import find_executable_batch_size

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)



class DefaultDataset(Dataset):
    def __init__(self,data,tokenizer):
        self.data=tokenizer(data,return_tensors='pt',padding=True)

    
    def __getitem__(self,idx):
        return {'input_ids':self.data['input_ids'][idx]}
    
    def __len__(self):
        return self.data['input_ids'].size(0)




class NiuInference:
    def __init__(self,model_dir,data,dtype=bfloat16,dataset=None,data_collator=None,output_path='niuinference.out',auto_batch_size=True,batch_size=1,generation_config=None):
        self.model_dir=model_dir
        self.dtype=dtype
        self.data=data
        self.dataset=dataset
        self.data_collator=data_collator
        self.output_path=output_path
        self.batch_size=batch_size
        self.auto_batch_size=auto_batch_size
        self.generation_config=generation_config
        
        
    def _load_model_and_tokenizer(self,device):
        print(self.dtype)
        config=AutoConfig.from_pretrained(self.model_dir)
        if config.model_type in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES:
            model=AutoModelForCausalLM.from_pretrained(self.model_dir,torch_dtype=self.dtype)
        else:
            model=AutoModelForSeq2SeqLM.from_pretrained(self.model_dir,torch_dtype=self.dtype)
        model.to(device)
        tokenizer=AutoTokenizer.from_pretrained(self.model_dir)
        return model,tokenizer
        

    def get_pred(self,rank,out_path,data,dict):
        batch_size=2
        
        try:
            device = torch.device(f'cuda:{rank}')
            model, tokenizer = self._load_model_and_tokenizer(device)
            if self.dataset is not None:
                dataset=self.dataset(data=data,tokenizer=tokenizer)
            else:
                dataset=DefaultDataset(data=data,tokenizer=tokenizer)

            if self.data_collator is not None:
                collator=self.data_collator(tokenizer,model=model,padding=True)
            else:
                collator= DataCollatorWithPadding(tokenizer)
            dataloader=DataLoader(dataset,batch_size,collate_fn=collator,pin_memory=True,num_workers=0)
            result=[]
            for input in tqdm(dataloader):
                input.to(device)
                print(input)
                output = model.generate(
                            input_ids=input['input_ids'],
                            attention_mask=input['attention_mask'],
                            num_beams=5,
                            do_sample=False,
                            temperature=1.0,
                            max_new_tokens=512,
                        )
                pred = tokenizer.batch_decode(output,skip_special_tokens=True)
                print(pred)
                result+=pred
            dict[f'{rank}']=result
        except Exception as e:
            print('error',device)
            raise
        
        
          
    
    def split_list(self,lst, n):
        avg = len(lst) / float(n)
        return [lst[int(avg * i):int(avg * (i + 1))] for i in range(n)]

    def run(self,):
    
        world_size = min(torch.cuda.device_count(),len(self.data)) # corner case， data<available GPU num
        
        data_subsets = self.split_list(self.data,world_size)
        print(data_subsets)
        processes = []
        manager = mp.Manager()
        record_dict = manager.dict()
        for rank in range(world_size):

            p = mp.Process(target=self.get_pred, args=(rank,self.output_path,data_subsets[rank],record_dict))
            p.start()
            processes.append(p)
        for p in processes:
            p.join()

        with open(self.output_path, "w", encoding="utf-8") as f:
            for rank in range(world_size):
                for r in record_dict[f'{rank}']:
                    f.write(r.replace('\n','\\n')+'\n')

  
if __name__=='__main__':
    mp.set_start_method('spawn')
    i=NiuInference(model_dir='Google/t5-v1_1-large',data=['hello,how is your day','my wish is that you happy','from scratch',])
    i.run()


```

### Versions

this is my environment
```shell
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2023 NVIDIA Corporation
Built on Tue_Feb__7_19:32:13_PST_2023
Cuda compilation tools, release 12.1, V12.1.66
Build cuda_12.1.r12.1/compiler.32415258_0

torch Version: 2.2.1
Python version: 3.10.13
transformers version: 4.39.0.dev0
```

cc @ptrblck @csarofeen @xwang233