应用 2:4 稀疏性并可选 FP8 量化
此脚本演示了如何使用 llm-compressor 库，对 Meta-Llama-3-8B-Instruct 模型应用2:4 结构化稀疏性（可选 FP8 量化）。压缩后的模型针对内存效率和支持 GPU 的推理速度进行了优化。
注意： FP8 动态精度计算支持 CUDA 计算能力 9.0 及以上的 Nvidia GPU。
安装
要开始，请安装 llm-compressor 库及其依赖项。
快速入门
使用以下命令运行脚本：
- 不带 FP8 量化
- 带 FP8 量化
该脚本使用以下方法压缩 Meta-Llama-3-8B-Instruct 模型：
- 2:4 结构化稀疏性： 应用结构化剪枝，将权重减少 50%。
- FP8 量化（可选）： 启用动态量化以进一步节省内存。
配置
- 模型：meta-llama/Meta-Llama-3-8B-Instruct
- 数据集：HuggingFaceH4/ultrachat_200k（训练集）
- 校准样本：512
- 最大序列长度：2048
运行步骤
- 选择模型、数据集和策略
模型和数据集已在脚本中预定义。策略会根据是否启用了 FP8 量化（--fp8 标志）动态调整。
示例策略
recipe = [
SparseGPTModifier(
sparsity=0.5,
mask_structure="2:4",
targets=[r"re:model.layers.\d*$"],
)
]
if fp8_enabled:
recipe.append(
QuantizationModifier(
targets=["Linear"],
ignore=["lm_head"],
scheme="FP8_DYNAMIC",
),
)
- 应用压缩 脚本使用 oneshot 函数应用压缩。
oneshot(
model=model,
dataset=ds,
recipe=recipe,
max_seq_length=2048,
num_calibration_samples=512,
)
保存压缩后的模型
压缩后的模型和 tokenizer 将被保存到输出目录。
输出目录：- 无 FP8：Meta-Llama-3-8B-Instruct-2of4-sparse - 带 FP8：Meta-Llama-3-8B-Instruct-2of4-W8A8-FP8-Dynamic-Per-Token
不进行稀疏压缩的保存
要将模型保存到磁盘而不进行稀疏压缩：
model.save_pretrained(save_dir, save_compressed=True, disable_sparse_compression=True)
tokenizer.save_pretrained(save_dir)
注意： 同时使用
save_compressed和disable_sparse_compression选项保存模型时，将使用量化压缩器压缩模型；但是，为了节省磁盘空间，将使用密集稀疏压缩器，而不是更节省磁盘空间的稀疏压缩器。dense稀疏压缩器按原样保存模型参数，不利用稀疏性来节省磁盘空间。这些选项仅影响模型在磁盘上的保存方式，而不影响实际的剪枝或量化过程。
验证
压缩后，脚本将通过生成示例输出来验证模型。
========== SAMPLE GENERATION ============
Hello my name is ...
=========================================
注意： - 确保您的 GPU 有足够的内存（至少约 25GB）来运行压缩脚本。 - 使用 --fp8 标志来启用 FP8 量化。
修改 MODEL_ID 和 DATASET_ID 变量以使用其他模型或数据集。
在 vLLM 中运行
使用 pip install vllm 安装 vLLM。
# run_model.py
import argparse
from vllm import LLM, SamplingParams
def run_inference(model_path, tensor_parallel_size, prompt="Hello my name is:"):
"""
Loads a model and performs inference using LLM.
"""
# Define sampling parameters
sampling_params = SamplingParams(
temperature=0.8,
top_p=0.95,
)
# Load the model
model = LLM(
model=model_path,
enforce_eager=True,
dtype="auto",
tensor_parallel_size=tensor_parallel_size,
)
# Generate inference
outputs = model.generate(prompt, sampling_params=sampling_params)
return outputs[0].outputs[0].text
def main():
""" Main function to handle CLI and process the model. """
# Argument parsing
parser = argparse.ArgumentParser(description="Run inference on a single model and print results.")
parser.add_argument(
"model_path",
type=str,
help="Path to the model to perform inference."
)
parser.add_argument(
"--tensor_parallel_size",
type=int,
default=1,
help="Tensor parallel size for the model. Default is 1."
)
args = parser.parse_args()
model_path = args.model_path
tensor_parallel_size = args.tensor_parallel_size
prompt = "Hello my name is:"
# Load model and perform inference
inference_result = run_inference(model_path, tensor_parallel_size)
print("="* 20)
print("Model:", model_path)
print(prompt, inference_result)
if __name__ == "__main__":
main()
运行模型的命令
示例