冷启动核心原因 #

首次推理时GPU需完成算子编译、显存分配、K/V缓存初始化等耗时操作，后续调用可复用这些资源，因此预热的核心是提前触发这些流程。

适配两种推理方式的预热实现 #

1. 基于pipeline()的预热

直接用短占位prompt调用pipeline，参数与真实推理保持一致，触发pipeline内部全链路初始化：

from transformers import pipeline

generator = pipeline('text-generation', model="tiiuae/falcon-7b-instruct", device=0)

def warm_up_pipeline():
    # 用极简prompt，参数和生产推理对齐
    generator("Hello", max_new_tokens=1, temperature=0.0)
    print("Pipeline预热完成")

# 模型加载后立即执行预热
warm_up_pipeline()

def generate_text(prompt):
    return generator(prompt, max_new_tokens=50)[0]['generated_text']

2. 基于model.generate()的预热

手动完成tokenization后调用generate()，同样用最小参数触发模型侧的初始化：

from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-7b-instruct")
model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-7b-instruct", device_map="auto")

def warm_up_model():
    inputs = tokenizer("Hello", return_tensors="pt").to(model.device)
    model.generate(**inputs, max_new_tokens=1, temperature=0.0)
    print("Model预热完成")

# 加载后执行预热
warm_up_model()

def generate_text(prompt):
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(**inputs, max_new_tokens=50)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

进阶优化要点 #

• 参数对齐：预热时的推理参数（如do_sample、top_p、max_new_tokens）必须和生产环境完全一致，避免算子重复编译。
• 模型编译：PyTorch 2.0+环境下，可对模型执行model = torch.compile(model)，预热时会提前生成优化后的算子，进一步压缩首token延迟。
• 显存优化：加载模型前执行torch.cuda.empty_cache()，提前预留足够显存，避免首次推理时因显存碎片化导致的分配耗时。

预热效果验证 #

通过统计首次与二次推理的时间差，确认预热生效：

import time

# 首次推理（预热后）
start = time.time()
generate_text("What is AI?")
print(f"首次推理耗时: {time.time() - start:.2f}秒")

# 二次推理
start = time.time()
generate_text("What is machine learning?")
print(f"二次推理耗时: {time.time() - start:.2f}秒")

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