在 NVIDIA Megatron Core 中实现 Falcon-H1 混合架构

在大型语言模型（LLM）开发迅速演进的格局中，NVIDIA Megatron Core 已成为大规模训练海量 Transformer 模型的基础框架。该开源库提供业界领先的并行技术与 GPU 优化性能。如今，该项目在 NVIDIA/Megatron-LM 仓库中采用 GitHub-first 开发模式，Megatron Core 正日益受到基础模型构建者贡献的塑造，使其成为面向开源 AI 模型更灵活、面向未来的引擎。

本文从技术角度概述了 Technology Innovation Institute (TII)（Falcon 模型家族的创建者）如何对 Megatron Core 和 Megatron Bridge 框架进行贡献与集成。第一部分探讨了 Falcon-H1 并行混合架构在 Megatron Bridge 中的实现，重点分析了在协调异构 Transformer 与 Mamba 层的同时结合不可学习 µP 乘数所面临的挑战。第二部分探讨了将 BitNet 集成至 Megatron Core 的过程，详细说明了如何使用三值参数对应层替换标准线性层，以及该方案对训练效率与可扩展性的影响。

这些贡献展示了 Megatron Core 用户如何扩展该框架，以支持其自定义模型架构和复杂的训练特性，并利用社区其他成员的工作成果。

Falcon-H1 混合架构在 Megatron Bridge 中的集成

在Megatron Bridge中实现Falcon-H1并行混合架构，凸显了协调异构Transformer层与Mamba层以及不可学习µP乘数所面临的挑战。该集成的具体细节将在以下章节中提供。

混合并行设计

TII 对 Megatron 的核心贡献在于 Falcon-H1 并行混合架构。该设计有别于其他近期混合模型中常见的串行层级结构。如图 1 所示，在每个处理块内，注意力机制与 SSM 并行运行，其输出在传入该块的输出投影层前会进行拼接。SSM 与注意力头的数量支持配置，可按需调整。

Falcon-H1 摒弃了传统的层叠结构，转而采用并行设计，使基于 Transformer 的注意力机制与 Mamba-2 状态空间模型（SSM）组件能够在每个核心处理块内同步处理输入。

注意力分支与 Mamba 分支的输出在投影前进行拼接，使模型得以融合 SSM 卓越的长上下文记忆能力与高效性，以及注意力机制的长程依赖建模能力。

模型中并行混合层、纯 Mamba 层、纯注意力层以及仅含多层感知机（MLP）层的比例可独立配置，从而实现灵活的架构探索。

双仓库集成

Falcon-H1 的支持跨越两个职责不同的仓库。在 Megatron Core（Megatron-LM）中，TII 贡献了：

• 基础的 ParallelHybridLayer，该层并行运行 Mamba 和注意力机制并汇总其输出
• 更新的层分配逻辑，在现有的 Mamba、注意力和 MLP 层类型旁引入了 PARALLEL 符号。

这也包括用于加载和保存并行混合模型的检查点转换工具。在 Megatron Bridge 中，TII 基于这些基础组件构建了完整的 Falcon-H1 模型：

• FalconH1Layer 扩展了并行设计，加入了 MLP 组件（从而形成完整的 Mamba 加注意力机制加 MLP 模块）
• FalconH1Bridge 提供 Hugging Face 与 Megatron 之间权重的双向转换功能，并为 Mamba 和注意力参数提供了专用映射
• FalconH1ModelProvider（包含针对 0.5B、1.5B-Deep、7B 和 34B 等不同规模的专属变体）封装了所有模型配置，其中包括前向传播的 µP 不可学习乘子

要将此混合设计集成至 Megatron 生态系统中，TII 必须通过几项关键架构创新来解决重大工程挑战，具体详见下文。

层规格统一

Megatron Core 使用 ModuleSpec 定义层配置。针对 Falcon-H1，这需要扩展 MambaStackSubmodules 以分别保存 mamba_layer、attention_layer、mlp_layer 以及新增的 parallel_hybrid_layer 的独立规范。MambaStack 模块会遍历层类型列表，并为每个位置构建相应的模块。

在 Megatron Bridge 中，对应的 FalconH1StackSubmodules 添加了一个 falconh1_layer 规范，将这三个组件集成在一起。这使得开发者能够在单一模型定义中灵活组合 Mamba 与 Transformer 组件。

检查点转换的权重映射

在 Megatron Bridge 中，将 Hugging Face 检查点转换为 Megatron 格式需要专门的参数映射。MambaInProjMapping 类负责处理复杂的拆分逻辑，将 Mamba 的 in_proj 权重拆分为 z、x、B、C 和 dt 组件。这些组件必须在保证数值正确性的前提下，准确地分布到各个张量并行秩上。

FalconH1Bridge 在单次处理中同时管理 Mamba 与 Attention 层的张量并行重分片，并结合用于融合独立 Q、K、V 投影的 QKVMapping，以及用于合并门控与上投影的 GatedMLPMapping。在 Megatron Core 中，检查点转换工具（loader_parallelhybrid 和 saver_parallelhybrid_hf）负责实现 Megatron 分布式格式与 Hugging Face FalconH1ForCausalLM 格式之间的转换。

SSM 层的张量并行

Mamba 层具有独特的张量并行需求。A_log、D 和 dt_bias 张量沿第 0 维切分，而 x_proj 沿第 1 维切分。对于 Mamba-2，in_proj 和 conv1d 层需要特殊处理，以便在各个 rank 上正确划分 z、x、B、C 和 dt 组件。

超越经典 μP

为优化 Falcon-H1 系列模型，TII 采用了定制化的最大更新参数化（μP）方法。经典 μP 植根于神经网络理论，旨在实现超参数从基础模型规模向更大模型的无缝迁移；而 Falcon-H1 在此基础上进一步扩展，通过对 μP 乘数本身进行调优，使每个组件都能以恰当的强度进行训练。

针对基于SSM的模型中常见的训练尖峰问题，通过在SSM模块内部应用衰减乘子予以缓解，从而带来更平滑的训练过程与更清晰的实验信号。

Falcon-H1 中的 µP 乘子以不可学习张量的形式存储。它们在前向传播过程中对激活值进行缩放，且不累积梯度。该方法在将内存开销降至最低的同时，实现了对学习动态的细粒度控制，涵盖嵌入、注意力、SSM 及 MLP 组件在内的 12 个独立缩放因子。

针对 Megatron Bridge，这需要在加载 Hugging Face 检查点时增加乘子提取步骤。该桥接模块从 HF 配置中读取乘子值，并将其应用于正确的前向传播位置。注意力与 Mamba 组件均会接收各自对应的缩放因子。

Megatron Core 中 Falcon Edge 的 BitNet 集成

Falcon Edge 是基于 BitNet 架构的一系列三值（1.58-bit）TII 语言模型。为实现 Falcon Edge 的大规模训练，TII 向 Megatron Core 贡献了面向类 GPT 架构的 BitNet 预训练支持。该集成是支持 1-bit LLM 可扩展预训练工作流的关键一步，同时完整保留了 Megatron 的并行架构与性能特性。

TII 引入了两个新的并行线性层：BitNetColumnParallelLinear 和 BitNetRowParallelLinear。这些层与现有的 Megatron 张量并行线性层架构相似，但融入了 BitNet 量化逻辑。通过将 BitNet 直接嵌入到层规范级别，该集成保持了对 Megatron 张量并行、流水线并行以及分布式训练基础设施的兼容性。

在底层实现上，该方案利用了 onebitllms 的 Triton 内核来实现高效的激活值与权重量化。

在前向传播过程中，BitNet 使用量化等效操作替换了全精度矩阵乘法：

• 权重通过绝对均值缩放被量化为三值 {−1, 0, +1}。具体而言，权重张量会乘以其绝对均值的倒数进行缩放，随后进行取整并限制至 {−1, 0, +1}。
• 激活值采用逐 token 绝对最大值缩放量化为 8 比特精度。对于每个 token，会计算其在隐藏维度上的最大绝对值，并以此将激活值缩放至 [−128, 127] 区间，最后将结果四舍五入为最接近的整数。
• 核心线性运算利用这些量化权重与激活值执行，并借助 onebitllms 提供的自定义 Triton 内核进行优化。
• 通过采用三值权重（1.58 位），该模型相比全精度模型显著降低了内存占用，并实现了更快的推理速度。

在反向传播过程中：

• 梯度将绕过不可微的量化函数，使反向传播得以继续进行，仿佛量化步骤为恒等函数。
• 权重梯度基于全精度权重进行计算。量化操作仅在正向传播阶段应用，从而确保优化器的更新保持高保真度。
• 激活梯度在通过量化感知层时遵循标准反向传播。

实现

Megatron Core 中集成 BitNet 引入的改动极小，同时保持了与现有并行策略及 Megatron Core 可扩展性的完全兼容。标准 Linear 层被替换为 BitNetLinear 变体，在维持 Megatron Core 层接口不变的前提下，实现了权重的三值量化。

激活值与权重量化算子已直接集成至 Megatron 计算流水线中。张量并行机制已扩展以支持分片量化权重，缩放因子在各分片中独立处理，以确保数值正确性。Megatron 融合算子与通信模式得以保留，确保三值量化在节省内存与带宽的同时，不牺牲吞吐量。

核心组件

• 自定义线性层：两个新类扩展了 Megatron 张量并行层：BitNetColumnParallelLinear 继承 ColumnParallelLinear，BitNetRowParallelLinear 继承 RowParallelLinear。
• 量化集成：两层均重写 _forward_impl，以使用 onebitllms 的 Triton 内核（weight_quant_triton 和 activation_quant_triton）实现三值权重量化与 8 位激活量化。
• 直通估计器（STE）：梯度通过 x_quantized = x + (quant(x) – x).detach() 模式绕过量化操作。

这使得在保持全精度权重更新的同时，能够通过不可微量化进行反向传播。

集成点

• 层规范系统：BitNet 层在 get_gpt_layer_local_spec 和 get_mlp_module_spec 中注册，可通过 --use-bitnet 标志启用
• 张量并行：权重划分后，量化独立应用于每个张量并行分片，从而在分布式计算中保持数值正确性
• 训练要求：BitNet 需要 --transformer-impl local 和 onebitllms 软件包。该实现直接复用现有的 Megatron 通信模式与融合内核，无需任何修改

该集成方案在大幅节省权重内存并提升带宽效率的同时，保持了对 Megatron 流水线并行、梯度累积及优化器基础设施的兼容性

开始使用 Megatron 构建基础模型

TII Falcon-H1 混合架构与 BitNet 三元训练支持展示了基础模型构建者如何为自身的架构与训练需求扩展 Megatron Core 和 Megatron Bridge。这些贡献现已可用。

如需上手 Megatron-LM，请查阅 BitNet 预训练与 ParallelHybrid 层支持的相关内容。如需上手 Megatron-Bridge，请查阅 Falcon-H1 检查点转换与 µP 乘数处理的相关内容。