一、训练核心准备：数据加载与Token处理
核心函数：tokenizing_distributed_data_loader
该函数用于从parquet文件流式读取预训练文本，进行Token化处理并生成训练批次，是训练数据准备的核心模块
def tokenizing_distributed_data_loader(B, T, split, tokenizer_threads=4, tokenizer_batch_size=128, device="cuda"):
"""从parquet文档读出流文本数据"""
assert split in ["train", "val"], "split must be 'train' or 'val'"
ddp, ddp_rank, ddp_local_rank, ddp_world_size = get_dist_info()
# +1 因为 target需要对inputs多一个token
needed_tokens = B * T + 1
# 获取 tokenizer 和 bos token
tokenizer = get_tokenizer()
bos_token = tokenizer.get_bos_token_id()
token_buffer = deque()
# 读出文档的batch
def document_batches():
while True:
# 获取train、val 的数据
for batch in parquets_iter_batched(split=split, start=ddp_rank, step=ddp_world_size):
# for the tokenizer we might want to go in usually smaller batches, e.g. 128 rows
for i in range(0, len(batch), tokenizer_batch_size):
yield batch[i:i+tokenizer_batch_size]
batches = document_batches()
batch_index = 0
while True:
# 获取指定量的token
while len(token_buffer) < needed_tokens:
doc_batch = next(batches)
# 文档内容转token，开头会添加bos_token
token_lists = tokenizer.encode(doc_batch, prepend=bos_token, num_threads=tokenizer_threads)
for tokens in token_lists:
# 存入token_buffer
token_buffer.extend(tokens)
batch_index += 1
# token_buffer取token到tokens
tokens = [token_buffer.popleft() for _ in range(needed_tokens)]
# CUDA卡则快速导入gpu，token转向量
scratch = torch.tensor(tokens, dtype=torch.int64, pin_memory=(device == "cuda"))
# 所有tokens作inputs，从下标1到结尾作targets
inputs_cpu = scratch[:-1].to(dtype=torch.int32)
targets_cpu = scratch[1:]
# 转成2D 便于并行gpu处理
inputs = inputs_cpu.view(B, T).to(device=device, dtype=torch.int32, non_blocking=True)
targets = targets_cpu.view(B, T).to(device=device, dtype=torch.int64, non_blocking=True)
yield inputs, targets
关键参数说明：
• B : 批次大小(batch size)，每个批次的序列数量
• T : 序列长度(sequence length)，每个序列的token数量
• split : 数据集分割，“train"或"val”
• tokenizer_threads : 分词器使用的线程数，默认为4
• tokenizer_batch_size : 分词器处理的批次大小，默认为128
• device : 目标设备，默认为"cuda"
• ddp : 是否启用分布式数据并行
• ddp_rank : 当前进程的排名
• ddp_local_rank : 本地排名
• ddp_world_size : 总进程数
核心处理逻辑：
1）Token数量计算：每个迭代需准备BT+1个Token，多1个Token是因为目标序列（targets）需比输入序列（inputs）多一个Token。
2）分词器与BOS Token获取：通过get_tokenizer()获取分词器，BOS（Beginning of Sequence）Token为序列起始标识，在该实现中定义为"<|bos|>"，通过encode_special方法获取其Token ID。
3）BOS Token的核心作用：
序列边界标记：明确标识新文档/序列的开始，帮助模型理解上下文边界。
4）训练稳定性：为模型提供一致的起始点，提升训练稳定性。
5）推理提示：推理时作为生成响应的起始信号。
6）数据流式读取与Token化：通过document_batches()函数从parquet文件流式读取数据，按分词器批次大小拆分后，调用分词器编码并在每个序列前添加BOS Token，将结果存入token_buffer。
7）批次生成：从token_buffer中提取所需数量的Token，转换为Tensor后拆分inputs（所有Token）和targets（从下标1开始的Token），最终reshape为2D张量（BT）并迁移至目标设备，生成训练批次。
BOS（Beginning of Sequence）token在语言模型训练中有几个重要作用：
1）序列边界标记：明确标识一个新文档/序列的开始，帮助模型理解上下文边界
2）训练稳定性：为模型提供一个一致的起始点，有助于提高训练稳定性
3）推理提示：在推理时，BOS token可以作为生成响应的起始信号
二、训练核心组件：优化器配置
优化器是训练过程中参数更新的核心，nanochat通过分组配置不同参数的学习率和优化器类型，保障训练效果和稳定性，核心函数为setup_optimizers。
def setup_optimizers(self, unembedding_lr=0.004, embedding_lr=0.2, matrix_lr=0.02, weight_decay=0.0):
model_dim = self.config.n_embd
ddp, rank, local_rank, world_size = get_dist_info()
# 将所有参数分为3组（矩阵参数、嵌入参数、语言模型头参数）
matrix_params = list(self.transformer.h.parameters())
embedding_params = list(self.transformer.wte.parameters())
lm_head_params = list(self.lm_head.parameters())
assert len(list(self.parameters())) == len(matrix_params) + len(embedding_params) + len(lm_head_params)
# dmodel_lr_scale 学习率缩放系数：随模型维度增大，学习率相对减小，保持训练稳定性和收敛性
# 结合预热和平滑衰减策略，减缓大维度模型初始化时的梯度振荡
dmodel_lr_scale = (model_dim / 768) ** -0.5
if rank == 0:
print(f"Scaling the LR for the AdamW parameters ∝1/√({model_dim}/768) = {dmodel_lr_scale:.6f}")
# AdamW优化器配置（用于嵌入层和语言模型头）
adam_groups = [
dict(params=lm_head_params, lr=unembedding_lr * dmodel_lr_scale),
dict(params=embedding_params, lr=embedding_lr * dmodel_lr_scale),
]
adamw_kwargs = dict(betas=(0.8, 0.95), eps=1e-10, weight_decay=weight_decay)
AdamWFactory = DistAdamW if ddp else partial(torch.optim.AdamW, fused=True)
adamw_optimizer = AdamWFactory(adam_groups, **adamw_kwargs)
# Muon优化器配置（用于线性层）
muon_kwargs = dict(lr=matrix_lr, momentum=0.95)
MuonFactory = DistMuon if ddp else Muon
muon_optimizer = MuonFactory(matrix_params, **muon_kwargs)
# 合并优化器并记录初始学习率
optimizers = [adamw_optimizer, muon_optimizer]
for opt in optimizers:
for group in opt.param_groups:
group["initial_lr"] = group["lr"]
return optimizers
核心设计思路：
1）参数分组：根据参数类型将模型参数分为三组——矩阵参数（transformer.h）、嵌入参数（transformer.wte）、语言模型头参数（lm_head），不同组采用不同优化策略，适配各组件特性。
2）学习率缩放：引入dmodel_lr_scale缩放系数，核心逻辑为“模型维度越大，学习率相对越小”，计算公式为(model_dim / 768) ** -0.5，结合预热和平滑衰减策略，减缓大维度模型初始化时的梯度振荡。
3）双优化器组合：
AdamW优化器：用于嵌入层（embedding_params）和语言模型头（lm_head_params），配置不同基础学习率并乘以缩放系数，设置betas=(0.8, 0.95)、eps=1e-10等超参数，分布式场景使用DistAdamW，非分布式场景使用融合优化的torch.optim.AdamW。
4）Muon优化器：用于线性层（matrix_params），基础学习率为0.02，动量0.95，分布式场景使用DistMuon，非分布式场景使用Muon。
5）初始学习率记录：为每个优化器的参数组记录初始学习率，用于后续学习率调度。
三、训练核心流程：迭代与参数更新
训练迭代过程是模型参数更新和性能优化的核心，包含评估、前向/后向传播、梯度处理、学习率调度和参数更新等关键步骤。
# 训练迭代
for step in range(num_iterations + 1):
last_step = (step == num_iterations)
# 每一定步数做评估
if last_step or step % eval_every == 0:
……
# 一定步数做核心评估或者最后一步评估
if core_metric_every > 0 and (last_step or (step > 0 and step % core_metric_every == 0)):
……
# 样本评估
if master_process and (last_step or (step > 0 and step % sample_every == 0)):
……
# 最后一步做权重保存
if master_process and last_step:
……
# 前向与后向传播，梯度计算
for micro_step in range(grad_accum_steps):
# 当使用CUDA设备时：启用自动混合精度训练（使用bfloat16）
# 当使用非CUDA设备（CPU/MPS）时：使用nullcontext()，即不进行任何特殊处理
autocast_ctx = torch.amp.autocast(device_type=device_type, dtype=torch.bfloat16) if device_type == "cuda" else nullcontext()
with autocast_ctx:
# 前向传播计算模型输出
loss = model(x, y)
train_loss = loss.detach() # for logging
# 梯度累加场景下归一化损失：避免梯度累积导致的梯度放大
loss = loss / grad_accum_steps
# 计算梯度并累积
loss.backward()
# 预取下一个批次（利用GPU前向/后向传播的空闲时间）
x, y = next(train_loader)
# 梯度裁剪：防止梯度爆炸
grad_clip_enabled = grad_clip > 0.0
if grad_clip_enabled:
grad_norm_tensor = torch.nn.utils.clip_grad_norm_(orig_model.parameters(), grad_clip)
grad_norm = grad_norm_tensor.item() # GPU tensor -> CPU float（注意：存在CPU-GPU同步）
# 学习率调度：根据步数调整学习率系数
lrm = get_lr_multiplier(step)
for opt in optimizers:
for group in opt.param_groups:
group["lr"] = group["initial_lr"] * lrm
# Muon优化器动量调度
muon_momentum = get_muon_momentum(step)
for group in muon_optimizer.param_groups:
group["momentum"] = muon_momentum
# 参数更新与梯度清零
for opt in optimizers:
opt.step() # 应用梯度更新参数
model.zero_grad(set_to_none=True) # 清零梯度准备下一轮训练
迭代核心步骤：
1）阶段性评估与保存：
定期评估：每eval_every步或最后一步执行常规评估；每core_metric_every步或最后一步执行核心指标评估。
2）样本评估：主进程（master_process）每sample_every步或最后一步执行样本评估，验证模型生成效果。
3）权重保存：主进程在最后一步保存模型权重，保障训练成果。
4）梯度累加与混合精度训练：
梯度累加：通过grad_accum_steps设置梯度累加步数，多次前向/后向传播后再更新参数，变相扩大批次大小，适配显存限制。
5）混合精度训练：CUDA设备下启用bfloat16自动混合精度训练，提升训练速度并减少显存占用；非CUDA设备使用默认精度。
6）损失归一化：梯度累加场景下，将损失除以grad_accum_steps，避免多次梯度累积导致的梯度放大，保障梯度规模合理性。
7）梯度处理：通过梯度裁剪（clip_grad_norm_）限制梯度范数，防止梯度爆炸，提升训练稳定性；梯度裁剪后将梯度张量转为CPU浮点数，用于训练监控（如wandb）。
8）学习率与动量调度：调用get_lr_multiplier(step)获取学习率乘数，调整各优化器参数组的学习率；调用get_muon_momentum(step)调整Muon优化器的动量参数，适配训练不同阶段需求。
参数更新与梯度清零：调用优化器step()方法应用梯度更新参数，随后调用model.zero_grad(set_to_none=True)清零梯度，为下一轮迭代做准备。
关键调度策略：学习率调度器：
学习率调度采用三阶段策略，根据训练步数动态调整学习率，平衡训练初期稳定性、中期收敛速度和末期精细收敛，核心函数为get_lr_multiplier。
def get_lr_multiplier(it):
warmup_iters = round(warmup_ratio * num_iterations)
warmdown_iters = round(warmdown_ratio * num_iterations)
if it < warmup_iters:
return (it + 1) / warmup_iters
elif it <= num_iterations - warmdown_iters:
return 1.0
else:
progress = (num_iterations - it) / warmdown_iters
return progress * 1.0 + (1 - progress) * final_lr_frac
get_lr_multipli er(it) 是一个三阶段学习率调度器，根据训练步数动态调整学习率。
1） 预热阶段 (Warmup) 持续时间：
学习率从0线性增长到初始学习率，避免训练初期梯度爆炸，帮助模型稳定收敛 warmup_ratio * num_iterations
2）稳定训练阶段 (Plateau)
持续时间：主要训练阶段学习率保持为初始学习率（乘数为1.0）模型进行主要的学习和参数更新
3） 冷却阶段 (Warmdown/Cosine Decay)
持续时间： warm down_ratio * num_iterations 步学习率从初始值线性衰减到 fina l_lr_frac 倍帮助模型在训练末期更精细地收敛
五、nanochat的训练分成预训练、中期预训练、微调、强化学习
-
其中预训练和中期训练在技术上都是使用无监督或自监督学习的方式进行，但它们不合并。原因：
1）预训练使用的数据量非常庞大（通常是TB级别）。中期训练使用的数据量相对较小，但质量要求极高，通常需要大量人工清洗和标注
2）预训练是一个极度昂贵的过程。一旦完成，开发者可以获得一个通用的基座模型。不同的下游任务（如对话、代码生成、医疗问诊）可以直接取用这个相同的基座模型，然后只针对自己的需求运行中期训练。这大大降低了进入特定领域的门槛和成本。
3）预训练的目标是“语言建模”，即预测下一个词。中期训练的目标是利用较低的学习率在不破坏原有通用知识的基础上，将模型的“知识分布”和“行为模式”向特定的领域或格式迁移。 -
预训练与sft微调训练的区别：
1）首先预训练（Pre-training）和有监督微调（Supervised Fine-Tuning, SFT），在技术上都是“预测下一个 Token”，但本质区别在于数据类型、学习目标和训练目的。
2）预训练是无监督/自监督，模型自己从连续文本中训练。目的是学习语言的统计分布、语法和世界知识。
3）有监督微调，数据经过人工筛选、标注或格式化，明确指定了输入和输出。目的是学习如何遵循指令、输出高质量、安全且有帮助的特定格式的回答。
4）另外计算损失的不同，SFT 的数据通常是[指令] + [目标回答]拼接在一起输入模型。但很多实现中，模型在计算损失时，会只计算[目标回答]部分的损失，而忽略[指令]部分的损失。
5）计算损失的相同，a. 模型的损失计算是逐 Token进行，每一次的损失计算都是使用交叉熵损失 (Cross-Entropy Loss)，它是衡量模型预测的 Token 概率分布与目标 Token 的"One-Hot" 分布之间差异的标准方法。b. 模型在计算后续 Token 的损失时，总是使用对应位置的“实际值”（目标 Token）作为输入上下文，而不是使用模型上一步的“预测值”，稳定训练。 -
强化学习RLHF 的损失计算机制：
1）不再是逐 Token 的交叉熵：RL 的损失不是将模型预测的 Token 和目标 Token 进行比较。
2）奖励是序列级别的：损失是基于整个生成序列获得的奖励分数。
3）同时优化与约束：RL 损失是两个项的组合：最大化奖励项（推动模型向人类偏好对齐）和 KL 散度约束项（与基础模型回复对比，防止模型遗忘基础语言能力）。
扫一扫
428
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
