Skip to content

optimizer.md

Latest commit

 

History

History
150 lines (125 loc) · 4.7 KB

optimizer.md

File metadata and controls

150 lines (125 loc) · 4.7 KB

Hugging Face 模型推理优化指南

本文档整合了多种优化策略，帮助你在硬件条件有限的情况下提升 Hugging Face 模型的推理效率，同时降低显存占用。主要内容包括混合精度计算、生成参数调整、使用 Pipeline、量化技术、内存管理、设备选择、批量处理等各个方面。

1. 混合精度计算（自动混合精度）

启用混合精度可以减少计算所需的显存，并提高推理速度。PyTorch 提供了 torch.cuda.amp.autocast() 来自动切换计算精度。
例如，在推理过程中：

from torch.cuda.amp import autocast

with autocast():
    generated_ids = model.generate(inputs['input_ids'], ...)

2. 调整生成参数

生成文本时，通过调整参数可以在保证效果的同时减少计算资源占用：

max_length：限制生成文本的最大长度。较短的生成长度会减少计算量。 temperature：调控生成的随机性，较低的值通常生成更保守的结果。 top_k 和 top_p：限制候选词的范围，减少计算负担。 例如：

generated_ids = model.generate(
    inputs['input_ids'], 
    max_length=100, 
    num_return_sequences=1, 
    temperature=0.7, 
    top_k=50, 
    top_p=0.9
)

3. 使用 Hugging Face Pipeline 简化代码

Hugging Face 的 pipeline 封装了模型加载、tokenization 以及推理的常见步骤，能大幅简化代码。

from transformers import pipeline

generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0)
response = generator("hi, how are you?", max_length=200, temperature=0.7)

4. 量化技术优化

量化可以将模型权重从高精度转换为低精度，从而降低显存占用并提升推理速度。不过量化可能会影响生成质量，需要权衡使用。

4.1. 启用 8-bit 量化 通过设置 load_in_8bit=True 可以启用 8-bit 量化。

model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.float16,
    device_map="auto",
    load_in_8bit=True  # 启用 8-bit 量化
)

4.2. 双重量化优化 在启用 8-bit 量化时，设置 bnb_4bit_use_double_quant=True 进一步优化显存占用。

model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.float16,
    device_map="auto",
    load_in_8bit=True,
    bnb_4bit_use_double_quant=True  # 启用双重量化优化
)

4.3. 使用 4-bit 量化 当显存非常紧张时，可以尝试 4-bit 量化，并配置相应参数：

model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.float16,
    device_map="auto",
    load_in_8bit=True,
    bnb_4bit_compute_dtype=torch.float16,  # 计算时使用 float16
    bnb_4bit_quant_type="nf4"              # 使用 nf4 量化类型
)

5. 内存管理优化

5.1. 清理显存缓存 在多次推理后，调用 torch.cuda.empty_cache() 可释放未使用的显存。

torch.cuda.empty_cache()

5.2. 优化 CUDA 内存分配策略 设置环境变量 PYTORCH_CUDA_ALLOC_CONF 来调整内存分配策略，例如：

os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"

6. 设备选择与模型加载优化

确保模型加载到合适的设备（例如 GPU），并利用 device_map="auto" 自动分配模型到可用设备。

device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)

model = AutoModelForCausalLM.from_pretrained(
    model_path,
    device_map="auto"  # 自动将模型分配到多个 GPU 或 CPU
)

7. 批量处理输入

将多个输入文本合并成一个批次（batch）处理，可以充分利用 GPU 的并行计算优势，提高整体推理效率。

input_texts = ["hi, how are you?", "What is the weather today?"]
inputs = tokenizer(input_texts, return_tensors="pt", padding=True, truncation=True).to(device)

with torch.no_grad():
    generated_ids = model.generate(
        inputs['input_ids'], 
        max_length=50, 
        num_return_sequences=1, 
        attention_mask=inputs['attention_mask']
    )
responses = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)

8. 其他优化策略

8.1. 延迟加载模型 使用 cache_dir 参数来指定缓存路径，避免每次都从远程加载模型。

model = AutoModelForCausalLM.from_pretrained(
    model_path,
    cache_dir="path/to/cache"
)

8.2. 调整束搜索参数 使用束搜索（beam search）并调整 num_beams 参数，可以在生成质量和计算开销之间取得平衡。

generated_ids = model.generate(
    inputs['input_ids'], 
    max_length=100, 
    num_beams=3, 
    no_repeat_ngram_size=2
)

8.3. 缓存 Tokenizer 结果 对于重复输入，缓存 tokenization 结果可以避免重复计算，提升整体效率。