关于ChatGLM3-6B 的简单非结构化剪枝尝试

Lao_666_

649人浏览 · 2026-03-19 19:27:45

Lao_666_ · 2026-03-19 19:27:45 发布

去年我曾在本地部署过 ChatGLM3-6B，并基于它做了一个 ASR→LLM→TTS 的简单实验。最近突然想深入了解一下大模型的剪枝技术，于是决定拿这个熟悉的模型开刀，用 PyTorch 官方的 torch.nn.utils.prune 库进行一次完整的剪枝实验，并记录下整个过程和发现。大模型的权重文件地址:

Huggingface https://huggingface.co/THUDM/chatglm3-6b

阿里云:https://www.modelscope.cn/models/ZhipuAI/chatglm3-6b

基本思路，先用原始的模型权重文件跑简单对话：

首先配置基础，我用的RTX5080,支持的是sm-120,去年的时候有些torch还不支持，想要知道怎么给 RTX5080配置sm-120的可以去看这篇博文:RTX50系列显卡SM_120 架构兼容性问题-CSDN博客

配置:

torch 2.11.0.dev20260105+cu128

pytorch-cuda 12.4

transformers 4.38.2

首先配置基础:

# ========== 基础配置 ==========
LOCAL_MODEL_PATH = r"C:\Users\k\models\ZhipuAI\chatglm3-6b"  # 你的模型路径
DEVICE = "cuda"                     # 固定为 GPU
QUANTIZE = 0

模型路径填自己的路径，如果没有下载也可以直接填官方网址会自动下。

然后加载模型，用pytorch的加载方式：

def init_model_and_tokenizer():
    """加载模型和分词器（CPU 优化）"""
    print("正在加载分词器...")
    tokenizer = AutoTokenizer.from_pretrained(
        LOCAL_MODEL_PATH,
        trust_remote_code=True,
        use_fast=True
    )
    
    print("正在加载模型（CPU 模式，可能需要几分钟）...")
    model = AutoModel.from_pretrained(
        LOCAL_MODEL_PATH,
        trust_remote_code=True,
        low_cpu_mem_usage=True,
        torch_dtype=torch.float16,      
    )
    
    model = model.to(DEVICE).eval()
    print("模型加载完成。")
    return tokenizer, model

这里torch_dtype在更高的transformer里面似乎改成了dtype,这里兼容我的transfotmer版本，保留了原写法。这里torch.float16用的16位浮点，是为了减少内存占用。试过，32位时候大概占用内存是25G左右，官方给出的是24G 多点，差不多。16FP的话只占用12G内存，并且可以加快速度效果非常明显。

接下来是简单的问答逻辑界面：

def chat_loop(tokenizer, model):
    """基础对话循环，无system prompt，无记忆"""
    history = []  # 初始为空，不包含任何system消息
    
    print("\n===== 原始模型对话（无预设角色）=====")
    print("输入 'exit' 退出对话\n")
    
    while True:
        user_input = input("👤 你: ").strip()
        if user_input.lower() == "exit":
            print("👋 再见！")
            break
        if not user_input:
            continue
        
        try:
            # 调用 model.chat 生成回复，不传入system参数
            response, history = model.chat(
                tokenizer,
                user_input,
                history=history,
                top_p=0.9,
                temperature=0.8
                # 注意：没有 system 参数
            )
            
            print(f"🤖 模型: {response}\n")
        
        except Exception as e:
            print(f"❌ 出错: {e}")

整体就是这样。模型的对话脚本就好了。

然后是剪枝。官方给出模型有28层，实际用layer输出也是28层对应，索引是0到27。

定义剪枝函数：


def prune_model_layer(model, amount, layer_index=19):
    target_prefix = f"transformer.encoder.layers.{layer_index}."
    parameters_to_prune = []
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear) and name.startswith(target_prefix):
            parameters_to_prune.append((module, 'weight'))
    
    if not parameters_to_prune:
        print(f"警告：未找到层 {layer_index} 中的 Linear 模块")
        return
    
    # 执行全局剪枝
    prune.global_unstructured(
        parameters_to_prune,
        pruning_method=prune.L1Unstructured,
        amount=amount,
    )
    
    # 使剪枝永久化，移除掩码，权重变为普通参数（含零值）
    for module, name in parameters_to_prune:
        prune.remove(module, name)
    
    print(f"层 {layer_index} 剪枝完成，剪枝比例: {amount*100:.1f}%")

这里只试验了一下剪枝第19层，剪枝度为0.5.

也可以用数组prune_model_index存储层的索引。然后主函数里for循环剪枝。

这里设置的是prune.L1Unsructured剪枝方式，也就是L1非结构化剪枝。L1剪枝对于输出差异大的模型剪枝效果会好一点。

剪枝好的模型放到工作目录文件下：

def save_pruned_model(model, save_path="pruned_model.pth"):
    """保存剪枝后的 state_dict（包含掩码）"""
    torch.save(model.state_dict(), save_path)
    print(f"模型已保存到: {save_path}")

掩码mask是这样一个概念，就是剪枝完之后不直接修改原模型权重，而是生成一个对应的mask矩阵，和原模型的参数矩阵耦合，让某些权重变成0。如果要固定，要调用prune.remove。

测试一下剪枝结果：

def load_pruned_model(model_class, model_path, prune_state_dict_path, device='cpu'):
    """
    加载原始模型结构，然后加载剪枝后的 state_dict
    """
    model = model_class.from_pretrained(
        model_path,
        trust_remote_code=True,
        low_cpu_mem_usage=True,
        torch_dtype=torch.float16,## 16位浮点数，减少内存占用和计算时间
    ).to(device)
    model.load_state_dict(torch.load(prune_state_dict_path, map_location=device))
    return model

def main():
    # 1. 加载原始模型（CPU 上，float32）
    print("正在加载原始模型...")
    model = AutoModel.from_pretrained(
        LOCAL_MODEL_PATH,
        trust_remote_code=True,
        low_cpu_mem_usage=True,
        torch_dtype=torch.float16,  # 加载为 float16 以节省内存
    )
    model.eval()

    # 可选：将模型移至 GPU 以加速剪枝（如果有）
    # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # model = model.to(device)

    # 2. 打印剪枝前的参数信息
    total_params_before = sum(p.numel() for p in model.parameters())
    print(f"剪枝前参数量: {total_params_before/1e9:.2f} B")

    # 3. 执行剪枝
    for layer_index in prune_model_index:
        prune_model_layer(model, amount=0.5, layer_index=layer_index)
    

    # 4. 打印剪枝后的参数信息（注意：非结构化剪枝不改变参数量，但部分权重变为零）
    total_params_after = sum(p.numel() for p in model.parameters())
    nonzero_params = sum(torch.count_nonzero(p) for p in model.parameters())
    print(f"剪枝后参数量: {total_params_after/1e9:.2f} B")
    print(f"非零参数数量: {nonzero_params/1e9:.2f} B")
    print(f"实际稀疏度: {1 - nonzero_params/total_params_after:.2%}")

    # 5. 保存剪枝后的模型
    save_pruned_model(model, "chatglm3_pruned.pth")

    # 6. （可选）演示加载剪枝后的模型

输出结果:

(llm_micro) C:\Users\k\LLM_Mciro\build>python model_prune.py
正在加载原始模型...
Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 23.48it/s]
剪枝前参数量: 6.24 B
层 19 剪枝完成，剪枝比例: 50.0%
剪枝后参数量: 6.24 B
非零参数数量: 6.14 B
实际稀疏度: 1.63%
模型已保存到: chatglm3_pruned.pth

降低了0.1B个参数。

测试输入，我想先测试情感类文本生成，所以问题是：写一首讲爱而不得的小诗。

加载模型的时候，剪枝过的模型的transformer和分词器从原来的地方加载，模型权重从工作目录加载。

LOCAL_MODEL_PATH = r"C:\Users\k\models\ZhipuAI\chatglm3-6b"  # 原始模型路径（用于加载配置和分词器）
PRUNED_MODEL_PATH = "chatglm3_pruned.pth"                    # 剪枝后的权重文件（位于工作目录）
DEVICE = "cuda"
def init_model_and_tokenizer():
    """加载分词器和剪枝后的模型"""
    print("正在加载分词器...")
    tokenizer = AutoTokenizer.from_pretrained(
        LOCAL_MODEL_PATH,
        trust_remote_code=True,
        use_fast=True
    )
    
    print("正在加载剪枝后的模型（CPU 模式，可能需要几分钟）...")
    # 1. 加载原始模型结构（不加载权重，仅用于构建）
    model = AutoModel.from_pretrained(
        LOCAL_MODEL_PATH,
        trust_remote_code=True,
        low_cpu_mem_usage=True,
        torch_dtype=torch.float16,      # CPU 必须用 float16
    )
    
    # 2. 加载剪枝后的权重
    state_dict = torch.load(PRUNED_MODEL_PATH, map_location="cuda", weights_only=True)
    model.load_state_dict(state_dict)
    
    model = model.to(DEVICE).eval()
    print("剪枝模型加载完成。")
    return tokenizer, model

输入写一首讲爱而不得的小诗后，剪枝后的模型在思考时间上明显超过非剪枝模型。

神奇的是，剪枝模型在第一次回答的质量上表现超过了非剪枝模型。

epoch 1
👤 你: 写一首讲爱而不得的小诗
🤖 模型: 爱而不得,心痛难忍,
看着你微笑,心中酸涩难言。
你的眼神,如明媚的春天,
而我却只能,远远地望着。

我们的距离,如海洋般遥远,
你的身影,在我心中飘散。
我想要的,是那份真挚的情感,
却只能,在心里默默思念。

或许有一天,你会有机会明白,
我的心意,早已深入骨髓。
但到那时,你可能已经,
 gone,离开我的世界。

爱而不得,让我孤独难耐,
我只能,将这份爱,藏在心底。
愿你能幸福快乐,找到你心中的他,
而我,也会默默地祝福你,永远。

其余几轮的输出基本水平大概是这样：

epoch 2
🤖 模型: 爱而不得，心伤难愈，
夜夜思量，泪如雨落。

繁星闪烁，似你眼眸，
寂静星空，唤不回的温柔。

情似网，编制的太紧，
心 general，如何破茧成蝶？

情可伤，心可碎，
小草绿，花儿败，

爱而不得，苦涩入喉，
只能沉默，忍受折磨。

望穿秋水，望断春晖，
情爱之间，难以取舍。

只能祝福，愿你幸福，
自己的痛苦，自己承受。

爱而不得，心伤难愈，
夜夜思量，泪如雨落。

出现了一些语法上的问题，但是剪枝与非剪枝的模型的表现相当。

接下来可以从ChatGLM3-6B 的官方github库下载微调的方法，准备之后再做：

zai-org/ChatGLM3: ChatGLM3 series: Open Bilingual Chat LLMs | 开源双语对话语言模型