Megatron-Bridge 深度剖析：MoE 大模型训练加速的工程之道

Motivation

假设你拿到了一个 DeepSeek-style 的 30B-A3B MoE 模型：128 个 expert，每个 token 激活 top-6，sigmoid routing。怎么训？

先算一笔账。128 个 expert，每个 expert 是一个 SwiGLU FFN（hidden_size=1856），单个 expert 参数量约 $3 \times 1856 \times 2688 \approx 15M$（gate_proj + up_proj + down_proj）。128 个 expert 共 $128 \times 15M \approx 1.9B$ 参数——光 expert 就占了总参数量 30B 的大头。加上 attention、embedding、shared expert 等非 expert 参数，BF16 下模型权重约 60 GB。

这只是权重。训练时还需要：

• 优化器状态：Adam 的 momentum + variance，FP32 下又是 $2 \times 60 = 120$ GB
• 梯度：与权重同规模，~60 GB
• 激活值：随 expert 数线性增长，因为每个 token 要经过 6 个 expert 的前向计算

总计轻松突破 300 GB——4 张 A100-80GB 都不够。而且这还是"小"模型，DeepSeek V3 有 671B 参数、256 个 expert。

显存只是第一道坎。MoE 训练面临三大工程挑战：

显存墙：128 个 expert 的参数量远超 dense 模型。即使 active 参数只有 3B，total 参数是 30B——你需要为所有 expert 分配显存，即使大部分 expert 在每个 batch 中并不被大多数 token 激活。

通信墙：Expert Parallel 的核心操作是 All-to-All dispatch——每个 GPU 需要把自己持有的 token 发送到对应 expert 所在的 GPU，计算完再收回来。通信量与 expert 数和 token 数成正比。当 128 个 expert 分布在 64 张 GPU 上时，All-to-All 的通信量可以轻松达到每步数 GB。

负载不均衡：token routing 天然不均匀。某些 expert 因为学到了更通用的特征而被频繁选中（“热门 expert”），而另一些几乎无人问津。这导致部分 GPU 过载、部分 GPU 空转——系统吞吐量由最慢的 GPU 决定。

这三个问题不是孤立的——显存约束限制了你的并行策略选择，通信开销受并行拓扑影响，负载均衡策略又会改变通信模式。MoE 训练本质上是一个联合优化问题：需要在并行策略、通信调度、计算融合、精度管理之间找到全局最优解。

Megatron-Core + Megatron-Bridge 是这套系统方案的工业级代表。Megatron-Core 提供了 TP/PP/DP/EP/CP 五维并行的底层引擎和 MoE 加速原语（Grouped GEMM、Flex Dispatcher、通信重叠）；Megatron-Bridge 则在此之上架起了通往 HuggingFace 生态的桥梁——让你能直接加载 HF 格式的预训练权重，用 Megatron 训练，再导出回 HF 格式部署。

本文将从系统视角出发，帮你理解：Megatron 如何组织五维并行、Bridge 如何实现权重双向转换、MoE 训练中的四层加速策略，以及如何用这套体系端到端地训练一个 30B MoE 模型。

先看一张全局架构图，后面逐一展开：

flowchart LR
    subgraph HF["HuggingFace 生态"]
        HF_CK["HF Pretrained Checkpoint<br/>(DeepSeek, Llama, ...)"]
    end

    subgraph MCore["Megatron 训练引擎 (Megatron-Core)"]
        PS["parallel_state<br/>TP/PP/DP/EP/CP 五维网格"]
        MOE["transformer.moe<br/>Router + Dispatcher<br/>Grouped GEMM"]
        ACC["加速原语<br/>DeepEP / HybridEP<br/>FP8 / MXFP8 / FP4<br/>TP Comm Overlap"]
    end

    HF_CK -- "AutoBridge<br/>权重映射+切分" --> PS
    PS --- MOE
    MOE --- ACC
    MCore -- "反向转换<br/>训练完成后导出回 HF 格式<br/>用于推理部署" --> HF_CK

    style HF fill:#fff3cd,stroke:#856404
    style MCore fill:#cce5ff,stroke:#004085

Megatron-Core 架构速览

在深入 Bridge 和 MoE 加速之前，先建立对 Megatron-Core 的整体认知。Megatron-Core 是 NVIDIA 的分布式训练引擎，为 LLM 训练提供了工业级的并行基础设施。

五维并行网格

Megatron-Core 的并行策略组织为五维网格：TP（Tensor Parallel）、PP（Pipeline Parallel）、DP（Data Parallel）、EP（Expert Parallel）、CP（Context Parallel）。每张 GPU 在五个维度上各属于一个并行组。

flowchart TD
    subgraph Grid["Megatron-Core 五维并行网格<br/>假设 64 张 GPU，配置：TP=4, PP=2, DP=4, EP=2, CP=1"]
        TP["<b>TP (4)</b><br/>张量切分<br/>all-reduce<br/>切分权重矩阵的行/列"]
        PP["<b>PP (2)</b><br/>流水线分层<br/>send/recv<br/>切分 Transformer 层"]
        DP["<b>DP (4)</b><br/>数据并行<br/>all-reduce<br/>切分梯度"]
        EP["<b>EP (2)</b><br/>专家分配<br/>all-to-all<br/>切分 MoE expert"]
        CP["<b>CP (1)</b><br/>上下文并行<br/>ring attention<br/>切分序列维度"]
    end

    TOTAL["总 GPU 数 = TP × PP × DP × EP × CP = 4 × 2 × 4 × 2 × 1 = 64<br/><i>注意：EP 和 DP 共享同一组 GPU，EP × DP = 数据并行维度的总 GPU 数</i>"]

    Grid --> TOTAL

    style Grid fill:#cce5ff,stroke:#004085
    style TOTAL fill:#fff3cd,stroke:#856404

parallel_state.initialize_model_parallel_group() 是并行组初始化的核心。它将所有 GPU 按五个维度划分成互不相交的通信组：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

def initialize_model_parallel(
    tensor_model_parallel_size: int = 1,
    pipeline_model_parallel_size: int = 1,
    expert_model_parallel_size: int = 1,
    context_parallel_size: int = 1,
    ...
):
    """
    将 world_size 个 GPU 划分为五维并行组。
    每个 rank 属于每个维度的一个组，用于后续通信。
    """
    world_size = torch.distributed.get_world_size()
    data_parallel_size = world_size // (
        tensor_model_parallel_size *
        pipeline_model_parallel_size *
        context_parallel_size
    )
    # EP 从 DP 维度中"借"GPU：
    # 实际 DP = data_parallel_size // expert_model_parallel_size
    # EP 组内的 GPU 持有不同的 expert，DP 组内的 GPU 持有相同的数据分片

关键洞察：EP 和 DP 是"竞争"关系——EP 越大，实际的 DP 越小。如果 64 张 GPU 配置 TP=4, PP=2, EP=8，那么 DP = 64 / (4 × 2) / 8 = 1——没有数据并行！这意味着 batch size 受限于单个数据并行组。生产中的 DeepSeek V3 用 TP=2, PP=16, EP=64，通过极大的 PP 来释放 EP 所需的 GPU 数量。

核心模块架构

graph TD
    subgraph GPT["megatron.core.models.gpt"]
        GPTModel["GPTModel"]
        EMB["embedding<br/>(VocabParallelEmbedding)<br/>TP 切分词表"]
        DEC["decoder (TransformerBlock)"]
        OUT["output_layer<br/>(ColumnParallelLinear)"]

        GPTModel --> EMB
        GPTModel --> DEC
        GPTModel --> OUT

        TL["layers[] (TransformerLayer)"]
        DEC --> TL

        SA["self_attention"]
        MLP["mlp"]
        TL --> SA
        TL --> MLP

        QKV["linear_qkv<br/>fused QKV, TP column parallel"]
        PROJ["linear_proj<br/>output proj, TP row parallel"]
        SA --> QKV
        SA --> PROJ

        DENSE["[Dense] linear_fc1 / linear_fc2"]
        MOE_SUB["[MoE] router + experts"]
        MLP --> DENSE
        MLP --> MOE_SUB

        ROUTER["router<br/>Top-K / Sigmoid 路由"]
        DISP["token_dispatcher<br/>All-to-All 调度"]
        EXPERTS["experts<br/>GroupedMLP / SequentialMLP"]
        SHARED["shared_experts<br/>共享专家 (可选)"]
        MOE_SUB --> ROUTER
        MOE_SUB --> DISP
        MOE_SUB --> EXPERTS
        MOE_SUB --> SHARED
    end

    subgraph DIST["megatron.core.distributed"]
        DDP["DistributedDataParallel<br/>gradient all-reduce"]
        DOPT["DistributedOptimizer<br/>ZeRO-style 优化器状态切分"]
    end

    subgraph TMOE["megatron.core.transformer.moe"]
        MOEL["MoELayer — MoE 层封装"]
        TKR["TopKRouter / SigmoidRouter"]
        A2A["AllToAllTokenDispatcher"]
        FLEX["FlexTokenDispatcher<br/>DeepEP / HybridEP 后端"]
        GMLP["GroupedMLP<br/>多 expert 融合计算"]
    end

    style GPT fill:#cce5ff,stroke:#004085
    style DIST fill:#d4edda,stroke:#155724
    style TMOE fill:#fff3cd,stroke:#856404

Megatron-Core 通过 get_gpt_decoder_block_spec() 定义 decoder 的层规格（layer spec）。对于 MoE 模型，MoE 层和 dense 层可以交替出现——这是 Nemotron-3-Nano 等 hybrid 架构的基础。

TransformerConfig 控制所有配置，包括 MoE 相关的关键参数：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

TransformerConfig(
    num_moe_experts=128,              # expert 总数
    moe_router_topk=6,                # 每个 token 激活的 expert 数
    moe_router_score_function="sigmoid",  # 路由评分函数
    moe_grouped_gemm=True,            # 是否用 Grouped GEMM 融合计算
    moe_token_dispatcher_type="flex",  # 调度器类型
    moe_shared_expert_overlap=True,   # 共享专家与路由专家并行
    expert_model_parallel_size=8,     # EP 大小
    ...
)

Megatron-Bridge 桥接机制

Megatron-Core 的训练性能是一流的，但它有一个实际门槛：你不能直接拿 HuggingFace 格式的预训练权重开始训练。Megatron 有自己的权重命名和存储格式——fused QKV、column/row parallel 切分、Grouped GEMM 的 expert 权重布局——与 HF 格式完全不同。

Megatron-Bridge 解决的就是这个问题：它在 HF 生态和 Megatron 训练引擎之间架起了一座双向桥梁。

数据流全景

flowchart TD
    A["① 加载 HF 权重<br/>HuggingFace Hub<br/><i>deepseek-ai/DeepSeek-V3</i>"]
    B["② AutoBridge 自动识别<br/>检测 HF config → 匹配 DeepSeekV3Bridge"]
    C["③ 配置翻译 (provider_bridge)<br/>HF config → TransformerConfig<br/>num_attention_heads → num_attention_heads<br/>num_key_value_heads → num_query_groups<br/>intermediate_size → ffn_hidden_size<br/>n_routed_experts → num_moe_experts"]
    D["④ 权重映射 (mapping_registry)<br/>Q, K, V 分离 → fused QKV<br/>gate_proj + up_proj → fused fc1<br/>per-expert 权重 → GroupedMLP 布局"]
    E["⑤ 分布式切分 (scatter_to_tp_ranks)<br/>根据 TP/PP/EP 配置，将权重切分到各 GPU"]
    F["⑥ Megatron GPTModel 就绪<br/>开始训练！"]
    G["⑦ 训练完成后：反向转换<br/>Megatron → gather_from_tp_ranks → unfuse → HF 格式<br/>导出到 HF Hub，用于推理部署"]

    A --> B --> C --> D --> E --> F
    F --> G
    G -.-> A

    style A fill:#fff3cd,stroke:#856404
    style F fill:#d4edda,stroke:#155724
    style G fill:#cce5ff,stroke:#004085

AutoBridge：一行代码加载任意 HF 模型

AutoBridge 是 Bridge 的用户入口。它通过注册机制自动匹配 HF 模型类型：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

from megatron.bridge import AutoBridge

# 一行代码：加载 HF 预训练模型，自动检测架构
bridge = AutoBridge.from_hf_pretrained(
    "deepseek-ai/DeepSeek-V3",
    trust_remote_code=True
)

# 转换为 Megatron provider，配置并行策略
provider = bridge.to_megatron_provider()
provider.tensor_model_parallel_size = 2
provider.pipeline_model_parallel_size = 16
provider.expert_model_parallel_size = 64
provider.finalize()  # 执行权重转换 + 分布式切分

# 获取分布式模型，开始训练
model = provider.provide_distributed_model(wrap_with_ddp=False)

自动检测机制：每个模型 bridge 通过装饰器注册自己支持的 HF 模型类：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

@MegatronModelBridge.register_bridge(
    source="DeepseekV3ForCausalLM",   # HF 模型类名
    target=GPTModel,                   # Megatron 目标模型
    provider=MLAModelProvider,         # 模型提供器
    model_type="deepseek_v3",
)
class DeepSeekV3Bridge(MegatronModelBridge):
    CONFIG_MAPPING = [
        ("num_attention_heads", "num_attention_heads"),
        ("num_key_value_heads", "num_query_groups"),
        ("n_routed_experts", "num_moe_experts"),
        ("num_experts_per_tok", "moe_router_topk"),
        # ... 更多字段映射
    ]

调用 AutoBridge.from_hf_pretrained() 时，它读取 HF config 中的 architectures 字段（如 "DeepseekV3ForCausalLM"），在注册表中查找对应的 bridge class，然后执行配置翻译和权重映射。

目前已注册的模型覆盖了主流 LLM 家族：Llama 2/3、DeepSeek V2/V3、Qwen 2/2.5/3、Gemma、Mistral、Mamba，以及 NVIDIA 自家的 Nemotron 系列。

权重映射策略

权重映射是 Bridge 的核心技术挑战。HF 和 Megatron 在以下方面存在本质差异：

flowchart LR
    subgraph HF["HuggingFace 格式"]
        HF_QKV["QKV 投影: 分离的<br/>q_proj, k_proj, v_proj"]
        HF_MLP["MLP 上投影: 分离的<br/>gate_proj, up_proj"]
        HF_EXP["Expert 存储:<br/>per-expert 独立参数"]
        HF_TP["TP 切分: 无<br/>(单卡完整权重)"]
        HF_PP["PP 切分: 无<br/>(所有层在一起)"]
        HF_EP["EP 切分: 无<br/>(所有 expert 在一起)"]
    end

    subgraph MG["Megatron-Core 格式"]
        MG_QKV["QKV 投影:<br/>fused linear_qkv"]
        MG_MLP["MLP 上投影:<br/>fused linear_fc1"]
        MG_EXP["Expert 存储:<br/>GroupedMLP 融合"]
        MG_TP["TP 切分:<br/>按行/列切分到多卡"]
        MG_PP["PP 切分:<br/>按层分配到不同 stage"]
        MG_EP["EP 切分:<br/>expert 分配到不同卡"]
    end

    HF_QKV -.-> MG_QKV
    HF_MLP -.-> MG_MLP
    HF_EXP -.-> MG_EXP
    HF_TP -.-> MG_TP
    HF_PP -.-> MG_PP
    HF_EP -.-> MG_EP

    style HF fill:#fff3cd,stroke:#856404
    style MG fill:#cce5ff,stroke:#004085

Bridge 通过一组 MegatronParamMapping 子类处理这些差异：

QKVMapping — 最复杂的映射之一。HF 存储分离的 Q、K、V 权重，Megatron 将它们融合为一个 linear_qkv 张量。融合时还要考虑 GQA（Grouped Query Attention）：Q 的 head 数可能是 K/V 的 N 倍，交错排列确保 TP 切分后每个 GPU 拿到完整的 Q-K-V head 组。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

# HF: 分离的 Q, K, V
# q_proj.weight: (num_heads * head_dim, hidden_size)
# k_proj.weight: (num_kv_heads * head_dim, hidden_size)
# v_proj.weight: (num_kv_heads * head_dim, hidden_size)

# Megatron: 融合的 QKV，交错排列以支持 TP
# linear_qkv.weight: ((num_heads + 2*num_kv_heads) * head_dim, hidden_size)
# 排列方式：[Q_group0, K0, V0, Q_group1, K1, V1, ...]

# TP 切分：沿第一维（output dim）均匀切分
# GPU 0 拿到前 1/TP 的 head 组，GPU 1 拿到下一组 ...

GatedMLPMapping — 处理 SwiGLU 等 gated activation 的 MLP。HF 分开存储 gate_proj 和 up_proj，Megatron 将它们拼接为 linear_fc1：

1
2
3
4
5
6

# HF: gate_proj.weight (ffn_hidden, hidden)  + up_proj.weight (ffn_hidden, hidden)
# Megatron: linear_fc1.weight (2 * ffn_hidden, hidden)  ← 拼接

# 对于 MoE expert，每个 expert 的 gate+up 独立融合
# GroupedMLP 模式下，所有 expert 的 fc1 存储为 weight0, weight1, ... weightN
# 支持 CUTLASS GroupedGEMM 一次性计算所有 expert

反向转换同样重要：训练完成后，megatron_to_hf 将 fused 权重拆分回 HF 格式。Bridge 的 gather_from_tp_ranks 先从各 TP rank 收集切片，然后 unfuse 回分离的权重。这使得 Megatron 训练的模型可以无缝用 HF 生态（vLLM、SGLang 等）部署。

Bridge 的价值

Bridge 消除了 Megatron 的生态锁定。在此之前，要用 Megatron 训练一个 HF 模型，需要手动编写转换脚本——不同模型架构的 QKV 排列方式、MoE expert 布局、attention 变体（MHA/GQA/MLA）各不相同，每个模型都需要定制转换逻辑。Bridge 将这些转换标准化为可复用的 mapping 策略，新增模型支持只需注册新的 bridge class 和 mapping。

MoE 加速深度剖析

这是本文的核心章节。MoE 训练的加速可以分为四个层次：并行策略、通信调度、计算融合、模型实例。我们逐层展开。

Expert Parallel 基础

第 2 篇已经介绍了 EP 的基本概念：将 expert 分配到不同 GPU，通过 All-to-All dispatch/combine 实现 token 路由。这里聚焦 EP 在实际大规模训练中的工程细节。

EP 的工作流分为三步：

flowchart TD
    subgraph S1["Step 1: Router 决策"]
        INPUT["Input tokens: t0, t1, t2, ..., t7"]
        ROUTER["Router(x) → TopK scores<br/>t0 → Expert 3,7,12,45,67,99 (top-6 sigmoid)<br/>t1 → Expert 1,5,23,44,88,101<br/>..."]
        INPUT --> ROUTER
    end

    subgraph S2["Step 2: All-to-All Dispatch (token 重分布)"]
        BEFORE["<b>Before:</b> GPU 按数据分片持有 token<br/>GPU 0: [t0, t1] &nbsp; GPU 1: [t2, t3] &nbsp; ..."]
        A2A1["═══ All-to-All ═══"]
        AFTER["<b>After:</b> GPU 按 expert 分组持有 token<br/>GPU 0 (E0-E15): 收到所有发给 E0-E15 的 token<br/>GPU 1 (E16-E31): 收到所有发给 E16-E31 的 token"]
        BEFORE --> A2A1 --> AFTER
    end

    subgraph S3["Step 3: Expert Compute + All-to-All Combine"]
        COMPUTE["各 GPU 对收到的 token 运行自己的 expert FFN<br/>GPU 0: GroupedGEMM(E0-E15, tokens)"]
        A2A2["═══ All-to-All ═══"]
        RESULT["token 结果送回原始 GPU<br/>恢复数据分片布局"]
        COMPUTE --> A2A2 --> RESULT
    end

    S1 --> S2 --> S3

    style S1 fill:#fff3cd,stroke:#856404
    style S2 fill:#cce5ff,stroke:#004085
    style S3 fill:#d4edda,stroke:#155724

与其他并行维度的组合是 MoE 训练的关键设计决策。以 DeepSeek V3（671B，256 expert）为例：

总 GPU 数：$2 \times 16 \times 64 \times 1 \times 1 = 2048$ 张。注意 DP=1 意味着每个 micro-batch 只在一组 GPU 上处理，全局 batch size 通过 gradient accumulation 实现。

DeepEP 与 Flex Dispatcher

标准的 All-to-All 通信有一个问题：它是同步的。所有 GPU 必须同时参与 dispatch 和 combine，期间 GPU 的计算单元闲置。对于 128 expert、EP=8 的配置，每个 MoE 层需要两次 All-to-All，而一个 30B 模型可能有 40+ 个 MoE 层——通信开销累积起来非常可观。

DeepEP（Deep Expert Parallelism）和 Flex Dispatcher 是 Megatron-Core 对 All-to-All 的优化方案。

flowchart TD
    subgraph STD["标准 All-to-All (AllToAllTokenDispatcher)"]
        direction LR
        S_DISP["dispatch"] --> S_WAIT1["等待通信完成<br/>(GPU 计算闲置)"] --> S_COMP["expert compute"]
        S_COMP --> S_COMB["combine"] --> S_WAIT2["等待通信完成"]
    end

    subgraph FLEX["Flex Dispatcher (FlexTokenDispatcher + DeepEP)"]
        direction LR
        F_DISP["dispatch 开始<br/>(异步通信)"] --> F_COMP["expert compute 开始<br/>(收到一部分就开始算)"]
        F_DISP --> F_DONE["dispatch 完成"]
        F_COMP --> F_ALL["全部 expert compute 完成"]
        F_ALL --> F_COMB["combine (异步)"]
    end

    KEY["关键：通信和计算可以部分重叠"]

    STD ~~~ FLEX
    FLEX ~~~ KEY

    style STD fill:#fff3cd,stroke:#856404
    style FLEX fill:#d4edda,stroke:#155724
    style KEY fill:#cce5ff,stroke:#004085

Megatron-Bridge 通过配置选择 dispatcher 后端：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

# 标准 All-to-All（默认）
cfg.model.moe_token_dispatcher_type = "alltoall"

# DeepEP 后端（Ampere/Hopper/Blackwell）
cfg.model.moe_token_dispatcher_type = "flex"
cfg.model.moe_flex_dispatcher_backend = "deepep"

# HybridEP 后端（GB200 NVL72 最优，也支持其他 GPU）
cfg.model.moe_token_dispatcher_type = "flex"
cfg.model.moe_flex_dispatcher_backend = "hybridep"
cfg.model.moe_hybridep_num_sms = 16  # 分配给通信的 SM 数量

DeepEP 的核心思想：利用 RDMA 和 NVLink 的异步传输能力，将 All-to-All 拆分为多个小批次，每个批次的 token 到达后立即开始 expert 计算，无需等待全部 token 到齐。这在 expert 计算量不均（负载不均衡）时尤其有效——热门 expert 的 GPU 可以尽早开始计算，不用等冷门 expert 的 GPU。

HybridEP 更进一步，专门为 NVIDIA GB200 NVL72 拓扑优化——NVL72 将 72 个 GPU 通过 NVSwitch 全互联，提供 1.8 TB/s 的对分带宽。HybridEP 根据 GPU 间是 NVLink 直连还是需要跨 switch，自动选择最优的通信路径。

重要约束：使用 Flex Dispatcher 时，moe_shared_expert_overlap 必须设为 False。因为 Flex Dispatcher 本身已经在做通信-计算重叠，再叠加 shared expert overlap 会导致 CUDA stream 竞争。

计算融合

通信优化解决的是"数据搬运"问题，计算融合解决的是"计算效率"问题。MoE 层的计算有三个维度的融合机会。

Grouped GEMM

这是 MoE 计算加速最重要的一项优化。

问题：标准 MoE 实现中，N 个 expert 各自独立做矩阵乘法。如果 EP=8、每 GPU 持有 16 个 expert，每个 MoE 层需要 16 次独立的 GEMM 调用。每次 GEMM 都有 kernel launch 开销和 GPU 利用率损失（因为单个 expert 处理的 token 数可能很少）。

解决方案：将所有 expert 的矩阵乘法合并为一次 Grouped GEMM 调用。CUTLASS 的 GroupedGEMM kernel 接受一组形状可能不同的矩阵乘法，在 GPU 上统一调度执行。

flowchart TD
    subgraph SEQ["SequentialMLP (每个 expert 独立)"]
        direction LR
        E0["GEMM E0<br/>20 tok"]
        E1["GEMM E1<br/>5 tok"]
        E2["GEMM E2<br/>35 tok"]
        EN["...<br/>GEMM E15<br/>8 tok"]
        E0 --> E1 --> E2 --> EN
    end
    SEQ_NOTE["16 次 kernel launch，GPU 利用率低"]

    subgraph GRP["GroupedMLP (融合)"]
        GGEMM["GroupedGEMM<br/>[E0:20tok, E1:5tok, E2:35tok, ...]<br/>1 次 kernel launch，GPU 利用率高"]
    end

    subgraph STORE["权重存储对比"]
        direction LR
        S1["Sequential:<br/>expert_0.fc1.weight<br/>expert_1.fc1.weight, ..."]
        S2["Grouped:<br/>linear_fc1.weight0<br/>linear_fc1.weight1, ...<br/>(连续存储，便于 CUTLASS 批量访问)"]
    end

    SEQ --- SEQ_NOTE
    SEQ_NOTE ~~~ GRP
    GRP ~~~ STORE

    style SEQ fill:#fff3cd,stroke:#856404
    style GRP fill:#d4edda,stroke:#155724
    style STORE fill:#cce5ff,stroke:#004085

在 Megatron-Core 中启用 Grouped GEMM 只需一个配置项：

1

cfg.model.moe_grouped_gemm = True  # 默认 False

开启后，MoE 层使用 GroupedMLP 而非 SequentialMLP，expert 权重以融合格式存储。Bridge 的 GatedMLPMapping 自动处理 HF 格式到 Grouped 格式的转换。

Permute Fusion

在 All-to-All dispatch 之后，token 需要按 expert 分组重排（permute），然后送入 expert 计算。标准实现中，permute 是一次独立的内存拷贝操作。

Permute Fusion 将 token 重排与后续的 GEMM 融合为一个 kernel——在读取 token hidden states 时"顺便"完成重排，避免额外的内存读写。

1

cfg.model.moe_permute_fusion = True  # 融合 permute 与 GEMM

这个优化看起来简单，但对于 hidden_size=2688、batch 内 token 数以万计的场景，减少一次全量 hidden states 的内存拷贝可以节省数 GB/s 的带宽。

Shared Expert Overlap

DeepSeek-style MoE 有两种 expert：所有 token 都要经过的 shared expert，和经过 router 选择的 routed expert。标准实现中两者串行执行。

Shared Expert Overlap 让 shared expert 和 routed expert 在不同 CUDA stream 上并行计算：

flowchart TD
    subgraph SERIAL["串行（默认）"]
        direction LR
        R_DISP["Routed Expert<br/>Dispatch"] --> R_GEMM["Routed<br/>GEMM"] --> R_COMB["Combine"] --> S_GEMM1["Shared<br/>GEMM"]
    end
    SERIAL_T["总时间 = T_dispatch + T_routed + T_combine + T_shared"]

    subgraph OVERLAP["并行（overlap=True）"]
        direction LR
        STREAM0["Stream 0:<br/>Routed Dispatch → Routed GEMM → Combine"]
        STREAM1["Stream 1:<br/>Shared GEMM (与 Stream 0 并行)"]
    end
    OVERLAP_T["总时间 = max(T_dispatch + T_routed + T_combine, T_shared)<br/>节省：T_shared 被完全隐藏（通常 shared expert 较小）"]

    SERIAL --- SERIAL_T
    SERIAL_T ~~~ OVERLAP
    OVERLAP --- OVERLAP_T

    style SERIAL fill:#fff3cd,stroke:#856404
    style OVERLAP fill:#d4edda,stroke:#155724
    style OVERLAP_T fill:#cce5ff,stroke:#004085

1
2

cfg.model.moe_shared_expert_overlap = True
# 注意：与 Flex Dispatcher 不兼容，二者只能选其一

其他融合优化

Megatron-Core 还提供了一系列通用计算融合：

1
2
3
4

cfg.model.cross_entropy_loss_fusion = True    # 融合 cross-entropy loss 计算
cfg.model.gradient_accumulation_fusion = True  # 梯度累积融合
cfg.model.bias_activation_fusion = True        # bias + activation 融合
cfg.model.bias_dropout_fusion = True           # bias + dropout 融合

这些融合减少了 kernel launch 次数和中间结果的显存占用，对 MoE 训练的影响虽然不如 Grouped GEMM 显著，但累积效果不可忽视。

Nemotron-3-Nano (30B-A3B) 实例

NVIDIA 的 Nemotron-3-Nano 是一个很好的 MoE 工程实例——它不仅是 MoE，还是 Hybrid Mamba + Transformer + MoE 架构。

模型架构

Hybrid 层模式：52 层中混合了三种层类型：

flowchart LR
  subgraph Pattern["Nemotron-3-Nano 52 层 Hybrid 模式"]
    direction TB
    ROW1["M E M E M * E M E M E M * E M E M E M * E M E M E M *"]
    ROW2["E M E M E M E M * E M E M E M E M E"]
    subgraph Legend["图例"]
      direction LR
      LM["M = Mamba (SSM)\n线性复杂度"]
      LE["E = MoE (专家层)\n稀疏激活"]
      LA["* = Attention\n标准注意力"]
    end
    NOTE["52 层中：Attention 仅 ~7 层\n大部分是 Mamba 和 MoE 交替"]
  end

  style LM fill:#d4edda,stroke:#155724
  style LE fill:#fff3cd,stroke:#856404
  style LA fill:#cce5ff,stroke:#004085
  style NOTE fill:#f8f9fa,stroke:#6c757d

这个设计的动机是：Attention 的计算复杂度是 $O(n^2)$，Mamba 是 $O(n)$。对于长序列（262K tokens），用少量 Attention 层捕获全局依赖，用 Mamba 层处理局部上下文，比纯 Transformer 在推理效率上有巨大优势——同时通过 MoE 保持了大参数量带来的模型容量。

并行配置

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

# 并行策略
cfg.model.tensor_model_parallel_size = 4   # TP=4
cfg.model.pipeline_model_parallel_size = 1  # 无流水线并行
cfg.model.expert_model_parallel_size = 8   # EP=8
cfg.model.sequence_parallel = True          # SP 开启

# MoE 加速
cfg.model.moe_grouped_gemm = True           # Grouped GEMM
cfg.model.moe_permute_fusion = True         # Permute 融合
cfg.model.moe_router_score_function = "sigmoid"  # Sigmoid 路由
cfg.model.moe_router_enable_expert_bias = True   # Expert bias 修正
cfg.model.moe_aux_loss_coeff = 0.0001       # 负载均衡 loss 系数

# 计算融合
cfg.model.cross_entropy_loss_fusion = True
cfg.model.gradient_accumulation_fusion = True
cfg.model.bias_activation_fusion = True
cfg.model.use_fused_weighted_squared_relu = True  # Nemotron 专用融合

128 个 expert，EP=8，每张 GPU 持有 16 个 expert。TP=4 意味着非 expert 参数（attention、embedding 等）在 4 张 GPU 间切分。假设用 32 张 GPU（4 节点 × 8 GPU），则 DP = 32 / (4 × 1) / 8 = 1——单数据并行，全局 batch size 通过 gradient accumulation 扩大。

与 DeepSeek V3 的对比

两者的 MoE 路由策略几乎相同（sigmoid scoring + expert bias 修正），但并行策略差异巨大——DeepSeek V3 的 EP=64 和 PP=16 是其 256 expert、61 层深度的必然选择。Nemotron-3-Nano 因为模型较小（52 层、128 expert）且有 Mamba 层替代部分 Attention，可以用更简单的 TP=4 + EP=8 配置。

通信隐藏与混合精度

上一节聚焦 MoE 层内部的加速。但 Transformer 模型不只有 MoE 层——attention、embedding、LayerNorm 等 dense 组件同样需要 TP 通信，而整个训练过程的精度策略直接影响显存和吞吐量。

TP 通信重叠

Tensor Parallelism 的每一层都需要通信：column parallel 层在前向传播后做 all-reduce（或 reduce-scatter），row parallel 层在前向传播前做 all-gather。对于一个 52 层的模型，每层 2 次 TP 通信（attention + MLP），每步就是 100+ 次通信操作。

**通信重叠（Comm Overlap）**的核心思想：不要等通信完成再开始下一步计算，而是让通信和计算在不同的硬件单元上同时进行——GEMM 用 Tensor Core，通信用 NVLink/PCIe 的 DMA 引擎。

Megatron-Core 提供三种重叠策略：

flowchart TD
    subgraph PIPE["1. Pipeline Overlap (流水线重叠)"]
        direction LR
        C0["Chunk 0: GEMM"] --> CC0["comm"]
        C1["Chunk 1: GEMM"] --> CC1["comm"]
        C2["Chunk 2: GEMM"] --> CC2["comm"]
    end
    PIPE_NOTE["通信和下一个 chunk 的 GEMM 在不同 SM 上并行<br/>适用于：前向传播的 fprop"]

    subgraph RING["2. Ring-Exchange Overlap (环形交换重叠)"]
        direction LR
        RS0["Step 0: recv chunk<br/>+ GEMM on local chunk"]
        RS1["Step 1: recv chunk<br/>+ GEMM on received chunk"]
        RS0 --> RS1
    end
    RING_NOTE["每步接收一个 chunk，同时计算上一步收到的 chunk<br/>适用于：前向传播，特别是 FP8 场景"]

    subgraph BULK["3. Bulk Overlap (批量重叠)"]
        direction LR
        B_GEMM["GEMM (大部分 SM)"]
        B_COMM["Comm (少量 SM)"]
    end
    BULK_NOTE["不拆分 GEMM，在专用 SM 上启动整个通信操作<br/>适用于：反向传播的 dgrad / wgrad"]

    PIPE --- PIPE_NOTE
    PIPE_NOTE ~~~ RING
    RING --- RING_NOTE
    RING_NOTE ~~~ BULK
    BULK --- BULK_NOTE

    style PIPE fill:#cce5ff,stroke:#004085
    style RING fill:#d4edda,stroke:#155724
    style BULK fill:#fff3cd,stroke:#856404

实际配置模式：Megatron-Bridge 提供了针对不同硬件的预配置 profile。以 H100 + TP=4 为例：

规律：前向传播倾向于 Pipeline 或 Ring-Exchange（GEMM 和通信量相当），反向传播倾向于 Bulk（GEMM 更大，通信可以被完全隐藏在 GEMM 之后）。FP8 场景下，GEMM 速度翻倍但通信量不变，因此通信成为更大的瓶颈，Ring-Exchange 的细粒度重叠更有优势。

FP8 训练

FP8（8-bit floating point）是 Hopper 及更新架构的核心加速能力。相比 BF16，FP8 的 Tensor Core 算力翻倍，显存带宽需求减半。但 FP8 的动态范围很小（E4M3 格式只有 ~480），需要精心管理缩放因子。

Megatron-Core 支持四种 FP8 策略：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

from megatron.bridge.training import MixedPrecisionConfig

cfg.mixed_precision = MixedPrecisionConfig(
    fp8="e4m3",                    # 启用 FP8 E4M3 格式
    fp8_recipe="tensorwise",        # 缩放策略
    fp8_param=True,                 # 参数也存为 FP8
    fp8_param_gather=True,          # all-gather 时用 FP8 通信
    fp8_wgrad=True,                 # 权重梯度用 FP8
    first_last_layers_bf16=True,    # 首尾层保持 BF16（稳定性）
    moe_router_padding_for_fp8=True,  # MoE router 输出对齐 FP8
)

对 MoE 的影响：FP8 对 MoE 的 Grouped GEMM 特别有益。128 个 expert 的矩阵乘法在 BF16 下已经是 memory-bound（每个 expert 处理的 token 少），FP8 将访存量减半，直接提升 roofline 利用率。fp8_param_gather=True 还减少了 FSDP all-gather 的通信量——参数用 FP8 传输，到目标 GPU 后再 upcast。

Activation Recomputation

MoE 模型的激活值显存随 expert 数线性增长。每个 token 经过 top-6 个 expert，每个 expert 产生独立的中间激活，128 expert 模型的激活值显存可以是 dense 模型的数倍。

Megatron-Core 的 activation recomputation 有两种模式：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

# Selective recomputation（推荐）：只重计算特定算子
cfg.model.recompute_granularity = "selective"
# 重计算 attention 的 softmax 和 dropout，保留 GEMM 结果
# MoE 场景下：重计算 router softmax、dispatch permute

# Full recomputation：整层重计算
cfg.model.recompute_granularity = "full"
cfg.model.recompute_method = "uniform"
cfg.model.recompute_num_layers = 1
# 每 1 层做一次完整重计算，显存节省最大但计算增加 ~33%

DeepSeek V3 的 32 节点配置就使用了 full recomputation——在 2048 张 GPU 上，显存是更紧张的约束（EP=64 意味着每 GPU 仍持有 4 个 expert 的完整参数和优化器状态），而额外的计算开销可以被大量 GPU 分摊。

端到端实战

前面介绍了架构、Bridge、加速原语和精度策略。把它们组合起来，端到端训练一个 MoE 模型是什么样的？以 Nemotron-3-Nano (30B-A3B) 为例。

Recipe 配置解读

Megatron-Bridge 通过 Python recipe 文件定义完整的训练配置。以下是 Nemotron-3-Nano 预训练 recipe 的关键参数逐项注释：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51

from megatron.bridge import AutoBridge
from megatron.bridge.training import TrainingConfig, MixedPrecisionConfig

def nemotron_3_nano_pretrain_config():
    cfg = TrainingConfig()

    # ─────────── 模型加载 ───────────
    cfg.model = AutoBridge.from_hf_pretrained(
        "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16"
    ).to_megatron_provider(load_weights=True)

    # ─────────── 并行策略 ───────────
    cfg.model.tensor_model_parallel_size = 4   # TP=4, 节点内 NVLink
    cfg.model.pipeline_model_parallel_size = 1  # 无 PP（模型够小）
    cfg.model.expert_model_parallel_size = 8   # 128 experts / 8 = 16 per GPU
    cfg.model.sequence_parallel = True          # 与 TP 配合的序列并行

    # ─────────── MoE 加速 ───────────
    cfg.model.moe_grouped_gemm = True          # Grouped GEMM（必开）
    cfg.model.moe_permute_fusion = True        # Permute 融合
    cfg.model.moe_token_dispatcher_type = "alltoall"  # 标准 All-to-All
    cfg.model.moe_shared_expert_overlap = True  # Shared expert 并行

    # ─────────── MoE 路由 ───────────
    cfg.model.moe_router_score_function = "sigmoid"
    cfg.model.moe_router_enable_expert_bias = True  # DeepSeek-style bias
    cfg.model.moe_aux_loss_coeff = 0.0001      # 负载均衡 loss

    # ─────────── 计算融合 ───────────
    cfg.model.cross_entropy_loss_fusion = True
    cfg.model.gradient_accumulation_fusion = True
    cfg.model.bias_activation_fusion = True
    cfg.model.use_fused_weighted_squared_relu = True

    # ─────────── 训练超参 ───────────
    cfg.training.max_steps = 100000
    cfg.training.global_batch_size = 256
    cfg.training.micro_batch_size = 1
    cfg.training.seq_length = 8192

    # ─────────── 优化器 ───────────
    cfg.optimizer.lr = 1e-4
    cfg.optimizer.min_lr = 1e-5
    cfg.optimizer.weight_decay = 0.1
    cfg.optimizer.adam_beta1 = 0.9
    cfg.optimizer.adam_beta2 = 0.95

    # ─────────── 混合精度 ───────────
    cfg.mixed_precision = MixedPrecisionConfig()  # 默认 BF16

    return cfg

几个关键决策的解读：

为什么 PP=1？ Nemotron-3-Nano 只有 30B 参数，TP=4 + EP=8 已经能让每张 GPU 的显存需求可控。PP 会引入流水线 bubble，对吞吐量有损，小模型没必要。

为什么 micro_batch_size=1？ MoE 的 All-to-All 通信量与 micro-batch 内的 token 数成正比。MBS=1 意味着每次前向/反向只处理 8192 个 token，限制了 All-to-All 的峰值通信量。全局 batch size 通过 gradient accumulation 扩大到 256。

为什么 moe_aux_loss_coeff=0.0001？ 负载均衡 loss 鼓励 token 均匀分布到各 expert，但系数过大会干扰主任务 loss。0.0001 是经验值——足够防止 expert 坍缩（所有 token 只选少数 expert），又不至于损害模型质量。

启动命令

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

# 4 节点 × 8 GPU = 32 GPU
# 使用 torchrun 启动，Megatron-Bridge 自动初始化并行组
torchrun --nproc_per_node=8 --nnodes=4 \
    --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT \
    -m megatron.bridge.training.train \
    --recipe nemotron_3_nano_pretrain_config

# 环境要求：
# - PyTorch >= 2.4 with CUDA 12.x
# - Megatron-Core (via 3rdparty/Megatron-LM submodule)
# - Transformer Engine (TE) >= 1.7
# - NCCL >= 2.18
# - 节点间 InfiniBand / RoCE RDMA 互联

SFT 场景差异

SFT（Supervised Fine-Tuning）与预训练在配置上有几个关键差异：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

# 并行策略简化：SFT 通常数据量较小，可以降低并行度
cfg.model.tensor_model_parallel_size = 1   # TP=1（不切分权重）
cfg.model.expert_model_parallel_size = 8   # EP 保持（expert 数不变）

# 学习率降低
cfg.optimizer.lr = 2e-5       # 预训练的 1/5
cfg.optimizer.min_lr = 2e-6

# 序列长度可能更长（SFT 数据含完整对话）
cfg.training.seq_length = 16384

# PEFT：LoRA/DoRA 只微调部分参数
cfg.peft.method = "lora"
cfg.peft.lora_rank = 16
cfg.peft.lora_alpha = 32
cfg.peft.target_modules = [
    "linear_qkv",    # Attention QKV
    "linear_proj",    # Attention output
    "linear_fc1",     # MLP (含 expert)
    "linear_fc2",     # MLP down
    "in_proj",        # Mamba in projection
    "out_proj",       # Mamba out projection
]

注意：LoRA 的 target_modules 覆盖了 Transformer 和 Mamba 两种层类型的关键投影——这是 Hybrid 架构特有的需求。只微调 Transformer 部分的 LoRA 会丢失 Mamba 层的适配能力。

常见调参陷阱

1. EP 与 DP 的平衡：EP 越大，每 GPU 持有的 expert 越少，显存越宽裕。但 EP 会"吃掉"DP 维度——如果 EP 过大导致 DP=1，全局 batch size 只能通过 gradient accumulation 实现，训练吞吐量可能下降（因为 GA 的 micro-step 之间没有计算-通信重叠）。

2. Grouped GEMM 必须开启：不开 moe_grouped_gemm，128 个 expert 意味着每 GPU 16 次独立 GEMM，kernel launch overhead 会吃掉大量 GPU 时间。这是 MoE 训练最容易忽略的性能杀手。

3. Router score function 要与预训练一致：如果预训练用的是 sigmoid routing，SFT 也必须用 sigmoid。切换为 softmax 会导致 token 分布突变，训练不稳定。

4. FP8 + MoE 的陷阱：FP8 训练时，MoE router 的输出需要对齐到 FP8 kernel 的要求。忘记设置 moe_router_padding_for_fp8=True 可能导致 CUDA kernel crash 或静默精度问题。

5. Shared expert overlap 与 Flex Dispatcher 互斥：不能同时开启 moe_shared_expert_overlap=True 和 moe_token_dispatcher_type="flex"。前者用 CUDA multi-stream 并行 shared/routed expert，后者也使用 multi-stream 做异步 dispatch，两者会产生 stream 竞争。

如何进一步加速？

MoE 训练的性能优化是一个不断演进的领域。以下是当前最有潜力的几个方向。

通信瓶颈分析

在做任何优化之前，先定位瓶颈。Megatron-Core 集成了 PyTorch 的 profiling 工具：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

# 使用 PyTorch Profiler 抓取 trace
with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ],
    schedule=torch.profiler.schedule(wait=5, warmup=2, active=3),
    on_trace_ready=torch.profiler.tensorboard_trace_handler('./log/trace'),
) as prof:
    for step, batch in enumerate(dataloader):
        train_step(model, batch)
        prof.step()

在 TensorBoard 中查看 trace，重点关注：

1. All-to-All 占比：如果 All-to-All dispatch + combine 占单步时间超过 30%，说明 EP 通信是瓶颈，可以考虑 Flex Dispatcher 或增大 EP（减少每次 All-to-All 的通信量）
2. GEMM 利用率：如果 expert GEMM 的 GPU 利用率低于 50%，说明每个 expert 处理的 token 太少，应该增大 micro-batch size 或减小 EP
3. 通信-计算重叠率：如果 TP comm 和 GEMM 没有重叠（在 trace 中表现为通信和计算交替出现），需要开启 comm overlap

Kernel Fusion 前沿

FlashAttention-3：针对 Hopper 架构的 warp-specialization 优化，进一步提升 attention 计算的 FLOPs 利用率。对 Nemotron-3-Nano 的 7 个 attention 层虽然影响有限，但对纯 Transformer MoE（如 DeepSeek V3 的 61 层 attention）效果显著。

MLA 专用 Kernel：DeepSeek V3 使用的 Multi-Latent Attention 将 KV 投影到低维空间，减少 KV Cache 显存。针对 MLA 的专用 CUDA kernel 可以将 latent 压缩/解压与 attention 计算融合，避免额外的内存读写。

通信-计算极致重叠

当前的 Shared Expert Overlap 和 Flex Dispatcher 各自做局部的通信-计算重叠，但它们之间是互斥的。未来的方向是全局调度：

gantt
    title 理想的 MoE 层执行流水线
    dateFormat X
    axisFormat %s

    section Pipeline
    Router           :r, 0, 3
    Dispatch (异步)  :d, 3, 10
    Shared Exp (与 dispatch 并行) :s, 10, 18
    Routed Exp (dispatch 完成后) :re, 18, 28
    Combine (异步)   :c, 28, 35
    Next Layer       :nl, 35, 39

EP dispatch 与上一层的 dense 计算流水线化，shared expert 与 dispatch 并行，combine 与下一层的 router 重叠——将空闲 bubble 压缩到极致。

新硬件适配

GB200 NVL72：72 个 GPU 通过第五代 NVSwitch 全互联，对分带宽 1.8 TB/s。这个拓扑对 MoE 的 All-to-All 是革命性的——NVL72 内的 All-to-All 带宽比当前 8-GPU NVLink 域提升 ~9 倍。HybridEP 已经为这个拓扑做了优化，区分域内（NVSwitch 直连）和域间（跨机架）的通信路径。

MXFP8 / FP4：Blackwell 架构引入了 Microscaling FP8（更细粒度的缩放因子）和 FP4（4-bit 浮点）。FP4 可以将 expert 权重的显存进一步减半，让单 GPU 能容纳更多 expert——这可能改变 EP 的最优配置，因为 EP 可以更小（更多 expert per GPU），从而释放更多 GPU 给 DP。

系统级优化

异步 Checkpoint：大规模训练中，checkpoint 保存可以花费数分钟。异步 checkpoint 在后台线程完成存储 I/O，训练不中断。对于 MoE 模型尤其重要——128 expert 的全量 checkpoint 可能达数十 GB，同步保存会严重拖慢训练。

Fault Tolerance：2048 张 GPU 的集群中，单卡故障是常态而非异常。弹性训练框架需要支持：检测故障 GPU → 重新分配并行组 → 从最近 checkpoint 恢复 → 继续训练。Megatron-Core 目前通过与 NVIDIA NeMo 的集成支持部分故障恢复能力。

Key Takeaways

1. MoE 训练的三大挑战：显存墙（128 expert 参数量巨大）、通信墙（All-to-All dispatch 开销与 expert 数成正比）、负载不均衡（token routing 天然不均匀）。这三者相互耦合，需要联合优化。

2. Megatron-Core 的五维并行：TP/PP/DP/EP/CP 构成完整的并行空间。EP 和 DP 竞争同一组 GPU——EP 越大，DP 越小。DeepSeek V3 用 TP=2, PP=16, EP=64 训练 671B 模型，需要 2048 张 GPU。

3. Megatron-Bridge 消除了生态锁定：AutoBridge 自动检测 HF 模型类型，QKVMapping 和 GatedMLPMapping 处理权重格式差异，训练完成后可无缝导出回 HF 格式用于推理部署。

4. MoE 加速的四个层次：

   • 并行策略：EP 分配 expert，与 TP/PP 组合
   • 通信调度：Flex Dispatcher（DeepEP/HybridEP）实现异步 All-to-All
   • 计算融合：Grouped GEMM 将 N 个 expert 合并为一次 kernel 调用
   • 精度优化：FP8 将 GEMM 算力翻倍、通信量减半

5. Nemotron-3-Nano 是 Hybrid 架构的代表：52 层中混合 Mamba（SSM）、MoE、Attention 三种层类型，用少量 Attention 捕获全局依赖，Mamba 处理局部上下文，MoE 提供大参数容量。

6. 实战中的关键配置：moe_grouped_gemm=True 是必开项；moe_shared_expert_overlap 与 moe_flex_dispatcher_backend 互斥；FP8 训练需要 moe_router_padding_for_fp8；SFT 的 router score function 必须与预训练一致。

推荐源码阅读路径

如果你想深入 Megatron-Core 和 Bridge 的源码，推荐以下阅读顺序：

1. megatron/core/parallel_state.py — 理解五维并行组的初始化逻辑
2. megatron/core/transformer/moe/moe_layer.py — MoE 层的完整工作流
3. megatron/core/transformer/moe/token_dispatcher.py — All-to-All dispatch/combine 实现
4. megatron/core/transformer/moe/grouped_mlp.py — Grouped GEMM 的权重布局和计算
5. megatron/bridge/models/conversion/auto_bridge.py — AutoBridge 自动检测机制
6. megatron/bridge/models/conversion/param_mapping.py — 权重映射基类和 QKV/GatedMLP 映射
7. megatron/bridge/training/comm_overlap.py — TP 通信重叠的三种策略配置

参考资料

1. Shoeybi et al., 2019. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism — Megatron 的开创性工作，定义了 TP 列切分/行切分范式

2. NVIDIA, 2025. Megatron-Bridge GitHub Repository — Megatron-Bridge 开源仓库

3. DeepSeek-AI, 2024. DeepSeek-V3 Technical Report — 671B MoE 模型的训练细节，包括 EP=64 的并行配置

4. Fedus et al., 2022. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity — MoE 在 Transformer 中的经典应用

5. Dai et al., 2024. DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models — Fine-grained expert + shared expert 的设计

6. NVIDIA, 2025. Nemotron-3-Nano Technical Blog — Hybrid Mamba + Transformer + MoE 架构解析

7. Gu & Dao, 2023. Mamba: Linear-Time Sequence Modeling with Selective State Spaces — Mamba SSM 架构，Nemotron-3-Nano 的层类型之一