Create convert_v2.py

convert_v2.py

@@ -0,0 +1,90 @@
+# code adapted from https://huggingface.co/fahadh4ilyas
+mport argparse, json
+from safetensors import safe_open
+from safetensors.torch import save_file
+from pathlib import Path
+
+parser = argparse.ArgumentParser(description="Convert original dbrx model into quantizable model")
+
+parser.add_argument("--model-dir", type=str, required=True, help="directory to the original dbrx model")
+parser.add_argument("--output-dir", type=str, required=True, help="directory for the converted dbrx model")
+args = parser.parse_args()
+
+model_dir = Path(args.model_dir)
+output_dir = Path(args.output_dir)
+output_dir.mkdir(parents=True, exist_ok=True)
+
+NUM_EXPERTS = 16
+HIDDEN_SIZE = 6144
+HEAD_DIM = 128
+NUM_KV_HEAD = 8
+FFN_HIDDEN_SIZE = 10752
+
+def change_tensor_attn(tensor):
+
+    return [x.contiguous() for x in tensor.split([HIDDEN_SIZE, NUM_KV_HEAD*HEAD_DIM, NUM_KV_HEAD*HEAD_DIM])]
+
+def change_attn(tensors):
+
+    keys = list(tensors.keys())
+    for k in keys:
+        if 'Wqkv' in k:
+            prefix = k.removesuffix('.Wqkv.weight')
+            tensor = tensors.pop(k)
+            output_tensor = change_tensor_attn(tensor)
+            for dtype,t in zip(['q_proj', 'k_proj', 'v_proj'], output_tensor):
+                tensors[f'{prefix}.{dtype}.weight'] = t
+    
+    return tensors
+
+def change_tensor_mlp(tensor, reverse=False):
+
+    output = [x.contiguous() if not reverse else x.t().contiguous() for x in tensor.reshape(NUM_EXPERTS, FFN_HIDDEN_SIZE, HIDDEN_SIZE)]
+
+    return output
+
+def change_mlp(tensors):
+
+    keys = list(tensors.keys())
+    for k in keys:
+        if any([x in k for x in ['w1', 'v1', 'w2']]):
+            prefix,dtype = k.rsplit('.', 1)
+            tensor = tensors.pop(k)
+            output_tensor = change_tensor_mlp(tensor, dtype=='w2')
+            for i,t in enumerate(output_tensor):
+                tensors[f'{prefix}.{i}.{dtype}.weight'] = t
+
+    return tensors
+
+for file in sorted(list(model_dir.glob('*.safetensors'))):
+    print(file)
+    tensors = {}
+    with safe_open(file, 'pt') as f:
+        metadata = f.metadata()
+        for k in f.keys():
+            tensors[k] = f.get_tensor(k)
+    tensors = change_attn(tensors)
+    tensors = change_mlp(tensors)
+    save_file(tensors, (output_dir / file.name).as_posix(), metadata)
+
+with open(model_dir / 'model.safetensors.index.json') as f:
+    weight_map = json.load(f)
+
+weight_keys = list(weight_map['weight_map'])
+for k in weight_keys:
+    if any([x in k for x in ['w1', 'v1', 'w2']]):
+        prefix,dtype = k.rsplit('.', 1)
+        value = weight_map['weight_map'].pop(k)
+        for i in range(NUM_EXPERTS):
+            weight_map['weight_map'][f'{prefix}.{i}.{dtype}.weight'] = value
+    elif 'Wqkv' in k:
+        prefix = k.removesuffix('.Wqkv.weight')
+        value = weight_map['weight_map'].pop(k)
+        for dtype in ['q_proj', 'k_proj', 'v_proj']:
+            weight_map['weight_map'][f'{prefix}.{dtype}.weight'] = value
+
+sorted_map = sorted(weight_map['weight_map'].items())
+weight_map['weight_map'] = dict(sorted_map)
+
+with open(output_dir / 'model.safetensors.index.json', 'w') as f:
+    json.dump(weight_map, f, indent=4)