[BUG] RuntimeError: 'weight' must be 2-D during inferencing after loading the model saved by shard

Describe the bug
I fine-tuned bloomz-176b model with zero-stage-3 successfully. Previously, I call save_zero_three_model function which is in turn to call torch.save(output_state_dict, output_model_file) to save the model parameter in a single huge file (355G). However, it's very very slow to load such a huge file into main memory by from_pretrained method and it causes OOM by ds script. So I modify the function save_zero_three_model by two ways:

1. replace torch.save with model_to_save.save_pretrained to save the model by shard
2. replace torch.save with model_ema.save_checkpoint
   Either way is OK during saving. However, after loading by from_pretrained, it cause the following runtime error:

** Starting to generate 100 tokens with bs=1
Generate args {'max_new_tokens': 100, 'do_sample': False}
*** Running generate
Traceback (most recent call last):
File "/home/ec2-user/xxx/bloom-accelerate-inference.py", line 183, in
generated = generate()
File "/home/ec2-user/xxx/bloom-accelerate-inference.py", line 168, in generate
print(tokenizer.decode(model.generate(inputs_prompt["input_ids"].to(torch.cuda.current_device()), attention_mask=inputs_prompt['attention_mask'].to(torch.cuda.current_device()), max_length=100)[0]))
File "/opt/conda/envs/optimusgptx/lib/python3.9/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
return func(*args, **kwargs)
File "/opt/conda/envs/optimusgptx/lib/python3.9/site-packages/transformers/generation/utils.py", line 1515, in generate
return self.greedy_search(
File "/opt/conda/envs/optimusgptx/lib/python3.9/site-packages/transformers/generation/utils.py", line 2332, in greedy_search
outputs = self(
File "/opt/conda/envs/optimusgptx/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
File "/opt/conda/envs/optimusgptx/lib/python3.9/site-packages/accelerate/hooks.py", line 165, in new_forward
output = old_forward(*args, **kwargs)
File "/opt/conda/envs/optimusgptx/lib/python3.9/site-packages/transformers/models/bloom/modeling_bloom.py", line 913, in forward
transformer_outputs = self.transformer(
File "/opt/conda/envs/optimusgptx/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
File "/opt/conda/envs/optimusgptx/lib/python3.9/site-packages/transformers/models/bloom/modeling_bloom.py", line 730, in forward
inputs_embeds = self.word_embeddings(input_ids)
File "/opt/conda/envs/optimusgptx/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
File "/opt/conda/envs/optimusgptx/lib/python3.9/site-packages/accelerate/hooks.py", line 165, in new_forward
output = old_forward(*args, **kwargs)
File "/opt/conda/envs/optimusgptx/lib/python3.9/site-packages/torch/nn/modules/sparse.py", line 162, in forward
return F.embedding(
File "/opt/conda/envs/optimusgptx/lib/python3.9/site-packages/torch/nn/functional.py", line 2210, in embedding
return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
RuntimeError: 'weight' must be 2-D

To Reproduce
Steps to reproduce the behavior:

1. Simple inference script to reproduce
   the inference commend is : CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python bloom-accelerate-inference.py

the script is here:

import argparse
import gc
import math
import os
import time

import torch
import torch.distributed as dist

from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

def get_args():
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", required=False, type=int, help="used by dist launchers")
parser.add_argument("--name", type=str, help="Name path", required=False)
parser.add_argument("--batch_size", default=1, type=int, help="batch size")
parser.add_argument("--benchmark", action="store_true", help="additionally run benchmark")
parser.add_argument("--greedy", action="store_true")
parser.add_argument("--top-k", type=int, default=0)
parser.add_argument("--top-p", type=float, default=0.0)
parser.add_argument("--dtype", type=str, help="float16 or int8", choices=["int8", "float16"], default="float16")

return parser.parse_args()

t_start = time.time()

num_tokens = 100

args = get_args()

local_rank = int(os.getenv("LOCAL_RANK", "0"))
world_size = torch.cuda.device_count()

rank = local_rank

def print_rank0(*msg):
if rank != 0:
return
print(*msg)

print_rank0(f"Using {world_size} gpus")
model_name = args.name
print_rank0(f"Loading model {model_name}")

tokenizer = AutoTokenizer.from_pretrained(model_name)

tokenizer = AutoTokenizer.from_pretrained("/home/ec2-user/xiaojhui/Models/JP_Vender_bloomz176b_lora_gc_bs8_ml512_lr1e6_e1_JPVender5k")

XXX: can't automatically derive dtype via config's from_pretrained

dtype = torch.bfloat16 if model_name in ["bigscience/bloom", "bigscience/bigscience-small-testing"] else torch.float16

print(get_max_memory_per_gpu_dict())

infer_dtype = args.dtype
if infer_dtype == "int8":
dtype = torch.int8

kwargs = dict(
device_map="auto",
)

def get_world_size() -> int:
if dist.is_initialized():
return dist.get_world_size()
else:
return 1

balanced_low_0 - because it allows a larger batch size with multiple GPUs

if get_world_size() > 1:
kwargs["device_map"] = "balanced_low_0"

if infer_dtype == "int8":
print_rank0("Using load_in_8bit=True to use quanitized model")
kwargs["load_in_8bit"] = True
else:
kwargs["torch_dtype"] = dtype

import accelerate

from accelerate.state import AcceleratorState

from transformers.utils import ContextManagers

def get_deepspeed_plugin():

if accelerate.state.is_initialized():

return AcceleratorState().deepspeed_plugin

else:

return None

def deepspeed_zero_init_disabled_context_manager():

"""

returns either a context list that includes one that will disable zero.Init or an empty context list

"""

deepspeed_plugin = get_deepspeed_plugin()

if deepspeed_plugin is None:

return []

return [deepspeed_plugin.zero3_init_context_manager(enable=False)]

with ContextManagers(deepspeed_zero_init_disabled_context_manager()):

model = AutoModelForCausalLM.from_pretrained("/home/ec2-user/xiaojhui/Models/deepspeedchat_try/4debug5/", **kwargs)

model = AutoModelForCausalLM.from_pretrained(model_name, **kwargs)

model = AutoModelForCausalLM.from_pretrained("/home/ec2-user/xiaojhui/Models/JP_Vender_bloomz176b_lora_gc_bs8_ml512_lr1e6_e1_JPVender5k", **kwargs)

model = AutoModelForCausalLM.from_pretrained("/home/ec2-user/xiaojhui/Models_external/bloomz-1b7/", **kwargs)

model = AutoModelForCausalLM.from_pretrained("/home/ec2-user/xiaojhui/Models_external/bloomz-7b1/", **kwargs)

model = AutoModelForCausalLM.from_pretrained("/home/ec2-user/xiaojhui/Models_external/bloomz/", **kwargs)

model = AutoModelForCausalLM.from_pretrained("/home/ec2-user/xiaojhui/Models/JP_Vender_bloomz7b1_loraOnly02_bs16_ml512_lr3e5_e10_JPVender5k", **kwargs)

orignal_model = AutoModelForCausalLM.from_pretrained("/home/ec2-user/xiaojhui/Models/deepspeedchat_try/4debug5/", **kwargs)

model = AutoModelForCausalLM.from_pretrained("/home/ec2-user/xiaojhui/Models/deepspeedchat_try/4debug3/", **kwargs)

model = AutoModelForCausalLM.from_pretrained("/home/ec2-user/xiaojhui/Models/deepspeedchat_try/4debug/", **kwargs)

import deepspeed

model = deepspeed.init_inference(

orignal_model,

mp_size=world_size,

base_dir="/home/ec2-user/xiaojhui/Models/deepspeedchat_try/4debug5/",

dtype=getattr(torch, infer_dtype),

**kwargs,

)

addddddd

model.eval()

if args.benchmark:
t_ready = time.time()

Generate

print_rank0(f"*** The loaded model is {model} ")

print_rank0(f"*** Starting to generate {num_tokens} tokens with bs={args.batch_size}")

input_sentences = [
"DeepSpeed is a machine learning framework",
"He is working on",
"He has a",
"He got all",
"Everyone is happy and I can",
"The new movie that got Oscar this year",
"In the far far distance from our galaxy,",
"Peace is the only way",
]

if args.batch_size > len(input_sentences):
# dynamically extend to support larger bs by repetition
input_sentences *= math.ceil(args.batch_size / len(input_sentences))

generate_kwargs = dict(max_new_tokens=num_tokens, do_sample=False)

generate_kwargs = dict(max_new_tokens=num_tokens, use_cache=False, do_sample=False)

generate_kwargs = dict(min_length=num_tokens, max_length=num_tokens, do_sample=False)

print_rank0(f"Generate args {generate_kwargs}")
inputs = input_sentences[: args.batch_size]

def generate():
"""returns a list of zipped inputs, outputs and number of new tokens"""

input_tokens = tokenizer.batch_encode_plus(inputs, return_tensors="pt", padding=True)
for t in input_tokens:
    if torch.is_tensor(input_tokens[t]):
        input_tokens[t] = input_tokens[t].to("cuda:0")

# input_tokens = tokenizer.batch_encode_plus(inputs, return_tensors="pt", padding=True)
# for t in input_tokens:
#     if torch.is_tensor(input_tokens[t]):
#         input_tokens[t] = input_tokens[t].to(torch.cuda.current_device())

prompt = "It was a dark and stormy night"
result_length = 50
inputs_prompt = tokenizer(prompt, return_tensors="pt")
print(tokenizer.decode(model.generate(inputs_prompt["input_ids"].to(torch.cuda.current_device()), attention_mask=inputs_prompt['attention_mask'].to(torch.cuda.current_device()),  max_length=100)[0]))

outputs = model.generate(**input_tokens, **generate_kwargs)

input_tokens_lengths = [x.shape[0] for x in input_tokens.input_ids]
output_tokens_lengths = [x.shape[0] for x in outputs]

total_new_tokens = [o - i for i, o in zip(input_tokens_lengths, output_tokens_lengths)]
outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True)

return zip(inputs, outputs, total_new_tokens)

print_rank0("*** Running generate")
t_generate_start = time.time()
generated = generate()
t_generate_span = time.time() - t_generate_start
for i, o, _ in generated:
print_rank0(f"{'-'*60}\nin={i}\nout={o}\n")

Benchmark

if args.benchmark:
# clear cache / free memory
torch.cuda.empty_cache()
gc.collect()

print_rank0("*** Running benchmark")
# warm up
for i in range(1):
    _ = generate()
torch.cuda.synchronize()

# benchmark
t0 = time.time()
cycles = 5
total_new_tokens_generated = 0
for i in range(cycles):
    generated = generate()
    total_new_tokens_generated += sum(new_tokens for _, _, new_tokens in generated)
torch.cuda.synchronize()
throughput = (time.time() - t0) / (total_new_tokens_generated)
print_rank0(
    f"""

*** Performance stats:
Throughput per token including tokenize: {throughput*1000:.2f} msecs
Start to ready to generate: {t_ready - t_start:.3f} secs
Tokenize and generate {total_new_tokens_generated} (bs={args.batch_size}) tokens: {t_generate_span:.3f} secs
Start to finish: {t_ready - t_start + t_generate_span:.3f} secs
"""
)

the saving function of save_zero_three_model is here:

def save_zero_three_model(model_ema, global_rank, save_dir, zero_stage=0):
zero_stage_3 = (zero_stage == 3)
os.makedirs(save_dir, exist_ok=True)
WEIGHTS_NAME = "pytorch_model.bin"
output_model_file = os.path.join(save_dir, WEIGHTS_NAME)

model_to_save = model_ema.module if hasattr(model_ema,
                                            'module') else model_ema
if not zero_stage_3:
    if global_rank == 0:
        torch.save(model_to_save.state_dict(), output_model_file)
else:
    output_state_dict = {}
    for k, v in model_to_save.named_parameters():

        if hasattr(v, 'ds_id'):
            with deepspeed.zero.GatheredParameters(_z3_params_to_fetch([v
                                                                        ]),
                                                   enabled=zero_stage_3):
                v_p = v.data.cpu()
        else:
            v_p = v.cpu()
        if global_rank == 0 and "lora" not in k:
            output_state_dict[k] = v_p
    # if global_rank == 0:
        #torch.save(output_state_dict, output_model_file)
        # model_to_save.save_pretrained (
        #     save_directory=save_dir,
        #     state_dict=output_state_dict,
        #     max_shard_size='5GB'
        # )
    model_ema.save_checkpoint(
        save_dir=save_dir,
        client_state=output_state_dict)
    del output_state_dict

I ever try to load the original bloomz and inference, everything seems OK.

4. What packages are required and their versions
   deepspeed, torch

5. How to run the script
   CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python bloom-accelerate-inference.py

6. ...

Expected behavior
Inference without error.

ds_report output
Please run ds_report to give us details about your setup.

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] sparse_attn requires a torch version >= 1.5 and < 2.0 but detected 2.0
[WARNING] using untested triton version (2.0.0), only 1.0.0 is known to be compatible
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 ............... ['/opt/conda/envs/optimusgptx/lib/python3.9/site-packages/torch']
torch version .................... 2.0.0+cu117
deepspeed install path ........... ['/opt/conda/envs/optimusgptx/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.8
deepspeed wheel compiled w. ...... torch 0.0, cuda 0.0

Screenshots
If applicable, add screenshots to help explain your problem.

System info (please complete the following information):

• OS: [e.g. Ubuntu 18.04]
• GPU count and types [one machines with x8 A100s 80G memory each]
• (if applicable) what DeepSpeed-MII version are you using
• (if applicable) Hugging Face Transformers (4.29.2) /Accelerate (0.19.0 )
• Python version: 3.9.16
• Any other relevant info about your setup

Docker context
Are you using a specific docker image that you can share?

Additional context
Add any other context about the problem here.

Start with the provided bloom-accelerate-inference.py reproduction and the save_zero_three_model function, then compare the sharded checkpoint loading path with the single-file path. Verify the loaded word-embedding weight shape before generation; done means the sharded model loads and generates without the 2-D weight RuntimeError.