Issues with running inference on quantized model

I have a question on MoQ usage and appreciate the help. I have run the example that train a bert-base model with MoQ, then trying to load it back for inference using the code snippet below.

import torch
import copy
from transformers.models.bert import BertLayer
from transformers import (
    AutoConfig,
    AutoModelForSequenceClassification,
    AutoTokenizer)
device = 'cuda'
model = AutoModelForSequenceClassification.from_pretrained('../DeepSpeedExamples/MoQ/output-8bits/qnli/checkpoint-9000/')
tokenizer = AutoTokenizer.from_pretrained('bert-base-cased')

inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
input_ids = inputs["input_ids"].to(device)
attention_mask = inputs["attention_mask"].to(device)
test_model = copy.deepcopy(model).to('cuda')
import deepspeed
import deepspeed.module_inject as module_inject
injection_policy={BertLayer:
                    module_inject.HFBertLayerPolicy}
ds_engine = deepspeed.init_inference(model,
                            mp_size=1,
                            dtype=torch.int8,
                            replace_method='auto',
                            quantization_setting=8,
                            injection_policy=injection_policy)
model = ds_engine.module
output = model(input_ids,attention_mask=attention_mask) # or output = model(input_ids,attention_mask)

This runs into an issue that complains about missing input_mask in the inputs, logs,

TypeError: compute_attention() missing 1 required positional argument: 'input_mask'

I was wondering if the setting for inference engine is missing something and what would be the best practice to run inference on a MOQ quantized model. Thanks.

--- Update, it seems like a bug that compute_attention(qkv_out) is missing input_mask, added and could pass the error.

However, I am not seeing any speed up (it actually degraded) compared to running a based-bert model with the inference engine. (quantized 6.62 ms vs bert-base 5.18 ms). I wonder if somehow should have injected quantized modules. I appreciate your help and seems an e2e tutorial for MoQ + inference would be very helpful for the community.

@RezaYazdaniAminabadi

System info (please complete the following information):

• OS: Ubuntu 18.04
• GPU count and types T4 (g4dn.2xlarge EC2)
• Interconnects (if applicable) [e.g., two machines connected with 100 Gbps IB]
• Python version 3.8
• Pytorch version 1.9
• Cuda version 11.1

Start with deepspeed/ops/transformer/inference/transformer_inference.py at the compute_attention call linked in the issue, then review the MoQ example and the referenced module_quantize.py. Reproduce the missing input_mask error and compare quantized and non-quantized inference benchmarks. Done means inference runs correctly and the expected quantized-module path or performance behavior is documented or fixed.