大模型加速与压缩之wanda代码详解

A SIMPLE AND EFFECTIVE PRUNING APPROACH FORLARGE LANGUAGE MODELS

github: https://github.com/locuslab/wanda

• 核心idea：基于权重和激活来计算权重重要性

这里我们以剪枝opt-125m为例，即main_opt.py中涉及的代码进行详细介绍

wanda剪枝算法步骤

主要剪的是模型layers的线性权重(主要占绝大多数参数，如nn.Linear，这里以opt-125m为例

可以看到opt-125m的模块为OPTDecoder构成包含

• emd_tokens–>Embedding(50272,768)–>词库的个数为50272
• embed_positions–>(2050, 768)–>seq_len+(2个其他添加的特殊标记字符token,如句子开始字符BOS,结束字符EOS)
• final_layer_norm–>(,768)
• layers包括12层OPTDecoderLayer
• lm_head–>nn.Linear()

第一步：加载预训练模型以及tokenizer

# 导入相关库
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from importlib.metadata import version

D:\Applications\anaconda\envs\llm\lib\site-packages\transformers\utils\hub.py:127: FutureWarning: Using `TRANSFORMERS_CACHE` is deprecated and will be removed in v5 of Transformers. Use `HF_HOME` instead.
  warnings.warn(

print('torch', version('torch'))
print('transformers', version('transformers'))
print('# of gpus: ', torch.cuda.device_count())

torch 1.13.1+cu116
transformers 4.42.4
# of gpus:  1

def get_llm(model_name, cache_dir="llm_weights"):
    model = AutoModelForCausalLM.from_pretrained(
        model_name, 
        torch_dtype=torch.float16, # 默认是torch.float32,但半精度float16足够并且可以减少一般显存大小
        cache_dir=cache_dir, # 模型缓存路径如果不指定模型缓存路径会通过huggingface自动下载到./cache/huggingface/hub下
        low_cpu_mem_usage=True, # 模型加载减少cpu内存的使用，默认为False
        device_map="auto"  # 自动根据是否具有gpu进行模型加载，如果有gpu，则模型加载会直接加载到gpu上即device为"cuda"
    )

    model.seqlen = model.config.max_position_embeddings 
    return model

model = get_llm('llm_weights/models--facebook--opt-125m/snapshots/opt-125m', cache_dir="llm_weights") # 直接加载已经缓存好的模型路径
model.eval() # 推理评估模式下进行否则dropout没有去掉,layer_norm没有根据训练的结果进行
tokenizer = AutoTokenizer.from_pretrained('llm_weights/models--facebook--opt-125m/snapshots/opt-125m', use_fast=False)

model.device # device_map不设置则会为cpu

device(type='cuda', index=0)

• 需要注意的是模型加载如果你已经缓存好了，以后可以直接加载缓存好的模型路径定位到snapshots下一个文件夹(一般地，下载缓存好的文件夹名是一段哈希字符，)这里重命名为opt-125m
• tokenizer中的use_fast是否采用分词器加速(Rust编程优化)

第二步：使用校验数据执行前向传播计算得到layers的输入

import numpy as np
import random
import torch
import torch.nn as nn
from datasets import load_dataset

• 通常剪枝设定的校验数据batch_size = 128
• 模型的seq_len = 2048 (一般根据预训练模型配置的最大嵌入维度获得，即model.config.max_position_embeddings),也就是说seq_len是定了的，你要剪枝什么模型，那么它之前训练时的seq_len就会固定不变
• 接着需要考虑的是隐藏层维度（由于大模型的基础架构是Transformer，这个隐藏层维度其实就是d_model,理论上不同的大模型，设置的d_model有所区别在opt-125m中，维度为768=256*3，即代码中model.config.hidden_size)

那么我们其实已经确定好了，模型的OPTDecoder的layers的输入数据维度应该是[128, 2048, 768]，那么我如何得到这个输入呢？这就需要校验数据，通过前向传播经过OPTDecoder的layers前的那几层计算（如emd_tokens,embed_positions等组成的嵌入层）得到输出即是OPTDecoder的layers的输入

获取校验数据

• 通常会选择c4（Colossal Clean Crawled Corpus，是一个大规模的、清洁的、从互联网上爬取的文本数据集）的一部分作为校验数据，主要是因为c4数据集规模大，多样性丰富
• 用它作为校验数据通用性比较强
• 那会选择多少c4的样本数据集去校验呢？
• 至少要确保c4某一样本分词后的长度必须大于seq_len（2048）,这样你才可以取得完整的句子序列

那么现在，你有一个重要的问题，可能要问：你为什么要使用校验数据？为什么不直接根据权重进行剪枝呢？因为wanda考虑了激活就必须需要校验数据才能获得激活X

def get_c4(nsamples=128, seed=0, seq_len=2048,tokenizer=tokenizer,
          data_path=None):
    '''
    这里我们修改了部分代码，但差距不大
    return: train_loader-->list-->length:128
    valenc: Tensor-->[1:,词个数]
    '''
    # 获取c4数据集
    # 直接下载
    if not data_path: # 需要联网（外网）下载数据
        print('you need open the network to download data')
        train_data = load_dataset('allenai/c4', 
                                 data_files={
   
   'train': 'en/c4-train.00000-of-01024.json.gz'},
                                 split='train')
        val_data = load_dataset('allenai/c4', 
                               data_files={
   
   'validation': 'en/c4-validation.00000-of-00008.json.gz'}, 
                               split='validation')
    else:
        print('load data from the local directory')
        train_data = load_dataset(data_path, split='train')
        val_data = load_dataset(data_path, split='validation')
    # 获取校验数据nsample个
    random.seed(seed)
    train_loader = []
    for _ in range(nsamples):
        while True:
            i = random.randint(0, len(train_data) - 1)
            # 对随机选择的样本进行分词，返回的是字典key-->'input_ids', 'attention_mask'->维度为[1,分词个数]
            trainenc = tokenizer(train_data[i]['text'], return_tensors='pt') 
            # 如果分词后的词的长度大于seq_len,则保留trainenc，并且跳出循环，否则继续寻找trainenc直到跳出循环
            if trainenc['input_ids'].shape[1] > seq_len:
                break
            else:
                continue
        # 对保留后的trainenc随机获取一个长度为2048的词序列
        j = random.randint(0, trainenc['input_ids'].shape[1] - seq_len-1)
        inp = trainenc.input_ids[:, j:(j+seq_len)] # 获取词序列
        # 获取预测任务的序列
        tar = inp.clone()
        tar[:, :-1] = -100
        # 添加(inp,tar)到train_loader中
        train_loader.append((inp, tar))

    # 获取评估验证集
    valenc = tokenizer(' '.join(val_data[:1100]['text']), return_tensors='pt')
    valenc = valenc.input_ids[:, :(256 * seq_len)]
    print("data loading complete") 
    return train_loader, valenc

data_loader,_ = get_c4(nsamples=128, seed=0, seq_len=2048,tokenizer=tokenizer,
                       data_path='./data/c4')

load data from the local directory
data loading complete

获取layers的输入

• 需要注意的点：
• 前向传播无需梯度传播–>必须with torch.no_grad()
• 无需使用kv缓存因为要重新计算一边，因此刚开始的use_cache为False，最后前向传播结束后在use_cache=True

def prepare_layers_input(model, data_loader, device='cuda'):
    """
    return inps, outs, attention_mask
    """
    use_cache = model.config.use_cache
    model.config.use_cache = False
    layers = model.model.decoder.layers # 一共12层
    if "model.embed_tokens" in model.hf_device_map: # model.hf_device_map-->字典，分配模块的gpu位置这里为{"":0}
        device = model.hf_device_map["model.embed_tokens"]
    dtype = next(iter(model.parameters())).dtype # torch.float16
    inps = torch.zeros((128, model.seqlen, model.config.hidden_size), dtype=dtype, device=device)
    inps.requires_grad = False # 单独设置的数据需要device和设置无grad
    cache = {
   
   'i': 0, 'attention_mask': None, "position_ids": None} # 用来记录前向传播

    # 捕获layers[0]的输入,一旦捕获就抛出异常，让异常模块直接pass掉
    class Catcher(nn.Module):
        def __init__(self, module):
            super()