[BUG] ds-model inference results go far away from that of original model (megatron ) attention-context_layer error

Describe the bug
I create a 8 megatron transformer layer model (ori-model) and transform it into deepspeed model (ds-model) by deepspeed.init_inference.
Then I compare the results of ori-model with that of ds-model.
I found results are quiet different.

To Reproduce
Steps to reproduce the behavior:

1. create a 8 megatron transformer layer model (4 layers are with moe )
2. load parameter from a trained model ( with parameters that are not inited results are more different )
3. get deepspeed model with deepspeed.init_inference
4. run both two model with the same input
5. now you find results are different.

part of the code:

args_defaults={'tokenizer_type': 'GPT2BPETokenizer'}

initialize_megatron(extra_args_provider=None,
                    args_defaults=args_defaults)

args = get_args()
tp_config=DeepSpeedTPConfig()
tp_config.tp_size=1

model = GPTModel(
            num_tokentypes=0,
            parallel_output=True,
            pre_process=True,
            post_process=True,
            return_moe_loss=False,
        )
######load ckpt for ori model
base_dir="xxx"
checkpoint_name=base_dir+"xxx.pt"
state_dict = torch.load(checkpoint_name, map_location='cpu')
state_dict=state_dict['module']
model.load_state_dict(state_dict, strict=False)

device = torch.cuda.current_device()
model.to(device)
model.half()
model.eval()
######get sample from dataloader or random ########
seg_len=2048
batch_size=1
tokens=(50254*torch.rand(batch_size,seg_len)).to(device).type(torch.int64)
position_ids=torch.linspace(0,seg_len-1,steps=seg_len).to(device).type(torch.int64).view(1,seg_len).expand_as(tokens)
attention_mask = torch.tril(torch.ones(
(1, seg_len, seg_len), device=tokens.device)).view(
        1, 1, seg_len, seg_len)
attention_mask = (attention_mask < 0.5)
#######run ori - model
args.ds_inference=False
outputs_ori=model(tokens, position_ids, attention_mask)
print(outputs_ori)
outputs_ori_ind = torch.argmax( outputs_ori, dim=-1)
args.ds_inference=True( I set all " args.ds_inference "  related variables  update every iter )


######deepspeed init
ds_engine = deepspeed.init_inference(model,
                                 max_out_tokens=4096,
                                 tensor_parallel=tp_config,
                                 dtype=args.params_dtype,
                                 checkpoint=checkpoint_path,
                                 replace_method='auto',
                                 replace_with_kernel_inject=True,)

model = ds_engine.module

#######run
val_times=1
with torch.no_grad():
    for i in range(val_times):
        outputs=model(tokens, position_ids, attention_mask) 
######print
outputs_ind = torch.argmax( outputs, dim=-1)
print(outputs)
print(((outputs_ori_ind==outputs_ind).sum())*1.0/(outputs_ind.shape[-1]))

Expected behavior
results are the same or quiet similar

ds_report output
JIT compiled ops requires ninja
ninja .................. [OKAY]
op name ................ installed .. compatible
[WARNING] async_io requires the dev libaio .so object and headers but these were not found.
[WARNING] async_io: please install the libaio-dev package with apt
[WARNING] If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
cpu_adagrad ............ [NO] ....... [OKAY]
cpu_adam ............... [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
[WARNING] please install triton==1.0.0 if you want to use sparse attention
sparse_attn ............ [NO] ....... [NO]
spatial_inference ...... [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]
utils .................. [NO] ....... [OKAY]
DeepSpeed general environment info:
torch install path ............... ['/home/xxx/miniconda3/envs/lwxds083/lib/python3.9/site-packages/torch']
torch version .................... 1.12.1
deepspeed install path ........... ['/home/xxx/miniconda3/envs/lwxds083/lib/python3.9/site-packages/deepspeed']
deepspeed info ................... 0.8.0, unknown, unknown
torch cuda version ............... 11.3
torch hip version ................ None
nvcc version ..................... 11.3
deepspeed wheel compiled w. ...... torch 1.12, cuda 11.3

System info (please complete the following information):

• OS: Ubuntu 18.04
• one RTX 2080ti
• Python version 3.9

Start by reproducing the mismatch between GPTModel and the model returned by deepspeed.init_inference with the stated Megatron, checkpoint, and attention-mask setup. Compare the outputs before and after kernel injection, then narrow the issue to the inference conversion or configuration; done means identifying the cause and confirming matching or acceptably similar results.