diff --git a/Dockerfile b/Dockerfile index f278e82..9436b1f 100644 --- a/Dockerfile +++ b/Dockerfile @@ -21,10 +21,8 @@ RUN pip install auto-gptq==0.7.1 --no-build-isolation RUN pip3 install /rknn-llm/rkllm-toolkit/packages/rkllm_toolkit-1.2.3-cp310-cp310-linux_x86_64.whl # Клонирование модели -WORKDIR /models - -# Копируем скрипт компиляции -COPY entrypoint.py /entrypoint.py +WORKDIR /work +VOLUME ["/models", "/output"] # Точка входа для компиляции -ENTRYPOINT ["python3", "/entrypoint.py"] \ No newline at end of file +CMD ["bash"] \ No newline at end of file diff --git a/build.sh b/build.sh index 3174f4f..1bdbea8 100644 --- a/build.sh +++ b/build.sh @@ -3,9 +3,16 @@ set -e MODELS_PATH="models" OUTPUT_DIR="compiled-models" -IMAGE_NAME="rkllm-builder" +IMAGE_NAME="git.home/drholy/builder_rkllm:latest" CONTAINER_NAME="rkllm-compile-$$" +MODEL_NAME="Qwen3-4B-Instruct-2507-Q4_0" +MODEL_PATH="/models/$MODEL_NAME/$MODEL_NAME.gguf" +MODEL_OUTPUT="/output/" +MODEL_FORMAT="gguf" +MODEL_QUANT_ENABLE="True" +MODEL_QUANT="w8a8" + mkdir -p "$MODELS_PATH" "$OUTPUT_DIR" echo "🏗️ Сборка образа для компиляции..." @@ -27,6 +34,8 @@ docker run --rm \ --name "$CONTAINER_NAME" \ -v "$(pwd)/$MODELS_PATH:/models" \ -v "$(pwd)/$OUTPUT_DIR:/output" \ - "$IMAGE_NAME" + -v "$(pwd)/entrypoint.py:/work/entrypoint.py:ro" \ + "$IMAGE_NAME" \ + python3 /work/entrypoint.py --model $MODEL_PATH --output $MODEL_OUTPUT --format $MODEL_FORMAT --enable-q $MODEL_QUANT_ENABLE --quant $MODEL_QUANT echo "✅ Модель сохранена в: $OUTPUT_DIR" \ No newline at end of file diff --git a/entrypoint.py b/entrypoint.py index aedb2e4..cc0a190 100644 --- a/entrypoint.py +++ b/entrypoint.py @@ -1,37 +1,91 @@ -from rkllm.api import RKLLM -from tqdm import tqdm -import torch -from torch import nn +#!/usr/bin/env python3 +import argparse import os +from rkllm.api import RKLLM -modelpath = '/models/Qwen3-4B-Instruct-2507-F16/Qwen3-4B-Instruct-2507-F16.gguf' -output = '/output/Qwen3-4B-Instruct-2507-F16-Q.rkllm' -os.makedirs("/output", exist_ok=True) +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() -print(f"Загрузка модели из: {modelpath}") +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") -llm = RKLLM() -ret = llm.load_gguf(model=modelpath) -# ret = llm.load_huggingface(model=modelpath, model_lora = None, device='cpu') - -if ret != 0: - print(f"❌ Ошибка загрузки модели: {ret}") - exit(ret) - -print("Компиляция для RK3588 (NPU)...") -ret = llm.build(do_quantization=True, optimization_level=1, quantized_dtype='w8a8', - quantized_algorithm='normal', target_platform='rk3588', num_npu_core=3, extra_qparams=None) - - -if ret != 0: - print(f"❌ Ошибка компиляции: {ret}") - exit(ret) - -# Export rkllm model -ret = llm.export_rkllm(output) -if ret != 0: - print(f"❌ Ошибка экспорта: {ret}") - exit(ret) - -print(f"✅ Модель скомпилирована: {output}") -print(f"Размер: {os.path.getsize(output) / 1024**3:.2f} GB") \ No newline at end of file +if __name__ == '__main__': + main() \ No newline at end of file