低硬件资源微调预训练Mamba模型的方法
深入探索Mamba模型架构与应用 - 商品搜索 - 京东
在前面的章节中，我们详细介绍了基于预训练模型Mamba的全参数微调方法，并展示了其在实际应用中的有效性。这种全参数微调确实是一种切实可行的模型训练与预测手段，通过对模型的所有参数进行调整，可以使其更好地适应特定的任务和数据集。然而，微调的方法并非仅限于此。
事实上，除全参数微调外，还存在多种灵活的微调策略，能够在满足我们需求的同时，完成对模型的精细化调整。例如，部分参数微调是一种更加聚焦的调整方式，只针对模型中的特定层或参数进行调整，而保持其他部分不变。这种方法既可以节省计算资源，又能有针对性地提升模型在特定任务上的性能。
此外，还有基于适配器的微调方法，这种方法通过在预训练模型中插入额外的适配器模块，并仅对这些模块进行训练来实现对模型的微调。这种方式能够在保留预训练模型大部分知识的同时，快速适应新的任务。
因此，除全参数微调外，还有多种微调方法可供选择，以便根据实际情况灵活调整模型，以实现最佳性能。在接下来的章节中，我们将深入探讨这些微调方法的具体实现和应用场景。
8.3.1 使用冻结模型参数的微调方法
我们的目标是利用预训练模型来完成模型微调，以实现更高的适应性和性能。一个直观且实用的策略是，在原有模型的基础上，冻结部分底层参数，而集中训练模型的高层参数。
通常模型的底层负责抽取基础且相对“粗糙”的特征，这些特征虽然对理解输入数据至关重要，但在不同的任务和领域具有一定的通用性。相对而言，模型的高层更专注于抽取“细粒度”的特征，这些特征对于特定任务的性能提升尤为关键，如图8-7所示。
图8-7 只在部分层上进行训练
基于这一特性，我们可以采用一种策略：冻结模型的底层参数，仅对高层参数进行训练。这种微调方式不仅可以减少训练过程中的计算负担，还能使模型更快地适应新任务，因为高层参数的调整能够更直接地影响模型的最终输出。
通过这种方法，我们可以在保留预训练模型强大特征抽取能力的同时，使模型更好地适应特定任务的需求。在接下来的实验中，我们将验证这种微调策略的有效性，并探索其对模型性能的具体影响。
在具体实现时，我们首先获取所有层的名称，此时PyTorch给我们提供了打印所有模型名称的函数，代码如下：
# 冻结所有层
for name,param in mamba_model.named_parameters():
print(name)
param.requires_grad = False可以看到，在打印层名称的同时，param.requires_grad=False这个函数帮助我们完成了对所有层的冻结。打印结果如下：
embedding.weight
layers.0.mixer.A_log
layers.0.mixer.D
...
layers.23.mixer.dt_proj.bias
layers.23.mixer.out_proj.weight
layers.23.norm.weight
norm_f.weight
可以看到，这里实际上是根据不同的层号对每个位置所处的层进行编号，核对原有的Mamba的config设置参数：
{
"d_model": 768,
"n_layer": 24,
"vocab_size": 50277,
"ssm_cfg": {},
"rms_norm": true,
"residual_in_fp32": true,
"fused_add_norm": true,
"pad_vocab_size_multiple": 8
}
可以很容易地发现，在预训练的Mamba层中，层是按照数字顺序从0~23排列的，下面我们仅根据名称，在使用名称的基础上对模型进行微调，代码如下：
model_dir = snapshot_download('AI-ModelScope/mamba-130m',cache_dir="./mamba/")
mamba_model = Mamba.from_pretrained("./mamba/AI-ModelScope/mamba-130m")
# 冻结除目标层外的所有层
for name,param in mamba_model.named_parameters():
print(name)
if "23" not in name:
param.requires_grad = False
BATCH_SIZE = 320
import all_config
device = all_config.device
model = mamba_model
model.to(device)
可以看到，此时我们冻结了除目标层外的所有层，而在具体代码编写上，我们通过传入特定字符的方式将无须冻结的层排除，同时冻结特定目标层[晓王1] ，此时完整代码如下：
import os
import math
from tqdm import tqdm
import torch
from torch.utils.data import DataLoader
from model import Mamba
from modelscope import snapshot_download,AutoTokenizer
model_dir = snapshot_download('AI-ModelScope/mamba-130m',cache_dir="./mamba/")
mamba_model = Mamba.from_pretrained("./mamba/AI-ModelScope/mamba-130m")
# 冻结所有层
for name,param in mamba_model.named_parameters():
print(name)
if "23" not in name:
param.requires_grad = False
BATCH_SIZE = 320
import all_config
device = all_config.device
model = mamba_model
model.to(device)
import get_data_emotion
train_dataset = get_data_emotion.TextSamplerDataset(get_data_emotion.token_list)
train_loader = (DataLoader(train_dataset, batch_size=BATCH_SIZE,shuffle=True))
save_path = "./saver/glm_text_generator.pth"
optimizer = torch.optim.AdamW(model.parameters(), lr = 2e-5)
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,T_max = 1200,eta_min=2e-7,last_epoch=-1)
criterion = torch.nn.CrossEntropyLoss()
for epoch in range(12):
pbar = tqdm(train_loader,total=len(train_loader))
for token_inp,token_tgt in pbar:
token_inp = token_inp.to(device)
token_tgt = token_tgt.to(device)
logits = model(token_inp)
loss = criterion(logits.view(-1, logits.size(-1)), token_tgt.view(-1))
optimizer.zero_grad()
loss.backward()
optimizer.step()
lr_scheduler.step() # 执行优化器
pbar.set_description(f"epoch:{epoch +1}, train_loss:{loss.item():.5f}, lr:{lr_scheduler.get_last_lr()[0]*1000:.5f}")
torch.save(model.state_dict(), save_path)
可以看到，在原有的架构上，我们没有进行修改，而是采用前面讨论的冻结部分层（编号23）的方式对模型进行微调。此时可以采用载入预测的方法对模型进行输出，这点请读者按8.2.4节的预训练模型输出方案对结果进行打印。部分结果如下：
酒店里没有讲住这的估计，服务人员的地理位置还好，也不如酒店的优越性价比
位置，但是在海有不合观，另外大型的服务怎么说，服务员也很不错。※房间干净、洗
-------------
酒店房间台总经理人员有称客意见，但是本人所说没有态度反映，所以让我先说没有意见，
位置很好。房间有宽带，但晚上睡觉很吵
-------------
酒店服务员态度、房间服务员态度都一般，而服务态度都好像不错。入住酒店就在市中心太
位置很高，在我们的五星级酒店，环境也很吵。但酒店环境不太好，所
-------------
酒店设施，否则在酒店服务态度很好
位置，没有预定的，不仅因我们自己没有
-------------
酒店了，这次我们还在还要去。服务员可以带热的水果
位置比较少。早餐还要10元的房价格也不符合星级标准酒店房间太
可以看到，此时的输出结果并不是很符合我们的输出要求，究其原因，一般是需要对更多的层进行微调，此时我们可以修改冻结部分的代码如下：
for name,param in mamba_model.named_parameters():
print(name)
if "23" not in name or "22" not in name:
param.requires_grad = False
这样即可通过修改冻结层的数目，从而设定可训练的参数数目，这点请读者基于自身的硬件设备进行处理。
8.3.2 通过替换制定层的方式完成微调
除冻结部分层而对某些特定层进行微调外，还有一种方法是通过替换一些关键的特定层，从而完成预训练模型的微调。首先打印完整的模型层名称：
from model import Mamba
from modelscope import snapshot_download,AutoTokenizer
model_dir = snapshot_download('AI-ModelScope/mamba-130m',cache_dir="./mamba/")
mamba_model = Mamba.from_pretrained("./mamba/AI-ModelScope/mamba-130m")
print(mamba_model)
打印结果如下：
Mamba(
(embedding): Embedding(50280, 768)
(layers): ModuleList(
(0-23): 24 x ResidualBlock(
(mixer): MambaBlock(
(in_proj): Linear(in_features=768, out_features=3072, bias=False)
(conv1d): Conv1d(1536, 1536, kernel_size=(4,), stride=(1,), padding=(3,), groups=1536)
(x_proj): Linear(in_features=1536, out_features=80, bias=False)
(dt_proj): Linear(in_features=48, out_features=1536, bias=True)
(out_proj): Linear(in_features=1536, out_features=768, bias=False)
)
(norm): RMSNorm()
)
)
(norm_f): RMSNorm()
(lm_head): Linear(in_features=768, out_features=50280, bias=False)
)
在这里，我们希望替换最后一层lm_head，并冻结其他所有层的内容，只保持对这一层的训练。因此，具体来看，我们需要根据lm_head的大小和架构样式，重新构建一个相同的层完成模型的设计。这部分代码如下：
for param in mamba_model.parameters():
param.requires_grad = False
lm_head_new = torch.nn.Linear(in_features=768, out_features=50280, bias=False,device=device)
model.lm_head = lm_head_new
train_params = (model.lm_head.parameters())
optimizer = torch.optim.AdamW(train_params , lr = 2e-5)
可以看到，首先我们对所有的层和参数设置了不再更新，之后使用了一个新的全连接层lm_head_new = torch.nn.Linear(in_features=768, out_features=50280, bias=False,device=device)替换原有Mamba模型的lm_head层。而在具体的torch.optim.AdamW更新时，制定了只对这一层的参数进行更新，从而完成模型的微调。完整训练代码如下：
import os
import math
from tqdm import tqdm
import torch
from torch.utils.data import DataLoader
from model import Mamba
from modelscope import snapshot_download,AutoTokenizer
model_dir = snapshot_download('AI-ModelScope/mamba-130m',cache_dir="./mamba/")
mamba_model = Mamba.from_pretrained("./mamba/AI-ModelScope/mamba-130m")
BATCH_SIZE = 512
import all_config
device = all_config.device
model = mamba_model
model.to(device)
import get_data_emotion
train_dataset = get_data_emotion.TextSamplerDataset(get_data_emotion.token_list)
train_loader = (DataLoader(train_dataset, batch_size=BATCH_SIZE,shuffle=True))
save_path = "./saver/glm_text_generator.pth"
model.load_state_dict(torch.load(save_path))
for param in mamba_model.parameters():
param.requires_grad = False
lm_head_new = torch.nn.Linear(in_features=768, out_features=50280, bias=False,device=device)
model.lm_head = lm_head_new
train_params = (model.lm_head.parameters())
optimizer = torch.optim.AdamW(train_params , lr = 2e-3)
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,T_max = 1200,eta_min=2e-7,last_epoch=-1)
criterion = torch.nn.CrossEntropyLoss()
for epoch in range(24):
pbar = tqdm(train_loader,total=len(train_loader))
for token_inp,token_tgt in pbar:
optimizer.zero_grad()
token_inp = token_inp.to(device)
token_tgt = token_tgt.to(device)
logits = model(token_inp)
loss = criterion(logits.view(-1, logits.size(-1)), token_tgt.view(-1))
loss.backward()
optimizer.step()
lr_scheduler.step() # 执行优化器
pbar.set_description(f"epoch:{epoch +1}, train_loss:{loss.item():.5f}, lr:{lr_scheduler.get_last_lr()[0]*1000:.5f}")
torch.save(model.state_dict(), save_path)
预测部分请读者自行参考上面的预测模块进行处理。
8.3.3 对模型参数进行部分保存和载入的方法
在完成模型的微调，并设定部分参数可参与梯度更新后，大部分原有预训练模型的参数通常不会发生变化，即只有部分模型参数的值被更新。在这种情况下，如果按照原有全部参数量进行保存和处理，显然有些多余。
PyTorch还为我们提供了仅仅对部分参数进行保存和读取的方法，当我们只需要对某个特定层的参数进行保存时，可以这样：
...
for param in mamba_model.parameters():
param.requires_grad = False
...
lm_head_new = torch.nn.Linear(in_features=768, out_features=50280, bias=False,device=device)
model.lm_head = lm_head_new
train_params = (model.lm_head.parameters())
...
torch.save([model.lm_head.state_dict()],"./saver/lm_head.pth")
即在选定特定层后，直接对参数输入新的地址并保存即可。这里我们只是显式地对model.lm_head.state_dict()的内容进行了保存，同时读者需要注意，在这里实际上可以对多个层的内容进行保存，作者使用list作为一个存储目标，有兴趣的读者可以仔细尝试。
而重新载入这部分参数只需要在预训练载入完成后，仅载入保存的新的部分存档即可，代码如下：
import model
from model import Mamba
import torch
from modelscope import snapshot_download,AutoTokenizer
model_dir = snapshot_download('AI-ModelScope/mamba-130m',cache_dir="./mamba/")
mamba_model = Mamba.from_pretrained("./mamba/AI-ModelScope/mamba-130m")
save_path = "./saver/lm_head.pth"
mamba_model.load_state_dict(torch.load(save_path) [0],strict=False)
tokenizer = AutoTokenizer.from_pretrained('./mamba/tokenizer')
for _ in range(10):
print(model.generate(mamba_model, tokenizer, '酒店'))
print(model.generate(mamba_model, tokenizer, '位置'))
print("-------------")
可以看到，当前的保存地址就是前面保存的部分参数，而此时对模型参数的载入也仅限于这部分更新后的参数。另外需要注意的是，当前保存的模型是以列表（list）的形式保存的，因此在载入时也需要按照要求对特定的目标进行载入。
