
如何把模型扩到2B参数train-llm-from-scratch的大模型扩展路径【免费下载链接】train-llm-from-scratchA straightforward method for training your LLM, from downloading data to generating text.项目地址: https://gitcode.com/GitHub_Trending/tr/train-llm-from-scratchtrain-llm-from-scratch 是一个从零训练大模型的开源项目纯 PyTorch 手写 Transformer不依赖 trl、peft、transformers完整覆盖数据下载、预训练、SFT 与 RLHF 对齐。本文以把基座模型从默认的 400M 扩到 2B 参数为主线给你一条可直接照做的大模型扩展路径改哪 4 个配置、显卡够不够、显存不够怎么省。为什么小模型仓库也能扩到 2B这个项目最打动人的地方是扩展模型 改数字。架构完全透明每个阶段只是复用同一个 Transformer 骨干换数据和换损失函数规格配置n_embed / n_head / n_blocks参数量入门128 / 8 / 1约 13M教程基座512 / 8 / 8约 77M后训练默认1024 / 16 / 24约 406M2B 目标2048 / 16 / 36约 2.0B参数量可以用一个公式快速估算≈ 12 × n_embed² × n_blocks 2 × 词表 × n_embed注意力占 4 倍、MLP 占 8 倍 n_embed²嵌入和 lm_head 各占一份。代入默认配置 1024/24 正好得到 406M说明这个公式可以直接拿来当扩参计算器。模型本体在 src/models/transformer.py由 src/post_training/utils.py 里的build_model_from_config从配置一键构建——你不需要碰任何模型代码。三步把配置扩到 2B ⚡第一步改 4 个架构参数所有阶段共享一份模型配置 configs/base.json改它一处预训练、SFT、DPO、PPO、GRPO 全线对齐不会阶段之间参数量对不上参数默认值2B 建议值作用n_embed10242048隐藏层宽度参数量增长的主引擎n_blocks2436Transformer 层数决定模型深度n_head1616注意力头数须能整除 n_embedcontext_length10241024上下文长度可再调大但注意力开销随长度平方增长配置解析遵循四层合并规则优先级从低到高dataclass 默认值 configs/base.json 阶段 JSON 命令行参数机制详见 config/loader.py 与 config/post_training_config.py。所以你也可以不改文件用 CLI 一次性覆盖python scripts/pretrain_base.py --n_embed 2048 --n_blocks 36第二步对一遍 GPU 显存预算2B 参数的 bf16 权重约 4GB但训练还要装优化器状态和激活值不能只看权重体积。以下是仓库 README 给出的官方参考2B LLM Training列GPU显存能否训练 2BA10040GB✔上限约 6–8BRTX 4090 / 309024GB✔上限约 4BRTX 509032GB✔V100 / T4 / RTX 408016GB✘上限约 1.5–2BRTX 40608GB✘上限约 1B第三步备好三件省显存工具bf16 自动混合精度scripts/pretrain_base.py 默认开启bf16autocast无需额外配置梯度累积grad_accum让微批可以很小有效批大小 batch_size × grad_accum × 卡数显存只按微批算旧版路径的可选开关scripts/train_transformer.py 支持--amp --grad-checkpointing --grad-accum 8其中梯度检查点用计算换激活内存是 OOM 时最有用的一个。跑一次完整的 2B 预训练 备足数据模型越大越需要更多 token多下载几个分片即可scripts/prepare_pretrain_data.py 的--num_shards参数启动训练单卡直接运行多卡用 DDP 数据并行同一命令换一行启动命令即可# 单卡 python scripts/pretrain_base.py --n_embed 2048 --n_blocks 36 # 双卡数据并行注意每张卡仍需装下完整模型 torchrun --standalone --nproc_per_node2 scripts/pretrain_base.py --n_embed 2048 --n_blocks 36长训必备断点续训脚本按save_every周期保存检查点中断后用--resume接着跑2B 级别的长训练几乎离不开它开工前验证加--print-config打印合并后的最终配置确认参数无误想先花几秒在 CPU 上验证整条流程可以切到 configs/smoke/ 里的微型配置。训练时盯住 loss 曲线它应从接近ln(词表)的均匀猜测水平约 10.8稳步下降。上图是 77M 基座在 2×L40 上的真实曲线11.14 → 3.73扩到 2B 后轨迹相同只是需要更多 token 才能到同等水平。学习率保持3e-4的余弦退火 预热是默认稳妥值细节见 docs/02_pretraining.md。扩完之后后训练链自动跟上 这是这个项目最省心的地方——configs/base.json被所有阶段共享所以 2B 的预训练检查点可以直接流入 SFT → 奖励模型 → DPO/PPO/GRPO 整条链中间零改造。只需注意各阶段的显存倍数DPO要冻结一份参考模型副本显存约为 2 倍模型PPO要同时装 actor 价值头 参考模型约为 3 倍是显存大户2B 放不下时优先降batch_size和rollout_len。每扩一档用 scripts/eval_post_training.py 在同一个 GSM8K 轴上对比各阶段仓库的结论是基座越大、预训练算力越多这条曲线整体越高。一页速查大模型扩展配置表 ✅目标规模推荐配置备注约 1.5Bn_embed1536, n_blocks4824GB 显卡可尝试约 2.0Bn_embed2048, n_head16, n_blocks3624GB 起步40GB 从容约 3.1Bn_embed2560, n_blocks36建议 40GB 或多卡遇到 OOM 的处理优先级开 bf16 → 加梯度累积 → 开梯度检查点 → 降 context_length / batch_size → 多卡 DDP。建议的路径和作者给出的建议一致先训 13M 看到它说出通顺的词再依次 77M、406M、2B把n_embed和n_blocks一路推高直到触到显卡上限。每一档的每个阶段都小到一坐就能读完而它们共享同一个模型——这大概就是从零二字的意义。更多配置说明可查 docs/howto/configs.md完整命令清单见 docs/howto/commands.md。【免费下载链接】train-llm-from-scratchA straightforward method for training your LLM, from downloading data to generating text.项目地址: https://gitcode.com/GitHub_Trending/tr/train-llm-from-scratch创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考