[BUG] Non-deterministic model output if `replace_with_kernel_inject=True` for GPT-neo-1.3B/Bloom 1b7 and others

Describe the bug
For the models GPT-neo-1.3B, Bloom 1b7, Pythia 1.4b, GPT2-xl, I get non-deterministic model outputs when using context length 1 and engine = deepspeed.init_inference(model, dtype=torch.float16, replace_with_kernel_inject=True).

The context length 1 may let this sound like a low-priority bug, BUT it may be not: When using transformer's model.generate, the context length of the ids may be cut down to 1 (because of the use of past_key_values to speed up model inference). In particular, the above bug is a blocker to me rewriting transformer's model.generate for my needs.

To Reproduce

# !pip install --upgrade torch==1.13.1
# !pip install --upgrade transformers==4.28.1
# !pip install --upgrade deepspeed==0.9.1

SEED = 42

from typing import Any

import random
random.seed(SEED)

import numpy as np
np.random.seed(SEED)

import torch
torch.manual_seed(SEED)
torch.cuda.manual_seed_all(SEED)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True

import deepspeed
from transformers import AutoTokenizer, AutoModelForCausalLM

ARCHITECTURE = "EleutherAI/gpt-neo-1.3B"  # error -> non-deterministic
# ARCHITECTURE = "EleutherAI/pythia-1.4b"  # error -> non-deterministic
# ARCHITECTURE = "bigscience/bloom-1b7"  # error -> non-deterministic
# ARCHITECTURE = "gpt2-xl"  # error -> slightly non-deterministic

DEVICE = "cuda"

model = AutoModelForCausalLM.from_pretrained(ARCHITECTURE).to(DEVICE).eval()
tokenizer = AutoTokenizer.from_pretrained(ARCHITECTURE, use_fast=True)

def test_if_model_is_deterministic(mod: Any, tok: Any, device: str) -> None:
    mod = mod.eval()
    with torch.inference_mode():
        throw_away_ids = tok.encode(" 1 2 3 4 5 6 7", return_tensors="pt").to(device)
        throw_away_output = mod(throw_away_ids)
        
        for _ in range(10):
            ids = tok.encode(" 4", return_tensors="pt").to(device)
            output = mod(ids)
            token_id = torch.argmax(output.logits[0][-1]).item()
            logit = output.logits[0][-1][token_id].item()
            token = tok.decode(token_id)
            print(f"{token=}, {token_id=}, {logit=}")

test_if_model_is_deterministic(mod=model, tok=tokenizer, device=DEVICE) 
# prints:
# token='.', token_id=13, logit=-0.9108065366744995
# token='.', token_id=13, logit=-0.9108065366744995
# token='.', token_id=13, logit=-0.9108065366744995
# token='.', token_id=13, logit=-0.9108065366744995
# token='.', token_id=13, logit=-0.9108065366744995
# token='.', token_id=13, logit=-0.9108065366744995
# token='.', token_id=13, logit=-0.9108065366744995
# token='.', token_id=13, logit=-0.9108065366744995
# token='.', token_id=13, logit=-0.9108065366744995
# token='.', token_id=13, logit=-0.9108065366744995

engine = deepspeed.init_inference(model, dtype=torch.float16, replace_with_kernel_inject=True)

test_if_model_is_deterministic(engine.module, tokenizer, device=DEVICE) 
# prints:
# token=' 5', token_id=642, logit=-2.376953125
# token=' 4', token_id=604, logit=-1.349609375
# token=' 4', token_id=604, logit=-0.7822265625
# token=' 4', token_id=604, logit=-0.662109375
# token=' 4', token_id=604, logit=-0.59765625
# token=' 4', token_id=604, logit=0.01448822021484375
# token='\n', token_id=198, logit=-3.369140625
# token='t', token_id=83, logit=-5.5
# token='t', token_id=83, logit=-7.88671875
# token='l', token_id=75, logit=-11.1484375

Expected behavior
The 10 print outs after test_if_model_is_deterministic(engine.module, tokenizer, device=DEVICE) should not change. They should be identical to the print outs after test_if_model_is_deterministic(engine.module, tokenizer, device=DEVICE), but they are not.

ds_report output

(pytorch_p39) sh-4.2$ ds_report
--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
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-devel package with yum
 [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/ec2-user/anaconda3/envs/pytorch_p39/lib/python3.9/site-packages/torch']
torch version .................... 1.13.1
deepspeed install path ........... ['/home/ec2-user/anaconda3/envs/pytorch_p39/lib/python3.9/site-packages/deepspeed']
deepspeed info ................... 0.9.1, unknown, unknown
torch cuda version ............... 11.7
torch hip version ................ None
nvcc version ..................... 11.7
deepspeed wheel compiled w. ...... torch 1.13, cuda 11.7

System info (please complete the following information):

• OS: Amazon Linux 2
• GPU count and types: 1 GPU NVIDIA A10G
• Python version: 3.9.15
• Any other relevant info about your setup: Run inside a Jupyter notebook (notebook instance in AWS)