Phind-CodeLlama-34B-v2模型加载与推理提速咨询

一、优化模型加载速度（解决4分钟加载耗时问题） #

• 使用4-bit/8-bit量化加载模型：34B参数的float16模型体积约68GB，远超A10G的24GB显存，导致大量参数需offload到CPU，加载时磁盘IO和跨设备数据传输耗时极长。改用4-bit量化后模型体积仅17GB，可大幅压缩加载时间，同时避免不必要的CPU offload。
  示例代码：

  from transformers import AutoTokenizer, LlamaForCausalLM
  import torch
  
  model_path = "Phind/Phind-CodeLlama-34B-v2"
  
  model = LlamaForCausalLM.from_pretrained(
      model_path,
      device_map="auto",
      torch_dtype=torch.float16,
      load_in_4bit=True,
      bnb_4bit_use_double_quant=True,
      bnb_4bit_quant_type="nf4",
      bnb_4bit_compute_dtype=torch.float16
  )
  tokenizer = AutoTokenizer.from_pretrained(model_path)
  

• 将模型存储到本地NVMe磁盘：g5.8xlarge实例自带900GB本地NVMe SSD，读写速度远高于EBS卷。将模型下载到本地磁盘（如/mnt/local_storage），可大幅降低磁盘IO耗时。

• 启用低内存加载参数：添加low_cpu_mem_usage=True参数，减少加载时的CPU内存占用，同时加快模型初始化速度：

  model = LlamaForCausalLM.from_pretrained(
      model_path,
      device_map="auto",
      torch_dtype=torch.float16,
      low_cpu_mem_usage=True,
      # 搭配量化参数使用
  )
  

二、优化推理速度（解决60秒+推理耗时问题） #

• 核心：用4-bit量化将模型完全放入GPU显存：如上述代码，4-bit量化后的34B模型约17GB，加上推理缓存可完全放入A10G的24GB显存，彻底避免GPU与CPU之间的频繁数据传输（这是当前推理慢的核心原因）。

• 启用Flash Attention 2：Flash Attention可大幅提升注意力计算速度并减少显存占用，需先安装flash-attn库，再在加载模型时启用：

  model = LlamaForCausalLM.from_pretrained(
      model_path,
      device_map="auto",
      torch_dtype=torch.float16,
      load_in_4bit=True,
      use_flash_attention_2=True
  )
  

• 调整生成参数减少计算量：

  • 若不需要随机性，将do_sample=False改为贪心解码，速度会显著提升；
  • 适当降低max_new_tokens（如不需要384个token可按需减少）；
  • 若保留采样，可简化top_p/top_k参数，减少计算开销。
    调整后的生成代码示例：

  def generate_one_completion(prompt: str):
      tokenizer.pad_token = tokenizer.eos_token
      inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=4096).to("cuda")
  
      # 贪心解码，速度最快
      generate_ids = model.generate(
          inputs.input_ids,
          max_new_tokens=384,
          do_sample=False
      )
      completion = tokenizer.batch_decode(generate_ids, skip_special_tokens=True)[0].replace(prompt, "")
      return completion.split("\n\n")[0]
  

• 使用专业推理框架：相比直接用transformers的generate方法，vLLM或Text Generation Inference（TGI）针对LLM推理做了深度优化（如PagedAttention、连续批处理），可将推理速度提升数倍。以vLLM为例，部署命令如下：

  pip install vllm
  python -m vllm.entrypoints.api_server --model Phind/Phind-CodeLlama-34B-v2 --quantization awq --gpu-memory-utilization 0.9
  

内容的提问来源于stack exchange，提问作者Abhijith