91 lines
3.9 KiB
Python
91 lines
3.9 KiB
Python
#!/usr/bin/env python3
|
||
import argparse
|
||
import os
|
||
from rkllm.api import RKLLM
|
||
|
||
def parse_args():
|
||
parser = argparse.ArgumentParser(description='Компиляция LLM для Rockchip NPU')
|
||
|
||
# Обязательные параметры
|
||
parser.add_argument('--model', type=str, required=True,
|
||
help='Путь к модели (Hugging Face или GGUF)')
|
||
parser.add_argument('--output', type=str, required=True,
|
||
help='Путь для сохранения .rkllm файла')
|
||
|
||
# Опциональные параметры с умными значениями по умолчанию
|
||
parser.add_argument('--format', type=str, choices=['huggingface', 'gguf'], default='huggingface',
|
||
help='Формат входной модели (по умолчанию: huggingface)')
|
||
parser.add_argument('--enable-q', type=str, choices=['True', 'False'], default='True',
|
||
help='Включить квантование (по умолчаниюЖ True)')
|
||
parser.add_argument('--quant', type=str, default='w8a8_g256',
|
||
choices=['w4a16', 'w4a16_g32', 'w4a16_g64', 'w4a16_g128',
|
||
'w8a8', 'w8a8_g128', 'w8a8_g256', 'w8a8_g512'],
|
||
help='Тип квантования (по умолчанию: w8a8_g256)')
|
||
parser.add_argument('--platform', type=str, default='rk3588',
|
||
choices=['rk3588', 'rk3576', 'rk3562', 'rv1126b'],
|
||
help='Целевая платформа (по умолчанию: rk3588)')
|
||
parser.add_argument('--npu-cores', type=int, default=3, choices=[1, 2, 3],
|
||
help='Количество ядер NPU (по умолчанию: 3)')
|
||
parser.add_argument('--max-context', type=int, default=8192,
|
||
help='Максимальная длина контекста (по умолчанию: 8192)')
|
||
|
||
return parser.parse_args()
|
||
|
||
def main():
|
||
args = parse_args()
|
||
os.makedirs(os.path.dirname(args.output) or '.', exist_ok=True)
|
||
|
||
print(f"🔧 Компиляция модели:")
|
||
print(f" Модель: {args.model}")
|
||
print(f" Формат: {args.format}")
|
||
print(f" Включено квантование: {args.enable_q}")
|
||
print(f" Квантование: {args.quant}")
|
||
print(f" Платформа: {args.platform}")
|
||
print(f" Ядра NPU: {args.npu_cores}")
|
||
print(f" Длина контекста: {args.max_context}")
|
||
|
||
llm = RKLLM()
|
||
|
||
# Загрузка в зависимости от формата
|
||
if args.format == 'huggingface':
|
||
ret = llm.load_huggingface(
|
||
model=args.model,
|
||
device='cpu',
|
||
trust_remote_code=True,
|
||
quantized_model=False
|
||
)
|
||
else: # gguf
|
||
ret = llm.load_gguf(model=args.model)
|
||
|
||
if ret != 0:
|
||
print(f"❌ Ошибка загрузки модели (код {ret})")
|
||
exit(ret)
|
||
|
||
# Определение алгоритма квантования
|
||
algo = 'grq' if 'w4a16' in args.quant else 'normal'
|
||
|
||
ret = llm.build(
|
||
do_quantization=args.enable_q,
|
||
optimization_level=1,
|
||
quantized_dtype=args.quant,
|
||
quantized_algorithm=algo,
|
||
target_platform=args.platform,
|
||
num_npu_core=args.npu_cores,
|
||
max_context=args.max_context
|
||
)
|
||
|
||
if ret != 0:
|
||
print(f"❌ Ошибка компиляции (код {ret})")
|
||
exit(ret)
|
||
|
||
ret = llm.export_rkllm(args.output)
|
||
if ret != 0:
|
||
print(f"❌ Ошибка экспорта (код {ret})")
|
||
exit(ret)
|
||
|
||
size_gb = os.path.getsize(args.output) / (1024**3)
|
||
print(f"\n✅ Успешно! Сохранено: {args.output}")
|
||
print(f"📦 Размер: {size_gb:.2f} GB")
|
||
|
||
if __name__ == '__main__':
|
||
main() |