调用AWS SageMaker Llama2推理端点遇内存分配错误求助

错误信息Allocation larger than expected: tag 'rms_qkv', requested size: 164831232, expected max size: '100663296'本质是模型推理时显存不足，rms_qkv是LLaMA2注意力模块的参数，说明输入/生成的token总量超出了当前端点实例的显存承载能力。

以下是具体解决方向：

一、减少输入输出的Token总量 #

• 降低生成token上限：把model_kwargs里的max_new_tokens从2048调小，减少模型生成内容的显存占用：

  llm = SagemakerEndpoint(
      endpoint_name=endpoint_name, 
      region_name=region, 
      model_kwargs={"max_new_tokens": 1024, "top_p": 0.9, "temperature": 0.1},  # 调整max_new_tokens
      endpoint_kwargs={"CustomAttributes": 'accept_eula=true'},
      content_handler=content_handler
  )
  

• 精简输入文档：当前用chain_type="stuff"会把所有docs内容拼接进prompt，文档过长会导致输入token暴增。可以换成map_reduce或refine链，避免一次性塞入所有文档：

  chain = load_qa_chain(llm, chain_type="map_reduce", memory=st.session_state['memory'], prompt=prompt_template)
  

• 限制对话历史长度：memory中的聊天记录会持续累积，增加输入token量。可以设置内存的token上限，或清理旧对话轮次：

  from langchain.memory import ConversationTokenBufferMemory
  # 限制历史记录最多占用1024个token
  memory = ConversationTokenBufferMemory(llm=llm, max_token_limit=1024)
  st.session_state['memory'] = memory
  

二、升级端点实例类型 #

当前使用的实例显存不足以支撑当前的token规模，换成带更大显存的GPU实例，比如：

• 从t2/t3系列CPU实例换成g4dn.xlarge（16GB显存）
• 或升级到g5.xlarge（24GB显存）、g5.2xlarge（48GB显存）等更高配置的实例

三、优化模型推理参数减少显存占用 #

• 启用模型量化：如果创建端点时选择了INT8/FP16量化版本的LLaMA2模型，可大幅降低显存占用。若支持自定义参数，可在model_kwargs中添加配置：

  model_kwargs={"max_new_tokens": 1024, "quantize": "int8"}
  

• 开启Flash Attention：部分LLaMA2的JumpStart版本支持Flash Attention，可减少注意力模块的内存开销，在model_kwargs中开启：

  model_kwargs={"max_new_tokens": 1024, "flash_attention": True}
  

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