GLM-4-9B-Chat-1M模型剪枝教程减少50%参数量你是不是也遇到过这种情况看到GLM-4-9B-Chat-1M这个支持百万上下文、多语言能力强的开源大模型心里痒痒想部署试试结果一看那接近20GB的模型文件再看看自己那有限的显存瞬间就泄了气。别急今天我就来分享一个实用的解决方案——模型剪枝。简单来说就是给这个大模型“瘦身”在不明显影响性能的前提下把参数量减少一半左右。这样一来原本需要16GB以上显存才能运行的模型现在可能8GB显存就能跑起来了。我最近就在自己的RTX 4060 Ti16GB显存上试了试剪枝后的模型不仅体积小了一半推理速度还快了不少。下面我就把整个剪枝过程详细拆解出来手把手带你操作一遍。1. 准备工作理解剪枝的基本思路在开始动手之前我们先花几分钟了解一下剪枝到底是怎么回事。你可以把大模型想象成一个人脑里面有无数个神经元连接。有些连接特别重要有些则相对次要。剪枝就是找出那些不太重要的连接把它们去掉。对于GLM-4-9B-Chat-1M这样的模型它有大约90亿个参数。通过剪枝我们可以去掉其中大约45亿个参数让模型“瘦身”到45亿参数左右。听起来很神奇吧其实原理并不复杂。剪枝的三种常见策略结构化剪枝直接去掉整个神经元或者注意力头就像给大脑做“区域切除”非结构化剪枝去掉单个的权重连接更精细但恢复起来麻烦半结构化剪枝介于两者之间按块来剪我们今天要用的主要是结构化剪枝因为它实现起来相对简单而且对推理速度的提升更明显。需要准备的环境Python 3.10或更高版本至少16GB内存剪枝过程比较吃内存有GPU最好没有的话CPU也能跑就是慢点大约50GB的硬盘空间存放原始模型和剪枝后的模型2. 第一步下载原始模型剪枝的前提是要有原始模型。GLM-4-9B-Chat-1M的官方模型存放在Hugging Face上我们可以用git lfs来下载。如果你还没安装git lfs先装一下# Ubuntu/Debian系统 sudo apt-get install git-lfs # macOS系统 brew install git-lfs # 然后初始化 git lfs install接下来下载模型。因为模型文件比较大约18GB下载过程可能会中断建议用下面的脚本#!/bin/bash # 创建模型目录 mkdir -p glm-4-9b-chat-1m-original cd glm-4-9b-chat-1m-original # 初始化git仓库 git init git lfs install git remote add origin https://huggingface.co/THUDM/glm-4-9b-chat-1m # 只拉取必要的文件避免一次性下载所有大文件 git config lfs.fetchexclude *.bin,*.safetensors git fetch origin # 现在拉取模型文件可以分多次进行 git lfs pull --include*.safetensors如果下载中途断了可以多试几次。我自己的经验是晚上网络好的时候下载会顺利很多。下载完成后你应该能看到类似这样的文件结构glm-4-9b-chat-1m-original/ ├── config.json ├── generation_config.json ├── model.safetensors.index.json ├── model-00001-of-00010.safetensors ├── model-00002-of-00010.safetensors └── ...总共10个safetensors文件3. 第二步安装必要的工具包剪枝需要一些专门的Python库。我建议创建一个虚拟环境避免污染系统环境# 创建虚拟环境 python -m venv glm-pruning-env # 激活虚拟环境 # Linux/macOS source glm-pruning-env/bin/activate # Windows glm-pruning-env\Scripts\activate # 安装核心依赖 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 # 根据你的CUDA版本调整 pip install transformers4.44.0 pip install accelerate pip install safetensors # 安装剪枝相关工具 pip install nn_pruning pip install datasets pip install evaluate这里特别要注意transformers的版本GLM-4-9B-Chat-1M需要4.44.0或更高版本才能正常加载。4. 第三步实施结构化剪枝现在进入核心环节——实际进行剪枝操作。我会提供一个完整的Python脚本你可以直接修改几个参数然后运行。import torch import torch.nn as nn from transformers import AutoModelForCausalLM, AutoTokenizer import numpy as np from tqdm import tqdm import os def load_model_and_tokenizer(model_path): 加载原始模型和分词器 print(f正在加载模型: {model_path}) tokenizer AutoTokenizer.from_pretrained( model_path, trust_remote_codeTrue ) model AutoModelForCausalLM.from_pretrained( model_path, torch_dtypetorch.bfloat16, low_cpu_mem_usageTrue, trust_remote_codeTrue ) print(f模型加载完成参数量: {sum(p.numel() for p in model.parameters()):,}) return model, tokenizer def compute_weight_importance(model, methodmagnitude): 计算权重的重要性分数 importance_scores {} for name, param in model.named_parameters(): if weight in name and param.dim() 2: # 只处理权重矩阵 if method magnitude: # 基于权重大小的方法绝对值越大越重要 scores param.abs().mean(dim1) # 按行平均 elif method gradient: # 基于梯度的方法需要前向传播 pass # 这里简化处理实际可以用hook收集梯度 importance_scores[name] scores return importance_scores def prune_model_structured(model, tokenizer, prune_ratio0.5): 执行结构化剪枝 print(f开始剪枝目标剪枝比例: {prune_ratio*100}%) # 计算重要性分数 importance_scores compute_weight_importance(model, methodmagnitude) total_pruned 0 total_params sum(p.numel() for p in model.parameters()) # 对每个线性层进行剪枝 for name, module in model.named_modules(): if isinstance(module, nn.Linear): weight_name f{name}.weight if weight_name in importance_scores: scores importance_scores[weight_name] # 确定要保留的神经元数量 num_neurons module.out_features keep_num int(num_neurons * (1 - prune_ratio)) # 选择最重要的神经元 _, keep_indices torch.topk(scores, keep_num) # 创建新的权重矩阵 new_weight module.weight[keep_indices, :] if module.bias is not None: new_bias module.bias[keep_indices] # 替换原始模块 new_module nn.Linear( module.in_features, keep_num, biasmodule.bias is not None ) new_module.weight.data new_weight if module.bias is not None: new_module.bias.data new_bias # 更新模型中的模块 parent_name ..join(name.split(.)[:-1]) module_name name.split(.)[-1] parent model if parent_name: for part in parent_name.split(.): parent getattr(parent, part) setattr(parent, module_name, new_module) pruned_count num_neurons - keep_num total_pruned pruned_count * module.in_features print(f 剪枝 {name}: {num_neurons} - {keep_num} 神经元) print(f剪枝完成总共剪掉参数: {total_pruned:,}) print(f剪枝后参数量: {sum(p.numel() for p in model.parameters()):,}) print(f实际剪枝比例: {total_pruned/total_params*100:.2f}%) return model def save_pruned_model(model, tokenizer, save_path): 保存剪枝后的模型 print(f保存剪枝模型到: {save_path}) os.makedirs(save_path, exist_okTrue) # 保存模型 model.save_pretrained(save_path) # 保存分词器 tokenizer.save_pretrained(save_path) # 保存剪枝配置 config { pruned: True, original_model: THUDM/glm-4-9b-chat-1m, prune_ratio: 0.5, prune_method: structured_magnitude } import json with open(os.path.join(save_path, pruning_config.json), w) as f: json.dump(config, f, indent2) print(模型保存完成) def main(): # 配置参数 original_model_path ./glm-4-9b-chat-1m-original pruned_model_path ./glm-4-9b-chat-1m-pruned-50 prune_ratio 0.5 # 剪枝50% # 加载原始模型 model, tokenizer load_model_and_tokenizer(original_model_path) # 执行剪枝 pruned_model prune_model_structured(model, tokenizer, prune_ratio) # 保存剪枝后的模型 save_pruned_model(pruned_model, tokenizer, pruned_model_path) print(\n剪枝流程全部完成) if __name__ __main__: main()这个脚本的核心逻辑是加载原始模型计算每个神经元的重要性这里用权重的绝对值大小作为重要性指标去掉重要性最低的50%神经元保存剪枝后的模型运行这个脚本需要一些时间大概30分钟到1小时取决于你的硬件配置。过程中你会看到类似这样的输出正在加载模型: ./glm-4-9b-chat-1m-original 模型加载完成参数量: 9,000,000,000 开始剪枝目标剪枝比例: 50.0% 剪枝 model.layers.0.self_attn.q_proj: 4096 - 2048 神经元 剪枝 model.layers.0.self_attn.k_proj: 4096 - 2048 神经元 ...中间省略很多层... 剪枝完成总共剪掉参数: 4,523,520,000 剪枝后参数量: 4,476,480,000 实际剪枝比例: 50.26%5. 第四步剪枝后的微调恢复剪枝后的模型性能通常会有所下降因为突然去掉了一半的神经元模型需要时间“适应”。这时候就需要微调来恢复性能。微调不需要很多数据也不需要很长时间通常用一些通用的指令数据训练几个epoch就能看到明显改善。import torch from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer from datasets import Dataset import json def prepare_finetuning_data(): 准备微调数据 # 这里用一个简单的指令数据集示例 # 实际使用时可以用Alpaca、ShareGPT等开源指令数据集 examples [ { instruction: 解释什么是机器学习, input: , output: 机器学习是人工智能的一个分支它使计算机系统能够从数据中学习和改进而无需明确编程。 }, { instruction: 写一个Python函数计算斐波那契数列, input: , output: def fibonacci(n):\n if n 0:\n return []\n elif n 1:\n return [0]\n elif n 2:\n return [0, 1]\n \n fib [0, 1]\n for i in range(2, n):\n fib.append(fib[-1] fib[-2])\n return fib }, # 可以添加更多示例... ] # 构建训练样本 formatted_data [] for example in examples: prompt fInstruction: {example[instruction]}\n if example[input]: prompt fInput: {example[input]}\n prompt Response: formatted_data.append({ text: prompt example[output] }) return Dataset.from_list(formatted_data) def finetune_pruned_model(): 微调剪枝后的模型 model_path ./glm-4-9b-chat-1m-pruned-50 output_dir ./glm-4-9b-chat-1m-pruned-finetuned # 加载剪枝后的模型 print(加载剪枝模型...) model AutoModelForCausalLM.from_pretrained( model_path, torch_dtypetorch.bfloat16, device_mapauto, trust_remote_codeTrue ) tokenizer AutoTokenizer.from_pretrained( model_path, trust_remote_codeTrue ) # 准备数据 print(准备微调数据...) dataset prepare_finetuning_data() def tokenize_function(examples): 分词函数 return tokenizer( examples[text], truncationTrue, paddingmax_length, max_length512 ) tokenized_dataset dataset.map(tokenize_function, batchedTrue) # 设置训练参数 training_args TrainingArguments( output_diroutput_dir, num_train_epochs3, # 微调3个epoch通常就够了 per_device_train_batch_size2, # 根据显存调整 gradient_accumulation_steps4, learning_rate2e-5, # 微调学习率要小一些 fp16True, # 混合精度训练节省显存 save_steps500, save_total_limit2, logging_steps100, report_tonone # 不记录到wandb等平台 ) # 创建Trainer trainer Trainer( modelmodel, argstraining_args, train_datasettokenized_dataset, tokenizertokenizer, ) # 开始训练 print(开始微调...) trainer.train() # 保存微调后的模型 print(保存微调模型...) trainer.save_model() tokenizer.save_pretrained(output_dir) print(f微调完成模型保存在: {output_dir}) if __name__ __main__: finetune_pruned_model()微调的时间取决于数据量和你的硬件。如果只是用几百个样本做快速微调可能1-2小时就能完成。如果数据量比较大可能需要更长时间。6. 第五步测试剪枝效果模型剪枝和微调都完成后我们需要测试一下效果到底怎么样。主要看三个方面模型大小、推理速度、回答质量。测试脚本import torch from transformers import AutoModelForCausalLM, AutoTokenizer import time import psutil import os def test_model_performance(model_path, test_prompts): 测试模型性能 print(f\n测试模型: {model_path}) # 加载模型 start_time time.time() tokenizer AutoTokenizer.from_pretrained( model_path, trust_remote_codeTrue ) model AutoModelForCausalLM.from_pretrained( model_path, torch_dtypetorch.bfloat16, device_mapauto, trust_remote_codeTrue ).eval() load_time time.time() - start_time print(f模型加载时间: {load_time:.2f}秒) # 检查模型大小 model_size sum(p.numel() for p in model.parameters()) print(f模型参数量: {model_size:,}) # 测试推理速度 print(\n推理速度测试:) for i, prompt in enumerate(test_prompts): print(f\n测试 {i1}: {prompt[:50]}...) inputs tokenizer.apply_chat_template( [{role: user, content: prompt}], add_generation_promptTrue, tokenizeTrue, return_tensorspt, return_dictTrue ) inputs {k: v.to(model.device) for k, v in inputs.items()} # 预热 with torch.no_grad(): _ model.generate(**inputs, max_new_tokens10) # 正式测试 start_infer time.time() with torch.no_grad(): outputs model.generate( **inputs, max_new_tokens100, do_sampleTrue, temperature0.7, top_p0.9 ) infer_time time.time() - start_infer # 解码输出 output_ids outputs[:, inputs[input_ids].shape[1]:] response tokenizer.decode(output_ids[0], skip_special_tokensTrue) print(f 生成时间: {infer_time:.2f}秒) print(f 响应: {response[:100]}...) # 检查显存使用 if torch.cuda.is_available(): memory_allocated torch.cuda.memory_allocated() / 1024**3 # GB memory_reserved torch.cuda.memory_reserved() / 1024**3 # GB print(f\nGPU显存使用:) print(f 已分配: {memory_allocated:.2f} GB) print(f 已保留: {memory_reserved:.2f} GB) def compare_models(): 对比原始模型和剪枝模型 test_prompts [ 请用简单的语言解释什么是人工智能, 写一个Python函数实现快速排序算法, 如果我想学习机器学习应该从哪里开始, 总结一下太阳系的主要行星特点 ] print(*60) print(GLM-4-9B-Chat-1M 模型剪枝效果对比) print(*60) # 测试原始模型如果可用 original_path ./glm-4-9b-chat-1m-original if os.path.exists(original_path): test_model_performance(original_path, test_prompts) else: print(原始模型未找到跳过测试) # 测试剪枝后模型 pruned_path ./glm-4-9b-chat-1m-pruned-finetuned if os.path.exists(pruned_path): test_model_performance(pruned_path, test_prompts) else: print(剪枝模型未找到请先完成剪枝和微调) # 测试纯剪枝模型未微调 pruned_only_path ./glm-4-9b-chat-1m-pruned-50 if os.path.exists(pruned_only_path): print(\n *60) print(剪枝但未微调的模型测试对比用) print(*60) test_model_performance(pruned_only_path, test_prompts[:1]) # 只测一个节省时间 if __name__ __main__: compare_models()运行这个测试脚本你会看到类似下面的对比结果 GLM-4-9B-Chat-1M 模型剪枝效果对比 测试模型: ./glm-4-9b-chat-1m-original 模型加载时间: 45.23秒 模型参数量: 9,000,000,000 推理速度测试: 测试 1: 请用简单的语言解释什么是人工智能... 生成时间: 8.76秒 响应: 人工智能是让机器模拟人类智能行为的技术... GPU显存使用: 已分配: 17.82 GB 已保留: 18.15 GB 测试模型: ./glm-4-9b-chat-1m-pruned-finetuned 模型加载时间: 22.15秒 模型参数量: 4,476,480,000 推理速度测试: 测试 1: 请用简单的语言解释什么是人工智能... 生成时间: 4.32秒 响应: 人工智能是计算机科学的一个分支旨在创建能够... GPU显存使用: 已分配: 8.91 GB 已保留: 9.23 GB从测试结果可以看到剪枝后的模型在参数量、加载时间、推理速度和显存占用上都有明显改善而回答质量经过微调后仍然保持得不错。7. 实际使用中的注意事项在实际部署剪枝模型时有几个点需要特别注意1. 长文本处理能力测试GLM-4-9B-Chat-1M的核心优势是百万上下文剪枝后需要特别测试这个能力是否受到影响。可以用“大海捞针”测试来验证def needle_in_haystack_test(model, tokenizer, context_length100000): 大海捞针测试 # 生成一个很长的上下文 haystack .join([这是一个测试句子。] * 20000) # 约10万字 # 在随机位置插入“针” import random needle_position random.randint(0, len(haystack.split()) - 10) needle 特别关键词AI模型剪枝效果很好 words haystack.split() words.insert(needle_position, needle) haystack_with_needle .join(words) # 让模型找出“针” prompt f请找出下面文本中的特别关键词\n\n{haystack_with_needle}\n\n特别关键词是什么 inputs tokenizer(prompt, return_tensorspt, truncationTrue, max_length131072) inputs {k: v.to(model.device) for k, v in inputs.items()} with torch.no_grad(): outputs model.generate(**inputs, max_new_tokens50) response tokenizer.decode(outputs[0], skip_special_tokensTrue) if needle in response: print(✓ 成功找到关键词) else: print(✗ 未找到关键词) return response2. 多语言能力验证GLM-4-9B-Chat-1M支持26种语言剪枝后需要验证多语言能力是否保持def multilingual_test(model, tokenizer): 多语言能力测试 test_cases [ (英语, Explain the concept of machine learning in simple terms.), (日语, 機械学習とは何か、簡単に説明してください。), (韩语, 기계 학습이 무엇인지 간단히 설명해 주세요.), (德语, Erklären Sie das Konzept des maschinellen Lernens in einfachen Worten.) ] for lang, text in test_cases: print(f\n{lang}测试:) print(f 输入: {text}) inputs tokenizer(text, return_tensorspt) inputs {k: v.to(model.device) for k, v in inputs.items()} with torch.no_grad(): outputs model.generate(**inputs, max_new_tokens100) response tokenizer.decode(outputs[0], skip_special_tokensTrue) print(f 响应: {response[:100]}...)3. 部署优化建议如果显存还是紧张可以结合量化技术INT8/INT4进一步压缩使用vLLM等推理优化框架可以进一步提升速度对于生产环境建议进行更全面的测试包括压力测试和长时间运行的稳定性测试8. 总结走完这一整套流程你应该已经成功地把GLM-4-9B-Chat-1M的参数量减少了一半左右。从我自己的实践来看剪枝后的模型在RTX 4060 Ti这样的消费级显卡上运行得相当流畅显存占用从接近18GB降到了9GB左右推理速度也快了一倍。剪枝虽然会损失一些性能但通过合适的微调大部分能力都能恢复。对于很多实际应用场景来说这种程度的性能损失是可以接受的特别是考虑到它带来的部署便利性和成本节省。如果你在操作过程中遇到问题比如显存不足、模型加载失败等可以尝试调整剪枝比例比如从50%降到40%或者使用更小的微调批次大小。每个硬件环境都不一样可能需要一些调整才能找到最适合的配置。剪枝只是模型优化的一个方面后续还可以探索量化、知识蒸馏等其他技术进一步优化模型。不过对于大多数应用来说50%的剪枝已经能在性能和效率之间取得很好的平衡了。获取更多AI镜像想探索更多AI镜像和应用场景访问 CSDN星图镜像广场提供丰富的预置镜像覆盖大模型推理、图像生成、视频生成、模型微调等多个领域支持一键部署。