| """Generic Hugging Face text model -> RKLLM w8a8 for RK3588 (run on the 3090 WSL box). |
| |
| python convert_rkllm.py MODEL_DIR OUT.rkllm [--max-context 4096] [--device cuda] [--fla] |
| |
| --fla: import flash-linear-attention first so `rwkv7` (fla-hub checkpoints) is registered with transformers. The fla |
| pinned by requirements_rwkv7.txt also registers names (bitnet, ...) that transformers 5.8 now ships, which raises |
| "already used"; those are skipped. |
| """ |
| import argparse |
| import os |
|
|
| from rkllm.api import RKLLM |
|
|
| ap = argparse.ArgumentParser() |
| ap.add_argument("model_dir") |
| ap.add_argument("out") |
| ap.add_argument("--max-context", type=int, default=4096) |
| ap.add_argument("--device", default="cuda") |
| ap.add_argument("--dtype", default="w8a8") |
| ap.add_argument("--fla", action="store_true") |
| args = ap.parse_args() |
| if args.fla: |
| from transformers.models.auto import auto_factory, configuration_auto |
|
|
| def skip_taken(orig): |
| def register(self, key, value, exist_ok=False): |
| try: |
| orig(self, key, value, exist_ok=exist_ok) |
| except ValueError as e: |
| if "already" not in str(e): |
| raise |
| return register |
|
|
| configuration_auto._LazyConfigMapping.register = skip_taken(configuration_auto._LazyConfigMapping.register) |
| auto_factory._LazyAutoMapping.register = skip_taken(auto_factory._LazyAutoMapping.register) |
| import fla.models |
| import transformers |
| |
| transformers.RWKV7ForCausalLM = fla.models.RWKV7ForCausalLM |
| llm = RKLLM() |
| assert llm.load_huggingface(model=args.model_dir, device=args.device) == 0, "load failed" |
| assert llm.build(do_quantization=True, optimization_level=1, quantized_dtype=args.dtype, quantized_algorithm="normal", |
| target_platform="rk3588", num_npu_core=3, max_context=args.max_context) == 0, "build failed" |
| assert llm.export_rkllm(args.out) == 0, "export failed" |
| print("EXPORTED", args.out, os.path.getsize(args.out)) |
|
|
