【Bug已解决】[Qwen3.5] Missing linear_attn entries in base_model_tp_plan causes OOM and shape error at TP>1 解决方案

一、现象长什么样

在 Qwen3.5 这类混合注意力架构（标准全注意力 self_attn + 线性注意力 linear_attn）上开启张量并行，当 world_size > 1 时，程序有两种不一致的表现，且都指向同一个根因：

# 表现 A：显存爆炸（OOM）
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.10 GiB
  (GPU 0; 79.15 GiB total capacity; 76.88 GiB already allocated)
  # 报错发生在第 18 层附近，而不是传统的 embedding / lm_head

# 表现 B：形状错误（shape error）
RuntimeError: mat1 and mat2 shapes cannot be multiplied (1024x4096 and 8192x4096)
  # 注意：输入的最后一维是 4096（完整隐藏维度），
  # 但权重被切成了 8192（因为相邻层按 tp_size=2 切了列），维度对不上

最让人困惑的地方是：不使用 TP 时一切正常；TP=1 时也正常；只要 TP >= 2 就必崩，而且崩溃点不固定——小模型倾向于先 OOM，大模型倾向于先报 shape error。这说明问题不在某个具体算子，而在"层与层之间的切分状态没有被统一规划"。

二、背景

HuggingFace Transformers 在 4.40+ 引入了一套声明式张量并行方案：模型在 base_model_tp_plan 里登记哪些子模块走列并行（"col_parallel"）、行并行（"row_parallel"）或不切（"local"）。引擎（如 torch.distributed.tensor 的 ColwiseParallel/RowwiseParallel，或 device_mesh 上的 parallelize_module）会按这张表把对应权重沿 tp 维度切分。

Qwen3.5 在结构上比 Qwen3 多了一族 linear_attn（线性/门控线性注意力）子层，位于每个 decoder block 内，与 self_attn 并列。问题在于：早期合并进主干的 base_model_tp_plan 只枚举了 self_attn.*、mlp.*、gate_*，遗漏了 linear_attn.* 这一族。于是这族层在 parallelize_module 遍历时被跳过，既没有挂上并行策略，也没有被显式标成 local——它们被当成"普通模块"原样保留。

三、根因

根因有两条，且会互相放大：

1. 遗漏登记导致不切分 → OOM。 linear_attn 内部包含数个 nn.Linear（输入投影、门控、输出投影），其中输出投影维度通常是 hidden_size（例如 4096）。因为没在 tp_plan 里，它不会被 ColwiseParallel 切分，于是每个 rank 都完整持有这组权重 + 完整激活，等于把"本应被 2/4/8 路均摊的显存"又重新全量复制了一份。层数越多，额外显存越大，直接触发 OOM。

2. 切分状态不一致导致形状错位 → shape error。 假设 self_attn 的输出投影被登记为 RowwiseParallel（沿输入维切分、输出维不切），它的输出在 all-reduce 后是完整 hidden_size；但紧邻的 linear_attn 既然没被切分，它内部的 col_parallel 风格投影却被"错误地"按全局 tp_size 初始化或读取了被切过的权重。当 linear_attn 的输入来自一个已被行并行收束的完整向量，而它内部却用"按 tp 切过的权重矩阵"去做乘加时，就出现 1024x4096 vs 8192x4096 这种维度对不上。

换句话说：tp_plan 必须是一个"闭包"——它要么显式切分某个线性层，要么显式声明 local，绝不能留空。留空等于"隐式全量复制"，既浪费显存又破坏相邻层的维度契约。

四、最小可运行复现

下面用一个可独立运行的 Python 片段复现"计划漏登 → 切分状态不一致"的核心逻辑（不依赖真实多卡，用纯逻辑模拟每个 rank 持有哪些权重维度）：

from dataclasses import dataclass, field
from typing import Dict, List

HIDDEN = 4096
TP = 2  # 假设 2 路张量并行

# 模型里每个 decoder 块实际存在的可切分层
REAL_LINEAR_MODULES = [
    "model.layers.0.self_attn.q_proj",
    "model.layers.0.self_attn.k_proj",
    "model.layers.0.self_attn.v_proj",
    "model.layers.0.self_attn.o_proj",
    "model.layers.0.mlp.gate_proj",
    "model.layers.0.mlp.up_proj",
    "model.layers.0.mlp.down_proj",
    # —— 线性注意力这一族，实际存在却被 tp_plan 漏登 ——
    "model.layers.0.linear_attn.in_proj",
    "model.layers.0.linear_attn.gate_proj",
    "model.layers.0.linear_attn.out_proj",
]

# 早期（有 bug）的 base_model_tp_plan：漏了 linear_attn.*
BUGGY_TP_PLAN: Dict[str, str] = {
    "model.layers.*.self_attn.q_proj": "col_parallel",
    "model.layers.*.self_attn.k_proj": "col_parallel",
    "model.layers.*.self_attn.v_proj": "col_parallel",
    "model.layers.*.self_attn.o_proj": "row_parallel",
    "model.layers.*.mlp.gate_proj": "col_parallel",
    "model.layers.*.mlp.up_proj": "col_parallel",
    "model.layers.*.mlp.down_proj": "row_parallel",
}


def plan_has(module: str, plan: Dict[str, str]) -> bool:
    """模拟 parallelize_module 用通配符匹配：返回该模块是否出现在计划里。"""
    for key, _ in plan.items():
        if key.replace(".*", ".0") == module:
            return True
    return False


def simulate_memory_and_shapes(plan: Dict[str, str]):
    unsharded_count = 0
    shape_mismatch = False
    for m in REAL_LINEAR_MODULES:
        if not plan_has(m, plan):
            # 没登记 = 隐式全量复制，每个 rank 持有完整 HIDDEN x HIDDEN
            unsharded_count += 1
        else:
            # 登记过的列并行层，单个 rank 持有 HIDDEN/TP x HIDDEN
            pass
    # 关键契约：linear_attn 输入来自相邻 row_parallel 收束后的完整向量(4096)
    # 但其 out_proj 若被当成"按 tp 切过的权重"去乘，就出现 8192 维权重
    if not plan_has("model.layers.0.linear_attn.out_proj", plan):
        # 漏登时，引擎有时仍按全局权重初始化 → 维度错位
        shape_mismatch = True
    return unsharded_count, shape_mismatch


unsharded, mismatch = simulate_memory_and_shapes(BUGGY_TP_PLAN)
print("未切分层数量:", unsharded)      # 输出 3（linear_attn 三族）
print("存在形状错位风险:", mismatch)   # 输出 True
assert unsharded > 0, "复现失败：应当存在未切分层"

运行后会看到 未切分层数量: 3、存在形状错位风险: True，正好对应 OOM 与 shape error 两条根因。

五、解决方案（第一层：最小直接修复）

最快的止血办法：在加载模型、调用 parallelize_module 之前，手动把 linear_attn 这一族补齐进 base_model_tp_plan。具体规则与 self_attn/mlp 对称即可：

def patch_qwen35_tp_plan(model):
    """第一层修复：把漏登的 linear_attn 子层补进 tp_plan。"""
    plan = dict(model.base_model_tp_plan)  # 复制一份，避免改到类属性

    # 与 self_attn 对称：q/k/v 风格投影走列并行，输出投影走行并行
    plan["model.layers.*.linear_attn.in_proj"] = "col_parallel"
    plan["model.layers.*.linear_attn.gate_proj"] = "col_parallel"
    plan["model.layers.*.linear_attn.out_proj"] = "row_parallel"
    # 如果某个变体里 linear_attn 还有独立的 value/feature 投影，也一并补上
    plan.setdefault("model.layers.*.linear_attn.feat_proj", "col_parallel")

    model.base_model_tp_plan = plan
    return model


# 使用示意
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3.5-xxx", torch_dtype="auto")
model = patch_qwen35_tp_plan(model)

# 之后再用官方 TP 工具切分，就不会再漏
from torch.distributed.tensor.parallel import parallelize_module
from torch.distributed.device_mesh import init_device_mesh
mesh = init_device_mesh("cuda", (TP,), mesh_dim_names=("tp",))
for name, child in model.named_children():
    if name == "model":
        parallelize_module(child, mesh, model.base_model_tp_plan)

这一层修复后，OOM 与 shape error 都会消失，因为 linear_attn 现在和其它层一样被统一规划，权重沿 tp 均摊，维度契约恢复一致。

六、解决方案（第二层：结构性改进）

每次换模型都要手写补丁容易遗漏。更稳妥的做法是把"计划完整性校验 + 自动补全"做成模型加载流程的一部分，用一个中心化的 TpPlanAuditor 来兜底：

from dataclasses import dataclass, field
from typing import Dict, List, Set

@dataclass
class TpPlanAuditor:
    """对 base_model_tp_plan 做闭包校验与自动补全。"""
    model_name: str
    parallelizable_prefixes: List[str] = field(default_factory=lambda: [
        "self_attn", "mlp", "linear_attn", "gate_attn",
    ])
    # 已知的"列并行候选"子模块名片段
    col_candidates: Set[str] = field(default_factory=lambda: {
        "in_proj", "q_proj", "k_proj", "v_proj",
        "gate_proj", "up_proj", "feat_proj",
    })
    row_candidates: Set[str] = field(default_factory=lambda: {
        "o_proj", "down_proj", "out_proj",
    })

    def audit(self, plan: Dict[str, str], layer_indices: List[int]) -> Dict[str, str]:
        new_plan = dict(plan)
        missing: List[str] = []
        for idx in layer_indices:
            for prefix in self.parallelizable_prefixes:
                base = f"model.layers.{idx}.{prefix}"
                for cand in self.col_candidates:
                    key = f"model.layers.*.{prefix}.{cand}"
                    if key not in new_plan:
                        new_plan[key] = "col_parallel"
                        missing.append(key)
                for cand in self.row_candidates:
                    key = f"model.layers.*.{prefix}.{cand}"
                    if key not in new_plan:
                        new_plan[key] = "row_parallel"
                        missing.append(key)
        if missing:
            print(f"[TpPlanAuditor] {self.model_name} 自动补全 {len(missing)} 项:")
            for m in missing[:5]:
                print("  +", m)
        return new_plan


# 使用
auditor = TpPlanAuditor(model_name="Qwen3.5")
patched = auditor.audit(model.base_model_tp_plan, layer_indices=list(range(36)))
model.base_model_tp_plan = patched

TpPlanAuditor 的语义是：凡是模型里真实存在的、可并行的线性子层，必须在计划中有明确归宿（切分或 local），否则就按命名约定自动补全。这样即便上游 transformers 版本又新增了一族注意力层，也不会再次因为漏登而崩。

七、解决方案（第三层：断言 / CI 守护）

把"计划必须覆盖所有线性层"固化成测试，防止回退。下面是一组 pytest，可在 CI 中常驻：

import pytest

# 假设有一个可在测试里构造的最小 Qwen3.5 配置
from transformers import Qwen3p5Config, Qwen3p5ForCausalLM

ALL_PARALLELIZABLE = ["self_attn", "mlp", "linear_attn"]


def _collect_linear_paths(model):
    paths = []
    for name, mod in model.named_modules():
        if name.endswith((".weight",)) and mod.__class__.__name__ == "Linear":
            paths.append(name.rsplit(".", 1)[0])
    return paths


def test_linear_attn_present_in_tp_plan():
    cfg = Qwen3p5Config(num_hidden_layers=2, hidden_size=256, num_attention_heads=4)
    model = Qwen3p5ForCausalLM(cfg)
    plan = dict(model.base_model_tp_plan)
    # 关键断言：linear_attn 这一族必须被登记
    assert any("linear_attn" in k for k in plan), \
        "tp_plan 漏登 linear_attn，会在 TP>1 时 OOM / shape error"


def test_no_unsharded_linear_attn():
    cfg = Qwen3p5Config(num_hidden_layers=2, hidden_size=256, num_attention_heads=4)
    model = Qwen3p5ForCausalLM(cfg)
    plan = dict(model.base_model_tp_plan)
    for key in plan:
        if "linear_attn" in key:
            assert plan[key] in ("col_parallel", "row_parallel", "local"), \
                f"linear_attn 子层 {key} 必须是显式切分或 local，不能是空"


def test_auditor_closes_gap():
    from tp_audit import TpPlanAuditor
    cfg = Qwen3p5Config(num_hidden_layers=2, hidden_size=256, num_attention_heads=4)
    model = Qwen3p5ForCausalLM(cfg)
    auditor = TpPlanAuditor(model_name="Qwen3.5")
    fixed = auditor.audit(model.base_model_tp_plan, layer_indices=[0, 1])
    assert any("linear_attn" in k for k in fixed)

在 CI 里加一行 pytest tests/test_tp_plan.py，以后只要有人改了 base_model_tp_plan 又漏掉 linear_attn，测试会立刻红灯，把问题拦在合并前。

八、排查清单

当你在混合注意力模型上遇到"TP>1 才崩、TP=1 正常"的情况，按下面顺序查：

1. 打印 model.base_model_tp_plan，确认是否覆盖所有 *_attn / *_mlp / linear_attn 子层。
2. 用 named_modules() 统计真实存在的 nn.Linear，与 tp_plan 里的键做差集，差集里的就是"隐式全量复制"的层。
3. OOM 优先看漏登层是不是输出维度大的（out_proj/down_proj 类），它们不切分最吃显存。
4. shape error 优先看相邻层的切分类型是否一致：行并行输出是完整维，列并行输入应是完整维，错位往往出在"一个被切、一个没被切"的边界。
5. 升级 transformers 后务必重跑上面的 pytest，因为 tp_plan 是随模型类一起演进的，新注意力族容易被漏。

九、小结

这个问题的本质不是算子 bug，而是张量并行计划的不完备：linear_attn 一族在 base_model_tp_plan 中漏登，导致它们既没被切分（显存 OOM）、又破坏了与相邻层之间的维度契约（shape error）。

• 第一层：加载前手动把 linear_attn.* 按 col_parallel/row_parallel 补齐进 tp_plan，立即止血。
• 第二层：用 TpPlanAuditor 做结构性兜底，自动补全所有可并行子层，杜绝未来新增注意力族再次漏登。
• 第三层：用 pytest 把"计划必须闭包覆盖所有线性层"固化成 CI 守护，防止回归。

记住一个原则：张量并行计划要么显式切分、要么显式标 local，绝不能有"留空"的层——留空就是隐式全量复制，迟早会在 TP>1 时以 OOM 或 shape error 的形式找回来。