在构建以大模型推理为主的应用流程时，出于推理性能与应用解耦等考量，通常会希望使用一些大模型推理引擎。比如在使用langchain时，自带的llm就提供了包括openai、tgi等接口。 

ModelArts文档中提到了tgi、vllm的昇腾适配，vllm包无法直接在该处获取，且华为工作人员表示我们的场景不支持这个包，故尝试部署tgi。以下记录一些个人部署经验，如有错误请指出。 

环境 

ModelArts云平台 

aarch64 cpu  

Ascend 910PremiumA计算卡 

待部署模型为llama结构 

部署过程 

非flash attention部署 

tgi的基本部署过程可参考上述ModelArts文档 。

几点注释： 

• 云平台没有root权限，rust可以conda安装 
• 最好严格按照文档中提供的包的版本安装，比如如果pytorch装的是2.1.0，会碰到ProcessGroupHCCL未定义的问题（虽然个人感觉那处代码注释掉影响不大），且调用昇腾atb加速库时会碰到symbol undefined的问题 
  • atb中以.so的形式提供了算子，应该是应用了TorchScript C++扩展的机制（TorchScript Custom Classes） ，版本不匹配无法正常读取。 

此时应当可以正常部署tgi： 

ASCEND_RT_VISIBLE_DEVICES=0 BIND_CPU=0 FLASH_ATTENTION=0 MAX_MEMORY_GB=28 text-generation-launcher --model-id <your_llama_path> --port 8080 

非flash attention的llama直接使用transformer库上的实现，正常情况下是能在昇腾平台上跑的。但可惜由于padding的问题，在同时执行多个推理请求时，continuous batching会导致精度丢失。这个问题在tgi的文档上有说明（huggingface.co/docs/text-generation-inference/main/en/basic_tutorials/launcher）： 

--max-batch-total-tokens <MAX_BATCH_TOTAL_TOKENS> 

          **IMPORTANT** This is one critical control to allow maximum usage of the available hardware. 

           

          This represents the total amount of potential tokens within a batch. When using padding (not recommended) this would be equivalent of `batch_size` * `max_total_tokens`. 

           

          However in the non-padded (flash attention) version this can be much finer. 

           

          For `max_batch_total_tokens=1000`, you could fit `10` queries of `total_tokens=100` or a single query of `1000` tokens. 

           

          Overall this number should be the largest possible amount that fits the remaining memory (after the model is loaded). Since the actual memory overhead depends on other parameters like if you're using quantization, flash attention or the model implementation, text-generation-inference cannot infer this number automatically. 

           

          [env: MAX_BATCH_TOTAL_TOKENS=] 

解决方案也提供了，换flash attention版本的实现。个人理解tgi在flash attention版本提供了各种优化手段，包括这里对我们最重要的padding free实现，这段代码确实有些复杂，个人认为单独抽出其中一些优化实有难度也不太合算。下面介绍flash attention版本的部署经验。 

 flash attention部署 

如果直接就能正常跑起来(910A)，我当然不会令开一节写这个事。 

首先介绍一下tgi与这个部分相关的代码结构（昇腾适配tgi版本为0.9.4） 

• server/text_generation_server/models/__init__.py 
  • 定义了get_model函数，从这里可以看到各种参数下实际模型的入口，比如非flash attention的llama模型入口在同文件夹下causal_lm.py，相对应的falsh attention在flash_causal_lm.py 
• server/text_generation_server/models/flash_llama.py 
  • flash attention版本的各模型在tgi单独实现，flash attention llama的实现的一些公共部分在这里，其他的在server/text_generation_server/models/custom_modeling/flash_llama_modeling_ascent.py等处 

我们看server/text_generation_server/models/flash_llama.py开头的模型实现选择： 

if ENV.use_ascend: 

    from text_generation_server.models.custom_modeling.flash_llama_modeling_ascend import ( 

        FlashLlamaForCausalLM, 

        LlamaConfig, 

    ) 

    from text_generation_server.models.custom_modeling.flash_llama_modeling_quant_ascend import ( 

        FlashLlamaForCausalLM as QuantFlashLlamaForCausalLM, 

        LlamaConfig as QuantLlamaConfig, 

    ) 

else: 

    from text_generation_server.models.custom_modeling.flash_llama_modeling import ( 

        FlashLlamaForCausalLM, 

        LlamaConfig, 

    ) 

Quant模型个人暂没有使用打算。这里注意flash_llama_modeling与flash_llama_modeling_ascend，flash_llama_modeling是tgi原始实现，昇腾替换了其中部分实现进行了适配，flash_llama_modeling_ascend是atb实现，与CANN transformer_llm包里example中的llama_pa类似。 

我们所使用的卡为ModelArts平台的910PremiumA，flash_llama_modeling_ascend虽能够跑起来，但输出完全异常，个人也没有调试二进制算子的经验，无法定位问题。 

flash_llama_modeling直接也没能跑起来，且没观察到有意义的报错，最后发现在取消亲和力优化的need_nz flag后能够跑通： 

server/text_generation_server/utils/npu.py:24: self.need_nz改为False 

 

启动命令： 

ASCEND_RT_VISIBLE_DEVICES=0 BIND_CPU=0 FLASH_ATTENTION=1 MAX_MEMORY_GB=28 text-generation-launcher --model-id <your_llama_path> --port 8080 

注意MAX_MEMORY_GB不要贴着卡上限设，prefill阶段可能会爆。 

 

目前这样的方案能够成功部署，推理精度不再由于continuous batching降低，但好像并不是很flash，比非flash attention还慢一点。 

其他部署经验 

top logprobs 

昇腾适配的tgi版本为0.9.4，相比最新实现落后不少版本。个人需要返回topn token以及logprobs的功能，自行进行了实现。tgi还挺复杂的，加个字段在三端python、protobuf、rust间来来回回折腾小几十处。十分希望昇腾能有更方便的跟踪最新库生态的方法。 

多卡部署 

tgi不直接支持数据并行的推理，可以同时在多卡上部署多个tgi实例，用nginx做个反向代理，此处提供一个脚本供参考： 

if [ -z $1 ] 

then 

    echo "pass num npu" 

else 

    for ((i=0; i<$1; i++)) 

    do 

ASCEND_RT_VISIBLE_DEVICES=$i BIND_CPU=0 FLASH_ATTENTION=1 MAX_MEMORY_GB=28 \ 

    nohup text-generation-launcher --model-id <your_llama_path> \ 

    --port $((8081+${i})) --shard-uds-path /tmp/text-generation-server-$i-0 >nohup.out.$i 2>&1 & 

    done 

 

    NGINX_CONF_PATH=nginx.self_rag.$1.conf 

    cat > "$NGINX_CONF_PATH" <<EOF 

events { 

    worker_connections 1024; 

} 

http { 

 

upstream backend { 

EOF 

 

    for ((i=0; i<(($1)); i++)) do 

        echo "    server 127.0.0.1:$((8081+${i}));" >> $NGINX_CONF_PATH 

    done 

 

    cat >> "$NGINX_CONF_PATH" <<EOF 

} 

 

server { 

    listen 8080; 

 

    location / { 

        proxy_pass http://backend; 

        proxy_set_header Host \$host; 

        proxy_set_header X-Real-IP \$remote_addr; 

        proxy_set_header X-Forwarded-For \$proxy_add_x_forwarded_for; 

        proxy_set_header X-Forwarded-Proto \$scheme; 

 

        proxy_connect_timeout 10m; 

        proxy_read_timeout 20m; 

        proxy_send_timeout 20m; 

    } 

} 

} 

EOF 

 

    nginx -c `pwd`/$NGINX_CONF_PATH 

 

fi 

注意--shard-uds-path做个区分，不要设一样的 

一些问题 

• atb的速度确实是快，但是无法正常输出也不易确定原因，希望能够得到解答，是不是910A就是用不了，还是需要装特定版本 
• 910A需要的need_nz做了什么工作，为什么导致flash_llama_modeling版本无法启动，且这个代码跑的比较慢是不是没做nz转换的问题 
• ascend-vllm是否910A确定无法使用 
• MindSpore有实现类似功能的大模型推理引擎吗，最好不是单点优化的实现，做研究目的用，所以希望能够达到一定程度的易用性