从零手搓中文大模型|Day04|模型参数配置和训练启动
今天我们来深入研究一下模型的配置。在基于 litgpt 框架进行大语言模型训练时,配置文件是连接数据、架构与训练逻辑的核心枢纽。
1. 训练配置文件概览
litgpt 使用的配置文件格式为 YAML,与 transformers 库的配置方式略有不同。其仓库中提供了一些预训练所用的 YAML 配置文件样例,主要用于需要自定义模型架构的场景。此外,litgpt 也内置了一些 Hugging Face 上的现成模型,可以直接拿来使用,但在从头训练或微调特定任务时,自定义配置更为常见。
以下是本次定义的一个完整配置文件,涵盖了模型架构、训练策略、数据加载及优化器设置等关键部分。
model_name: microstories
model_config:
name: microstories
hf_config: {}
scale_embeddings: false
block_size: 512
padded_vocab_size: 65024
vocab_size: 64798
n_layer: 6
n_head: 6
n_query_groups: 6
n_embd: 512
head_size: 48
rotary_percentage: 1.0
parallel_residual: false
bias: false
norm_class_name: RMSNorm
mlp_class_name: LLaMAMLP
intermediate_size: 768
out_dir: /home/puppyapple/Server/BigAI/Chinese_LLM_From_Scratch/Experiments/Output/pretrain/microstories
precision: bf16-mixed
initial_checkpoint_dir:
resume: true
data:
class_path: litgpt.data.LitData
init_args:
data_path: Chinese_LLM_From_Scratch/Data/TinyStoriesChinese/processed_data
split_names:
- train
- val
train:
save_interval: 1000
log_interval: 1
global_batch_size: 512
micro_batch_size: 32
lr_warmup_steps: 1000
max_tokens: 3000000000000
max_seq_length: 512
tie_embeddings: true
max_norm: 1.0
min_lr: 0.0
eval:
interval: 2000
max_new_tokens:
max_iters: 100
initial_validation: false
final_validation: false
optimizer:
class_path: torch.optim.AdamW
init_args:
lr: 0.0005
weight_decay: 0.1
betas:
- 0.9
- 0.95
devices: auto
num_nodes: 1
tokenizer_dir: Chinese_LLM_From_Scratch/References/chatglm3-6b
logger_name: wandb
seed: 42
2. 模型架构参数详解 (model_config)
model_config 部分定义了神经网络的具体拓扑结构。对于构建一个轻量级的 Micro 模型,以下参数的调整至关重要。
2.1 基础维度
block_size: 即最大序列长度(max_seq_length),这里设置为 512。这决定了模型一次能处理多少个 token。较小的值有助于降低显存占用,适合资源受限的训练环境。
vocab_size & padded_vocab_size: 词汇表大小。vocab_size 取自实际 tokenizer 的词汇量(64798),而 padded_vocab_size 通常向上取整到特定倍数(如 65024)以优化矩阵运算效率。
n_layer & n_head: 层数和头数。这里构建了 6 层 6 头的 Transformer。较浅的网络适合快速验证想法或处理简单任务,但表达能力有限。
n_embd: 隐藏层维度,设为 512。这是模型内部特征向量的长度,影响模型的容量。
head_size: 每个注意力头的维度,计算方式为 n_embd / n_head,此处为 48。
2.2 注意力机制与归一化
n_query_groups: 分组查询注意力(GQA)的参数。当 n_query_groups 等于 n_head 时,退化为多头注意力(MHA)。GQA 通过共享 key/value 头来减少内存访问,提升推理速度,同时保持较好的性能。
rotary_percentage: 旋转位置编码(RoPE)的比例。设置为 1.0 表示对所有维度应用 RoPE,这有助于模型更好地处理长序列和相对位置信息。
parallel_residual: 残差连接模式。设置为 false 表示使用非并行残差(类似原始 Transformer 设计),而 true 则对应 Pre-Norm 结构。当前主流大模型多采用 Pre-Norm 以提升训练稳定性。
bias: 控制 Linear 层的偏置是否存在。现代大模型通常将 bias 设为 false 以减少参数量并简化计算。
norm_class_name: 归一化类名,这里使用 RMSNorm。相比 LayerNorm,RMSNorm 去除了均值中心化步骤,仅保留均方根缩放,计算更高效且效果相当。
mlp_class_name: MLP 模块名称,使用 LLaMAMLP,通常包含 SwiGLU 激活函数,比传统 ReLU 具有更好的表现。
intermediate_size: MLP 中间层大小,这里是 768。通常设置为隐藏维度的 2~4 倍之间。
按照上述配置,得到的模型参数量约为 44M(0.044B)。根据微软 TinyStories 论文的结论,10-80M 级别的模型能在小故事生成这种简单的语言任务上达到不错的效果,足以理解基本的语法和语义。
2.3 Embedding 缩放
scale_embeddings: 控制是否在 forward 函数中对 embedding 进行缩放。如果为 True,会对输入嵌入向量乘以 sqrt(hidden_size)。虽然理论上对收敛有帮助,但在实践中区别不大,通常可设为 False。
3. 训练参数详解 (train)
训练参数决定了模型如何学习以及学习的节奏。
global_batch_size: 全局批次大小,跨所有数据并行副本的样本总数,这里为 512。较大的 batch size 有助于梯度估计更稳定。
micro_batch_size: 单卡/单设备批次大小,这里为 32。由于显存限制,通常无法直接设置全局 batch size,需要通过梯度累积来实现。
lr_warmup_steps: 学习率预热步数,设为 1000。在训练初期逐步增加学习率,避免初始梯度爆炸导致模型发散。
max_tokens: 总训练 Token 数上限。这是一个软约束,用于控制训练时长。
tie_embeddings: 是否绑定输入输出 Embedding 权重。设为 True 可以减少参数量,并在某些情况下提升泛化能力。
max_norm: 梯度裁剪的最大范数,防止梯度爆炸。
4. 优化器与精度
optimizer: 使用 torch.optim.AdamW,这是目前最主流的优化器。
lr: 学习率 0.0005,针对小模型较为合适。
weight_decay: 权重衰减 0.1,用于正则化,防止过拟合。
betas: Adam 的动量参数 (0.9, 0.95)。
precision: 精度设置为 bf16-mixed。BF16(Brain Floating Point 16)保留了 FP32 的动态范围,同时减少了显存占用,是现代 GPU 训练大模型的首选。
5. 数据与 Tokenizer
data: 指定了数据类路径 litgpt.data.LitData。LitData 是 Lightning 提供的高效数据加载工具,支持流式读取和压缩存储,能显著提升 IO 效率。
tokenizer_dir: 分词器路径。这里使用了 ChatGLM3-6B 的 tokenizer,因为其在中文分词方面表现良好,且兼容性好。
6. 启动训练与监控
预训练启动的命令非常简单,只需指定配置文件的路径即可。
litgpt pretrain --config Experiments/configs/microstories.yaml
执行后,训练过程将自动开始。系统会根据配置保存 Checkpoint 到 out_dir 目录,并通过日志记录器发送指标。
6.1 监控指标
我们选择了 wandb 作为日志记录器。在训练过程中,可以通过 WandB 网页界面实时查看 Loss 曲线、Learning Rate 变化、GPU 利用率等关键指标。这对于调试超参数和调整训练策略非常有帮助。
6.2 常见问题排查
- 显存不足:如果 OOM,尝试减小
micro_batch_size 或 block_size。
- Loss 不下降:检查学习率是否过大,或确认数据预处理是否正确(如 Padding 是否一致)。
- 训练中断:利用
resume: true 配置,可以从最近的 Checkpoint 继续训练,无需从头开始。
7. 小结
本文详细介绍了 litgpt 的预训练模型配置文件,重点解析了模型架构参数(如 GQA、RMSNorm)、训练超参数(Batch Size、LR)以及数据配置。通过合理的参数组合,即使是 Micro 级别的模型也能有效运行。
训练启动后,建议密切观察 WandB 图表,确保 Loss 呈下降趋势。后续章节将探讨如何进行推理测试以及如何评估模型质量。积少成多,每次进步一点点就够了。
注:本文内容基于开源项目 LitGPT 文档整理,旨在分享技术实践细节。