波仔
🦙 第5章 基于 Llama 2 定制行业大模型
📄 一句话概括

加载模型这件事，说到底就是三个数：参数量、字节数、显存余量。 三个数算清，多卡怎么切、能不能跑、能扛几个并发，全都有答案。 这一节真读一份 config，算出 67.38 亿参数、fp16 12.6 GiB，再把三十二层分到两张卡上。

📝 本章和前面几章的关系

上一节备好了数据，这一节把模型立起来。 第 2 章 2.1 节讲过三笔显存账的算法， 这一节把它落到一份真的配置文件上——从配置字段一路算到「每张卡还剩多少」。

一、读配置：先看四个数

hidden_size 是隐藏维（4096），num_hidden_layers 是层数（32）， intermediate_size 是前馈中间维（11008），vocab_size 是词表大小（32000）。 这四个数决定 95% 的参数量。剩下那些——头数、最大位置、rope 参数——只影响怎么算，不影响多大。 参数量按块算：一层注意力是 4 个 d×d，一层前馈是 3 个 d×F，加上词嵌入和输出头，就是全部。

演示里算出来 67.38 亿，官方公布 67.39 亿，差 0.01%。 这点差来自偏置、归一化层的参数，以及输出头是否共享权重——算到百分之一以内，就够做显存规划了。

二、算体积：三样东西

权重：参数量乘字节数。fp16 是两个字节，67.38 亿就是 12.6 GiB。 KV 缓存：2 乘层数乘 KV 头数乘每头维度乘序列长度乘字节数。7B 在 4096 上下文下是 2.00 GiB。 激活与框架开销：一到两 GiB，跟批大小和解码方式有关。这一项最难算准，所以留余量要留够。

三项里只有第一项是固定的。后两项跟着你的用法变——上下文翻倍，缓存翻倍；并发翻倍，缓存再翻倍。

三、切分：贪心装箱

多卡放不下整个模型时，就按层切。device_map='auto' 做的是贪心：每块放进当前占用最小的那张卡。 演示里两种切法对比：把嵌入层和输出头都塞给卡 0、剩下按层数平均，结果是 6.52 对 6.03， 不均衡度 1.08；按字节贪心则是 6.28 对 6.28，1.00。 这个差看着不大，但模型越大、卡越多，差得越多。 70B 八十层切八张卡，嵌入层那一份就能让某张卡提前 OOM。

不均衡还有第二个代价：切了就有跨卡通信。 每层算完要把激活传给下一张卡，卡间带宽比卡内显存低一个数量级， 所以「能放下」的切法未必是「跑得快」的切法。

四、初始化：四个开关

torch_dtype 决定运行精度。不写这一项，很多仓库默认 fp32，显存直接翻倍。 device_map 决定放哪。单卡就写 cuda:0，多卡写 auto。 pad_token 要跟 eos_token 设成同一个。 很多模型没有专门的填充符，不设会在批处理时报错，或者悄悄用 0 号 token 填充，把结果带歪。

use_fast 建议设成 False。 快速分词器基于另一套实现，在个别字符上跟原版不一致——微调时这种不一致会被放大成「模型学歪了」。

五、动手 Demo：读配置、算体积、把三十二层分到两张卡上

造一份 config.json 真落盘、真读回来，按 Llama 2 7B 的规格。 然后分项算参数量，再算两种切法下的显存占用。 全部是算术，不需要显卡。算对了，买卡和排期才有依据。

d, L, F, V = 4096, 32, 11008, 32000

attn = 4 * d * d                  # q k v o 四路
ffn = 3 * d * F                   # gate up down 三块
tot = V * d + L * (attn + ffn) + d + V * d
print(tot / 1e8, '亿', tot * 2 / GB, 'GiB')

# 贪心装箱：每块放进当前占用最小的卡
def pack(blocks, cards):
    load = [0.0] * cards
    for b in blocks:
        i = load.index(min(load))
        load[i] += b
    return load

# KV 缓存：2 x 层数 x KV头数 x 每头维度 x 序列长度 x 字节
kv = 2 * L * 32 * 128 * 4096 * 2 / GB

实跑输出：

--------------------------------------------------------------------
读配置、算体积、把三十二层分到两张卡上
--------------------------------------------------------------------


  ① 配置是真读进来的
  --------------------------------------------------------------------
| 字段 | 值 |
|---|---|
| hidden_size | 4096 |
| intermediate_size | 11008 |
| max_position_embeddings | 4096 |
| num_attention_heads | 32 |
| num_hidden_layers | 32 |
| torch_dtype | float16 |
| vocab_size | 32000 |

  这份 config 是照着 Llama 2 7B 的形状造的，字段名跟官方一致。
  ! 加载模型时第一步就是读它。
  维度对不上会在第一层就报错，而报错信息只说形状不匹配，不说哪个字段错了——所以出问题先看这份配置。

  ② 参数量：分项算
  --------------------------------------------------------------------
| 部分 | 参数量 |
|---|---:|
| 一层注意力（q k v o） | 0.67 亿 |
| 一层前馈（gate up down） | 1.35 亿 |
| 三十二层合计 | 64.76 亿 |
| 词嵌入 + 输出头 | 2.62 亿 |
| 总计 | 67.38 亿 |

  fp16 存一份是 12.6 GiB。
  ! 词嵌入加输出头占了 2.6 亿，接近总量的 8%。
  这部分在小模型上占比更高，所以小模型压缩词表往往比剪层更划算。

  ③ 两种切法：按层平均，还是按字节均衡
  --------------------------------------------------------------------
| 切法 | 两张卡的权重 GiB | 不均衡度 |
|---|:---:|:---:|
| 按层平均 | 6.52 / 6.03 | 1.08 |
| 按字节均衡 | 6.28 / 6.28 | 1.00 |

  两种切法的总量一样，差别只在「谁多背一点」。不均衡度越接近 1 越好。
  ! 真实加载时，device_map='auto' 做的就是这个贪心装箱，
  但它还要算上激活和 KV 缓存的余量——所以它给出的切法常常不是「数学上最均衡」，
  而是「最不容易 OOM」。

  ④ 加上 KV 缓存和激活，够不够
  --------------------------------------------------------------------
| 卡 | 权重 | KV 缓存 | 激活 | 剩余 |
|---|---:|---:|---:|---:|
| 卡 0 | 6.28 | 2.00 | 1.00 | 14.72 |
| 卡 1 | 6.28 | 2.00 | 1.00 | 14.72 |

  KV 缓存 2.00 GiB 是按 4096 上下文、单条请求算的。并发四条就是四倍。
  ! 所以「权重放得下」不等于「能跑」。
  决定并发上限的不是权重，是 KV 缓存——它是唯一随请求数线性涨的那一项。

  ⑤ 真机上的加载
  --------------------------------------------------------------------

  [真机代码 · 本沙箱不执行]
  | from transformers import AutoModelForCausalLM, AutoTokenizer
  | 
  | tok = AutoTokenizer.from_pretrained(ckpt, use_fast=False)
  | model = AutoModelForCausalLM.from_pretrained(
  |     ckpt, torch_dtype=torch.float16, device_map='auto')
  要 transformers 与权重，本沙箱没有，不执行。

  device_map='auto' 会把逐层的分配结果打在日志里。加载完看一眼那张表，就知道有没有哪张卡被塞满了。
  ! 一个常被忽略的点：切分了就有跨卡通信。
  每层都要把激活传给下一张卡，两张卡之间的带宽比卡内显存低一个数量级，
  所以「能放下」的切法未必是「跑得快」的切法。

①配置是真落盘真读回来的，七个字段。字段名跟官方一致，所以这份 config 换成真权重也能用。 出问题时先看这份配置：维度对不上会在第一层就报错，而报错只说形状不匹配，不说哪个字段错了。 ②分项算：一层注意力 0.67 亿，一层前馈 1.35 亿，三十二层合计 64.76 亿，加上词嵌入和输出头 2.62 亿， 总计 67.38 亿。 词嵌入加输出头占 8%。 这个比例在小模型上更高——所以小模型压词表往往比剪层划算，大模型则相反，压层更有效。 ③两种切法：按层平均是 6.52 对 6.03（不均衡度 1.08），按字节贪心是 6.28 对 6.28（1.00）。 差别来自嵌入层和输出头都被塞给了卡 0。 这两块加起来 524 MB，对 7B 不算什么，但对 70B 八十层切八张卡就很致命。 ④加上 KV 缓存和激活：每张卡还剩 14.72 GiB。看着很宽裕，但这是「单条请求、4096 上下文」的前提。 并发四条，KV 缓存就是 8 GiB；上下文翻到 8k，再翻一倍。决定并发上限的从来不是权重，是缓存。 ⑤真机加载那三行，重点在最后两个参数：torch_dtype 不写会按 fp32 加载，显存直接翻倍； device_map='auto' 会把逐层分配打在日志里，加载完看一眼那张表，就知道哪张卡被塞满了。

30 秒实验

  1. 把 vocab_size 从 32000 改成 128000（换更大的词表），重算总参数量和体积。 想一想：词表翻四倍，代价是什么？
  2. 把上下文从 4096 改成 8192，重算 KV 缓存。四张卡上还能扛几个并发？
  3. 把卡数从 2 改成 4，用同一个贪心函数装箱。看每卡占用变成多少，不均衡度有没有变好。

常见误区

  1. 不写 torch_dtype。很多仓库默认 fp32，加载完显存翻倍，然后你以为是模型太大。
  2. 以为「权重放得下就能跑」。 KV 缓存和激活还没算，一到长上下文或高并发就 OOM，而 OOM 发生在生成到一半的时候，最难查。
  3. 忘了设 pad_token。 批处理时要么报错，要么用 0 号 token 填充——后者更坏，因为它不报错，只是让结果慢慢变歪。
自测三问（点开对答案） 1. 为什么算到 67.38 亿就够用了，不用算到个位？ —— 因为这是用来做显存规划的，差 0.01% 只差十几 MB，不影响「能不能放下」的判断。 要精确数字去看官方公布的参数量，那里的口径包含偏置和归一化层。 2. 为什么说多卡切分「能放下」不等于「跑得快」？ —— 因为切了就有跨卡通信。 每层算完要把激活传给下一张卡，卡间带宽比卡内显存低一个数量级， 层数多的时候通信开销能吃掉一半加速比。 3. `pad_token` 设错会出什么问题？ —— 批处理时不同长度的序列要用填充符补齐。 没有专门的填充符，框架可能用 0 号 token 顶替， 而 0 号是个真实字符——模型会把填充当成内容读进去。

模型立起来了，下一节决定怎么改它：继续预训练补知识，还是指令微调学格式。

相关笔记