Error when using GPT2 `model.forward` method with DeepSpeed inference

Hi,

I am trying to use DeepSpeed with GPT-2/Neo for inference, but encountering an error when using model.forward method directly. It seems that the error occurs if I don’t provide an attention_mask.

import deepspeed
from transformers import GPT2Tokenizer, GPT2LMHeadModel
import torch

tokenizer = GPT2Tokenizer.from_pretrained("gpt2-medium")
model = GPT2LMHeadModel.from_pretrained("gpt2-medium")
model = model.cuda()

model = deepspeed.init_inference(model,
                                 mp_size=1,
                                 dtype=torch.float,
                                 replace_method='auto')

ids = tokenizer.encode("A valley full of unicorns was discovered",
                       add_special_tokens=False, return_tensors="pt").cuda()
# attn = torch.ones((1, 8), dtype=torch.long).cuda()
output = model(
    input_ids=ids,
    # attention_mask=attn, # <-- need to provide this
    return_dict=True,
    use_cache=True
)
print(output)

Traceback (most recent call last):
  File "deepspeed_attn.py", line 25, in <module>
    use_cache=True
  File "torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "deepspeed/inference/engine.py", line 222, in forward
    return self.module(*inputs, **kwargs)
  File "torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "transformers/models/gpt2/modeling_gpt2.py", line 954, in forward
    return_dict=return_dict,
  File "torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "transformers/models/gpt2/modeling_gpt2.py", line 797, in forward
    output_attentions=output_attentions,
  File "torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "deepspeed/ops/transformer/inference/transformer_inference.py", line 611, in forward
    self.norm_b)
  File "torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "deepspeed/ops/transformer/inference/transformer_inference.py", line 393, in forward
    self.qkv_merging)
  File "deepspeed/ops/transformer/inference/transformer_inference.py", line 141, in forward
    while len(input_mask.shape) < 4:
AttributeError: 'NoneType' object has no attribute 'shape'

When I use model.generate, as in the example scripts, I am not getting any errors as generate takes care of creating attention_mask (but in general I think model.forward should work even without providing attention_mask):

https://github.com/huggingface/transformers/blob/1ed2ebf60d87ef12bd063c7c58e484e19189c754/src/transformers/generation_utils.py#L395

[Express50]

Thanks! Just tested it on my end as well and it works both with and without attention_mask.

[RezaYazdaniAminabadi]

I see. I will fix it soon. Thanks for trying it out 👍