赤兔Chitu:清華大學團隊開源的高性能大語言模型推理框架
赤兔(Chitu)是什么?
赤兔(Chitu)是由清華大學團隊開源的一個高性能大語言模型推理框架,專注于提升推理效率、降低成本,并支持多種硬件平臺和部署場景。它可以讓DeepSeek推理成本降一半,性能翻番。赤兔(Chitu)的目標是為企業(yè)和開發(fā)者提供一個高效、靈活且易于部署的推理引擎,加速大語言模型(LLM)在實際應(yīng)用中的落地。
赤兔(Chitu)功能特點
多元算力適配:支持 NVIDIA 最新旗艦到舊款的多系列產(chǎn)品,并為國產(chǎn)芯片提供優(yōu)化支持。
全場景可伸縮:從純 CPU 部署、單 GPU 部署到大規(guī)模集群部署,提供靈活的可擴展解決方案。
長期穩(wěn)定運行:適用于實際生產(chǎn)環(huán)境,能夠穩(wěn)定承載并發(fā)業(yè)務(wù)流量。
高效推理性能:在 A800 集群測試中,相比部分國外開源框架,GPU 使用量減少 50% 的情況下推理速度提升 3.15 倍。
降低成本:通過優(yōu)化硬件資源使用,降低了企業(yè)部署 ai 模型的門檻和成本。
赤兔(Chitu)應(yīng)用場景
企業(yè)級 AI 應(yīng)用:需要高性能、低延遲和高吞吐量的推理服務(wù)。
大規(guī)模集群部署:需要在多 GPU 或多節(jié)點環(huán)境中高效運行的場景。
資源受限的環(huán)境:需要在有限的硬件資源下實現(xiàn)高效推理的場景。
國產(chǎn)芯片適配:需要在國產(chǎn)硬件平臺上運行大語言模型的場景。
赤兔(Chitu)使用方法
1. 源碼安裝
git clone --recursive https://github.com/thu-pacman/chitu && cd chitu pip install -r requirements-build.txt pip install -U torch --index-url https://download.pytorch.org/whl/cu124 # 根據(jù) CUDA 版本調(diào)整 TORCH_CUDA_ARCH_LIST=8.6 CHITU_SETUP_JOBS=4 MAX_JOBS=4 pip install --no-build-isolation
2.單 GPU 推理
torchrun --nproc_per_node 8 test/single_req_test.py request.max_new_tokens=64 models=DeepSeek-R1 models.ckpt_dir=/data/DeepSeek-R1 infer.pp_size=1 infer.tp_size=8
3.混合并行 (TP+PP)
torchrun --nnodes 2 --nproc_per_node 8 test/single_req_test.py request.max_new_tokens=64 infer.pp_size=2 infer.tp_size=8 models=DeepSeek-R1 models.ckpt_dir=/data/DeepSeek-R1
4.啟動服務(wù)
export WORLD_SIZE=8 torchrun --nnodes 1 \ --nproc_per_node 8 \ --master_port=22525 \ example/serve.py \ serve.port=21002 \ infer.stop_with_eos=False \ infer.cache_type=paged \ infer.pp_size=1 \ infer.tp_size=8 \ models=DeepSeek-R1 \ models.ckpt_dir=/data/DeepSeek-R1 \ keep_dtype_in_checkpoint=True \ infer.mla_absorb=absorb-without-precomp \ infer.soft_fp8=True \ infer.do_load=True \ infer.max_reqs=1 \ scheduler.prefill_first.num_tasks=100 \ infer.max_seq_len=4096 \ request.max_new_tokens=100 \ infer.use_cuda_graph=True
5.性能測試
python benchmarks/benchmark_serving.py \ --model "deepseek-r1" \ --iterations 10 \ --seq-len 10 \ --warmup 3 \ --base-url http://localhost:21002
GitHub倉庫:https://github.com/thu-pacman/chitu