如何构建参数规模小于7B的LLaMA v2因果解码器预训练模型？

需求背景 #

预训练参数规模小于7B的Decoder架构因果语言模型，因7B及以上参数模型训练过程不稳定，需保障预训练顺利、减少人工维护；选择LLaMA v2作为基础架构，需实现两个核心目标：

• 初始化参数更少的LLaMA v2架构（如减小特征宽度、减少Transformer层数）；
• 对模型进行随机初始化。

实现步骤与代码 #

1. 自定义小参数量LLaMA v2配置

从transformers库的LlamaConfig继承，修改关键结构参数来控制模型总参数量。常见可调整的参数包括：

• hidden_size：Transformer层的特征维度（LLaMA-7B为4096，调小到2048/1024可大幅减少参数）
• num_layers：Transformer decoder层数（LLaMA-7B为32，调小到24/16）
• num_attention_heads：注意力头数（需与hidden_size匹配，比如hidden_size=2048时设为16）
• max_position_embeddings：最大序列长度（可选调整）

2. 基于配置随机初始化模型

使用AutoModelForCausalLM.from_config()方法直接生成随机初始化的模型，无需加载预训练权重。

完整可运行代码

from transformers import AutoModelForCausalLM, LlamaConfig, AutoTokenizer

# 1. 定义小参数量的LLaMA v2配置
class SmallLlamaConfig(LlamaConfig):
    def __init__(self, **kwargs):
        # 调用父类构造方法，覆盖关键参数
        super().__init__(
            hidden_size=2048,          # 特征维度，从4096缩小到2048
            num_layers=24,             # Decoder层数，从32减少到24
            num_attention_heads=16,    # 注意力头数，与hidden_size匹配（2048/16=128）
            max_position_embeddings=4096,
            vocab_size=32000,
            **kwargs
        )

# 2. 初始化配置并创建随机初始化模型
config = SmallLlamaConfig()
model = AutoModelForCausalLM.from_config(config)

# 验证模型结构与参数量
print("模型结构：")
print(model)
print("\n总参数量：", sum(p.numel() for p in model.parameters())/1e9, "B")

# 加载对应的LLaMA v2 tokenizer
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf", use_auth_token=True)
tokenizer.pad_token = tokenizer.eos_token  # LLaMA tokenizer默认无pad token，需手动设置

关键说明 #

• 调整参数时需保证hidden_size能被num_attention_heads整除（注意力头的维度=hidden_size/num_attention_heads），否则会报错；
• 若需要更精确控制参数量，可以通过调整intermediate_size（FFN中间层维度，LLaMA默认是hidden_size*4）进一步压缩；
• 代码中使用transformers官方库的LlamaConfig，避免自定义库的依赖问题；
• 随机初始化的模型可直接用于预训练流程，无需额外权重加载步骤。

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