[BUG] : Inconsistent results on model save and load using stage 3

Describe the bug
The model gives different results after saving and loading the model again.
The primary objective is to train the model and save it and load it back for inference. I tried using huggingface save_pretrained, even that is not giving me exact generation.

To Reproduce
file1.py

from transformers import AutoTokenizer, AutoModelForCausalLM
import deepspeed
import os
import torch

world_size = int(os.getenv("WORLD_SIZE", "1"))

train_batch_size = 8 * 1
ds_config = ds_config = {
    "fp16": {
        "enabled": False
    },
    "bf16": {
        "enabled": False
    },
    "gradient_accumulation_steps": 1,
    "zero_optimization": {
        "stage": 3,
        "offload_param": {
            "device": "cpu",
            "pin_memory": True
        },

    },
    "steps_per_print": 2000,
    "train_batch_size": train_batch_size,
    "train_micro_batch_size_per_gpu": 1,
    "wall_clock_breakdown": False
}


model_name = "google/gemma-2b"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
text_in = "Is this review positive or negative? Review: this is the best cast iron skillet you will ever buy"
inputs = tokenizer.encode(text_in, return_tensors="pt")


ds_engine = deepspeed.initialize(model=model, config_params=ds_config)[0]
ds_engine.module.eval()
torch.save(ds_engine.module.state_dict(), 'pytorch.bin')

with torch.no_grad():
    outputs = ds_engine.module.generate(inputs.to('cuda'), max_new_tokens=32)
text_out = tokenizer.decode(outputs[0], skip_special_tokens=True)

Now in a file2.py

from transformers import AutoTokenizer, AutoModelForCausalLM
import deepspeed
import os
import torch

world_size = int(os.getenv("WORLD_SIZE", "1"))

train_batch_size = 8 * 1
ds_config = ds_config = {
    "fp16": {
        "enabled": False
    },
    "bf16": {
        "enabled": False
    },
    "gradient_accumulation_steps": 1,
    "zero_optimization": {
        "stage": 3,
        "offload_param": {
            "device": "cpu",
            "pin_memory": True
        },

    },
    "steps_per_print": 2000,
    "train_batch_size": train_batch_size,
    "train_micro_batch_size_per_gpu": 1,
    "wall_clock_breakdown": False
}


model_name = "google/gemma-2b"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
text_in = "Is this review positive or negative? Review: this is the best cast iron skillet you will ever buy"
inputs = tokenizer.encode(text_in, return_tensors="pt")


ds_engine = deepspeed.initialize(model=model, config_params=ds_config)[0]
ds_engine.module.eval()
ds_engine.module.load_state_dict(torch.load('pytorch.bin', map_location=torch.device('cpu')))

with torch.no_grad():
    outputs = ds_engine.module.generate(inputs.to('cuda'), max_new_tokens=32)
text_out = tokenizer.decode(outputs[0], skip_special_tokens=True)

Expected behavior
The generated text should be the same.

System info (please complete the following information):

• OS: [e.g. Ubuntu 18.04]: 22.04.3
• GPU count and types [e.g. two machines with x8 A100s each] : 8 x A100
• Python version : 3.9

Start by running the file1.py and file2.py reproducer with the listed DeepSpeed ZeRO stage 3 configuration, then compare the generated text before saving and after loading pytorch.bin. Check the save_pretrained path mentioned in the report as well; done means inference produces the same generated text after reload.