在Day2 中,我们梳理了trt-llm对于TinyLLama的调用,在Day3,我们也熟悉了一下Trt-llm常规的三步流程。
这里其实有个问题,在针对tiny-llama的部署中,其实没有显式的进行模型转换,那麽其推理接口中到底包含了什么?
#demo.pyfrom tensorrt_llm import LLM, SamplingParams
prompts = ["Hello, my name is","The president of the United States is","The capital of France is","The future of AI is",
]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)model_path = "/data/TinyLlama-1.1B-Chat-v1.0"# add break point
# import pdb;pdb.set_trace()
llm = LLM(model= model_path)outputs = llm.generate(prompts, sampling_params)
经过PDB的断点,可以看到下列函数被调用。
model_loader = CachedModelLoader(self.args,
298 mpi_session=self.mpi_session,
299 workspace=self.workspace,
300 llm_build_stats=self.llm_build_stats)
301 self._engine_dir = model_loader()
其中, 在model_loader()的调用中,实现了模型的加载和engine的构造。而在step_forward函数中,则实现了每次的step。
这里附上step_forward的代码
主要包含两步,在第二步中,主要包含
> /usr/local/lib/python3.10/dist-packages/tensorrt_llm/builder.py(1059)build()
-> builder = Builder()
(Pdb)
在准备好参数后,进行模型和engine的构造。
其中,主要的时间都用于模型的构造了。
> /usr/local/lib/python3.10/dist-packages/tensorrt_llm/builder.py(1203)build()-><tensorrt_llm...x7f119bb08a60>
-> return Engine(engine_config, engine, managed_weights)
在构造后,模型被存储在了一个临时目录,如下
ModelLoader.save(runtime_context, self.llm_args.model_dir, engine_dir)
其中,打开engine_dir 可以看到路径为
/tmp/tmpeefho9klllm-workspace/tmp.engine
基于下列代码,我们实现了executor的构造。
-> self._executor = self._executor_cls.create((Pdb) p self._executor
<tensorrt_llm.executor.ExecutorBindingsWorker object at 0x7f0ee555fdc0>
最后完成了模型的推理。
今天算是仔细过了一遍模型的生成,可以看到,其中包含了
1. 原始权重文件的的读取,加载,转换,优化
2,. engine文件的落盘
3. 基于engine文件实现推理