From d7d6e5d77e3299091df2d05cbcd1b95322bfb78c Mon Sep 17 00:00:00 2001 From: Danny Semiat Date: Thu, 14 Mar 2024 12:45:26 +0200 Subject: [PATCH] Added hqt unify_measurement script for synapse1.15 --- .../quantization_tools/unify_measurements.py | 171 ++++++++++++++++++ 1 file changed, 171 insertions(+) create mode 100644 examples/text-generation/quantization_tools/unify_measurements.py diff --git a/examples/text-generation/quantization_tools/unify_measurements.py b/examples/text-generation/quantization_tools/unify_measurements.py new file mode 100644 index 0000000000..75ae329a44 --- /dev/null +++ b/examples/text-generation/quantization_tools/unify_measurements.py @@ -0,0 +1,171 @@ +import argparse +import json +import os +import sys + +import numpy as np + + +def find_measurement_path(measurement, measurements_dir_path, scales, group_size): + measurment_card = measurement + "_" + str(group_size) + for measurment_file in os.listdir(measurements_dir_path): + filename = os.fsdecode(measurment_file) + if not filename.endswith(".json") or "_mod_list" in filename or measurment_card not in filename: + continue + if scales: + if "MAXABS" in filename: + return os.path.join(measurements_dir_path, measurment_file) + else: + if "MAXABS" not in filename: + return os.path.join(measurements_dir_path, measurment_file) + + +def unify_measurements(measurement_group, measurements_dir_path, output_path, scales=False): + measurements_paths = [] + group_name = "" + + # save all the jsons paths in the given measurement group + for measurement in measurement_group: + measurement_path = find_measurement_path(measurement, measurements_dir_path, scales, len(measurement_group)) + measurements_paths.append(measurement_path) + group_name += measurement + + # save all the jsons content in the given measurement group + measurements_jsons = [] + for measurement_path in measurements_paths: + with open(measurement_path, "r") as f: + js = json.load(f) + measurements_jsons.append(js["Nodes"]) + + # create a name for the unified json that will be created for this measurement group + unified_json_name = ( + find_measurement_path(measurement_group[0], measurements_dir_path, scales, len(measurement_group)) + .split("/")[-1] + .replace("_" + measurement_group[0] + "_" + str(len(measurement_group)), "") + ) + unified_json_path = os.path.join(output_path, unified_json_name) + + # open a unified json file + with open(measurements_paths[0], "r") as origin, open(unified_json_path, "w") as copy: + copy.write(origin.read()) + with open(unified_json_path, "r") as json_file: + unified_json = json.load(json_file) + + # iterate all unified json nodes + for node_name, node_values in unified_json["Nodes"].items(): + max_inputs = node_values["inputs"] + max_outputs = None + if node_values.get("outputs") is not None: + max_outputs = node_values["outputs"] + max_weight = None + if node_values.get("params") is not None and node_values["params"].get("weight") is not None: + max_weight = node_values["params"]["weight"] + + # iterate over all the measurment group and take the maximum for each tensor and its channel + if scales: + for measurement_json in measurements_jsons: + max_inputs[0] = max(measurement_json[node_name]["inputs"][0], max_inputs[0]) + if max_outputs is not None: + max_outputs = max(measurement_json[node_name]["outputs"], max_outputs) + if max_weight is not None: + max_weight = max(measurement_json[node_name]["params"]["weight"], max_weight) + else: + for measurement_json in measurements_jsons: + for i in range(0, len(max_inputs)): + for j in range(0, len(max_inputs[i])): + max_inputs[i][j][0] = max(measurement_json[node_name]["inputs"][i][j][0], max_inputs[i][j][0]) + if max_outputs is not None: + for i in range(0, len(max_outputs)): + max_outputs[i][0] = max(measurement_json[node_name]["outputs"][i][0], max_outputs[i][0]) + if max_weight is not None: + for i in range(0, len(max_weight)): + max_weight[i][0] = max(measurement_json[node_name]["params"]["weight"][i][0], max_weight[i][0]) + + # update the maximum in the unified json + if scales: + unified_json["Nodes"][node_name]["inputs"][0] = max_inputs[0] + if max_outputs is not None: + unified_json["Nodes"][node_name]["outputs"] = max_outputs + if max_weight is not None: + unified_json["Nodes"][node_name]["params"]["weight"] = max_weight + else: + for i in range(0, len(max_inputs)): + for j in range(0, len(max_inputs[i])): + unified_json["Nodes"][node_name]["inputs"][i][j][0] = max_inputs[i][j][0] + if max_outputs is not None: + for i in range(0, len(max_outputs)): + unified_json["Nodes"][node_name]["outputs"][i][0] = max_outputs[i][0] + if max_weight is not None: + for i in range(0, len(max_weight)): + unified_json["Nodes"][node_name]["params"]["weight"][i][0] = max_weight[i][0] + global_rank = None + local_rank = None + mode = "" + layers = {} + with open(unified_json_path, "w") as json_file: + json.dump(unified_json, json_file) + mode = unified_json["Mode"] + nodes = unified_json["Nodes"] + + # create unified npz file from the unified json + unified_npz_path = os.path.join(output_path, unified_json_name.replace(".json", ".npz")) + for layer, dlayer in nodes.items(): + layers[layer] = {} + layers[layer]["inputs"] = [np.array(x) for x in dlayer["inputs"]] + if dlayer.get("outputs") is not None: + layers[layer]["outputs"] = np.array(dlayer["outputs"]) + if dlayer.get("params") is not None and dlayer["params"].get("weight") is not None: + layers[layer]["params"] = {} + layers[layer]["params"]["weight"] = np.array(dlayer["params"]["weight"]) + df = {"GlobalRank": global_rank, "LocalRank": local_rank, "Mode": mode, "Nodes": layers} + with open(unified_npz_path, "w"): + np.savez(unified_npz_path, df) + + +def parse_args(args): + parser = argparse.ArgumentParser( + description="Run the measurements parser", formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + parser.add_argument( + "-m", "--measurements", type=str, help="path to the directory of the measurements that will be unified" + ) + parser.add_argument( + "-g", + "--groups", + type=list, + nargs="+", + help="the groups of cards that are going to be unified- e.g. 01 23 45 67", + ) + parser.add_argument( + "-o", + "--out", + type=str, + default=os.getcwd(), + help="path to the directory where the unified measurements will be written", + ) + return parser.parse_args(args) + + +def main(args): + args = parse_args(args) + output_path = args.out + if not os.path.exists(output_path): + os.mkdir(output_path) + measurements_path = args.measurements + groups = args.groups + + num_jsons = 0 + for path in os.listdir(measurements_path): + if path.endswith(".json"): + num_jsons += 1 + assert os.path.isdir(measurements_path) and (num_jsons % len(groups)) == 0 + + for group in groups: + unify_measurements(group, measurements_path, output_path, scales=False) + unify_measurements(group, measurements_path, output_path, scales=True) + + print("finished measurement unifier script") + + +if __name__ == "__main__": + main(sys.argv[1:])