152 lines
5.7 KiB
Python
152 lines
5.7 KiB
Python
# Merges a LoRA checkpoint in PyTorch format (.pth) into an rwkv.cpp model file.
|
|
# Usage: python merge_lora_into_ggml.py C:\rwkv.cpp-169M.bin C:\my-lora.pth 32 C:\rwkv.cpp-169M-with-my-lora.bin
|
|
# LoRA format is compatible with https://github.com/Blealtan/RWKV-LM-LoRA
|
|
# You need to know lora_alpha value to perform the merge.
|
|
# Source model must be in either FP16 or FP32 format. Quantization can be performed after merging.
|
|
|
|
import argparse
|
|
import struct
|
|
import torch
|
|
import numpy as np
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(description='Merge a PyTorch LoRA checkpoint (.pth) into an rwkv.cpp model file')
|
|
parser.add_argument('src_path', help='Path to source rwkv.cpp model')
|
|
parser.add_argument('lora_path', help='Path to LoRA checkpoint in PyTorch format')
|
|
parser.add_argument('lora_alpha', type=int, help='Value of lora_alpha parameter used when training this LoRA checkpoint')
|
|
parser.add_argument('dest_path', help='Path to destination rwkv.cpp model, will be overwitten with the merged model')
|
|
return parser.parse_args()
|
|
|
|
def write_parameter(out_file, key: str, parameter: torch.Tensor) -> None:
|
|
assert parameter.dtype == torch.float32 or parameter.dtype == torch.float16
|
|
|
|
key_encoded: bytes = key.encode('utf-8')
|
|
|
|
out_file.write(struct.pack(
|
|
'=iii',
|
|
len(parameter.shape),
|
|
len(key_encoded),
|
|
1 if parameter.dtype == torch.float16 else 0
|
|
))
|
|
|
|
# Dimension order is reversed here:
|
|
# * PyTorch shape is (x rows, y columns)
|
|
# * ggml shape is (y elements in a row, x elements in a column)
|
|
# Both shapes represent the same tensor.
|
|
for dim in reversed(parameter.shape):
|
|
out_file.write(struct.pack('=i', dim))
|
|
|
|
out_file.write(key_encoded)
|
|
|
|
parameter.numpy().tofile(out_file)
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
|
|
print(f'Reading {args.lora_path}')
|
|
|
|
lora_state_dict: dict[str, torch.Tensor] = torch.load(args.lora_path, map_location='cpu')
|
|
|
|
print(f'Merging')
|
|
|
|
with open(args.src_path, 'rb') as in_file, open(args.dest_path, 'wb') as out_file:
|
|
# noinspection PyTypeChecker
|
|
header: tuple[int, int, int, int, int, int] = struct.unpack('=iiiiii', in_file.read(6 * 4))
|
|
|
|
assert header[0] == 0x67676d66, 'Invalid magic value'
|
|
assert header[1] == 100, 'Invalid version number'
|
|
assert header[5] == 0 or header[5] == 1, 'Only FP32 and FP16 models are supported'
|
|
|
|
out_file.write(struct.pack('=iiiiii', *header))
|
|
|
|
while True:
|
|
parameter_header_bytes: bytes = in_file.read(3 * 4)
|
|
|
|
if len(parameter_header_bytes) == 0:
|
|
break
|
|
|
|
dim_count, key_length, data_type = struct.unpack('=iii', parameter_header_bytes)
|
|
|
|
# noinspection PyTypeChecker
|
|
shape: tuple[int] = struct.unpack('=' + 'i' * dim_count, in_file.read(dim_count * 4))
|
|
# ggml order to PyTorch
|
|
shape: list[int] = [d for d in reversed(shape)]
|
|
|
|
key: str = in_file.read(key_length).decode('utf-8')
|
|
|
|
print(f'* {key} {shape}')
|
|
|
|
assert data_type == 0 or data_type == 1, 'Only FP32 and FP16 models are supported'
|
|
|
|
element_count: int = 1
|
|
|
|
for dim in shape:
|
|
element_count *= dim
|
|
|
|
parameter_np: np.ndarray = np.frombuffer(
|
|
in_file.read((2 if data_type == 1 else 4) * element_count),
|
|
dtype=(np.half if data_type == 1 else np.single)
|
|
)
|
|
|
|
parameter: torch.Tensor = torch.tensor(parameter_np).view(shape)
|
|
|
|
if key in lora_state_dict:
|
|
replacement: torch.Tensor = lora_state_dict[key].float()
|
|
|
|
# Same processing as in convert_pytorch_to_ggml.py
|
|
if '.time_' in key:
|
|
# (1, 1, n_embed) -> (n_embed)
|
|
replacement = replacement.squeeze()
|
|
|
|
if '.time_decay' in key:
|
|
replacement = -torch.exp(replacement)
|
|
|
|
if parameter.dtype == torch.float16:
|
|
replacement = replacement.half()
|
|
|
|
assert replacement.shape == parameter.shape, f'Parameter {key} has shape {parameter.shape} in model file ' \
|
|
f'and shape {replacement.shape} in LoRA file'
|
|
|
|
parameter = replacement
|
|
|
|
print(f'Replaced parameter {key}')
|
|
|
|
del lora_state_dict[key]
|
|
|
|
for suffix in ['.weight', '']:
|
|
lora_A_key: str = key.replace('.weight', '') + '.lora_A' + suffix
|
|
lora_B_key: str = key.replace('.weight', '') + '.lora_B' + suffix
|
|
|
|
if lora_A_key in lora_state_dict:
|
|
lora_A: torch.Tensor = lora_state_dict[lora_A_key]
|
|
lora_B: torch.Tensor = lora_state_dict[lora_B_key]
|
|
|
|
assert lora_B.shape[1] == lora_A.shape[0], f'Invalid shape of LoRA matrices for {key}: ' \
|
|
f'{lora_A.shape}, {lora_B.shape}'
|
|
|
|
lora_R: int = lora_B.shape[1]
|
|
|
|
replacement: torch.Tensor = parameter + lora_B @ lora_A * (args.lora_alpha / lora_R)
|
|
|
|
if parameter.dtype == torch.float16:
|
|
replacement = replacement.half()
|
|
|
|
parameter = replacement
|
|
|
|
print(f'Merged LoRA into parameter {key}, lora_r = {lora_R}')
|
|
|
|
del lora_state_dict[lora_A_key]
|
|
del lora_state_dict[lora_B_key]
|
|
|
|
break
|
|
|
|
write_parameter(out_file, key, parameter)
|
|
|
|
for key in lora_state_dict:
|
|
print(f'WARNING: Unused parameter in LoRA state dict {key}')
|
|
|
|
print('Done')
|
|
|
|
if __name__ == "__main__":
|
|
main()
|