Oneflow-Inc / libai

LiBai(李白): A Toolbox for Large-Scale Distributed Parallel Training
https://libai.readthedocs.io
Apache License 2.0
391 stars 55 forks source link

Llama device #548

Closed ShawnXuan closed 1 month ago

ShawnXuan commented 2 months ago

This PR introduces the ability to specify the target device (e.g., CUDA, NPU, XPU) for Llama inference. The changes allow users to select the desired device through configuration, improving flexibility across different hardware platforms.

Please update the projects/Llama/configs/llama_config.py file to configure the model path and tokenizer path.

python projects/Llama/pipeline.py