[BUG] DeepSpeed non-deterministic inference with HF transformers when replace_with_kernel_inject=True

Describe the bug DeepSpeed non-deterministic inference with HF transformers when replace_with_kernel_inject=True.

To Reproduce

import os
import deepspeed
import torch
from transformers import pipeline, set_seed

local_rank = int(os.getenv("LOCAL_RANK", "0"))
world_size = int(os.getenv("WORLD_SIZE", "1"))
generator = pipeline("text-generation", model="gpt2", device=local_rank)
prompt = "In a shocking finding, scientist discovered a herd of unicorns living in a remote, previously unexplored valley, in the Andes Mountains. Even more surprising to the researchers was the fact that the unicorns spoke perfect English."

hf_greedy = generator(prompt, do_sample=False, max_length=128)

set_seed(42)
hf_sample = generator(prompt, do_sample=True, max_length=128)

generator.model = deepspeed.init_inference(
    generator.model,
    mp_size=world_size,
    dtype=torch.float,
    replace_method="auto",
    replace_with_kernel_inject=True, # Four assertions pass when it is False
)

set_seed(42)
ds_sample_1 = generator(prompt, do_sample=True, max_length=128)
set_seed(42)
ds_sample_2 = generator(prompt, do_sample=True, max_length=128)
assert hf_sample == ds_sample_1
assert ds_sample_1 == ds_sample_2

ds_greedy_1 = generator(prompt, do_sample=False, max_length=128)
ds_greedy_2 = generator(prompt, do_sample=False, max_length=128)
assert hf_greedy == ds_greedy_1
assert ds_greedy_1 == ds_greedy_2

Expected behavior The four assertions should all pass. I found that sometime some of the assertions would pass.

ds_report output

--------------------------------------------------
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
--------------------------------------------------
cpu_adam ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
 [WARNING]  please install triton==1.0.0 if you want to use sparse attention
sparse_attn ............ [NO] ....... [NO]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
 [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]
utils .................. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/path/miniconda3/envs/deepspeed/lib/python3.9/site-packages/torch']
torch version .................... 1.11.0
torch cuda version ............... 11.3
torch hip version ................ None
nvcc version ..................... 11.1
deepspeed install path ........... ['/path/miniconda3/envs/deepspeed/lib/python3.9/site-packages/deepspeed']
deepspeed info ................... 0.6.4, unknown, unknown
deepspeed wheel compiled w. ...... torch 1.11, cuda 11.3

System info (please complete the following information):

• OS: Ubuntu 20.04.4
• GPU count and types: one machine with 1 NVIDIA TITAN RTX
• Interconnects (if applicable): n/a
• Python version: Python 3.9.7
• Any other relevant info about your setup

Launcher context deepspeed launcher and jupyter notebook

Docker context n/a

[jeffra]

Done 😃 0.6.5 is out now on pypi

[jeffra]

@shijie-wu, can you give it a try now with this PR linked above? I think we have this fixed now. Tested on A100, A6000, V100, and P40.