Llama2/Falcon在AWS Inf2/G4dn实例推理可行性及选型咨询

问题背景 #

我们在AWS G5.8xl实例上成功完成Falcon/Falcoder大模型的推理实验，但将相同代码迁移到Inf2.8xlarge实例时运行失败。已知Inf2实例采用AWS Neuron加速器而非NVIDIA GPU，我们尝试通过torch-neuronx库适配Neuron核心能力，但仍出现报错。

未使用torch-neuronx的生成代码及报错 #

generation_output = model.generate(
input_ids = input_ids,
attention_mask = attention_mask,
generation_config = generation_config,
return_dict_in_generate = True,
output_scores = False,
max_new_tokens = max_new_tokens,
early_stopping = True
)
#print("generation_output")
#print(generation_output)
s = generation_output.sequences[0]
output = tokenizer.decode(s)

对应错误栈：未适配Neuron加速器时，直接调用model.generate()出现与GPU依赖相关的运行错误。

使用torch-neuronx适配后的代码及报错 #

生成样例输入函数 #

def generate_sample_inputs(tokenizer, sequence_length):
    dummy_input = "dummy"
    embeddings = tokenizer(dummy_input, max_length=sequence_length,      
                           padding="max_length",return_tensors="pt")
    return tuple(embeddings.values())

模型编译函数 #

def compile_model_inf2(model, tokenizer, sequence_length, num_neuron_cores):
    #use only one neuron core
    os.environ["NEURON_RT_NUM_CORES"] = str(num_neuron_cores)
    import torch_neuronx
    payload = generate_sample_inputs(tokenizer, sequence_length)
    return torch_neuronx.trace(model, payload)

编译调用代码 #

model = compile_model_inf2(model, tokenizer, sequence_length=512, num_neuron_cores=1)

对应错误栈：使用torch-neuronx.trace()编译模型时，出现算子不支持、编译流程失败相关错误。

疑问与核心诉求 #

• 疑问：关于text-generation-webui适配Neuron加速器的GitHub issue是否能解决当前报错？
• 核心诉求：
  1. 确认Llama2/Falcon模型在Inf2.8xlarge、G4dn.8xlarge实例上是否支持推理；
  2. 若上述实例不支持，推荐高性价比的替代实例。

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