
llama2.c 最小模型 stories260K 完全指南训练配置、2X 多查询注意力与 C/Python 双端采样实践【免费下载链接】llama2.cInference Llama 2 in one file of pure C项目地址: https://gitcode.com/GitHub_Trending/ll/llama2.c导读本文以 doc/stories260K.md 为核心系统讲解 llama2.c 仓库中用于测试的最小模型 stories260K仅 260K 参数——它如何在 TinyStories 数据集上从零训练、为何采用8 头注意力 4 个 KV 头的 2X 多查询multiquery架构、如何配套训练 512 token 的自定义 tokenizer以及如何分别在纯 C 推理引擎 run.c 与 Python 脚本 sample.py 中完成贪心temperature0.0与 top-ptemperature1.0, p0.9两种采样。读完本文你将掌握用 train.py 复现微型 Llama 2 训练的完整参数含义并理解 C 与 Python 两条推理链路在终止逻辑上的关键差异。stories260K 是什么为测试而生的 260K 参数微型模型stories260K 是 llama2.c 训练并托管的一组 TinyStories 模型系列中的最小成员。根据 README.md 中models一节的模型对照表260K 模型的规格如下模型dimn_layersn_headsn_kv_heads最大上下文长度参数量val loss260K64584512260K1.297它的定位非常明确不是用来产出高质量故事而是作为工程验证的最小载体。因为参数极少模型文件只有约 2MB可以在几秒钟内完成 C 前向推理因此仓库的测试体系test_all.py直接依赖它来验证 run.c 与 PyTorch 前向结果的一致性详见下文测试验证一节。用文档中的原话说你不能对一个 260K 参数的模型期待太多它的价值在于让开发者以最低成本跑通训练 → 导出 → C 推理全链路。训练配置逐项拆解一行命令复现 260K 模型原文档给出了完整的训练命令这是全文的核心实操骨架逐项复现如下python train.py \ --out_diroutmini \ --batch_size128 \ --max_seq_len512 \ --gradient_accumulation_steps1 \ --vocab_sourcecustom \ --vocab_size512 \ --dim64 \ --n_layers5 \ --n_heads8 \ --n_kv_heads4 \ --multiple_of4 \ --learning_rate1e-3 \ --dropout0.05 \ --weight_decay0.01 \ --max_iters100000 \ --beta20.99 \ --warmup_iters1000 \ --eval_interval2000 \ --eval_iters100 \ --compileTrue这些参数在 train.py 中均有对应的默认值与注释逐一解读如下--out_diroutminicheckpoint 输出目录训练循环会在每个 eval 节点把ckpt.pt与model.bin写入该目录见 train.py其中model_export负责把 PyTorch 权重导出为 C 可读的 .bin 格式。默认值为out。--batch_size128单次迭代的微批大小。train.py 默认值为 128注释明确说明如果 gradient_accumulation_steps 1这是微批大小。260K 模型很小显存完全不是瓶颈所以可以开满 128。--max_seq_len512最大序列长度。在 model.py 的ModelArgs中默认是 2048Llama 2 官方取值此处降为 512。训练时每次从预分词数据中切出max_seq_len 1个 token 作为一组(x, y)见 tinystories.py序列越短单个样本覆盖的故事片段越小。--gradient_accumulation_steps1梯度累积步数。train.py 默认值是 4这里设为 1 意味着每一微批就是一个完整更新。结合batch_size × max_seq_len可算出单次更新的总 token 数128 × 512 65536约 65K tokens符合 README.md 训练指南中小规模应用可低于 100K的建议。--vocab_sourcecustom 与 --vocab_size512使用自定义词表而非 Llama 2 官方的 32000 token 词表。train.py 中有两个硬性断言train.pyvocab_source只能取llama2或custom且当不是 custom 时vocab_size必须等于 32000。因此训练 260K 模型必须二者成对指定。512 token 的词嵌入表512 × 64非常小这正是参数量能压到 260K 的关键原因之一。--dim64Transformer 的隐藏维度也是每个 token 嵌入向量的维度。260K 模型取 64是仓库所有模型中最小的一档对比15M 模型为 28842M 为 512110M 为 768。--n_layers5TransformerBlock 层数model.py 据此构建ModuleList循环堆叠 5 个 block。--n_heads8 与 --n_kv_heads4注意力头数 8、KV 头数 4构成 2X 多查询multiquery配置——这是本模型最值得注意的架构点详见下一节。--multiple_of4MLP 隐藏层尺寸向上取整到该值的倍数。model.py 的FeedForward中先按4*dim计算候选隐藏维度再乘 2/3最后用multiple_of * ((hidden_dim multiple_of - 1) // multiple_of)对齐。dim64 时4*64*2/3 ≈ 170取 4 的倍数得到 172。--learning_rate1e-3 与 --beta20.99AdamW 优化器的峰值学习率与 beta2。train.py 默认learning_rate5e-4, beta10.9, beta20.95train.py。文档中特别强调非常小的网络可以用较大的学习率1e-3 甚至更高因为模型小、优化曲面简单更大的 LR 能更快收敛。--dropout0.05 与 --weight_decay0.01dropout 只作用于注意力输出与 FFN 输出model.pyweight_decay 只施加于 2D 权重张量matmul 与嵌入层1D 参数RMSNorm scale 等不衰减model.py。--max_iters100000 与 --warmup_iters1000总训练迭代数与学习率线性预热步数。学习率调度采用线性 warmup 余弦衰减train.pylr_decay_iters默认等于max_iters。--eval_interval2000 与 --eval_iters100每 2000 步在 train/val 两个 split 上各用 100 个 batch 估算损失train.pyval loss 有提升时保存 checkpoint。--compileTrue启用 PyTorch 2.0torch.compile加速训练train.py。架构亮点8 头注意力配 4 个 KV 头的 2X 多查询设计原文档明确指出n_kv_heads是 4 而n_heads是 8所以每两个头共享一对 key、value 投影即这个模型是 2X multiquery。这一设计在 model.py 的Attention类中有完整实现。关键逻辑如下头数换算self.n_rep self.n_local_heads // self.n_local_kv_headsmodel.py。当 n_heads8、n_kv_heads4 时n_rep 2即每个 KV 头被复制 2 份供 2 个 query 头共享。投影矩阵维度差异wq输出n_heads * head_dim维而wk、wv只输出n_kv_heads * head_dim维model.py。当head_dim dim / n_heads 64/8 8时Q 投影是64×64而 K/V 投影仅为64×32KV 投影参数量直接减半。共享展开repeat_kvmodel.py在注意力计算前把 KV 头沿 head 维复制n_rep份等价于复制后计算但训练时只保留一套 KV 权重。相比标准 MHA8 个 KV 头2X MQA 把 KV 投影参数缩减 2 倍缓存KV cache占用也随之减半相比 1 个 KV 头的极端 MQA2X 又保留了一定的表达能力。由于 260K 模型里 KV 投影只占很小一部分这一设计更多是验证仓库对n_kv_heads ≠ n_heads分组多查询注意力 GQA的完整支持——这也是 run.c 需要读取n_kv_heads字段、export.py 需要在 header 中写出n_kv_heads的原因export.py。另外值得注意模型词嵌入与输出层共享权重weight tyingmodel.py以及使用 RoPE 相对位置编码model.py 预计算旋转频率freqs_cis在 export.py 中被序列化进 .bin 文件。这些都是 Llama 2 架构的标配260K 模型一应俱全。配套自定义 tokenizer512 token 的 BPE 词表260K 模型使用vocab_sourcecustomvocab_size512意味着它不依赖 Llama 2 的 32000 词表而是为 TinyStories 语料专门训练了一个 512 token 的词表。仓库 tinystories.py 的train_vocab阶段封装了这一流程内部调用 sentencepiece关键训练设置包括model_typebpe、character_coverage1.0、split_digitsTrue、byte_fallbackTrue、normalization_rule_nameidentity。在 README.md 的custom tokenizers一节中作者给出了训练自定义词表的通用三步流程tinystories.py 的 CLI 帮助信息与之完全一致python tinystories.py download python tinystories.py train_vocab --vocab_size4096 python tinystories.py pretokenize --vocab_size4096对 260K 模型只需把--vocab_size换成 512 即可。train_vocab会从 TinyStories 数据集中抽取前 10 个 shard 拼成临时文本data/tiny.txt用于训练tinystories.py产物保存在data/tok512.model随后pretokenize阶段按该词表把全部语料转成 uint16 整数序列存为 .bin 文件。训练时 train.py 只关心词表大小以正确初始化嵌入层512 × 64的嵌入矩阵不关心 token 本身内容。原文档还特别提到512 这个数字本身就是小而美的体现——README 中指出针对 TinyStories 专门训练的 4096 词表在压缩率上已能媲美通用 32000 词表而 512 词表则把词嵌入参数压到极致同时让推理时每步的 softmax 分类面大幅缩小速度更快。采样实践一C 端贪心解码与命令行参数训练完成后可用 run.c 以温度 0.0确定性贪心 argmax采样$ ./run stories260K/stories260K.bin -z stories260K/tok512.bin -t 0.0原文档给出了该命令的确定性输出开头部分Once upon a time, there was a little girl named Lily. She loved to play outside in the park. One day, she saw a big, red ball. She wanted to play with it, but it was too high. ...后略这条命令涉及两个关键点-z stories260K/tok512.bin显式指定自定义 tokenizer 的 .bin 格式文件。run.c 的-z参数注释为 optional path to custom tokenizerrun.c。若省略run.c 会回退到默认 Llama 2 tokenizer而模型本身是用 512 词表训练的——即使模型能生成正确的整数序列也会被错误词表翻译成乱码文本所以-z在这里是必须的。-t 0.0温度设为 0。run.c 中温度 0 走的是确定性分支——直接取 logits 最大的 tokenargmax无任何随机性因此输出可完全复现。run.c 完整支持的采样相关参数run.c如下表参数含义默认值-t float采样温度取值 [0, inf]0 为贪心1.0-p floattop-p核采样阈值取值 [0,1]1.0 表示关闭0.9-s int随机种子time(NULL)-n int生成步数0 表示跑到 max_seq_len256-i string输入 prompt空-z string自定义 tokenizer 路径无-m string模式generate / chatgenerate其中-s种子参数保证了随机采样结果可复现——同一模型、同一种子、同一温度与 top-p产出完全一致的故事。采样实践二Python 端复现与 BOS 终止差异原文档给出 Python 端等价的贪心采样命令$ python sample.py --checkpointstories260K/stories260K.pt --tokenizerstories260K/tok512.model --temperature0.0 --max_new_tokens257注意两点差异文件格式不同sample.py读取 PyTorch 的.ptcheckpointtokenizer 也需指定 sentencepiece 的.model文件C 端则是.bin格式。这两个文件在 huggingface 上以同名不同后缀的形式成对提供。--max_new_tokens257是硬编码的。原文档解释了原因sample.py目前不会像run.c那样在特殊的 BOS token 处终止生成所以必须手动限制长度。从 sample.py 的生成循环可见它调用model.generate()并固定生成max_new_tokens个 token 后直接打印没有任何停止符检查而 model.py 的generate同样只按max_new_tokens迭代。相比之下run.c 在解码循环中遇到 BOStoken id1会提前结束输出这就是 C/Python 两端行为不一致的根源也是 257 这个数字的来历——预分词时每段故事前都会加 BOStinystories.py生成 257 个 token 恰好能覆盖到下一个 BOS 附近与 C 端输出对齐。采样实践三温度 1.0 top-p 0.9 的多样化采样贪心输出虽然稳定但故事会陷入简单重复原文档中的 Lily 样例明显在句子层面打转。改用带随机性的采样能显著提升文本多样性$ ./run stories260K/stories260K.bin -z stories260K/tok512.bin -t 1.0 -p 0.9 -s 133742参数含义温度保持 1.0不改变 logits 分布但走多项式采样分支而非 argmax-p 0.9开启核采样nucleus sampling——把概率累加到 0.9 的最小 token 集合之外的低概率 token 全部裁剪为 0 概率从剩余集合中采样-s 133742固定种子保证可复现。从 run.c 的实现看[run.c](https://link.gitcode.com/i/c5a035ffeb783c874bd7e6d91ec88a17#L625, L709) 附近的注释与逻辑top-p 采样需要把候选 token 按概率排序后做累积裁剪比贪心多一步排序开销。README 的采样建议与此一致追求最佳效果推荐-t 1.0 -p 0.9因为 top-p 能避免手气差采样到极小概率 token降低生成中途脱轨的概率。更一般的原则是调节多样性时要么只动温度-t在 0~1 间变化并把-p 0关掉 top-p要么只动 top-p-p在 0~1 间变化并保持-t 1不要同时调两者。值得一提的是固定种子下用温度 1.0 采样会得到不同于贪心的故事——这也直观演示了-t 1.0与-t 0.0在 run.c 中走的是完全不同的两条代码路径多项式采样 vs topk argmax。训练结果与模型边界原文档记录了 260K 模型的训练结果训练时长约 10 分钟在作者的单张 A100 上原文档以问号标注存疑验证损失达到 1.2968README 表格中记录为 1.297二者一致文本质量贪心采样能产出语法基本正确、但内容高度重复的儿童故事temperature 1.0 top-p 0.9 时故事更合理但逻辑跳跃依然明显。原作者的原话值得引用你不能对一个 260K 参数的模型期待太多我甚至对能走到这一步感到有点惊讶。 260K 参数放在今天的大模型语境里几乎可以忽略不计但它在 TinyStories 这种领域极窄的合成数据集上依然学会了基本的英语句法、词汇搭配和简单的叙事模式——这正是 README 开篇强调的观点只要把领域收窄到足够程度极小的 LLM 也能展现出令人惊讶的表现。测试验证stories260K 是 C/Python 一致性测试的基石stories260K 最大的工程价值在于测试。test_all.py 中的test_runc用例会自动下载stories260K.bin、stories260K.pt、tok512.bin、tok512.model四个文件到临时test目录仅约 2MBtest_all.py执行./run test/stories260K.bin -z test/tok512.bin -t 0.0 -n 200把输出与expected_stdout中预置的已知正确输出逐字节比对test_all.py。也就是说前文展示的那段 Lily 故事正是测试套件的黄金标准输出。由于-t 0.0是确定性贪心run.c 的输出必须精确复现预期文本任何前向计算错误矩阵维度、RoPE、RMSNorm、注意力实现等都会导致断言失败。运行方式pip install pytest pytest这与原文档 This is a tiny model used for testing 的定位完全吻合一个几十 MB 的模型跑完整测试太慢260K 模型让整个测试套件在几秒内完成同时仍能严格校验 C 推理引擎的正确性。扩展阅读从 260K 到 110M 的模型系列如果你觉得 260K 的故事太碎README 的模型表给出了同系列更大的预训练模型均为 TinyStories 上从零训练模型dimn_layersn_headsn_kv_heads最大上下文长度参数量val loss260K64584512260K1.297OG28866625615M1.07242M512888102442M0.847110M7681212121024110M0.760观察这张表可以更深刻地理解 260K 模型的特殊性它是唯一采用n_kv_heads n_heads的模型其余均为标准 MHA也是唯一使用自定义 512 词表的模型。从 260K 到 110Mval loss 从 1.297 一路降到 0.760——模型越大、故事越连贯这条曲线本身就是缩放定律在小尺度上的直观演示。110M 模型在参数量上相当于 GPT-1 / GPT-2 small 的量级其训练配置dim 768、n_layers 12、n_heads 12、seq len 1024、batch 16、grad accum 8、LR 4e-4、dropout 0.1、max_iters 200K在 README.md 训练指南中有完整记录可作为继续调大模型的参考模板。【免费下载链接】llama2.cInference Llama 2 in one file of pure C项目地址: https://gitcode.com/GitHub_Trending/ll/llama2.c创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考