docs/interview/pretrain.md

# 面试：预训练（Pretrain）深度

> 本仓库 `src/trainers/lm/pretrain.py`，对应 config：`configs/lm/lm_pretrain.yaml`

## 0. 整体流程

```
数据集 (pretrain_t2t_mini.jsonl)
    │  1,270,238 条纯文本
    ▼
PretrainDataset
    ├── 读取 jsonl，每条取 text 字段
    ├── tokenizer.encode() → input_ids
    ├── labels = input_ids 克隆（完整序列监督）
    └── 截断 / 填充到 max_seq_len
    │
    ▼
SkipBatchSampler → DataLoader
    │  batch_size 个样本 / 组
    │  SkipBatchSampler 支持断点续训跳过前 N 步
    ▼
Forward: model(input_ids, labels)
    ├── TokenEmbed → N×Block → RMSNorm → lm_head
    └── CE Loss(input=logits[..., :-1, :], target=labels[..., 1:])
    │
    ▼
Backward → 梯度累积 → 梯度裁剪 → Optimizer.step()
    │
    ▼
  Cosine LR 调度 → 定期保存 checkpoint
```

---

## Q1. PretrainDataset 怎么处理数据？

### 代码（`src/dataset/pretrain.py`）

```python
class PretrainDataset(Dataset):
    def __init__(self, data_path, tokenizer, max_length=340):
        self.data = self.load_data(data_path)
        self.tokenizer = tokenizer
        self.max_length = max_length

    def load_data(self, data_path):
        with open(data_path, 'r') as f:
            return [json.loads(line)['text'] for line in f]

    def __getitem__(self, index):
        text = self.data[index]
        input_ids = self.tokenizer.encode(text, max_length=self.max_length,
                                          truncation=True, padding='max_length',
                                          return_tensors='pt')[0]
        labels = input_ids.clone()
        return input_ids, labels
```

### 与 SFT 数据集的关键区别

| | PretrainDataset | SFTDataset |
|---|---|---|
| 数据格式 | 纯文本 | 多轮对话 JSON |
| label 构造 | `labels = input_ids.clone()`（全监督） | prompt 段置 `-100`，只标 assistant |
| 损失计算 | 所有 token 都参与 | 只有 assistant 回复参与 |
| 训练目标 | 下一个 token 预测 | 指令跟随回复生成 |

> 面试点：预训练时所有 token 都计算损失 vs SFT 只算 assistant 段，为什么？→ 预训练目标是学习语言分布，所有 token 都提供统计信息；SFT 是学习指令跟随能力，只需要拟合回复

---

## Q2. 损失函数具体怎么算？

### Next-Token Prediction + 标签平移

```python
# src/models/lm/model.py:63-70
def forward(self, input_ids, labels=None):
    hidden_states = self.model(input_ids)
    logits = self.lm_head(hidden_states)          # [B, S, V]

    if labels is not None:
        shift_logits = logits[..., :-1, :].contiguous()       # [B, S-1, V]
        shift_labels = labels[..., 1:].contiguous()           # [B, S-1]
        loss = F.cross_entropy(
            shift_logits.view(-1, shift_logits.size(-1)),
            shift_labels.view(-1)
        )
```

### 为什么平移？

- `logits[t]` 预测的是 `labels[t+1]`（第 t 个位置的输出预测第 t+1 个位置的 token）
- 不平移的话，模型学到的是恒等映射：`logits[t] ≈ labels[t]`
- 训练时看似 loss 很低，但生成时完全无法产生新内容

> 面试点：平移的本质是什么？→ 将语言模型建模为条件概率 P(xₜ|x₁,...,xₜ₋₁)，第 t 个位置的输出对应第 t+1 个 token 的概率分布

### 辅助损失：MoE aux_loss

```python
loss = res.loss + res.aux_loss
```

当 `use_moe=1` 时，MoE 层还会产生一个 auxiliary load-balancing loss，鼓励专家负载均衡。非 MoE 模式下 `aux_loss = 0`。

---

## Q3. 训练超参如何影响模型？

### 核心参数详解

| 参数 | 默认值 | 作用 | 调大影响 | 调小影响 |
|---|---|---|---|---|
| `hidden_size` | 512 | 每 token 的表示维度 | 容量↑，速度↓ | 容量↓，速度↑ |
| `num_hidden_layers` | 8 | Transformer 层数 | 深度↑，梯度传播难 | 深度↓，表达弱 |
| `batch_size` | 32 | 每步样本数 | 梯度稳，显存↑ | 梯度噪，显存↓ |
| `accumulation_steps` | 8 | 梯度累积步数 | 等效 batch↑，速度不变 | 等效 batch↓ |
| `max_seq_len` | 340 | 最大序列长度 | 上下文↑，显存↑ | 上下文↓ |
| `learning_rate` | 5e-4 | 初始学习率 | 收敛快，可能不稳 | 收敛慢，更稳 |
| `dtype` | bfloat16 | 混合精度类型 | 精度高，速度中 | 精度低，速度快 |

### 等效 Batch Size

```
等效 batch = batch_size × accumulation_steps
```

本仓库 pretrain 默认 `batch_size=32, accumulation_steps=8` → 等效 batch = 256。

### 梯度累积实现

```python
# src/trainers/lm/pretrain.py:36-47
loss = loss / args.accumulation_steps     # 归一化
scaler.scale(loss).backward()             # 累积梯度

if step % args.accumulation_steps == 0:   # 累积够 N 步后
    scaler.unscale_(optimizer)
    clip_grad_norm_(model.parameters(), grad_clip)
    scaler.step(optimizer)                # 更新参数
    scaler.update()
    optimizer.zero_grad(set_to_none=True) # 清空梯度
```

> 面试点：为什么 loss 要除以 accumulation_steps？→ 因为梯度是链式累积的，不归一化的话等效学习率会放大 accumulation_steps 倍。除以步数后，每步梯度相当于独立 batch 梯度的平均，保持学习率语义不变

---

## Q4. 学习率调度策略？

### Cosine Decay

```python
# src/utils/training.py:82-83
def get_lr(current_step, total_steps, lr):
    return lr * (0.1 + 0.45 * (1 + math.cos(math.pi * current_step / total_steps)))
```

### 调度曲线

```
lr
↑
│    lr * 0.55 ─────────── 余弦下降 ──→ lr * 0.1
│    (初始)                          (最终)
└─────────────────────────────────────→ step
```

- 初始 LR = `lr * 0.55`（cos 从 0 开始，`1 + cos(0) = 2` → `0.1 + 0.45*2 = 1.0`... 不对）
  仔细看：`current_step=0` 时 `cos(0)=1` → `0.1 + 0.45*2 = 1.0` → `lr * 1.0 = lr`
  所以在 step 0 时 LR 正好等于设置的 learning_rate。
- 最终 LR = `lr * 0.1`（cos(π) = -1 → `0.1 + 0.45*0 = 0.1`）
- 即 LR 从 `lr` 余弦衰减到 `0.1 * lr`

这种"不下到 0"的调度比标准 cosine 更好，保持模型在训练后期仍有适度更新能力。

---

## Q5. 混合精度训练怎么做？

### AMP (Automatic Mixed Precision)

```python
# src/trainers/lm/pretrain.py:121-123
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = torch.cuda.amp.autocast(dtype=dtype)
```

### bf16 vs fp16

| | bf16 | fp16 |
|---|---|---|
| 指数位 | 8 位（同 fp32） | 5 位 |
| 尾数位 | 7 位 | 10 位 |
| 数值范围 | 同 fp32（~3.4e38） | 有限（~6.5e4） |
| 是否需要 GradScaler | ❌ 不需要 | ✅ 需要 |
| 精度 | 低精度（7bit 尾数） | 高精度（10bit 尾数） |
| 硬件要求 | Ampere+（3090/A100 等） | 几乎所有 GPU |

### 为什么 bf16 不需要 GradScaler？

bf16 的指数范围和 fp32 一样，不会发生梯度下溢。fp16 的指数范围只有 5 位，小梯度会直接变 0，需要用 GradScaler 放大梯度再缩小。

### 精度保留技巧

```python
# src/core/norm.py:7-8
x = x.float()                              # RMSNorm 内部转 fp32
x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
return (self.weight * x).type_as(self.weight)  # 转回 bf16/fp16
```

即使训练是 bf16，归一化层内部转 fp32 计算再转回，避免归一化精度损失。

---

## Q6. 训练时显存都花在哪了？

### 显存分布（以 30M 参数模型为例）

```
模型参数 (fp32):      120 MB
梯度 (fp32):          120 MB
Adam 状态 (fp32×2):   240 MB
─────────────────────────────
模型状态总计:          480 MB

激活值 (bf16, 8层):   ~3-5 GB  ← 大头
输入数据:              <10 MB
CUDA 上下文:           ~200 MB
─────────────────────────────
总计:                  ~4-6 GB（取决于 batch_size 和 seq_len）
```

### 为什么激活值占这么多？

- 反向传播需要存储每层的中间激活值
- 每层存储 Q/K/V、注意力输出、MLP 中间结果等
- 数量级：`O(batch_size × seq_len × hidden_size × num_layers × k)`，k ≈ 15-34

### 没有激活检查点

本仓库**没有**使用 `torch.utils.checkpoint`（梯度检查点）。代价是激活值全存，好处是不需要重计算，训练速度更快。

> 面试点：激活检查点 trade-off 是什么？→ 节省显存（存部分激活，反向时重算），增加约 15-20% 计算时间

---

## Q7. 分布式训练是怎么做的？

### DDP (DistributedDataParallel)

```python
# src/utils/distributed.py
def init_distributed_mode():
    if int(os.environ.get("RANK", -1)) == -1:
        return 0                              # 单卡模式
    dist.init_process_group(backend="nccl")   # 多卡模式
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    return local_rank
```

### 启动方式

```bash
# 单卡
python -m trainers.lm.pretrain --config configs/lm/lm_pretrain.yaml

# 多卡（torchrun）
torchrun --nproc_per_node=4 -m trainers.lm.pretrain --config configs/lm/lm_pretrain.yaml
```

### DDP 原理

- 每个 GPU 一张完整的模型副本
- 前向/反向独立计算
- 反向传播后通过 `allreduce` 同步梯度
- 每个 GPU 独立执行 optimizer.step()

### 本仓库没有使用

- ❌ ZeRO（DeepSpeed）
- ❌ FSDP
- ❌ 张量/序列并行
- ❌ torch.compile（默认关闭，`use_compile: 0`）

> 面试点：为什么小模型不用 ZeRO/FSDP？→ 模型仅 30M 参数，单卡就能装下，DDP 的梯度同步开销也很小。ZeRO/FSDP 的通信量更大，对小模型反而可能更慢

---

## Q8. Checkpoint 如何保存和恢复？

### 保存（`train_epoch` 内）

```python
# src/trainers/lm/pretrain.py:59-67
if (step % args.save_interval == 0 or step == iters) and is_main_process():
    ckp = f'{args.save_dir}/{args.save_weight}_{hidden_size}.pth'
    torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)
    lm_checkpoint(lm_config, weight=args.save_weight, model=model,
                  optimizer=optimizer, epoch=epoch, step=step, ...)
```

### 保存两个文件

| 文件 | 内容 | 用途 |
|---|---|---|
| `pretrain_512.pth` | 模型权重 (fp16) | 推理/下游微调 |
| `pretrain_512_resume.pth` | 权重 + 优化器 + epoch + step | 断点续训 |

### 恢复训练 `from_resume=1`

```python
ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints')
if ckp_data:
    model.load_state_dict(ckp_data['model'])
    optimizer.load_state_dict(ckp_data['optimizer'])
    start_epoch = ckp_data['epoch']
    start_step = ckp_data.get('step', 0)
```

### 权重初始化 `from_weight`

```yaml
from_weight: none      # 从头训练（随机初始化）
from_weight: pretrain  # 加载 pretrain_512.pth 继续训练
from_weight: full_sft  # 加载 full_sft_512.pth 继续训练
```

`init_model` 会根据 `from_weight` 在 `save_dir` 下查找对应文件：

```python
weight_path = f'{save_dir}/{from_weight}_{hidden_size}.pth'
weights = torch.load(weight_path, map_location=device)
model.load_state_dict(weights, strict=False)
```

> 面试点：`strict=False` 意味着什么？→ 允许加载的权重和模型结构部分不匹配（如只加载 encoder 不加载 lm_head），LoRA 等场景常用

---

## Q9. SkipBatchSampler 如何实现断点续训？

```python
# src/utils/training.py:177-200
class SkipBatchSampler(Sampler):
    def __init__(self, sampler, batch_size, skip_batches=0):
        self.sampler = sampler                  # 原始索引（或 DistributedSampler）
        self.batch_size = batch_size
        self.skip_batches = skip_batches        # 跳过的步数

    def __iter__(self):
        batch = []
        skipped = 0
        for idx in self.sampler:
            batch.append(idx)
            if len(batch) == self.batch_size:
                if skipped < self.skip_batches:  # 跳过前 N 个 batch
                    skipped += 1
                    batch = []
                    continue
                yield batch
                batch = []
```

### 为什么需要跳过？

- 断点续训时已经从 checkpoint 恢复了优化器状态
- 但 DataLoader 从头开始迭代的话，会重复处理之前的数据
- `SkipBatchSampler` 通过 `skip_batches` 跳过已处理的 batch

### 使用场景

```python
skip = start_step if (epoch == start_epoch and start_step > 0) else 0
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
```

---

## Q10. 随机种子和数据打乱

```python
# src/trainers/lm/pretrain.py:160
setup_seed(42 + epoch)
indices = torch.randperm(len(train_ds)).tolist()
```

### 每 epoch 重新打乱

- 每个 epoch 用不同的 `indices` 顺序
- 种子 = `42 + epoch`，保证可复现
- `DistributedSampler` 模式下，每个 GPU 拿到不同但确定性的分片

### 为什么 seed 要 + epoch？

- 每个 epoch 的数据顺序不同
- 同一 epoch 在不同运行间可复现
- 分布式下每个 rank 拿到不同的子集

---

## Q11. YAML 配置是如何生效的？

### 配置优先级

```
CLI 参数 > YAML 默认值 > Python argparse 默认值
```

### 实现机制

```python
# src/utils/training.py:25-38
def apply_config(parser, default_config=None):
    pre, _ = parser.parse_known_args()
    config_path = getattr(pre, 'config', None) or default_config
    if config_path and os.path.exists(config_path):
        defaults = _load_yaml_config(config_path)
        parser.set_defaults(**defaults)     # YAML 值设置为 argparse 默认值
    return parser.parse_args()              # CLI 显式传参仍可覆盖
```

### YAML 映射

```yaml
model:
  hidden_size: 512    → args.hidden_size
train:
  batch_size: 32      → args.batch_size
paths:
  save_dir: ...       → args.save_dir
```

YAML 的三个 section（model/train/paths）被扁平化为 argparse 参数后注入，之后任何 CLI 传参都可以覆盖。

---

## Q12. 为什么选 SwiGLU 作为激活函数？

### 公式对比

| 激活函数 | 公式 | 参数量 |
|---|---|---|
| ReLU FFN | `ReLU(xW₁)W₂` | 2×d×d_ff |
| SwiGLU | `SiLU(xW_g) * (xW_u) * W_d` | 3×d×d_ff |

SwiGLU 用三个投影（gate/up/down）替代两个，参数量增加 50%，但同等参数下效果更好。

### 仓库实现

```python
# src/core/mlp.py:16-17
def forward(self, x):
    return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
```

### intermediate_size 的奇怪计算

```python
# src/models/lm/config.py:22
self.intermediate_size = math.ceil(hidden_size * math.pi / 64) * 64
```

为什么用 π？这是一个取整技巧：`hidden_size × π / 64` 取整再 ×64，保证 intermediate_size 是 64 的倍数，有利于 GPU 内存对齐和 Tensor Core 加速。

对于 `hidden_size=512`：`intermediate_size = ceil(512 × π / 64) × 64 = ceil(25.13) × 64 = 1664`

---

## Q13. 训练过程中 loss 的变化规律

典型 pretrain loss 曲线：

```
loss
↑
8.0 │  █
7.0 │    ██
6.0 │      ███
5.0 │        ████
4.0 │          █████
3.0 │            ██████
    └────────────────────→ step
```

### 特征

1. **快速下降期**（前 5-10% 步数）：loss 从 ~8.5 降到 ~5.0，模型学到基础词法/语法模式
2. **平稳下降期**：loss 稳步下降，学习更复杂的语义/知识
3. **接近收敛**：loss 下降变缓，接近理论下界

### 预估最终 loss

对于 vocab_size=6400 的随机初始化模型：
- 初始 loss ≈ log(6400) ≈ 8.76（均匀分布的交叉熵）
- 训练后 loss ≈ 2.5-3.5（取决于模型大小和数据量）
- 理论下限 ≈ 0（完美拟合数据分布，但实际上达不到）

---

## 面试高频题汇总

### 基础

1. **预训练和 SFT 的区别？** → 预训练从零学习语言分布（全 token 监督），SFT 学习指令跟随（只监督回复）
2. **为什么用 bf16 而不是 fp16？** → bf16 指数范围同 fp32，无需 GradScaler，训练更稳定
3. **梯度累积的作用？** → 显存不足时用计算换显存，等效增大 batch_size
4. **Cosine 学习率调度的优缺点？** → 平滑衰减，早期快速学习后期精细调优，但可能过早衰减

### 进阶

5. **如何估计训练时间？** → `总步数 = ceil(样本数 / batch_size) × epochs`，`总时间 = 总步数 × 每步时间`
6. **为什么 loss 除以 accumulation_steps？** → 保持梯度期望不变，等效于 "平均 N 个 mini-batch 的梯度"
7. **DDP 和 DP 的区别？** → DDP 每个 GPU 独立前反向 + allreduce 梯度，DP 是单进程多线程（GIL 限制性能）
8. **weight tying 的作用？** → 共享 embedding 和 lm_head 的权重矩阵，减少参数量（本项目默认开启）
9. **GQA 和 MHA 的区别？** → GQA 减少 KV head 数量，降低 KV cache 大小，推理时更省显存
10. **Flash Attention 为什么省显存？** → tiling 计算注意力矩阵，不显式存储 O(n²) 的 score 矩阵