如何为AWS Inferentia编译支持灵活输入尺寸的HuggingFace模型？

问题背景 #

使用HuggingFace的bigscience/bloom-560m模型，核心推理代码如下：

import torch
from transformers import BloomTokenizerFast, BloomForCausalLM

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
device

tokenizer = BloomTokenizerFast.from_pretrained("bigscience/bloom-560m")
model = BloomForCausalLM.from_pretrained("bigscience/bloom-560m").to(device)

text = tokenizer.encode(seed)
inputs, past_key_values = torch.tensor([text[0]]), None

with torch.no_grad():
    while #condition met:
        model_out = model(input_ids=inputs.to(device), past_key_values=past_key_values)
        ...
        # Generate new inputs and go back to the start

尝试部署到AWS SageMaker Inf1实例时，使用以下编译代码：

from sagemaker.pytorch.model import PyTorchModel

pytorch_model = PyTorchModel(
    model_data=model_path,
    role=role,
    entry_point="my_entry_point_file.py",
    framework_version="1.5.1",
    py_version="py3",
)

neo_model = pytorch_model.compile(
    target_instance_family="ml_inf1",
    input_shape={"input0": [1, 3, 224, 224]},
    output_path=compiled_model_path,
    framework="pytorch",
    framework_version="1.5.1",
    role=role,
    job_name=compilation_job_name,
)

出现编译错误：

UnexpectedStatusException: Error for Compilation job bloom-compiled-inf-inf1-202304-1921-4203: Failed. Reason: ClientError: CompilationError: Unable to compile model for ml_inf1:', 'No operations were successfully partitioned and compiled to neuron for this model - aborting trace!')

核心问题：输入input_ids尺寸为[1,1]，但past_key_values是长度为24的元组，每个元素是包含两个张量（尺寸[16, 64, 6]和[16, 6, 64]）的元组，复杂结构导致Neo无法正确识别输入形状。

可行解决方案 #

方向一：配置包含完整输入结构的input_shape #

SageMaker Neo编译需要明确所有输入的形状，包括past_key_values的每个子张量。需要将复杂的元组结构拆分为可识别的命名输入，再在推理代码中重新组装。

1. 修改编译代码的input_shape参数：

neo_model = pytorch_model.compile(
    target_instance_family="ml_inf1",
    input_shape={
        "input_ids": [1, 1],
        # 遍历24层的past_key_values，定义每个子张量的形状
        **{f"past_key_values_{i}_0": [16, 64, 6] for i in range(24)},
        **{f"past_key_values_{i}_1": [16, 6, 64] for i in range(24)}
    },
    output_path=compiled_model_path,
    framework="pytorch",
    framework_version="1.13.1",  # 升级到兼容Neuron的新版本
    role=role,
    job_name=compilation_job_name,
)

2. 在my_entry_point_file.py的predict_fn中组装输入结构：

def predict_fn(input_data, model):
    # 解析输入的input_ids
    input_ids = input_data["input_ids"]
    
    # 重新组装past_key_values元组
    past_key_values = []
    for i in range(24):
        key_tensor = input_data[f"past_key_values_{i}_0"]
        value_tensor = input_data[f"past_key_values_{i}_1"]
        past_key_values.append((key_tensor, value_tensor))
    past_key_values = tuple(past_key_values)
    
    # 执行模型推理
    with torch.no_grad():
        outputs = model(input_ids=input_ids, past_key_values=past_key_values)
    
    # 返回推理结果和更新后的past_key_values
    result = {"logits": outputs.logits}
    for i in range(24):
        result[f"past_key_values_{i}_0"] = outputs.past_key_values[i][0]
        result[f"past_key_values_{i}_1"] = outputs.past_key_values[i][1]
    return result

方向二：修改推理逻辑，取消past_key_values传递（适合短文本生成） #

如果生成任务不需要长上下文，可以改为每次输入完整文本，避免传递past_key_values，简化输入结构。

1. 修改编译代码的input_shape：

MAX_SEQ_LENGTH = 200  # 根据需求设置最大文本长度
neo_model = pytorch_model.compile(
    target_instance_family="ml_inf1",
    input_shape={"input_ids": [1, MAX_SEQ_LENGTH]},
    output_path=compiled_model_path,
    framework="pytorch",
    framework_version="1.13.1",
    role=role,
    job_name=compilation_job_name,
)

2. 修改my_entry_point_file.py的推理逻辑：

from transformers import BloomTokenizerFast

def predict_fn(input_data, model):
    tokenizer = BloomTokenizerFast.from_pretrained("bigscience/bloom-560m")
    seed_text = input_data["text"]
    MAX_GENERATE_LENGTH = 50  # 设置生成的最大token数
    
    # 初始输入编码
    inputs = tokenizer.encode(seed_text, return_tensors="pt")
    
    with torch.no_grad():
        for _ in range(MAX_GENERATE_LENGTH):
            outputs = model(input_ids=inputs)
            # 获取下一个token
            next_token = torch.argmax(outputs.logits[:, -1, :], dim=-1)
            inputs = torch.cat([inputs, next_token.unsqueeze(0)], dim=-1)
            # 遇到结束token则停止
            if next_token == tokenizer.eos_token_id:
                break
    
    generated_text = tokenizer.decode(inputs[0], skip_special_tokens=True)
    return {"generated_text": generated_text}

额外优化建议 #

• 升级PyTorch版本：你当前使用的1.5.1版本过于老旧，与Neuron SDK兼容性差，建议升级到1.12+版本（如1.13.1），能大幅降低编译兼容性问题。
• 预追踪模型：使用torch.neuron.trace提前追踪模型，明确输入示例，再上传到S3部署，比直接使用SageMaker compile更灵活：

import torch
import torch.neuron
from transformers import BloomForCausalLM

# 加载模型
model = BloomForCausalLM.from_pretrained("bigscience/bloom-560m")
# 准备示例输入
input_ids = torch.tensor([[1]])
# 生成示例past_key_values
with torch.no_grad():
    outputs = model(input_ids=input_ids)
past_key_values = outputs.past_key_values

# 追踪模型，适配Neuron
neuron_model = torch.neuron.trace(model, example_inputs=(input_ids, past_key_values))
# 保存追踪后的模型
neuron_model.save("bloom_neuron.pt")

将保存的bloom_neuron.pt上传到S3，再通过SageMaker部署即可。

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