PetraAI/Nashmi

1from exllamav2 import ExLlamaV2, ExLlamaV2Config, ExLlamaV2Tokenizer2import argparse, os3import sys4import json5from conversion.tokenize import tokenize6from conversion.quantize import embeddings, measure_quant, quant7from conversion.optimize import optimize8from conversion.compile import compile_model9 10# import tracemalloc11# tracemalloc.start()12 13parser = argparse.ArgumentParser(description = "Convert model to ExLlamaV2")14parser.add_argument("-i", "--in_dir", type = str, help = "Input directory", default = "")15parser.add_argument("-o", "--out_dir", type = str, help = "Output directory")16parser.add_argument("-c", "--cal_dataset", type = str, help = "Calibration dataset (.parquet file)", default = "")17parser.add_argument("-r", "--dataset_rows", type = int, default = 100, help = "Number of rows to apply from dataset")18parser.add_argument("-mr", "--measurement_rows", type = int, default = 16, help = "Number of rows to apply from dataset when measuring")19parser.add_argument("-gr", "--gpu_rows", type = int, default = 16, help = "Threshold for paging hidden state to CPU")20parser.add_argument("-l", "--length", type = int, default = 2048, help = "Max no. tokens per sample")21parser.add_argument("-ml", "--measurement_length", type = int, default = 2048, help = "Max no. tokens per sample when measuring")22parser.add_argument("-b", "--bits", type = float, default = 4.156, help = "Target bits per weight")23parser.add_argument("-hb", "--head_bits", type = int, default = 6, help = "Target bits per weight (head layer)")24parser.add_argument("-m", "--measurement", type = str, help = "Reuse previous measurement")25 26args = parser.parse_args()27 28# Arguments29 30in_dir = None if args.in_dir == "" else os.path.abspath(args.in_dir)31out_dir = os.path.abspath(args.out_dir)32cal_dataset = None if args.cal_dataset == "" else os.path.abspath(args.cal_dataset)33dataset_rows = args.dataset_rows34measurement_rows = args.measurement_rows35gpu_rows = args.gpu_rows36length = args.length37measurement_length = args.measurement_length38bits = args.bits39head_bits = args.head_bits40reuse_measurement = args.measurement41 42if not os.path.exists(out_dir):43    print(f" ## Error: Directory not found: {out_dir}")44    sys.exit()45 46# Create model without loading weights47 48config = ExLlamaV2Config()49config.model_dir = in_dir50config.prepare()51 52model = ExLlamaV2(config)53model.load(lazy = True)54 55tokenizer = ExLlamaV2Tokenizer(config)56 57# Job file58 59job_file = os.path.join(out_dir, "job.json")60 61# Create new job62 63def save_job():64    global job_file, job65    with open(job_file, "w") as f:66        f.write(json.dumps(job, indent = 4))67 68if not os.path.exists(job_file):69 70    print(f" -- Beginning new job")71 72    if len(os.listdir(out_dir)) != 0:73        print(f" !! Warning: Output directory is not empty: {out_dir}")74 75    if in_dir is None:76        print(f" ## Error: No input directory specified")77        sys.exit()78 79    if cal_dataset is None:80        print(f" ## Error: No calibration dataset specified")81        sys.exit()82 83    job = { "in_dir": in_dir,84            "out_dir": out_dir,85            "cal_dataset": cal_dataset,86            "dataset_rows": dataset_rows,87            "measurement_rows": measurement_rows,88            "gpu_rows": gpu_rows,89            "length": length,90            "measurement_length": measurement_length,91            "bits": bits,92            "head_bits": head_bits,93            "progress": "begin",94            }95 96    if reuse_measurement is not None:97 98        with open(reuse_measurement, "r") as f:99 100            imp_measurement = json.load(f)101            job["measurement"] = imp_measurement["measurement"]102            job["last_module_idx"] = imp_measurement["last_module_idx"]103            job["base_perplexity"] = imp_measurement["base_perplexity"]104            job["reuse_measurement"] = reuse_measurement105 106    save_job()107 108# Resume existing job109 110else:111 112    print(f" -- Resuming job")113    print(f" !! Note: Overriding options with settings from existing job")114 115    with open(job_file, "r") as f:116        job = json.load(f)117 118    if "invalid" in job:119        print(" ** Error: Corrupted job")120        sys.exit()121 122    job["out_dir"] = out_dir123 124# Feedback125 126print(f" -- Input: {job['in_dir']}")127print(f" -- Output: {out_dir}")128print(f" -- Calibration dataset: {job['cal_dataset']}, {job['dataset_rows']} / {job['measurement_rows']} ({job['gpu_rows']}) rows, {job['length']} tokens per sample")129print(f" -- Target bits per weight: {job['bits']} (decoder), {job['head_bits']} (head)")130 131# Make sure subfolders exist132 133out_tensor_dir = os.path.join(job["out_dir"], "out_tensor")134if not os.path.exists(out_tensor_dir):135    os.makedirs(out_tensor_dir)136 137# Do the things138 139while True:140 141    progress = job["progress"]142 143    if progress == "begin":144 145        if "reuse_measurement" in job:146 147            print(f" -- Reusing measurement: {job['reuse_measurement']}")148            job["progress"] = "optimize"149            save_job()150 151        else:152 153            print(f" -- Tokenizing samples (measurement)...")154            tokenize(job, save_job, tokenizer, measure = True)155            job["progress"] = "initial_embeddings"156            save_job()157 158    if progress == "initial_embeddings":159 160        print(f" -- Token embeddings (measurement)...")161        embeddings(job, save_job, model)162        job["progress"] = "measure_quant"163        save_job()164 165    if progress == "measure_quant":166 167        print(f" -- Measuring quantization impact...")168        measure_quant(job, save_job, model)169        job["progress"] = "optimize"170        save_job()171 172    if progress == "optimize":173 174        print(f" -- Optimizing...")175        optimize(job, save_job)176        job["progress"] = "tokens_cal"177        save_job()178 179    if progress == "tokens_cal":180 181        print(f" -- Tokenizing samples...")182        tokenize(job, save_job, tokenizer)183        job["progress"] = "embeddings"184        save_job()185 186    if progress == "embeddings":187        print(f" -- Token embeddings again...")188        embeddings(job, save_job, model)189        job["progress"] = "quant"190        save_job()191 192    if progress == "quant":193 194        print(f" -- Quantizing...")195        quant(job, save_job, model)196        job["progress"] = "compile"197        save_job()198 199    if progress == "compile":200 201        print(f" -- Compiling output file...")202        compile_model(job, save_job, model)203        job["progress"] = "finished"204        save_job()205 206    if progress == "finished": break207 208print(f" -- Finished")