Falcon-H1混合架构与BitNet在Megatron中的实现

原创

用户11764306

发布于 2026-04-14 13:18:30

280

在大型语言模型（LLM）开发快速演进的背景下，NVIDIA Megatron Core已成为大规模训练Transformer模型的基础框架。该开源库提供业界领先的并行能力和GPU优化性能。Megatron Core目前在NVIDIA/Megatron-LM仓库中以GitHub优先的方式开发，并越来越多地受益于基础模型构建者的贡献，使其成为面向开放AI模型的更灵活、更具前瞻性的引擎。

本文提供了一项技术概述，介绍Falcon模型家族创建者——某机构（TII）——如何为Megatron Core和Megatron Bridge框架做出贡献并进行集成。第一部分探讨了在Megatron Bridge中实现Falcon-H1并行混合架构，重点阐述协调异构Transformer和Mamba层以及非可学习μP乘数所面临的挑战。第二部分阐述了BitNet在Megatron Core中的集成，详细说明了用三值参数层替代标准线性层的过程及其对训练效率和可扩展性的影响。

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

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

在Megatron Bridge中实现Falcon-H1并行混合架构，凸显了协调异构Transformer和Mamba层以及非可学习μP乘数的挑战。以下各节提供了该集成的细节。

混合并行设计

TII对Megatron贡献的核心是Falcon-H1并行混合架构。该设计不同于近期其他混合模型中的顺序分层方式。如图1所示，在每个块内部，注意力机制和SSM并行运行，它们的输出在通过块的输出投影之前被拼接在一起。SSM和注意力头的数量是可配置的，可根据需要调整。

图1. Falcon-H1混合架构有助于在每个核心处理块内同时处理输入，从而加速性能

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分量的复杂拆分。这些分量必须正确分布在张量并行 ranks 上，同时保持数值正确性。

FalconH1Bridge在一次遍历中管理Mamba和注意力层的张量并行重分片，同时使用QKVMapping融合独立的Q、K、V投影，使用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层需要特殊处理，以便在各 ranks 上正确划分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位）TII语言模型。为了规模化训练Falcon Edge，TII向Megatron Core贡献了对类似GPT架构的BitNet预训练支持。这一集成是在保留Megatron并行性和性能特征的同时，实现1位LLM可扩展预训练工作流的关键一步。

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

在底层，该实现利用onebitllms的Triton内核进行高效的激活和权重量化。

在前向传播过程中，BitNet用量化后的等效运算替换了全精度矩阵乘法：

权重使用绝对均值缩放量化到三值{−1, 0, +1}。权重张量先乘以其绝对均值的倒数，然后四舍五入并裁剪到{−1, 0, +1}。
激活值使用每token的absmax缩放量化到8位精度。对于每个token，计算隐藏维度上的最大绝对值，用于将激活值缩放到−128, 127范围，并将结果四舍五入到最接近的整数。
核心线性运算使用这些量化后的权重和激活值，并利用onebitllms提供的定制Triton内核进行优化。

通过使用三值权重（1.58位），与全精度模型相比，模型显著减少了内存占用并实现了更快的推理速度。

在反向传播过程中：

梯度绕过不可微的量化函数，使得反向传播如同量化步骤是恒等函数一样进行。
权重梯度在全精度权重上计算。量化仅在前向传播期间应用，确保优化器更新保持高保真度。
激活梯度通过量化感知层遵循标准反向传播。

实现

Megatron Core中的BitNet集成引入了最小的更改，同时保持与现有并行策略和Megatron Core可扩展性的完全兼容。标准线性层被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流水线并行、梯度累积和优化器基础设施的兼容性。