Gemini：异构显存 + CPU 卸载

Gemini 是 ColossalAI 自家的招牌组件，对标 DeepSpeed 的 ZeRO-3 + ZeRO-Offload， 但实现思路不同：它先把参数切成 chunk，再用一套显式调度器决定每个 chunk 此刻应该放 GPU 还是 CPU。本章把它拆开讲清楚。

5.1痛点：为什么不直接用 ZeRO-3 + Offload

DeepSpeed 的 ZeRO-3 + Offload 是按"逐张量"卸载的：每个 parameter 单独决定要不要 CPU。这有两个问题：

• 小张量太多：LLaMA-7B 有 ~290 个 parameter tensor，每个都要发起一次 CPU↔GPU 拷贝，PCIe overhead 巨大；
• GPU 显存碎片：每次 gather 用完释放，下一个又申请，造成大量 cudaMalloc 碎片。

Gemini 的破解思路：

1. 把多个相邻的小 tensor 打包成一个 chunk（默认 32MB / chunk），统一拷贝；
2. 用StateTensor 跟踪每个 chunk 的"位置 + 状态"，由统一调度器决定换入换出；
3. 预分配 chunk 内存池，复用空间，避免碎片。

5.2核心抽象：Chunk + ChunkManager + GeminiManager

看 colossalai/zero/gemini/：

colossalai/zero/gemini/
├── chunk/
│   ├── chunk.py            ★ Chunk 类
│   ├── manager.py          ★ ChunkManager
│   └── utils.py            按模型预切 chunk 的工具
├── gemini_ddp.py            ★ GeminiDDP（顶层模型 wrapper）
├── gemini_optimizer.py      ★ GeminiOptimizer
├── gemini_mgr.py            ★ GeminiManager（调度器）
├── placement_policy.py      ★ 放置策略
├── memory_tracer/            预跑一遍统计每层显存峰值
└── ...

Chunk：参数的"打包单位"

一个 chunk 由若干个相邻 tensor 组成，看 chunk.py::Chunk：

class Chunk:
    def __init__(self, chunk_size, dtype, ...):
        self.chunk_size = chunk_size          # 默认 32MB
        self.shard_size = chunk_size // dp_size  # 单卡那份
        # 实际只持有 shard
        self.cuda_shard: Optional[Tensor] = None
        self.cpu_shard: Optional[Tensor] = None

    def append_tensor(self, tensor): ...      # 把 tensor 塞进 chunk
    def shard_move(self, device): ...          # CPU↔GPU 整 shard 搬
    def access_chunk(self): ...                # all-gather 重组完整 chunk

注意 chunk 内部只持有自己那份 shard（即 ZeRO-3 切完后的本地分片），需要算时再 all-gather 拼成完整 chunk。

ChunkManager：管所有 chunk 的内存账

chunk/manager.py 维护：

• 所有 chunk 的列表 + 每个 chunk 当前的设备；
• GPU / CPU 上的容量预算；
• 当显存不够时，决定 evict 哪些 chunk 到 CPU。

GeminiManager：调度器

gemini_mgr.py 是 Gemini 的"大脑"，每次 forward / backward 触发参数访问时，由它决定：

1. 把要用的 chunk 从 CPU 搬到 GPU；
2. 预测下一步要用什么，提前搬；
3. 显存满了就把最久未用的 chunk 搬走（LRU-like）。

5.3三种 PlacementPolicy

placement_policy.py 定义了三种放置策略：

5.4使用：GeminiPlugin 全参数

看 colossalai/booster/plugin/gemini_plugin.py，关键参数：

plugin = GeminiPlugin(
    # —— 精度 ——
    precision="bf16",                  # bf16 / fp16
    # —— Chunk ——
    chunk_config_dict=None,            # 显式指定每层 chunk 大小（高阶）
    chunk_init_device="cuda",
    search_range_m=128,                # 自动搜 chunk size 的最大 MB
    # —— 放置策略 ——
    placement_policy="static",
    shard_param_frac=1.0,               # 多少比例参数走 ZeRO-3 切
    offload_optim_frac=1.0,             # 多少比例 optimizer state 卸到 CPU
    offload_param_frac=0.0,             # 多少比例 parameter 卸到 CPU
    # —— 通信 ——
    enable_gradient_accumulation=False,
    initial_scale=2**16,                # AMP 初始 scale
    # —— 进程组 ——
    enable_async_reduce=True,
    extra_dp_size=1,
)

显存预算公式

假设模型参数量 $P$、optimizer state 量 $S$、GPU 数 $N$，使用 Gemini 后单卡占用约为：

$$M_\text{gpu} = \underbrace{(1 - f_p) \cdot \frac{P}{N}}_{\text{保留在 GPU 的 param shard}} + \underbrace{(1 - f_o) \cdot \frac{S}{N}}_{\text{GPU 的 optim shard}} + M_\text{activation} + M_\text{buffer}$$

其中 $f_p = \text{offload\_param\_frac}$、$f_o = \text{offload\_optim\_frac}$。 当 $f_p=1$、$f_o=1$ 时，单卡显存只剩 activation + buffer，可以训远大于显存的模型（代价是吞吐受 PCIe 限制）。

5.5实战：8 张 24GB 卡训 13B

13B 模型 fp32 占 ~52GB，bf16 占 ~26GB。8 卡 24GB（共 192GB）理论能 ZeRO-3 切下来，但 activation + optimizer 还要占空间，纯 ZeRO-3 也可能 OOM。Gemini 配置：

from colossalai.booster.plugin import GeminiPlugin

plugin = GeminiPlugin(
    precision="bf16",
    placement_policy="auto",
    shard_param_frac=1.0,         # 全切
    offload_optim_frac=1.0,        # optimizer 全卸 CPU
    offload_param_frac=0.3,        # 30% param 也卸 CPU（不够用时换入）
    initial_scale=2**16,
)
booster = Booster(plugin=plugin)

跑起来后用 nvidia-smi + top 同时看：

• GPU 显存：每张应稳定在 ~18-20GB（不是 24GB，给 activation 留余地）；
• CPU 内存：会涨 ~80GB（13B × 2 bytes optimizer m + v + master weights）；
• PCIe 流量：可以用 nvidia-smi dmon -s p 看到持续 GB/s。

5.6Gemini vs DeepSpeed ZeRO 性能对照

同样 LLaMA-7B、8×A100-40GB、bs=4、seqlen=2048，bf16，约 50 step 平均：

结论：显存够时优先 HybridParallel；显存不够时 Gemini 比 DeepSpeed Offload 更快。

5.7常见踩坑

5.8调参建议

1. 第一次跑用 placement_policy="auto"，先把流程跑通。
2. 记下稳态显存使用率（nvidia-smi 平均值）。如果远低于 GPU 总显存（如只用 60%），说明 offload_param_frac 太激进，可以减小；反之如果 OOM，加大它。
3. 跑通后切到 placement_policy="static"，固定 offload_frac，吞吐会更稳。
4. 训练前期 chunk size 可以让框架自动搜（search_range_m=128），稳定后导出 chunk_config_dict 写死。

5.9Gemini 源码定位

5.10这章你需要带走的

• Gemini = chunk 化 + 显式调度 版的 ZeRO-3 + offload；
• 核心三件套：Chunk / ChunkManager / GeminiManager；
• 三种放置策略：static / auto / const；
• 四个关键 frac：shard_param_frac、offload_optim_frac、offload_param_frac、extra_dp_size；
• Gemini 比 DeepSpeed ZeRO-3 + offload 通常快 1.5–2×（chunk 化让 PCIe transfer 次数大幅下降）；
• 显存够时优先 HybridParallel，不够时 Gemini 是次优选择。