377 lines
13 KiB
Plaintext
377 lines
13 KiB
Plaintext
{
|
||
"cells": [
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "aba635c1-cdcd-4d00-a123-00045edbd10f",
|
||
"metadata": {},
|
||
"source": [
|
||
"# 3-Pretrain\n",
|
||
"\n",
|
||
"预训练是模型经历的第一个阶段，在该阶段，模型将会吸收知识，学习尽可能正确的下一词语预测范式。在这个笔记本中，我们仅对预训练的训练流程进行展示和学习，因此只给出必要的代码片段，如 wandb 和 ddp 不会在此笔记本中涉及。\n",
|
||
"\n",
|
||
"此笔记本的完整实现见主仓库 `/minimind/train_pretrain.py`"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 1,
|
||
"id": "98cb706f-2e03-4d8e-8e45-3663c9ab3b64",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"# 导入依赖\n",
|
||
"import os\n",
|
||
"import platform\n",
|
||
"import argparse\n",
|
||
"import time\n",
|
||
"import math\n",
|
||
"import warnings\n",
|
||
"import pandas as pd\n",
|
||
"import torch\n",
|
||
"import torch.distributed as dist\n",
|
||
"from torch import optim, nn\n",
|
||
"from torch.nn.parallel import DistributedDataParallel\n",
|
||
"from torch.optim.lr_scheduler import CosineAnnealingLR\n",
|
||
"from torch.utils.data import DataLoader, DistributedSampler\n",
|
||
"from contextlib import nullcontext\n",
|
||
"\n",
|
||
"from transformers import AutoTokenizer\n",
|
||
"\n",
|
||
"from model.model import MiniMindLM\n",
|
||
"from model.LMConfig import LMConfig\n",
|
||
"from model.dataset import PretrainDataset"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 2,
|
||
"id": "c3b6f015-0e5c-4e88-bb0d-32092993933f",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"warnings.filterwarnings('ignore')"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "a46c9b5b-12eb-4cad-a859-c5101cdcba82",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 可选参数设置\n",
|
||
"\n",
|
||
"首先，查看训练的可选参数，这些参数在实际使用时通过解析命令行进行导入，我们用 class 进行包装."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 3,
|
||
"id": "3f07e381-961e-4a4f-9f37-ff2f0bc2918c",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"class args:\n",
|
||
" # out_dir: str = \"out\" # pytorch 格式权重文件保存位置 我们只展示训练过程 所以不使用\n",
|
||
" epochs: int = 1 # 训练轮数\n",
|
||
" batch_size: int = 2 # pretrain 数据集仅两个样本，设置 batch 为 2\n",
|
||
" learning_rate: float = 5e-4 # 学习率\n",
|
||
" device: str = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
|
||
" dtype: str = 'bfloat16' # 16 bit 浮点数：8 bit 指数 + 7 bit 尾数\n",
|
||
" # use_wandb: bool = False # 是否使用 wandb 我们不使用\n",
|
||
" wandb_project: str = 'MiniMind-Notebook'\n",
|
||
" num_workers: int = 1 # 工作进程数\n",
|
||
" # ddp：bool = False # 单机多卡\n",
|
||
" accumulation_steps: int = 1 # 梯度累积步数\n",
|
||
" grad_clip: float = 1.0 # 梯度剪裁\n",
|
||
" warmup_iters: int = 0 # 学习率热启动\n",
|
||
" log_interval: int = 1 # 每一步打印日志 仅用于观察\n",
|
||
" # save_interval: int = 100 # checkpoint 保存点 我们不使用\n",
|
||
" local_rank: int = 1 # device 设备号\n",
|
||
" dim: int = 512 # 词嵌入维度 模型超参数\n",
|
||
" n_layers: int = 2 # MiniMind Block 数量 模型超参数\n",
|
||
" max_seq_len: int = 512 # 序列长度阈值\n",
|
||
" use_moe: bool = False # 是否启用混合专家\n",
|
||
" data_path: str = './toydata/pretrain_data.jsonl' # 数据集路径"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 4,
|
||
"id": "8f73286e-48a0-4a5b-b7c4-8c5fefea5f14",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"查看工作设备 cuda\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"print(f'查看工作设备 {args.device}')"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "7e4b0ea8-0f41-47fe-8125-93629d03d451",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 初始化训练\n",
|
||
"\n",
|
||
"\n",
|
||
"接下来，我们对一些重要模块进行初始化，我们已经了解过，分词器，模型和数据集是大模型的基本组件，我们对其进行初始化。"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 5,
|
||
"id": "9818e376-7f1e-4347-b8c3-c8e30f7d587f",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"def init_model(lm_config):\n",
|
||
" tokenizer = AutoTokenizer.from_pretrained('../model/minimind_tokenizer')\n",
|
||
" model = MiniMindLM(lm_config).to(args.device)\n",
|
||
" print(f'LLM总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')\n",
|
||
" return model, tokenizer"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 6,
|
||
"id": "24d15001-b670-4333-9fbe-b224ab1993c5",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"LLM总参数量：8.915 百万\n",
|
||
"模型位于设备：cuda:0, 词表长度：6400, DataLoader：<torch.utils.data.dataloader.DataLoader object at 0x000001C6BF33F280>\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"lm_config = LMConfig(dim=args.dim, n_layers=args.n_layers, max_seq_len=args.max_seq_len, use_moe=args.use_moe)\n",
|
||
"model, tokenizer = init_model(lm_config)\n",
|
||
"\n",
|
||
"train_ds = PretrainDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len)\n",
|
||
"\n",
|
||
"train_loader = DataLoader(\n",
|
||
" train_ds,\n",
|
||
" batch_size=args.batch_size,\n",
|
||
" pin_memory=True,\n",
|
||
" drop_last=False,\n",
|
||
" shuffle=False,\n",
|
||
" num_workers=args.num_workers,\n",
|
||
")\n",
|
||
"\n",
|
||
"print(f'模型位于设备：{model.device}, 词表长度：{tokenizer.vocab_size}, DataLoader：{train_loader}')"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 7,
|
||
"id": "60c8c970-f448-4ea0-b21b-8f929e03f6fa",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"打印一个 iter 的数据:\n",
|
||
"[tensor([[ 1, 46, 46, ..., 0, 0, 0],\n",
|
||
" [ 1, 5349, 1619, ..., 0, 0, 0]]), tensor([[ 46, 46, 47, ..., 0, 0, 0],\n",
|
||
" [5349, 1619, 2875, ..., 0, 0, 0]]), tensor([[1, 1, 1, ..., 0, 0, 0],\n",
|
||
" [1, 1, 1, ..., 0, 0, 0]])]\n",
|
||
"\n",
|
||
"数据集大小：2, DataLoader 大小：1\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"loader = iter(train_loader)\n",
|
||
"print(f'打印一个 iter 的数据:\\n{next(loader)}\\n')\n",
|
||
"print(f'数据集大小：{len(train_ds)}, DataLoader 大小：{len(loader)}')"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "0d91904d-7ddf-4105-a5f5-65b84315a73b",
|
||
"metadata": {},
|
||
"source": [
|
||
"我们发现，train loader 的每一个 iter 都包含一个长度为 3 的张量列表，这是因为 train_dataset 每一次取数据都会返回三个张量，分别为:\n",
|
||
"\n",
|
||
"- 样本 X: 包含 \\<bos> 在内的输入 content\n",
|
||
"- 标签 Y: 包含 \\<eos> 在内的输出 content\n",
|
||
"- 掩码 loss_mask: 指示需要计算损失的 token 位置\n",
|
||
"\n",
|
||
"由于我们的数据集只有两条数据，而 batch size 设置为 2，因此我们的 dataloader 只有一个 iter."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "40764309-9c98-45d5-ada5-e31c1aaa02c5",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 启动训练\n",
|
||
"\n",
|
||
"训练一个深度学习模型，还涉及到了优化器，损失函数和学习率调度。接下来，我们查看 MiniMind 训练部分的代码，并进行一轮简单的训练。"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 8,
|
||
"id": "4f70b9cf-fcff-4242-bd96-ce34a5da55e4",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"# 学习率调度方面 采用余弦退火学习率\n",
|
||
"def get_lr(current_step, total_steps, lr):\n",
|
||
" return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps))\n",
|
||
"\n",
|
||
"# 优化器方面 选择 AdamW 优化器 并在混精度场景下创建 scaler 进行梯度缩放避免数值下溢\n",
|
||
"scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))\n",
|
||
"optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)\n",
|
||
"\n",
|
||
"device_type = \"cuda\" if \"cuda\" in args.device else \"cpu\"\n",
|
||
"ctx = nullcontext() if device_type == \"cpu\" else torch.cuda.amp.autocast() # 在 cuda 上启动混精度训练，否则空白上下文"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "c101ffb4-1a97-4e29-8612-caf40a8aa05d",
|
||
"metadata": {},
|
||
"source": [
|
||
"接下来，我们来看看 MiniMind 的训练函数"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 9,
|
||
"id": "03043d8e-e90b-42a8-ba4d-de6187f6eb81",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"def train_epoch(epoch):\n",
|
||
" loss_fct = nn.CrossEntropyLoss(reduction='none') # 损失函数 采用交叉熵损失\n",
|
||
" start_time = time.time()\n",
|
||
" for step, (X, Y, loss_mask) in enumerate(train_loader):\n",
|
||
" X = X.to(args.device)\n",
|
||
" Y = Y.to(args.device)\n",
|
||
" loss_mask = loss_mask.to(args.device)\n",
|
||
"\n",
|
||
" lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch, args.learning_rate)\n",
|
||
" for param_group in optimizer.param_groups:\n",
|
||
" param_group['lr'] = lr\n",
|
||
"\n",
|
||
" with ctx:\n",
|
||
" res = model(X) # 前向推理\n",
|
||
" loss = loss_fct(\n",
|
||
" res.logits.view(-1, res.logits.size(-1)),\n",
|
||
" Y.view(-1)\n",
|
||
" ).view(Y.size()) # 取生成的最后一个 token 的 logits 计算损失\n",
|
||
" loss = (loss * loss_mask).sum() / loss_mask.sum()\n",
|
||
" loss += res.aux_loss # 若为混合专家 则将 MOE 辅助损失纳入考虑\n",
|
||
" loss = loss / args.accumulation_steps # 梯度累积\n",
|
||
"\n",
|
||
" scaler.scale(loss).backward() # 梯度缩放\n",
|
||
"\n",
|
||
" if (step + 1) % args.accumulation_steps == 0:\n",
|
||
" scaler.unscale_(optimizer)\n",
|
||
" torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) # 梯度剪裁\n",
|
||
"\n",
|
||
" scaler.step(optimizer)\n",
|
||
" scaler.update()\n",
|
||
"\n",
|
||
" optimizer.zero_grad(set_to_none=True) # 单步更新\n",
|
||
"\n",
|
||
" if step % args.log_interval == 0:\n",
|
||
" spend_time = time.time() - start_time\n",
|
||
" print(\n",
|
||
" 'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.12f} epoch_Time:{}min:'.format(\n",
|
||
" epoch + 1,\n",
|
||
" args.epochs,\n",
|
||
" step,\n",
|
||
" iter_per_epoch,\n",
|
||
" loss.item() * args.accumulation_steps,\n",
|
||
" optimizer.param_groups[-1]['lr'],\n",
|
||
" spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))\n",
|
||
"\n",
|
||
" # 到达指定保存步数时，save as PyTorch\n",
|
||
" # if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):\n",
|
||
" # model.eval()\n",
|
||
" # moe_path = '_moe' if lm_config.use_moe else ''\n",
|
||
" # ckp = f'{args.save_dir}/pretrain_{lm_config.dim}{moe_path}.pth'\n",
|
||
"\n",
|
||
" # if isinstance(model, torch.nn.parallel.DistributedDataParallel):\n",
|
||
" # state_dict = model.module.state_dict()\n",
|
||
" # else:\n",
|
||
" # state_dict = model.state_dict()\n",
|
||
"\n",
|
||
" # torch.save(state_dict, ckp)\n",
|
||
" # model.train()"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "7792bf9c-95a1-4c0d-bb53-fac2a1132ee0",
|
||
"metadata": {},
|
||
"source": [
|
||
"准备完毕，我们尝试一轮长度 1 个 iter 的训练."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 10,
|
||
"id": "17f51549-2009-4322-86c8-7aa783169e0b",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"Epoch:[1/1](0/1) loss:8.959 lr:0.000550000000 epoch_Time:0.0min:\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"iter_per_epoch = len(train_loader)\n",
|
||
"for epoch in range(args.epochs):\n",
|
||
" train_epoch(epoch)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 11,
|
||
"id": "6f257b9e-3a1e-4686-9516-3f9eeb7d0654",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"del model"
|
||
]
|
||
}
|
||
],
|
||
"metadata": {
|
||
"kernelspec": {
|
||
"display_name": "Python 3 (ipykernel)",
|
||
"language": "python",
|
||
"name": "python3"
|
||
},
|
||
"language_info": {
|
||
"codemirror_mode": {
|
||
"name": "ipython",
|
||
"version": 3
|
||
},
|
||
"file_extension": ".py",
|
||
"mimetype": "text/x-python",
|
||
"name": "python",
|
||
"nbconvert_exporter": "python",
|
||
"pygments_lexer": "ipython3",
|
||
"version": "3.10.14"
|
||
}
|
||
},
|
||
"nbformat": 4,
|
||
"nbformat_minor": 5
|
||
}
|