[RFC] Module Family：高性能模块的低精与并行自动适配

背景

HyperParallel 已支持通过 plan_overrides[].replace_module._target_ 将原始模型模块替换为 GQAAttention、SwiGLUMLP、GroupedExperts 等高性能 Module。

当前高性能 Module、低精实现和并行适配分别配置。用户选择高性能 Module 后，仍需额外指定 CP inner_wrapper._target_ 和 EP local_compute_fn._target_ 及其实现参数。

本方案引入 Module Family。用户只选择高性能 Module，并设置低精开关和并行规模；框架自动选择对应实现，并将现有 TP、CP、EP 路径适配到所选高性能 Module。

用户配置

accelerator:
  tp_size: 2
  cp_size: 2
  ep_size: 2

plan_overrides:
  - match: "*.self_attn"
    module_type: package.modeling.SourceAttention
    replace_module:
      _target_: hyper_parallel.components.modules.GQAAttention
      use_mxfp8: true
      use_fused_mxfp8: true

  - match: "*.mlp.experts"
    module_type: package.modeling.SourceExperts
    replace_module:
      _target_: hyper_parallel.components.modules.GroupedExperts
      use_mxfp8: true

选择 Module Family target 后，用户不再额外配置高性能 Module 对应的 CP wrapper、EP compute target 或 use_grouped_gemm。未选择 Family target 时保持当前行为。用户显式 plan 可以覆盖自动结果。

Module Family

Module Family 保存 baseline Module、低精实现和与现有并行路径的适配规则：

GQA_ATTENTION = ModuleFamily(
    name="gqa_attention",
    baseline=GQAAttention,
    variants={
        "mxfp8": MXFP8GQAAttention,
    },
)

GROUPED_EXPERTS = ModuleFamily(
    name="grouped_experts",
    baseline=GroupedExperts,
    variants={
        "mxfp8": MXFP8GroupedExperts,
    },
)

• use_mxfp8: false 使用 baseline Module；
• use_mxfp8: true 使用 MXFP8 Module；
• use_fused_mxfp8 由 MXFP8 Module 内部选择 fused forward，不增加新的 Family target；
• Family 解析完成后，再将现有 TP、CP、EP 路径作用到最终 Module。

Module Family 不选择或注册模型专属的 CP wrapper 和 EP compute。模型现有的并行 plan 保持不变；Family 只做高性能 Module 所需的调整：CP wrapper 调用最终 Module 原有的 attention_interface，EP compute 自动加入 use_grouped_gemm=true。

低精实现

低精只改变 Module 内部计算，不改变模型外部调用接口和并行通信语义。

GQAAttention
  └─ MXFP8GQAAttention

GroupedExperts
  └─ MXFP8GroupedExperts

普通和 MXFP8 Attention 保持相同的 attention_interface 调用位置。普通和 MXFP8 Experts 保持相同的 forward_expert_major() 调用方式。use_fused_mxfp8 只控制 MXFP8 Module 内部是否使用融合低精算子。

TP 实现

TP 不新增高性能 Module 专属开关或 forward。

Module replacement 在 ShardingPlanner 之前完成。替换后，Planner 直接根据最终 Module 的参数名、参数角色和维度布局生成 TP plan：

GQAAttention.linear_qkv  → FUSED_QKV
GQAAttention.o_proj      → ROWWISE

GroupedExperts.gate_up_proj → MOE_EXPERT
GroupedExperts.down_proj    → MOE_EXPERT

执行流程：

选择 baseline / MXFP8 Module
  → 完成 Module replacement
  → ShardingPlanner 读取替换后的参数
  → 使用现有 TP placement 和通信

如果新增实现改变参数命名或权重维度，必须在 Family 中补充对应 sharding rule；当前参数布局可被现有 Planner 识别时无需额外配置。

CP 实现

CP 继续使用现有 wrapper 和 flex_cp_allgather，不新增完整的 CPGQAAttention，也不新增一套 Attention forward。

GQAAttention.forward() 继续负责：

QKV projection
  → Q/K Norm
  → RoPE
  → KV Cache
  → attention_interface
  → output projection

当 cp_size > 1 且选择 GQAAttention Family 时，框架自动安装该 Family 对应的现有 CP wrapper。wrapper 保存 Module 原来的 attention_interface，然后完成：

local Q/K/V
  → flex_cp_allgather(K, V)
  → mask/layout 处理
  → 调用原来的 attention_interface

核心形式：

original_attention = target_module.attention_interface

def cp_attention_interface(
    module,
    query_states,
    key_states,
    value_states,
    attention_mask,
    **kwargs,
):
    key_states, value_states = flex_cp_allgather(
        key_states,
        value_states,
        2,
        cp_mesh,
    )
    return original_attention(
        module,
        query_states,
        key_states,
        value_states,
        attention_mask,
        **kwargs,
    )

现有 CP wrapper 继续处理对应的 mask 和 layout，并通过 Module 原有的 attention_interface 调用最终 Attention 实现。因此 GQAAttention 和 MXFP8GQAAttention 可以复用同一个 CP wrapper。

选择高性能 Module 不要求模型增加新的 CP wrapper。现有 CP plan 中的 wrapper 继续使用，只需通过 attention_interface 调用最终选中的 Module 实现。

EP 实现

EP 继续使用模型现有的 Router、token dispatch、All-to-All 和 combine，不在 GroupedExperts 内重新实现通信。

当 ep_size > 1 且选择 GroupedExperts Family 时，框架直接将现有 EP compute target 的 use_grouped_gemm 设置为 true：

ep_compute = existing_ep_compute.replace(
    use_grouped_gemm=True,
)

local_compute_fn 继续沿用现有 EP plan。Family resolver 只自动补充 use_grouped_gemm: true。

执行流程：

existing_ep_compute
  → Router
  → All-to-All dispatch
  → GroupedExperts.forward_expert_major()
  → All-to-All combine
  → routing weight 聚合

GroupedExperts 只负责本地专家计算。MXFP8GroupedExperts 保持相同的 forward_expert_major() 调用方式，因此使用同一套 EP 通信，仅替换本地专家计算实现。

选择高性能 Experts 不要求模型增加新的 EP factory；Family resolver 只修改现有 EP target 的 use_grouped_gemm 参数。

自动组合规则

if tp_size > 1:
    # 对最终选中的 Module 使用现有 ShardingPlanner。
    pass

if cp_size > 1 and selected_family is GQA_ATTENTION:
    # 安装现有 CP wrapper；wrapper 调用 flex_cp_allgather，
    # 然后委托 Module 原有 attention_interface。
    install_attention_cp_wrapper()

if ep_size > 1 and selected_family is GROUPED_EXPERTS:
    # 保留原 EP Router 和通信，只切换本地专家计算。
    ep_compute = ep_compute.replace(use_grouped_gemm=True)

对应关系：

Module Family	TP	CP	EP
GQAAttention	现有 Planner	自动安装 Attention CP wrapper	无额外适配
MXFP8GQAAttention	现有 Planner	复用同一 CP wrapper	无额外适配
GroupedExperts	现有 Planner	无额外适配	自动设置 use_grouped_gemm=true
MXFP8GroupedExperts	现有 Planner	无额外适配	复用原 EP 通信和高性能本地计算入口

执行流程

YAML replace_module target
  → target 是否属于 Module Family
      ├─ 否：保持当前 Target / plan 行为
      └─ 是：继续解析
  → 根据 use_mxfp8 选择 baseline / MXFP8 Module
  → 根据 use_fused_mxfp8 设置 Module 内部实现
  → 读取现有 CP wrapper / EP compute plan
  → 应用 Module replacement
  → TP：使用现有 ShardingPlanner
  → CP：GQA Family 自动安装现有 CP wrapper
  → EP：GroupedExperts Family 自动设置 use_grouped_gemm=true
  → 合并用户显式 plan
  → ShardingPlanner / applier

Resolver 最终仍生成现有对象：

• ModuleReplacementSpec：选择普通或低精高性能 Module；
• ModuleShardingSpec：保存 TP contract、CP wrapper 和 EP compute factory。

现有 Planner、forward rewriter、CP/EP collective 和 applier 保持不变。

当前代码入口

• GQAAttention.forward() 与 attention_interface
• flex_cp_allgather
• 通用 CP wrapper
• 通用 EP compute 与 use_grouped_gemm
• GroupedExperts.forward_expert_major()
• CP/EP Target 解析

Start with distributed/_builder/rule_resolver.py and trace how ModuleReplacementSpec and ModuleShardingSpec are built. Read the listed GQAAttention, GroupedExperts, context-parallel wrapper, collective, and expert-parallel recipe entry points before defining the Module Family resolution flow. Done means baseline or MXFP8 selection, existing TP/CP/EP adaptation, and explicit-plan overrides work without changing existing communication paths.