核心概念：Booster / Plugin / ShardFormer

ColossalAI 区别于 Megatron / DeepSpeed / FSDP 的根本设计，就是这三层抽象： Booster（顶层封装）+ Plugin（并行策略）+ ShardFormer（自动切模型）。 本章把这三个名词讲清楚，让你看任意 ColossalAI 代码都知道每一层在做什么。

4.1三层抽象总览

4.2Booster：顶层"加速器"

Booster 是用户唯一每次都用到的对象。看 colossalai/booster/booster.py 的接口：

class Booster:
    def __init__(self, plugin, mixed_precision=None, ...):
        self.plugin = plugin
        ...

    def boost(self, model, optimizer=None, criterion=None,
              dataloader=None, lr_scheduler=None):
        """
        把用户的纯 PyTorch 对象包装成分布式版本。
        返回 5 元组：包装后的 (model, optimizer, criterion, dataloader, lr_scheduler)
        """
        return self.plugin.configure(...)

    def backward(self, loss, optimizer):
        """触发 plugin 自家的反向（ZeRO scatter / Gemini swap / PP schedule）"""
        self.plugin.backward(loss, optimizer)

    def save_model(self, model, checkpoint, shard=True, ...):
        """plugin-aware 的 checkpoint 保存"""
        ...

    def load_model(self, model, checkpoint, ...):
        ...

    def execute_pipeline(self, data_iter, model, criterion, optimizer):
        """PP plugin 才用，跑一次 PP forward-backward schedule"""
        ...

关键设计：Booster 自己不实现任何并行逻辑，只是把所有请求转发给 plugin。 这种"strategy pattern"让 plugin 之间不互相干扰。

4.3Plugin：并行策略的实现

ColossalAI 自带 7 个 plugin，每个对应一种并行策略。看 colossalai/booster/plugin/：

colossalai/booster/plugin/
├── plugin_base.py              ★ 所有 plugin 的基类
├── dp_plugin_base.py            DP 风格 plugin 的中间基类
├── pp_plugin_base.py            PP 风格 plugin 的中间基类
├── torch_ddp_plugin.py          PyTorch DDP
├── torch_fsdp_plugin.py         PyTorch FSDP
├── low_level_zero_plugin.py     ZeRO-1/2
├── gemini_plugin.py             ★ Gemini（ZeRO-3 + 异构内存）
├── hybrid_parallel_plugin.py    ★ 3D 并行（TP/PP/DP/ZeRO 一锅炖）
└── moe_hybrid_parallel_plugin.py MoE 专用

Plugin 的接口

所有 plugin 继承 Plugin 抽象基类（看 plugin_base.py）：

class Plugin(ABC):
    @abstractmethod
    def configure(self, model, optimizer, criterion, dataloader, lr_scheduler):
        """把 PyTorch 对象转成分布式版本"""

    @abstractmethod
    def backward(self, loss, optimizer):
        """反向"""

    @abstractmethod
    def supported_devices(self) -> List[str]:
        """如 ["cuda", "npu"]"""

    @abstractmethod
    def support_no_sync(self) -> bool:
        """梯度累积是否支持 no_sync context"""

    # 还有 control_checkpoint_io, prepare_dataloader 等

关键 plugin 对照

4.4ShardFormer：自动按 layer 切模型

TP（张量并行）的痛点：传统 Megatron 要把模型从头重写一遍（写一份 ParallelLlamaAttention）。 ColossalAI 的 ShardFormer 用替换层的方式让任意 HF 模型自动支持 TP。

原理

policy 是什么

看 colossalai/shardformer/policies/，每个模型有自己的"policy 文件"：

colossalai/shardformer/policies/
├── llama.py            # LlamaPolicy
├── qwen2.py            # Qwen2Policy
├── gemma.py
├── mistral.py
├── mixtral.py
├── chatglm2.py
├── bert.py
├── gpt2.py
├── falcon.py
└── ...

每个 policy 描述哪几层要换成什么：

# 简化版（实际 LlamaPolicy 复杂得多）
class LlamaPolicy(Policy):
    def module_policy(self):
        return {
            LlamaAttention: ModulePolicyDescription(
                sub_module_replacement=[
                    SubModuleReplacementDescription(
                        suffix="q_proj",
                        target_module=Linear1D_Col,    # 列并行
                    ),
                    SubModuleReplacementDescription(
                        suffix="k_proj",
                        target_module=Linear1D_Col,
                    ),
                    SubModuleReplacementDescription(
                        suffix="v_proj",
                        target_module=Linear1D_Col,
                    ),
                    SubModuleReplacementDescription(
                        suffix="o_proj",
                        target_module=Linear1D_Row,    # 行并行
                    ),
                ],
            ),
            LlamaMLP: ModulePolicyDescription(
                sub_module_replacement=[
                    SubModuleReplacementDescription(suffix="gate_proj", target_module=Linear1D_Col),
                    SubModuleReplacementDescription(suffix="up_proj",   target_module=Linear1D_Col),
                    SubModuleReplacementDescription(suffix="down_proj", target_module=Linear1D_Row),
                ],
            ),
        }

核心是"按 module 类型替换 nn.Linear"，把 attention 的 QKV 列并行、out 行并行；MLP 的 gate/up 列并行、down 行并行。 这是 Megatron 经典切法（看 Megatron-LM ch07）的自动化版本。

4.5ShardFormer 的 layer 库

替换用的 layer 在 colossalai/shardformer/layer/：

4.6Plugin × ShardFormer 怎么配合

ShardFormer 自己不能用，必须被 plugin 调起来。这种 plugin 包括 HybridParallelPlugin 和 MoeHybridParallelPlugin：

plugin = HybridParallelPlugin(
    tp_size=2,                        # ★ tp_size>1 自动启用 ShardFormer
    pp_size=2,
    zero_stage=1,
    precision="bf16",
    enable_fused_normalization=True,   # ShardFormer 内的 fused norm
    enable_flash_attention=True,
)

booster = Booster(plugin=plugin)
model, ... = booster.boost(model, ...)
# 此时 model 已被 ShardFormer 改造，q_proj/k_proj 等已是 Linear1D_Col

Gemini / LowLevelZero plugin 不用 ShardFormer（它们走 ZeRO 风格，不需要切张量本身）。

4.7支持的模型清单

看 colossalai/shardformer/policies/，目前内置 policy 覆盖：

不在列表里的模型 → 走 Gemini / ZeRO 路径（不用 ShardFormer 也能训）；或自己写一份 policy。

4.8三层抽象的代码体感

把第 3 章那份 30 行脚本对应到三层抽象：

colossalai.launch_from_torch()                # 底层 PyTorch dist init

plugin = GeminiPlugin(...)                     # ← Plugin 层
booster = Booster(plugin=plugin)               # ← Booster 层
                                                 # (Gemini 不用 ShardFormer)

model, optimizer, _, dl, _ = booster.boost(model, optimizer, None, dl, None)
# ↑ booster.boost() 调用 plugin.configure()
# ↑ plugin 把 model 包装成 GeminiDDP / LowLevelZeroOptimizer / ChunkManager

for batch in dl:
    loss = model(**batch).loss
    booster.backward(loss, optimizer)          # ← Booster.backward() → plugin.backward()
    optimizer.step()                            # plugin 优化器自己知道怎么 step

4.9这种抽象的优劣

4.10Booster + 9 个 Plugin 源码定位

4.11ShardFormer 内部

Llama policy 范本（policies/llama.py）：QKV / gate / up 列切，o / down 行切，FFN 内零通信。要加新模型 = 拷贝一份 policy 改 module name。

4.12这章你需要带走的

• Booster + Plugin + ShardFormer 是 ColossalAI 全部架构的核心；
• Booster 是用户唯一接触的入口，自己不实现并行，只转发给 plugin；
• Plugin 决定走哪种并行（DDP / ZeRO / Gemini / 3D / FSDP / MoE）；
• ShardFormer 自动把 HF 模型切成 TP 版（只在 HybridParallel / MoE plugin 内启用）；
• 支持哪些模型 = colossalai/shardformer/policies/ 里有哪些 policy 文件；
• 不在 policy 列表的模型 → 走 Gemini / ZeRO 路径仍然能训。