ShardFormer：自动按 layer 切模型

Megatron 的 TP 要重写模型；ColossalAI 想让你直接用 HuggingFace 的 LlamaForCausalLM 也能 TP。 ShardFormer 就是这件事的实现。本章把它的内部机制拆给你看。

7.1问题：HF 模型的"原版 Linear" + Megatron 的"并行 Linear"

看一段 HuggingFace 原版 LLaMA：

# transformers/models/llama/modeling_llama.py
class LlamaAttention(nn.Module):
    def __init__(self, config):
        ...
        self.q_proj = nn.Linear(d, d, bias=False)
        self.k_proj = nn.Linear(d, d_kv, bias=False)
        self.v_proj = nn.Linear(d, d_kv, bias=False)
        self.o_proj = nn.Linear(d, d, bias=False)
        ...

它用的是 普通 nn.Linear，单卡跑没问题，但要在多卡 TP 上跑就缺三件事：

1. 把 weight 沿合适维度切到多张卡；
2. forward 时插入 all-reduce / all-gather / scatter；
3. backward 时插入对应的反向通信。

这就是 Megatron 当年要把模型整个重写一遍的原因——他们的 ColumnParallelLinear、RowParallelLinear 把这三件事都内置了。

7.2ShardFormer 的思路：替换不重写

ShardFormer 不要求你重写模型，而是在运行时把 nn.Linear 替换成 Colossal 自家的 ParallelLinear。整个流程：

7.3核心入口：ShardFormer.optimize()

看 colossalai/shardformer/shard/sharder.py：

class ModelSharder:
    def __init__(self, model, policy, shard_config):
        self.model = model
        self.policy = policy            # 该模型的 Policy 实例
        self.shard_config = shard_config

    def shard(self):
        # 1) 把 policy.preprocess_module() 跑一遍（注入 config）
        self.policy.preprocess()
        # 2) 拿 module_policy()，决定怎么替换
        policy_map = self.policy.module_policy()
        # 3) 递归遍历 model.named_modules()
        for name, module in self.model.named_modules():
            if type(module) in policy_map:
                self._replace_sub_module(module, policy_map[type(module)])
        # 4) 调 policy.postprocess()（如绑定 weight）
        self.policy.postprocess()

这就是 "按 module type 匹配 + 子模块替换" 的全部核心，剩下都是细节。

7.4Policy 的标准模板

每个支持的模型在 colossalai/shardformer/policies/ 下都有一个 policy。看 LLaMA 的简化版：

# colossalai/shardformer/policies/llama.py
class LlamaPolicy(Policy):

    def config_sanity_check(self):
        pass

    def preprocess(self):
        # 把 num_attention_heads / num_key_value_heads 按 tp_size 切
        if self.shard_config.tensor_parallel_size > 1:
            assert self.model.config.num_attention_heads % tp == 0
            self.model.config.num_attention_heads //= tp
            self.model.config.num_key_value_heads //= tp

    def module_policy(self) -> Dict[nn.Module, ModulePolicyDescription]:
        return {
            LlamaAttention: ModulePolicyDescription(
                attribute_replacement={...},
                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,
                    ),
                ],
                method_replacement={
                    "forward": llama_attention_forward,  # 自定义 forward
                },
            ),
            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),
                ],
            ),
            LlamaDecoderLayer: ModulePolicyDescription(
                sub_module_replacement=[
                    SubModuleReplacementDescription(
                        suffix="input_layernorm",
                        target_module=FusedRMSNorm,
                    ),
                    SubModuleReplacementDescription(
                        suffix="post_attention_layernorm",
                        target_module=FusedRMSNorm,
                    ),
                ],
            ),
        }

    def postprocess(self):
        return self.model

三个核心字段

7.5替换用的 layer 库

看 colossalai/shardformer/layer/，主要类：

7.6列并行 vs 行并行：为什么这样配对

看 LLaMA MLP（gate / up / down）的 forward：

$$y = \text{down}\bigl(\text{silu}(\text{gate}(x)) \odot \text{up}(x)\bigr)$$

当 TP=2 时分两份切：

核心原理："Col → element-wise → Row" 这条路径中间不需要通信，只在最后 Row 出来 all-reduce 一次。 attention（q/k/v Col + o Row）同理。

7.7method_replacement：什么时候必须改 forward

有些情况下"换 Linear"不够，整个 forward 都得改写。比如 attention 加上 RingAttention（沿 seq 切 + 环通信）时，原版 SDPA 不支持，必须替换：

# colossalai/shardformer/modeling/llama.py
def llama_attention_forward(self, hidden_states, ...):
    q = self.q_proj(hidden_states)           # 已是 ColumnParallel 输出
    k = self.k_proj(hidden_states)
    v = self.v_proj(hidden_states)
    ...
    if shard_config.enable_sequence_parallelism:
        attn_output = RingAttention.attention(q, k, v, ...)
    else:
        attn_output = flash_attn_func(q, k, v, ...)
    return self.o_proj(attn_output)

这是 method_replacement 真正发力的地方：让框架可以在不动 HF 源码的情况下插入自定义算子。

7.8支持模型清单

colossalai/shardformer/policies/ 当前 (v0.5.0) 提供：

未列出的模型不能用 TP，但可以走 Gemini / ZeRO / FSDP plugin。

7.9给一个新模型写 Policy（高阶）

假设你要支持 Yi 这种新模型，步骤：

1. 在 colossalai/shardformer/policies/yi.py 建文件，继承 Policy；
2. 实现 preprocess()（切 num_heads）和 module_policy()（替换 q/k/v/o/gate/up/down）；
3. 在 colossalai/shardformer/shard/shard_config.py 的 policy 列表里注册类型；
4. 跑 tests/test_shardformer/test_model/test_shard_yi.py 验证；
5. 提 PR。

7.10这章你需要带走的

• ShardFormer = "按 module type 替换 nn.Linear" 让 HF 模型自动支持 TP；
• 核心数据结构：Policy、ModulePolicyDescription、SubModuleReplacementDescription；
• 替换 layer：Linear1D_Col（输出维切） + Linear1D_Row（输入维切），attention/MLP 都靠这对组合；
• 不只换 Linear，还能用 method_replacement 替换 forward（RingAttention 必备）；
• policy 文件只在 HybridParallelPlugin 和 MoeHybridParallelPlugin 里启用；
• 给新模型加 TP = 写一个 policy 文件即可。