Support cpu inference with scaffolding

[WeiHaocheng]

Many community contributors find it difficult to find GPU based development environments and prefer CPU based scaffolding workers to complete some development tasks.
Potential options include pytorch, ollama, llama.cpp.

[amemov]

@WeiHaocheng, I'm interested in working on this issue - can this be assigned to me?

[amemov]

@WeiHaocheng

[WeiHaocheng]

@amemov Great, assign to you~
Feel free to left comment if you have any idea about this issue.
The introduction and code of scaffolding is here:
https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/scaffolding/README.md

We would like to develop a Worker which can support inference on cpu.

The code for Worker is here:
https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/scaffolding/worker.py

[amemov]

@amemov Great, assign to you~ Feel free to left comment if you have any idea about this issue. The introduction and code of scaffolding is here: https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/scaffolding/README.md

We would like to develop a Worker which can support inference on cpu.

The code for Worker is here: https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/scaffolding/worker.py

@WeiHaocheng , just wanted to double check that I'm on the right track: I looked at the code more in-depth and if I understood correctly I need to write a class for worker.py. I looked into the potential options that you mentioned above - they all seem pretty straightforward. For now I will focus on writing implementation with Pytorch.

About testing: I see that there is tests/unittest/scaffolding/test_worker.py ; I suppose my test should be similar to test_trtllm_worker_generation?

[WeiHaocheng]

@amemov

Yes, your understanding is right. We need a Worker likes "PytorchWorker" and it can handle and process the GenerationTask in CPU.

And you can read the README and just add a directory on "tensorrt_llm/scaffolding/contrib" and add an example of running the worker on "example/scaffolding/contrib". I would migrate it into worker.py and add the tests in later. We plan to use this approach to ensure flexible and scalable for contributors.

[WeiHaocheng]

Hi @amemov ~ May I ask if your progress is still going smoothly? If you encounter some difficult problems during the building, maybe you can try downloading the pre-compiled image directly from my dockerhub.

https://hub.docker.com/r/weihaocheng2/tensorrt_llm

Feel free to contact me if you got any problem.

[amemov]

Hi @WeiHaocheng - I'm still working on it. I had to take a long break due to personal issues. I think I will have a PR draft ready closer to the end of this week.

[WeiHaocheng]

Hi @WeiHaocheng - I'm still working on it. I had to take a long break due to personal issues. I think I will have a PR draft ready closer to the end of this week.

Thank you. Feel free to tell me if you need any help~