从零手搓中文大模型 Day04：模型参数配置和训练启动

从零手搓中文大模型｜Day04｜模型参数配置和训练启动

今天我们来深入研究一下模型的配置。在基于 litgpt 框架进行大语言模型训练时，配置文件是连接数据、架构与训练逻辑的核心枢纽。

1. 训练配置文件概览

litgpt 使用的配置文件格式为 YAML，与 transformers 库的配置方式略有不同。其仓库中提供了一些预训练所用的 YAML 配置文件样例，主要用于需要自定义模型架构的场景。此外，litgpt 也内置了一些 Hugging Face 上的现成模型，可以直接拿来使用，但在从头训练或微调特定任务时，自定义配置更为常见。

以下是本次定义的一个完整配置文件，涵盖了模型架构、训练策略、数据加载及优化器设置等关键部分。

# The name of the model to pretrain. Choose from names in ``litgpt.config``.
model_name: microstories

# A ``litgpt.Config`` object to define the model architecture.
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

# Directory in which to save checkpoints and logs.
out_dir: /home/puppyapple/Server/BigAI/Chinese_LLM_From_Scratch/Experiments/Output/pretrain/microstories

# The precision to use for pretraining.
precision: bf16-mixed

# Optional path to a checkpoint directory to initialize the model from.
initial_checkpoint_dir:

# Path to a checkpoint directory to resume from.
resume: true

# Data-related arguments
data:
  class_path: litgpt.data.LitData
  init_args:
    data_path: Chinese_LLM_From_Scratch/Data/TinyStoriesChinese/processed_data
    split_names:
      - train
      - val

# Training-related arguments
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

# Evaluation-related arguments
eval:
  interval: 2000
  max_new_tokens:
  max_iters: 100
  initial_validation: false
  final_validation: false

# Optimizer-related arguments
optimizer:
  class_path: torch.optim.AdamW
  init_args:
    lr: 0.0005
    weight_decay: 0.1
    betas:
      - 0.9
      - 0.95

# How many devices/GPUs to use
devices: auto
num_nodes: 1

# Optional path to the tokenizer dir
tokenizer_dir: Chinese_LLM_From_Scratch/References/chatglm3-6b

# The name of the logger to send metrics to
logger_name: wandb

# The random seed to use for reproducibility
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 文档整理，旨在分享技术实践细节。

从零手搓中文大模型 Day04：模型参数配置和训练启动