Mistral 7B Instruct部署AWS Sagemaker后令牌上限不符问题求助

原因分析 #

AWS Sagemaker的Hugging Face推理容器默认内置了max_total_tokens=4096的令牌数验证规则，该限制与模型本身的上下文窗口规格无关，因此即便部署的是支持32k上下文的Mistral 7B Instruct v0.3，也会被这个默认规则拦截。

解决方法 #

1. 自定义推理脚本覆盖验证逻辑 #

默认推理脚本会触发容器内置的令牌验证，你可以自定义inference.py脚本，直接基于模型配置调整令牌上限：

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

def model_fn(model_dir):
    tokenizer = AutoTokenizer.from_pretrained(model_dir)
    model = AutoModelForCausalLM.from_pretrained(
        model_dir,
        torch_dtype=torch.bfloat16,
        device_map="auto"
    )
    # 绑定模型原生上下文窗口大小
    model_max_length = model.config.max_position_embeddings
    tokenizer.model_max_length = model_max_length
    return model, tokenizer

def predict_fn(data, model_and_tokenizer):
    model, tokenizer = model_and_tokenizer
    inputs = data.pop("inputs", data)
    parameters = data.pop("parameters", {})
    
    # 将max_total_tokens默认值设为模型上下文窗口
    parameters["max_total_tokens"] = parameters.get("max_total_tokens", model.config.max_position_embeddings)
    
    inputs_tokenized = tokenizer(inputs, return_tensors="pt").to("cuda")
    outputs = model.generate(**inputs_tokenized, **parameters)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

2. 设置环境变量覆盖默认限制 #

在创建Sagemaker模型或端点时，添加HF_MODEL_MAX_LENGTH环境变量，值设为模型对应的上下文窗口大小（如Mistral 7B Instruct v0.3设为32768）：

from sagemaker.huggingface import HuggingFaceModel

huggingface_model = HuggingFaceModel(
    model_data="s3://your-model-path/model.tar.gz",
    role=your_iam_role,
    transformers_version="4.37",
    pytorch_version="2.1",
    py_version="py310",
    env={
        "HF_MODEL_MAX_LENGTH": "32768",
        "SAGEMAKER_CONTAINER_LOG_LEVEL": "20"
    }
)

该环境变量会被Hugging Face推理容器读取，自动调整令牌验证的上限值。

3. 调用端点时显式指定参数 #

发送请求到Sagemaker端点时，直接在请求体中设置max_total_tokens为模型支持的最大值，或调整max_new_tokens确保输入令牌数 + max_new_tokens不超过模型上下文窗口：

{
    "inputs": "你的输入文本内容",
    "parameters": {
        "max_new_tokens": 20000,
        "max_total_tokens": 32768
    }
}

注意：需保证输入令牌数 + max_new_tokens不超过模型实际支持的上下文窗口，否则模型会抛出自身的令牌溢出错误。

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