使用transformers.pipeline()编译Llama模型推理无日志问题排查求助

你尝试用transformers.pipeline()结合torch.compile()运行Llama模型推理，代码如下：

model = LlamaForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-70B",use_cache=True,device_map='auto')
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-70B",use_cache=True,truncation=True,padding="max_length",max_length=64,return_tensors="pt")
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "left"

model = torch.compile(model)
pipeline = transformers.pipeline(
                     "text-generation",
                       model=model,
                       model_kwargs={"torch_dtype": torch.bfloat16},
                       tokenizer = tokenizer,
                       device_map="auto",
                        )
            
generation_config = {
                        "num_beams":1,
                        "max_new_tokens":32,
                        "do_sample":True,
                        "use_cache":True,
                        }
outputs = pipeline(input_prompt,**generation_config)

预期会看到torch.compile()的编译日志，但无论用model = torch.compile(model)还是pipeline.model = torch.compile(pipeline.model)，都没有编译日志输出，也无报错，需要解决代码问题并掌握正确用法。

问题分析与修正步骤 #

1. 核心问题：编译时机与pipeline的冲突 #

当你把已编译的模型传入pipeline时，device_map="auto"会触发Transformers的模型重分配逻辑，直接覆盖掉之前的编译结果；同时model_kwargs={"torch_dtype": torch.bfloat16}会再次修改模型 dtype，破坏已编译的模型结构，导致编译完全失效。

另外，PyTorch默认不会输出编译日志，需要手动开启调试配置才能看到编译过程信息。

2. 正确实现代码 #

import torch
from transformers import LlamaForCausalLM, AutoTokenizer, pipeline
import logging

# 1. 加载模型与tokenizer，提前设置好dtype，暂不分配设备
model = LlamaForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-70B",
    use_cache=True,
    torch_dtype=torch.bfloat16  # 提前指定dtype，避免后续修改破坏编译
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-70B")
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "left"

# 2. 先编译模型，再分配设备
model = torch.compile(model, mode="reduce-overhead")  # 生成任务推荐用reduce-overhead模式
model = model.to("cuda")  # 手动分配设备，避免pipeline的device_map干扰

# 3. 创建pipeline，不传入device_map和model_kwargs
text_gen_pipeline = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer
)

# 4. 开启编译日志输出
logging.basicConfig(level=logging.INFO)
torch._logging.set_logs(compile=True)

# 5. 执行推理
generation_config = {
    "num_beams": 1,
    "max_new_tokens": 32,
    "do_sample": True,
    "use_cache": True,
}
outputs = text_gen_pipeline("你的输入prompt", **generation_config)
print(outputs)

3. 关键注意事项 #

• 编译顺序：必须先设置好模型的dtype，再执行torch.compile()，最后分配设备，避免后续操作覆盖编译结果。
• pipeline参数：创建pipeline时不要传入device_map="auto"和model_kwargs，否则会重新初始化模型，导致编译失效。
• 日志开启：通过torch._logging.set_logs(compile=True)和调整logging级别，才能看到编译相关的日志输出。
• 编译模式：文本生成属于循环推理场景，用mode="reduce-overhead"比默认模式更适配，能减少编译带来的额外开销。
• 可选优化：可以配合model.to_bettertransformer()调整模型结构，让其更适配PyTorch编译，进一步提升生成效率。

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