[BUG] Roberta Inference Engine RuntimeError: CUDA error: an illegal memory access was encountered

[shuyingsunshine21]

Describe the bug
A clear and concise description of what the bug is.

Encountered Illegal memory access when using inference engine for roberta model for long sequence (e.g. 512).

For shorter sequence (e.g. 20), it works.

Machine Specifications

• A6000 (1x)
• CUDA 11.3
• transformers==4.19.1
• deepspeed==0.7.2

To Reproduce

import torch, deepspeed
import json

from transformers import AutoConfig, AutoModel
import os

model_config = {
    "attention_probs_dropout_prob": 0.1,
    "bos_token_id": 0,
    "classifier_dropout": None,
    "eos_token_id": 2,
    "gradient_checkpointing": False,
    "has_spacev_lang_tokens": True,
    "hidden_act": "gelu",
    "hidden_dropout_prob": 0.1,
    "hidden_size": 384,
    "initializer_range": 0.02,
    "intermediate_size": 1536,
    "layer_norm_eps": 1e-05,
    "max_position_embeddings": 610,
    "model_type": "xlm-roberta",
    "num_attention_heads": 12,
    "num_hidden_layers": 12,
    "output_past": True,
    "pad_token_id": 1,
    "position_embedding_type": "absolute",
    "transformers_version": "4.11.0",
    "type_vocab_size": 1,
    "use_cache": True,
    "vocab_size": 250049
}

os.makedirs('/tmp/roberta_model_test/', exist_ok = True)
with open('/tmp/roberta_model_test/config.json', 'w+') as f:
    json.dump(model_config, f)

auto_config = AutoConfig.from_pretrained(
    pretrained_model_name_or_path="/tmp/roberta_model_test",
    params={
        "num_hidden_layers": 12,
        "hidden_size": 384,
    },
)
auto_model = AutoModel.from_config(config=auto_config)

sequence_length = 512

input_ids = [[i for i in range(sequence_length)]]
attention_mask = [[1 for _ in range(sequence_length)]]
    

input_ids = torch.Tensor(
    input_ids
).to(dtype=torch.long).to(torch.device("cuda:0"))

attention_mask = torch.Tensor(
    attention_mask
).to(dtype=torch.long).to(torch.device("cuda:0"))


auto_model.to(torch.device("cuda:0"))
auto_model.eval()

vallina_model_output = auto_model(input_ids=input_ids, attention_mask=attention_mask)

vallina_model_output_pooled = vallina_model_output["pooler_output"]
print(
    f"vallina_model_output_pooled is {vallina_model_output_pooled}."
)


deepspeed_infer_engine = deepspeed.init_inference(
    auto_model,
    mp_size=1,
    dtype=torch.float,
    replace_method="auto",
    replace_with_kernel_inject=True,
)

deepspeed_output = deepspeed_infer_engine.module(input_ids=input_ids, attention_mask=attention_mask)

deepspeed_output_pooled = deepspeed_output["pooler_output"]
print(f"deepspeed_output_pooled is: {deepspeed_output_pooled}.")

Expected behavior
A clear and concise description of what you expected to happen.

Expect to return inferenced output, as vallina model without deepspeed infer engine runs fine.

Current Behavior

Traceback (most recent call last):
  File "/home/core/miniconda3/envs/base/lib/python3.8/site-packages/transformers/models/roberta/modeling_roberta.py", line 847, in forward
    encoder_outputs = self.encoder(
  File "/home/core/miniconda3/envs/base/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/core/miniconda3/envs/base/lib/python3.8/site-packages/transformers/models/roberta/modeling_roberta.py", line 523, in forward
    layer_outputs = layer_module(
  File "/home/core/miniconda3/envs/base/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/core/miniconda3/envs/base/lib/python3.8/site-packages/deepspeed/ops/transformer/inference/transformer_inference.py", line 858, in forward
    output = self.mlp(attention_output, input, inp_norm, self.attention.attn_ob)
  File "/home/core/miniconda3/envs/base/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/core/miniconda3/envs/base/lib/python3.8/site-packages/deepspeed/ops/transformer/inference/transformer_inference.py", line 717, in forward
    return DeepSpeedMLPFunction.apply(input,
  File "/home/core/miniconda3/envs/base/lib/python3.8/site-packages/deepspeed/ops/transformer/inference/transformer_inference.py", line 638, in forward
    output = vector_matmul_func(intermediate, output_w, False)
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.

Screenshots
If applicable, add screenshots to help explain your problem.

[lokoppakmsft]

PR #2212 fixes the issue.