如何增大Mistral 7B模型中隐藏Linear4bit层的宽度？

1. 安装依赖 #

!pip install -U bitsandbytes
!pip install -U transformers
!pip install -U peft
!pip install -U accelerate
!pip install -U trl

2. 加载Mistral-7B模型 #

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig,HfArgumentParser,TrainingArguments,pipeline, logging
from peft import LoraConfig, PeftModel, prepare_model_for_kbit_training, get_peft_model
from datasets import load_dataset
from trl import SFTTrainer
from bitsandbytes.nn import Linear4bit

import torch

bnb_config = BitsAndBytesConfig(  
    load_in_4bit= True,
    bnb_4bit_quant_type= "nf4",
    bnb_4bit_compute_dtype= torch.bfloat16,
    bnb_4bit_use_double_quant= False,
)

base_model="mistralai/Mistral-7B-v0.1"

model = AutoModelForCausalLM.from_pretrained(
        base_model,
        quantization_config=bnb_config,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        trust_remote_code=True,
)
model.config.use_cache = False # 消除警告
model.config.pretraining_tp = 1
model.gradient_checkpointing_enable()

tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)
tokenizer.padding_side = 'right'
tokenizer.pad_token = tokenizer.eos_token
tokenizer.add_eos_token = True
tokenizer.add_bos_token, tokenizer.add_eos_token

3. 查看原始模型架构 #

执行：

>>> model

输出：

MistralForCausalLM(
  (model): MistralModel(
    (embed_tokens): Embedding(32000, 4096)
    (layers): ModuleList(
      (0-31): 32 x MistralDecoderLayer(
        (self_attn): MistralAttention(
          (q_proj): Linear4bit(in_features=4096, out_features=4096, bias=False)
          (k_proj): Linear4bit(in_features=4096, out_features=1024, bias=False)
          (v_proj): Linear4bit(in_features=4096, out_features=1024, bias=False)
          (o_proj): Linear4bit(in_features=4096, out_features=4096, bias=False)
          (rotary_emb): MistralRotaryEmbedding()
        )
        (mlp): MistralMLP(
          (gate_proj): Linear4bit(in_features=4096, out_features=14336, bias=False)
          (up_proj): Linear4bit(in_features=4096, out_features=14336, bias=False)
          (down_proj): Linear4bit(in_features=14336, out_features=4096, bias=False)
          (act_fn): SiLU()
        )
        (input_layernorm): MistralRMSNorm()
        (post_attention_layernorm): MistralRMSNorm()
      )
    )
    (norm): MistralRMSNorm()
  )
  (lm_head): Linear(in_features=4096, out_features=32000, bias=False)
)

目标将模型隐藏维度从4096增加800至4896，新增节点采用随机初始化，最终架构如下：

MistralForCausalLM(
  (model): MistralModel(
    (embed_tokens): Embedding(32000, 4896)
    (layers): ModuleList(
      (0-31): 32 x MistralDecoderLayer(
        (self_attn): MistralAttention(
          (q_proj): Linear4bit(in_features=4896, out_features=4896, bias=False)
          (k_proj): Linear4bit(in_features=4896, out_features=1024, bias=False)
          (v_proj): Linear4bit(in_features=4896, out_features=1024, bias=False)
          (o_proj): Linear4bit(in_features=4896, out_features=4896, bias=False)
          (rotary_emb): MistralRotaryEmbedding()
        )
        (mlp): MistralMLP(
          (gate_proj): Linear4bit(in_features=4896, out_features=14336, bias=False)
          (up_proj): Linear4bit(in_features=4896, out_features=14336, bias=False)
          (down_proj): Linear4bit(in_features=14336, out_features=4896, bias=False)
          (act_fn): SiLU()
        )
        (input_layernorm): MistralRMSNorm()
        (post_attention_layernorm): MistralRMSNorm()
      )
    )
    (norm): MistralRMSNorm()
  )
  (lm_head): Linear(in_features=4896, out_features=32000, bias=False)
)

直接修改模型各层参数，扩展维度并初始化新增部分：

# 定义维度参数
original_dim = 4096
new_dim = 4896
delta_dim = new_dim - original_dim

# 1. 扩展嵌入层
old_embedding = model.model.embed_tokens
new_embedding = torch.nn.Embedding(old_embedding.num_embeddings, new_dim)
new_embedding.weight.data[:, :original_dim] = old_embedding.weight.data
torch.nn.init.normal_(new_embedding.weight.data[:, original_dim:], mean=0, std=0.02)
model.model.embed_tokens = new_embedding.to(model.device)

# 2. 定义Linear4bit层扩展函数
def expand_linear4bit(layer, in_dim_new, out_dim_new, original_in_dim, original_out_dim):
    new_layer = Linear4bit(
        in_features=in_dim_new,
        out_features=out_dim_new,
        bias=layer.bias is not None,
        compute_dtype=layer.compute_dtype
    )
    # 复制原有权重
    old_weight = layer.weight.data
    new_layer.weight.data[:original_out_dim, :original_in_dim] = old_weight
    # 随机初始化新增权重
    if out_dim_new > original_out_dim:
        torch.nn.init.normal_(new_layer.weight.data[original_out_dim:, :original_in_dim], mean=0, std=0.02)
    if in_dim_new > original_in_dim:
        torch.nn.init.normal_(new_layer.weight.data[:, original_in_dim:], mean=0, std=0.02)
    # 处理bias
    if layer.bias is not None:
        new_bias = torch.zeros(out_dim_new, dtype=layer.bias.dtype)
        new_bias[:original_out_dim] = layer.bias.data
        if out_dim_new > original_out_dim:
            torch.nn.init.normal_(new_bias[original_out_dim:], mean=0, std=0.02)
        new_layer.bias = torch.nn.Parameter(new_bias).to(model.device)
    return new_layer.to(model.device)

# 3. 遍历所有DecoderLayer，扩展各Linear4bit层
for layer in model.model.layers:
    # 处理注意力层
    layer.self_attn.q_proj = expand_linear4bit(layer.self_attn.q_proj, new_dim, new_dim, original_dim, original_dim)
    layer.self_attn.k_proj = expand_linear4bit(layer.self_attn.k_proj, new_dim, 1024, original_dim, 1024)
    layer.self_attn.v_proj = expand_linear4bit(layer.self_attn.v_proj, new_dim, 1024, original_dim, 1024)
    layer.self_attn.o_proj = expand_linear4bit(layer.self_attn.o_proj, new_dim, new_dim, original_dim, original_dim)
    
    # 处理MLP层
    layer.mlp.gate_proj = expand_linear4bit(layer.mlp.gate_proj, new_dim, 14336, original_dim, 14336)
    layer.mlp.up_proj = expand_linear4bit(layer.mlp.up_proj, new_dim, 14336, original_dim, 14336)
    layer.mlp.down_proj = expand_linear4bit(layer.mlp.down_proj, 14336, new_dim, 14336, original_dim)

# 4. 更新RMSNorm层维度
from transformers.models.mistral.modeling_mistral import MistralRMSNorm
model.model.norm = MistralRMSNorm(new_dim, eps=model.model.norm.eps).to(model.device)
# 同时更新各DecoderLayer中的Norm层
for layer in model.model.layers:
    layer.input_layernorm = MistralRMSNorm(new_dim, eps=layer.input_layernorm.eps).to(model.device)
    layer.post_attention_layernorm = MistralRMSNorm(new_dim, eps=layer.post_attention_layernorm.eps).to(model.device)

# 5. 扩展lm_head
old_lm_head = model.lm_head
new_lm_head = torch.nn.Linear(new_dim, old_lm_head.out_features, bias=old_lm_head.bias is not None)
new_lm_head.weight.data[:, :original_dim] = old_lm_head.weight.data
torch.nn.init.normal_(new_lm_head.weight.data[:, original_dim:], mean=0, std=0.02)
if old_lm_head.bias is not None:
    new_lm_head.bias.data = old_lm_head.bias.data
model.lm_head = new_lm_head.to(model.device)

# 6. 更新模型配置
model.config.hidden_size = new_dim
model.config.n_embd = new_dim

执行完上述代码后，再次查看模型架构即可确认维度已扩展。

内容的提问来源于stack exchange，提问作者alvas