减少空间占用的生成模型实战与推理资源消耗量化对比
《DeepSeek原生应用与智能体开发实践》【摘要 书评 试读】- 京东图书
我们使用缓存对自回归生成模型进行计算时，最核心的目标是降低模型推理时的资源消耗。通过精心设计的缓存机制，我们得以高效地存储并复用中间计算结果，从而避免了在每次推理步骤中重复相同的计算。这种优化方法不仅大幅提升了模型的推理速度，还减少了对计算资源的依赖，使自回归生成模型在更多场景下都能实现高效运用。
在本节中，我们将从头开始训练一个自回归模型，并采用不同的推理方法，以便对资源消耗进行量化对比。我们希望通过这一对比实验，更直观地展示缓存机制在降低资源消耗方面的实际效果，进一步验证缓存优化策略的有效性。同时，我们也希望通过这些数据，为未来的模型优化和缓存设计提供更具体的指导和参考。通过这样的实证研究，我们不仅可以更深入地理解缓存机制在自回归模型中的作用，还能为相关领域的研究和实践提供有价值的经验和启示。
首先需要完成模型参数的配置。在这里我们简化了配置方法，采用类的形式对所有的参数进行管理，代码如下所示：
class GPT2Config:hidden_size = 384vocab_size = 4000num_attention_heads = 6assert hidden_size % num_attention_heads == 0, 'hidden_size must be divisible by num_head'intermediate_size = hidden_size * 4dropout = 0.1layer_norm_eps = 1e-12n_layers = 6is_cause = Truedevice = "cuda"max_length = 48
接下来，我们还是希望使用前面进行评论生成的数据进行本章的学习，在这里我们可以直接使用第5章的分词模型以及评论数据集，代码如下所示：
import sentencepiece as spm
class Tokenizer:def __init__(self,spm_path = './vocab_new/spm_model.model'):super().__init__()self.sp = spm.SentencePieceProcessor()self.sp.Load(spm_path)self.end_id = 4def encode(self,text):token = self.sp.EncodeAsIds(text)return tokendef decode(self,token):_text = self.sp.DecodeIds(token)return (_text)def vocab_size(self):return len(self.sp)
数据集的处理和准备：
import random,torch
from tqdm import tqdm
import sentencepiece as spmimport config
tokenizer_emo = config.Tokenizer()
print(tokenizer_emo.vocab_size())token_list = []
with open("../dataset/ChnSentiCorp.txt", mode="r", encoding="UTF-8") as emotion_file:for line in tqdm(emotion_file.readlines()):line = line.strip().split(",")text = "".join(line[1:]) + '<|end▁of▁sentence|>'token = tokenizer_emo.encode(text)for id in token:token_list.append(id)
token_list = torch.tensor(token_list * 4)class TextSamplerDataset(torch.utils.data.Dataset):def __init__(self, data = token_list, seq_len = 48):super().__init__()self.data = dataself.seq_len = seq_lendef __getitem__(self, index):rand_start = torch.randint(0, self.data.size(0) - self.seq_len, (1,))full_seq = self.data[rand_start : rand_start + self.seq_len + 1].long()return full_seq[:-1],full_seq[1:]def __len__(self):return self.data.size(0) // self.seq_len
在训练带有缓存的生成模型时，我们可以完全将其视为普通的生成模型进行处理。这是因为缓存机制主要影响的是模型的推理阶段，而非训练阶段。在训练过程中，模型需要学习的是如何生成合理的序列，而缓存的引入并不会改变这一学习目标。因此，我们可以按照标准的生成模型训练流程进行，无须对训练过程进行特殊调整。
相对于前面完成的生成模型，这里我们增加了训练次数，这一点请读者注意。完整的生成模型训练代码如下所示：
import torch
from tqdm import tqdm
import torch
import config
import gpt2_cachedimport get_dataset
from torch.utils.data import Dataset, DataLoaderdevice = torch.device("cuda" if torch.cuda.is_available() else "cpu")gpt2config = config.GPT2Config()
tokenizer = config.Tokenizer()model = gpt2_cached.GPT2().to(device)
model.load_state_dict(torch.load("./saver/model.pth"),strict=False)seq_len = 64
# 获取训练数据集
train_dataset = get_dataset.TextSamplerDataset(get_dataset.token_list,seq_len=seq_len)# 初始化 DataLoader
data_trainer = DataLoader(dataset=train_dataset,batch_size=640,shuffle=True)opt = torch.optim.AdamW(model.parameters(),lr=2e-4)
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(opt,T_max = 1200,eta_min=2e-6,last_epoch=-1)
# 损失函数
criterion = torch.nn.CrossEntropyLoss()for epoch in range(128):model.train() # 确保模型在训练模式下pbar = tqdm(data_trainer, total=len(data_trainer))for tok, lab in pbar:tok = tok.to(device)lab = lab.to(device)logits,kv_caches = model(tok)# 调整 logits 和 lab的维度logits = logits.view(-1, logits.size(-1)) # [batch_size * sequence_length, num_classes]lab = lab.view(-1) # [batch_size * sequence_length]# 计算损失loss = criterion(logits, lab)# 反向传播和优化opt.zero_grad() # 清除梯度loss.backward() # 计算梯度opt.step() # 更新参数lr_scheduler.step() # 执行优化器# 更新进度条上的描述pbar.set_description(f"Epoch {epoch + 1}, Loss: {loss.item():.4f}, lr:{lr_scheduler.get_last_lr()[0]*1000:.5f}")if epoch % 5 == 0:print("model saved")torch.save(model.state_dict(), "./saver/model.pth")torch.save(model.state_dict(), "./saver/model.pth")
值得注意的是，虽然在训练阶段缓存并不直接参与，但考虑到模型在实际应用中的推理效率，我们在设计模型结构时，仍然需要预留出与缓存机制相兼容的接口。这样做的好处是，一旦模型训练完成，我们可以轻松地整合缓存功能，从而在实际应用中实现更高效的推理。
此外，尽管缓存不在训练阶段直接使用，但了解并优化缓存机制对于提升模型的整体性能至关重要。因此，在训练过程中，我们也应不断思考如何更好地结合缓存策略，以便在后续的推理阶段达到最佳效果。通过这种方式，我们不仅可以确保模型在训练阶段学习到有效的生成策略，还能为其在实际应用中的高效推理奠定坚实基础。
