在大型语言模型（LLM）开发领域快速演进的背景下，NVIDIA Megatron Core已成为大规模训练Transformer模型的基础框架。该开源库提供了业界领先的并行能力和GPU优化性能。某机构为Falcon模型家族的创造者，对Megatron Core和Megatron Bridge框架做出了贡献并进行了集成。
Falcon-H1混合架构在Megatron Bridge中的集成
Falcon-H1并行混合架构在Megatron Bridge中的实现，突显了协调异构Transformer层和Mamba层以及非学习型μP乘数所面临的挑战。
混合并行设计
Falcon-H1采用并行设计，在每个核心处理块内，基于Transformer的注意力机制和状态空间模型（SSM）同时处理输入。注意力分支和Mamba分支的输出在投影之前进行拼接，使模型能够融合SSM的优越长上下文记忆和效率与注意力机制的远程依赖建模能力。模型中并行混合层、纯Mamba层、纯注意力层和纯多层感知机（MLP）层的比例可独立配置，支持灵活架构探索。
双仓库集成
Falcon-H1的支持横跨两个职责不同的仓库。在Megatron Core中贡献了：
- 基础性的
ParallelHybridLayer，该层并行运行Mamba和注意力并求和其输出 - 更新的层分配逻辑，在现有Mamba、注意力和MLP层类型旁引入
PARALLEL符号
这还包括用于加载和保存并行混合模型的检查点转换工具。在Megatron Bridge中，基于这些原语构建了完整的Falcon-H1模型：
FalconH1Layer扩展并行设计以包含MLP组件（构成完整的Mamba+注意力+MLP块）FalconH1Bridge提供双向的Hugging Face到Megatron权重转换，带有针对Mamba和注意力参数的特殊映射FalconH1ModelProvider（包含针对0.5B、1.5B-Deep、7B和34B的特定尺寸变体）封装所有模型配置，包括前向μP非学习型乘数
层规范统一
Megatron Core使用ModuleSpec定义层配置。对于Falcon-H1，这需要扩展MambaStackSubmodules以容纳mamba_layer、attention_layer、mlp_layer和新的parallel_hybrid_layer的独立规范。在Megatron Bridge中，相应的FalconH1StackSubmodules添加了一个捆绑所有三个组件的falconh1_layer规范。
检查点转换的权重映射
在Megatron Bridge中，将Hugging Face检查点转换为Megatron格式需要专门的参数映射。MambaInProjMapping类处理将Mamba的in_proj权重复杂拆分为z、x、B、C和dt分量。FalconH1Bridge在单次传递中管理Mamba和注意力层的张量并行重分片。
SSM层的张量并行
Mamba层具有独特的张量并行要求。A_log、D和dt_bias张量沿维度0拆分，而x_proj沿维度1拆分。对于Mamba-2，in_proj和conv1d层需要特殊处理以在各计算秩间正确划分分量。
超越经典μP
Falcon-H1系列采用定制的最大更新参数化（μP）。Falcon-H1通过调整μP乘数本身来扩展经典μP，使每个组件能以正确的强度进行训练。通过应用阻尼乘数来解决基于SSM的模型中常见的训练尖峰问题，实现更平滑的训练和更清晰的实验信号。
Falcon-H1中的μP乘数存储为非学习型张量，在前向传播中缩放激活而不累积梯度。这种方法在保持最小内存开销的同时，实现了对覆盖嵌入、注意力、SSM和MLP组件的12个不同缩放因子的学习动态的细粒度控制。
Falcon Edge的BitNet在Megatron Core中的集成
Falcon Edge是基于BitNet架构的一系列三值（1.58-bit）语言模型。为了大规模训练Falcon Edge，向Megatron Core贡献了针对GPT类架构的BitNet预训练支持。此次集成是在保留Megatron并行性和性能特性的同时，实现1-bit LLM可扩展预训练工作流的关键一步。
引入了两个新的并行线性层：BitNetColumnParallelLinear和BitNetRowParallelLinear。这些层镜像了现有的Megatron张量并行线性层，但融入了BitNet量化逻辑。
实现细节如下：
- 权重量化：使用权重的绝对值平均缩放，将权重量化为三值{−1, 0, +1}。权重张量先乘以其绝对值平均的倒数，然后四舍五入并限制到{−1, 0, +1}。
- 激活量化：使用每token的absmax缩放将激活量化为8位精度。对于每个token，计算隐藏维度上的最大绝对值，用于将激活缩放到[−128, 127]范围，结果四舍五入到最近整数。
- 核心计算：使用这些量化后的权重和激活，借助
onebitllms提供的自定义Triton内核进行核心线性运算。
反向传播过程中，梯度绕过不可微分的量化函数（使用直通估计器STE），使反向传播如同量化步骤是恒等函数一样进行。权重梯度在全精度权重上计算，确保优化器更新保持高保真度。
核心组件
- 自定义线性层：
BitNetColumnParallelLinear扩展ColumnParallelLinear，BitNetRowParallelLinear扩展RowParallelLinear - 量化集成：两个层都重写
_forward_impl，使用onebitllms的Triton内核（weight_quant_triton和activation_quant_triton）应用三值权重量化和8位激活量化 - 直通估计器：通过模式
x_quantized = x + (quant(x) – x).detach()实现梯度绕过量化
集成点
- 层规范系统：BitNet层在
get_gpt_layer_local_spec和get_mlp_module_spec中注册，通过--use-bitnet标志激活 - 张量并行：量化在权重被分区后独立应用于每个张量并行分片
- 训练要求：BitNet需要
--transformer-impl local和onebitllms包FINISHED
