146 lines
4.8 KiB
Python
146 lines
4.8 KiB
Python
import torch
|
|
import os
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def tensor_shape_to_string(tensor):
|
|
'''
|
|
convert a tensor shape to string description
|
|
'''
|
|
int_list = list(tensor.shape)
|
|
str_list = [str(num) for num in int_list]
|
|
str_shape = "x".join(str_list)
|
|
return str_shape
|
|
|
|
|
|
def save_prompt_token_ids(prompt_input_ids, args):
|
|
'''
|
|
save prompt_token_id
|
|
Args:
|
|
prompt_input_ids: prompt input_id assiged to llm.generate
|
|
args: arguments from main
|
|
'''
|
|
if args.dump_prompt_token_ids is not True:
|
|
return
|
|
output_dir = os.path.join(args.output_dir, "prompt_input_ids")
|
|
if not os.path.exists(output_dir):
|
|
os.makedirs(output_dir)
|
|
data_len = len(prompt_input_ids)
|
|
for data_index in range(data_len):
|
|
tensor = prompt_input_ids[data_index]
|
|
str_shape = tensor_shape_to_string(tensor)
|
|
file_path = os.path.join(output_dir, f"prompt_input_ids_{data_index}_{str_shape}.pt")
|
|
torch.save(tensor, file_path)
|
|
logger.info(f"Saved input_ids[{data_index}] to {file_path}")
|
|
|
|
|
|
def save_input_ids(input_ids, args):
|
|
'''
|
|
save input_ids
|
|
Args:
|
|
input_ids: input of qkv with layer0
|
|
args: arguments from main
|
|
'''
|
|
id_len = len(input_ids)
|
|
if args.dump_input_ids is not True or id_len == 0:
|
|
return
|
|
output_dir = os.path.join(args.output_dir, "input_ids")
|
|
if not os.path.exists(output_dir):
|
|
os.makedirs(output_dir)
|
|
for data_index in range(id_len):
|
|
tensor = input_ids[data_index]
|
|
str_shape = tensor_shape_to_string(tensor)
|
|
file_path = os.path.join(output_dir, f"input_ids_{data_index}_{str_shape}.pt")
|
|
torch.save(tensor, file_path)
|
|
logger.info(f"Saved input_ids[{data_index}] to {file_path}")
|
|
|
|
|
|
def save_act_range(act_range, args):
|
|
'''
|
|
save act_range
|
|
Args:
|
|
act_range: save act_range collected when model running
|
|
args: arguments from main
|
|
'''
|
|
if args.dump_act_range is not True:
|
|
return
|
|
output_dir = os.path.join(args.output_dir, "act_range")
|
|
if not os.path.exists(output_dir):
|
|
os.makedirs(output_dir)
|
|
|
|
for layer_name, layer_scale in act_range.items():
|
|
for tensor_key, tensor_value in layer_scale.items():
|
|
if isinstance(tensor_value, torch.Tensor):
|
|
str_shape = tensor_shape_to_string(tensor_value)
|
|
file_name = f'{layer_name}_{tensor_key}_{str_shape}.pt'
|
|
file_path = os.path.join(output_dir, file_name)
|
|
torch.save(tensor_value, file_path)
|
|
logger.info(f"Saved act_range[{layer_name}][{tensor_key}] to {file_path}")
|
|
|
|
|
|
def save_weights(weights, args):
|
|
'''
|
|
save hugging face weights
|
|
Args:
|
|
weights: hugging face weights merged with llm model named parameters
|
|
args: arguments from main
|
|
'''
|
|
if args.dump_weights is not True:
|
|
return
|
|
output_dir = os.path.join(args.output_dir, "weights")
|
|
if not os.path.exists(output_dir):
|
|
os.makedirs(output_dir)
|
|
|
|
for tensor_key, tensor_value in weights.items():
|
|
str_shape = tensor_shape_to_string(tensor_value)
|
|
file_name = f'{tensor_key}_{str_shape}.pt'
|
|
file_path = os.path.join(output_dir, file_name)
|
|
torch.save(tensor_value, file_path)
|
|
logger.info(f"Saved weights[{tensor_key}] to {file_path}")
|
|
|
|
|
|
def save_generate_weights(weights, args):
|
|
'''
|
|
save quantizated weights
|
|
Args:
|
|
weights: quantized weights of smoothquant or weightonly
|
|
args: arguments from main
|
|
'''
|
|
if args.dump_generate_weights is not True:
|
|
return
|
|
output_dir = os.path.join(args.output_dir, "generate_weights")
|
|
if not os.path.exists(output_dir):
|
|
os.makedirs(output_dir)
|
|
|
|
for tensor_key, tensor_value in weights.items():
|
|
str_shape = tensor_shape_to_string(tensor_value)
|
|
file_name = f'{tensor_key}_{str_shape}.pt'
|
|
file_path = os.path.join(output_dir, file_name)
|
|
torch.save(tensor_value, file_path)
|
|
logger.info(f"Saved generate weights[{tensor_key}] to {file_path}")
|
|
|
|
|
|
def dump_save_x_y(name, x, y, index):
|
|
'''
|
|
dump x, y when inferrence
|
|
output_dir need to modify by your self
|
|
'''
|
|
output_dir = "output_dir"
|
|
x_output_dir = os.path.join(output_dir, "x_tensor")
|
|
y_output_dir = os.path.join(output_dir, "y_tensor")
|
|
if not os.path.exists(x_output_dir):
|
|
os.makedirs(x_output_dir)
|
|
if not os.path.exists(y_output_dir):
|
|
os.makedirs(y_output_dir)
|
|
|
|
x_file_name = os.path.join(x_output_dir, f"{name}_x_{index}.pt")
|
|
y_file_name = os.path.join(y_output_dir, f"{name}_y_{index}.pt")
|
|
if isinstance(x, tuple):
|
|
x = x[0]
|
|
if not os.path.exists(x_file_name):
|
|
torch.save(x.cpu(), x_file_name)
|
|
if not os.path.exists(y_file_name):
|
|
torch.save(y.cpu(), y_file_name)
|