Persistent Compiled Model for Intel NPU Acceleration Library

[毕 成]

Problem Description

When using the Intel NPU Acceleration Library to accelerate large models like speech recognition networks or UNet(≥1GB), the compilation time can be excessively long, even though the acceleration effect is good. The library currently lacks the ability to export compiled binary files or binary caches, which means that every time a PyTorch model is loaded, it needs to be loaded into the CPU first and then compiled, which can take a long time. This can be especially problematic when distributing to other machines with lower configurations, where factors like insufficient memory or CPU power limitations can make the compilation time nearly impossible to complete.

Proposed Solution

Provide a persistence mechanism to decouple the model loading step from PyTorch in certain environments. Possible approaches:

1. Export an intermediate representation (IR) file similar to OpenVINO's IR format. This would avoid the need to use PyTorch for subsequent runs, even if some internal compilation still needs to be done (but at least it would complete your version of the static graph trace and weight saving, rather than relying on ONNX).

2. Provide a mechanism similar to OpenVINO's UMD Dynamic Model Caching. Perhaps the underlying work is the same as yours. Consider providing a UID or directly exposing the cache file, so that it can be hit without having to load the model into the CPU each time using a third-party library (PyTorch) before compiling.

Alternatives Considered

Exporting OpenVINO's IR file directly was considered. However, it is clear that the mechanisms are different, as this would require ONNX conversion, which may reduce optimization possibilities. Therefore, this is not an ideal approach.

Additional Context

The long compilation times can be a significant bottleneck, especially when distributing to machines with lower configurations. Introducing a persistence mechanism to save the compiled model and avoid relying on PyTorch for model loading could greatly improve the usability and efficiency of the library in various environments.

[Alessandro Palla]

I agree that is something that can improved. Models can be saved using torch.save (an example here: https://github.com/intel/intel-npu-acceleration-library/pull/20/files#diff-f1b1bbe273678edddc310741907c4962c7656a87b1975c9fd29b5398f48b367cR18) and loaded later but it is not a fool proof solution and I do not like it. For larger model still there is the weight load time that really takes a lot. This is a good issue and we'll definitively add a proper model saving mechanism

[Alessandro Palla]

I'm working on this issue in the following PR: #39

The reason is that now with neural compressor quantization scheme it is much slower than the older naive method in compiling quantized model. So now there are those classes, NPUModelForXXX.from_pretrained(...) that you can use and automatically handle the caching of the model. Any feedback is well accepted :)

Example: first time it compile the model and save it (disable the saving with .from_pretrained(export=False) ) and the other times it load that directly

from transformers import AutoTokenizer, TextStreamer
from intel_npu_acceleration_library import NPUModelForCausalLM
import torch

model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"

model = NPUModelForCausalLM.from_pretrained(
    model_id, use_cache=True, dtype=torch.int8
).eval()
tokenizer = AutoTokenizer.from_pretrained(model_id, use_default_system_prompt=True)
tokenizer.pad_token_id = tokenizer.eos_token_id
streamer = TextStreamer(tokenizer, skip_special_tokens=True)


query = input("Ask something: ")
prefix = tokenizer(query, return_tensors="pt")["input_ids"]


generation_kwargs = dict(
    input_ids=prefix,
    streamer=streamer,
    do_sample=True,
    top_k=50,
    top_p=0.9,
    max_new_tokens=512,
)

print("Run inference")
_ = model.generate(**generation_kwargs)