ChatGLM-6B轻量化LLM实战：如何在普通GPU上高效运行文本生成与问答任务

1次阅读

共计 2130 个字符，预计需要花费 6 分钟才能阅读完成。

大语言模型如 GPT-3、ChatGLM 等在 NLP 任务中表现出色，但它们的参数量通常达到数十亿甚至上千亿级别。在普通 GPU（如 RTX 3090 24GB）上部署这些模型时，开发者常遇到以下问题：

显存不足（OOM）：全精度模型参数占用显存远超消费级 GPU 容量
计算延迟高 ：大矩阵运算导致响应时间难以满足实时交互需求
能耗过大 ：高功耗不适合边缘设备部署

ChatGLM-6B 通过轻量化技术，将模型压缩到可在 16GB 显存的 GPU 上流畅运行，同时保持 90% 以上的原始模型性能。

方案	参数量	压缩方法	适用场景
ChatGLM-6B	6B	量化 + 剪枝	通用 NLP 任务
TinyBERT	14M	知识蒸馏	移动端文本分类
DistilGPT	82M	层数减半 + 蒸馏	英文文本生成

ChatGLM-6B 的优势在于：

保留完整的解码器架构，支持生成式任务
中文优化效果更好
支持动态批处理

结构化剪枝 ：移除整个注意力头或 FFN 层

# 示例：移除 50% 的注意力头
for layer in model.transformer.layers:
    layer.attention.num_heads = layer.attention.num_heads // 2

非结构化剪枝 ：按权重绝对值剪枝

import torch.nn.utils.prune as prune
parameters_to_prune = [(layer, 'weight') for layer in model.transformer.layers]
prune.global_unstructured(parameters_to_prune, pruning_method=prune.L1Unstructured, amount=0.3)

from transformers import AutoModelForCausalLM
import torch

# FP16 量化
model = AutoModelForCausalLM.from_pretrained("THUDM/chatglm-6b", torch_dtype=torch.float16)

# INT8 动态量化（需要安装 intel_extension_for_pytorch）model = torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8
)

梯度检查点 ：用计算换显存
```
model.gradient_checkpointing_enable()
```

动态加载 ：分片加载模型参数

model = AutoModel.from_pretrained("THUDM/chatglm-6b", 
                                device_map="auto",
                                offload_folder="offload")

from transformers import AutoTokenizer, pipeline
import torch

# 1. 加载量化模型
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    "THUDM/chatglm-6b",
    torch_dtype=torch.float16,
    device_map="auto"
)

# 2. 创建生成管道
generator = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device=0,
    max_new_tokens=256
)

# 3. 执行推理
input_text = "如何学习 Python 编程?"
output = generator(input_text, do_sample=True, temperature=0.7)
print(output[0]["generated_text"])

测试环境：RTX 3060 12GB