算力需求指南

大模型多卡训练笔记

1 GB = 1024 MB = 1024×1024 KB = 1024×1024×1024 Byte（字节） = 1024×1024×1024×8 Bit

参数类型与大小：

单卡哪怕有 80 GB 显存，也扛不住数十亿、上百亿参数的大模型全量训练。
权重只是一部分，梯度、优化器状态、激活都需要显存。

假设模型有 N 参数量（例如 2B = 20 亿）：

权重 W
- 存储方式：BF16 (2 Byte)
- 显存： $W = N \times 2$ Byte
- 例：2B 参数 → ≈ 4 GB
梯度 G
- 存储方式：BF16 (2 Byte)
- 显存： $G = N \times 2$ Byte
- 例：2B 参数 → ≈ 4 GB
优化器状态（Adam）
- 包含动量项 V、平方梯度 S，各自 FP32 (4 Byte)
- 每份 ≈ 8 GB，合计 ≈ 16 GB
  其中：优化器显存还需要的开销（权重梯度在训练中可能会复制很多份儿）
1. 权重 W
  - 模型本身的参数，BF16/FP16 存 2 Byte。
2. 梯度 G
  - 反向传播临时存储，BF16/FP16 2 Byte。
3. 优化器状态（不同优化器差别大）：
  - SGD：通常只需要梯度本身，0份复制。
  - SGDM（带动量的 SGD）：需要一个动量向量（FP32，4 Byte）。一份复制
  - Adam/AdamW：
  - 一阶动量 (V)：FP32（4 Byte）。
  - 二阶动量 (S)：FP32（4 Byte）。
  - 所以是 2 份状态。
4. Master 权重 ( $W^A$ )
  - 混合精度训练时常见：虽然前向/反向用 BF16，但优化器更新需要 FP32 精度 → 所以再存一份 FP32 权重。
Activations
- 依赖 batch size、seq_len、实现细节
- 粗略估计：≈ 0.7–1.0 × 权重大小

说明：ZeRO-3 显存更省，但通信和 I/O 开销更大。

架构：d_model ≈ 4096，ffn_dim ≈ 14336，每层 8 个专家，32 层，SwiGLU（gate/up/down 三个线性层）。
单个专家参数量：
$4096×14336×2 + 14336×4096 = 176,160,768$ ≈ 1.76×10^8
→ BF16 权重 ≈ 352 MB/专家。
一层 8 专家 ≈ 2.82 GB/层
32 层合计 ≈ 90 GB（仅专家权重）
→ 全专家全参训练在 44 GB 显存上不可能。

例：只训底部 6 层，每层 2 专家，同时训路由。

可训练参数数：
$6 × 2 × 176,160,768 = 2.114B$
权重（BF16）：≈ 4.23 GB
梯度（BF16）：≈ 4.23 GB
Adam 状态（V+S，FP32）：≈ 16.9 GB
Master 权重（FP32）：≈ 8.46 GB
合计（持久内存 + 梯度）：≈ 33.8 GB
再加冻结权重占位、激活开销，44 GB 卡上需：
- batch=1–2
- seq_len ≤ 1024
- use_cache=False
- gradient_checkpointing=True
可行，但需要严格控制。

在专家 / 路由上挂 LoRA（r=16）。

单专家 LoRA 参数：
$r × (4096+14336 + 4096+14336 + 14336+4096) = r × 55296$
→ r=16 → 0.885M/专家
每层 8 专家：7.08M
32 层：226.6M LoRA 参数
显存开销：
- 权重 ≈ 0.45 GB
- 梯度 ≈ 0.45 GB
- Adam + Master ≈ 2.72 GB
- 合计 ≈ 3.6 GB
远低于 44 GB 显存

显存碎片化：
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
（需在 import torch 前设置）
通信库：
NCCL > Gloo > MPI（除非特殊环境）
DDP：必须同步 random seed
eval：
- eval_mb_size 可大
- 训练 batch 小 + 梯度累积
- 关掉 model.config.use_cache
device = "auto": HF 会自动根据你机器的显存大小，把模型的不同部分切片分布,对于大模型，比如 7B，在单卡 44 GB 上：通常 attention + embedding + 部分 FFN 会放 GPU,冻结的模块、或者用不到的专家 (MoE inactive experts) 可以被放到 CPU。这个参数对推理来说自动分配还是很好的，但是训练的话需要梯度的参数必须常驻 GPU，否则每次 forward/backward 都要把参数搬上来，通信成本爆炸。所以，device_map=auto 对训练不总是安全，因为它可能把你要训练的层塞到 CPU 去，导致速度慢甚至无法训练。

作者：Yang Lewis
非商业转载请标明出处。
商业转载请联系作者：840691168ly@gmail.com