企业官网建设流程全解析

• 人工智能
• 大模型
• 强化学习
• AI Agent
• 微调

本文基于Megatron-LM/examples/mamba/README.md展开，结合仓库内train.sh、推理脚本、MambaModel/MambaStack源码与混合层分配实现，系统讲解如何在 Megatron-LM 中训练、部署基于 Mamba-2 的纯 SSM 与混合架构语言模型。读完本文，你将掌握从 Docker 环境搭建、800M/8B 规模预训练脚本的每个参数含义，到混合层比例配置、文本生成服务器启动，以及混合检查点在张量/流水线并行维度间的转换方法。

背景：Mamba-2 语言模型与混合架构

Megatron-LM/examples/mamba目录是 NVIDIA 技术报告《An Empirical Study of Mamba-based Language Models》所对应代码的入口点。该技术报告研究的是以结构化状态空间模型（SSM）为主体的语言模型——尤其是 Mamba-2 架构——并将其与 Transformer 架构在同等规模下进行对比。技术报告描述的 8B 规模混合模型架构，正是本目录脚本所演示的核心配置。

仓库在 Hugging Face 上发布了报告中部分模型的权重参数。需要特别注意：当前main分支的代码不再兼容Mamba2-*命名的历史检查点，如需加载这类旧权重，应当使用技术报告当时的固定代码快照（报告快照也支持运行最初的 Mamba 版本）。这意味着本文所有操作都应以当前main分支代码为准。

从仓库结构看，Mamba 相关的核心实现集中在 megatron/core/ssm（mamba_mixer.py、mamba_block.py、mamba_layer.py、mamba_hybrid_layer_allocation.py等）与 megatron/core/models/mamba（mamba_model.py、mamba_layer_specs.py）。模型入口为 pretrain_mamba.py，它通过mamba_builder（见 mamba_builders.py）构建MambaModel，并要求以--use-mcore-models运行——源码中mamba_builder明确断言args.use_legacy_models is False，即 Mamba 仅在 Megatron-Core 路径下受支持。

环境准备：通过 Dockerfile 构建运行容器

目录下的 Dockerfile 基于nvcr.io/nvidia/pytorch:24.01-py3镜像构建，核心步骤是：

1. 卸载自带 Triton，安装triton==2.1.0、sentencepiece==0.1.99、flask-restful；
2. 从源码编译causal-conv1d（固定v1.2.2.post1标签，设置CAUSAL_CONV1D_FORCE_BUILD=TRUE）；
3. 从源码编译mamba（固定v2.0.3标签，设置MAMBA_FORCE_BUILD=TRUE）。

注释里说明了为何要本地编译：NGC 基础镜像内嵌的是较老版本的 PyTorch（如2.2.0.dev231106），而causal-conv1d与mamba-ssm这两个包通常只为新版 PyTorch 发布 wheel，直接pip install会因版本不匹配产生 Python import 错误。因此对较老的 NGC PyTorch 容器，需要从源码构建这两个包。如果遇到MambaSSM is not installed之类的ImportError，说明mamba-ssm未正确安装或与 PyTorch 版本不兼容——这是 mamba_mixer.py 中MambaMixer构造时显式抛出的检查（该模块还要求einops已安装）。

构建并启动容器：

docker build -t your_image_name:your_tag . docker run --gpus all -it --rm \ -v /path/to/megatron:/workspace/megatron \ -v /path/to/dataset:/workspace/dataset \ -v /path/to/checkpoints:/workspace/checkpoints \ -w /workspace/megatron/examples/mamba \ your_image_name:your_tag

其中三个-v分别将 Megatron-LM 仓库、训练数据集、检查点目录挂载进容器，-w把工作目录设为examples/mamba，方便直接运行下文脚本。

预训练：train.sh 全参数拆解

train.sh是单节点（8 卡）预训练示例脚本，用法为：

./train.sh <data-path> <tokenizer-path>

通过MODEL_SCALE变量在 800M 与 8B 两档规模间切换，脚本用case语句为每档规模预设了张量并行度与模型尺寸：

脚本开头导出若干 NCCL/CUDA 环境变量（NCCL_IB_SL=1、CUDA_DEVICE_MAX_CONNECTIONS=1、NCCL_IB_TIMEOUT=19、NCCL_IB_QPS_PER_CONNECTION=4），并创建checkpoints、data-cache、tensorboard三个目录。Triton 缓存被导向本地目录，且指定了并行文件缓存管理器：

export TRITON_CACHE_DIR="./triton-cache/" export TRITON_CACHE_MANAGER="megatron.core.ssm.triton_cache_manager:ParallelFileCacheManager"

序列长度 4096，训练样本数按"300B tokens / 4096"折算为TRAIN_SAMPLES=73242188，学习率预热 50000 样本、衰减至TRAIN_SAMPLES - LR_WARMUP_SAMPLES = 73192188。最终通过torchrun --nproc_per_node 8 ../../pretrain_mamba.py启动训练。

模型与并行参数

• --tensor-model-parallel-size：按档位取 1 或 4；配合--sequence-parallel使用序列并行。
• --pipeline-model-parallel-size 1：示例为单节点单流水线。
• --use-distributed-optimizer、--overlap-param-gather、--overlap-grad-reduce：启用分布式优化器与参数收集/梯度归约重叠，降低通信开销。
• --untie-embeddings-and-output-weights：输入嵌入与输出 logits 权重解耦。对应MambaModel中share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights（见 mamba_builders.py）。
• --position-embedding-type none：Mamba 本身不依赖位置编码，混合架构中注意力层的相对位置信息由 SSM 的时序结构提供；若使用 RoPE 则需要--position-embedding-type rope并配合--rotary-percent。
• --group-query-attention --num-query-groups 8：混合架构中的注意力层采用 GQA，8 个 query group。

混合层比例

--hybrid-attention-ratio 0.08 --hybrid-mlp-ratio 0.5是技术报告中 8B 混合模型的标志性配置（48/56 层中约 8% 为注意力层、50% 为 MLP 层，其余为 Mamba 层）。这是本文后续"Hybrid Options"一节的核心，也是理解 8B 混合模型架构的钥匙。

数据、分词与训练超参

• 数据：--data-path、--data-cache-path、--split 99,1,0（99% 训练 / 1% 验证 / 0% 测试）；--no-create-attention-mask-in-dataloader表明混合模型在数据加载阶段不再需要显式构造注意力掩码。
• 分词器：--tokenizer-type GPTSentencePieceTokenizer --tokenizer-model <path>，需要与数据集预处理时一致。
• 训练策略：--micro-batch-size 4、--global-batch-size（按档位 32/8）；学习率--lr 2.5e-4、--min-lr 2.5e-5、cosine 衰减；--weight-decay 0.1、--clip-grad 1.0；Adam 的--adam-beta1 0.9、--adam-beta2 0.95。
• 数值与正则：--bf16；--attention-dropout 0.0 --hidden-dropout 0.0；--disable-bias-linear；归一化--normalization RMSNorm（MambaMixer 默认rmsnorm=True）。
• 日志与保存：--log-interval 10、--save-interval 2000、--eval-interval 2000、--eval-iters 32、--tensorboard-dir。

模型规格参数

--use-mcore-models \ --spec megatron.core.models.mamba.mamba_layer_specs mamba_stack_spec

mamba_stack_spec定义于 mamba_layer_specs.py，它把栈内四种层类型分别绑定到具体模块实现：Mamba 层用MambaLayer+MambaMixer（in_proj=TELayerNormColumnParallelLinear、out_proj=TERowParallelLinear），注意力层用带 causal mask 的SelfAttention（TEDotProductAttention），MLP 层用MLP，另外还预留了 MoE 层（get_moe_module_spec）。这就是混合模型"一层栈、四种层型"的模块化基础。

在 mamba_block.py 中，MambaStack.__init__调用allocate_layers(num_layers, hybrid_attention_ratio, hybrid_mlp_ratio, hybrid_override_pattern)生成layer_type_list，随后按符号逐层构建：Mamba 层走MambaLayer分支，注意力/MLP/MoE 层走TransformerLayer分支。前向时对TransformerLayer统一传attention_mask与rotary_pos_emb，对MambaLayer则直接调用。这也解释了为什么混合模型能把三种异构层放进同一个栈里训练。

文本生成：两种 8B 模型的推理服务器

目录提供两个文本生成脚本，均以torchrun单进程（--nproc_per_node 1，master 端口 6000）启动，用法为：

./run_text_gen_server_8b.sh <checkpoint-path> <tokenizer-path>

启动后客户端通过python ../../tools/text_generation_cli.py <URL-provided-by-server>发起生成请求。

run_text_gen_server_8b.sh：8B 混合模型

run_text_gen_server_8b.sh按技术报告中的 8B 混合模型配置启动，--tensor-model-parallel-size 1（单卡推理），参数与训练脚本保持一致：

• 架构：--num-layers 56 --hidden-size 4096 --num-attention-heads 32 --group-query-attention --num-query-groups 8；
• 混合比例：--hybrid-attention-ratio 0.08 --hybrid-mlp-ratio 0.5；
• 其余：--untie-embeddings-and-output-weights、--attention-dropout 0.0 --hidden-dropout 0.0、--disable-bias-linear、--normalization RMSNorm、--position-embedding-type none、--seq-length 4096 --max-position-embeddings 4096；
• 分词与精度：--tokenizer-type GPTSentencePieceTokenizer、--bf16、--micro-batch-size 1；
• 模型规格：--use-mcore-models --spec megatron.core.models.mamba.mamba_layer_specs mamba_stack_spec；
• 分布式：--distributed-backend nccl --distributed-timeout-minutes 1440，固定--seed 42。

脚本通过 tools/run_mamba_text_generation_server.py 以model_type="mamba"进入通用生成服务器入口（见 tools/run_text_generation_server.py 的main(model_type)分支），在该入口中按model_type选择mamba_builder构建模型并加载检查点。

重要限制：如果换用其他并行配置或架构差异的检查点，脚本参数必须相应修改。例如要推理 8B 纯 Mamba-2 模型，把--hybrid-attention-ratio与--hybrid-mlp-ratio改为0.0（或直接删掉这两个参数）。同理，若检查点使用不同的 TP/PP 配置，需同步调整--tensor-model-parallel-size等参数——这就是下文"检查点格式"一节强调的"推理配置必须与检查点严格匹配"。

run_text_gen_server_8b_gpt3.sh：8B 参考 Transformer

run_text_gen_server_8b_gpt3.sh用于启动 8B 参考 Transformer 检查点的推理服务器，作为对比基线。其参数明显不同：

• --num-layers 32 --hidden-size 4096 --num-attention-heads 32（标准 GPT 结构，无 GQA 与 hybrid 比例）；
• --use-flash-attn --apply-layernorm-1p、--squared-relu；
• 位置编码采用 RoPE：--position-embedding-type rope --rotary-percent 0.5；
• --transformer-impl local、--use-mcore-models；
• 同样固定--seed 42。

该脚本走的是 tools/run_text_generation_server.py 的默认model_type="gpt"路径，即gpt_builder。两个脚本放在同一目录，方便在完全相同的硬件与数据条件下对比 Mamba-2 混合模型与 Transformer 基线的生成效果。

检查点格式与并行规模转换

推理时模型配置必须与检查点文件严格一致，包括混合层配置（各层类型分布）与并行配置（TP/PP 大小）。如果需要一个混合检查点在不同的张量并行或流水线并行规模下使用，必须先用转换脚本处理：tools/checkpoint/hybrid_conversion.py。

运行前需要把 Megatron-LM 仓库根目录加入PYTHONPATH：

export PYTHONPATH=<path-to-megatron>:PYTHONPATH

脚本文件末尾（if __name__ == "__main__"）给出了示例命令：

python hybrid_conversion.py \ --load-dir mamba2-840m-test/checkpoints/ \ --save-dir mamba2-840m-test-conversion/checkpoints/ \ --target-pp-size 1 \ --target-tp-size 1

其核心逻辑（见 hybrid_conversion.py）会把检查点内记录的args.tensor_model_parallel_size/pipeline_model_parallel_size改写为--target-tp-size/--target-pp-size指定的目标值，并按 Mamba 混合模型特有的张量切分方式（Mamba 层的 in-proj/out-proj、因果卷积、SSM 状态等按 TP 维度分片）完成权重合并或再切分。其余可选参数包括：

• --reset-iterations：重置迭代计数；
• --d-model（默认 4096）：模型隐藏维度；
• --mamba-version（默认 2）：目标 Mamba 版本；
• --mamba-d-state（默认 128）、--mamba2-n-groups（默认 8）、--mamba2-head-dim（默认 64）：Mamba 层状态维度、分组数与头维度，需与模型实际配置一致。

需要说明：脚本本身不负责混合层模式的迁移——若源与目标检查点属于不同架构（如注意力比例不同），仍需先确保模型定义与检查点一致。

Hybrid Options：混合层比例的配置规则

--hybrid-attention-ratio与--hybrid-mlp-ratio是控制"注意力 / MLP / Mamba"三类层在总层数中占比的核心开关（定义于 megatron/training/arguments.py，默认均为 0.0，取值范围[0.0, 1.0]）。

--hybrid-attention-ratio ATT：目标注意力层数占总层数的比例。例如 48 层中 4 层为注意力层，写作--hybrid-attention-ratio 0.08。

--hybrid-mlp-ratio MLP：目标 MLP 层数占总层数的比例。例如 48 层中 24 层为 MLP 层，写作--hybrid-mlp-ratio 0.5。

约束与边界语义：

• ATT + MLP必须 ≤ 1.0；
• 1.0 - ATT - MLP即混合 Mamba 比例（Mamba 层占总层数的比例）；
• ATT = MLP = 0：纯 Mamba 模型；
• ATT = MLP = 0.5：Transformer 模型（所有层均为注意力 + MLP 的标准结构）。

当ATT或MLP任一大于 0.0，或指定了--hybrid-override-pattern时，日志会打印实际使用的混合层模式信息，包括每种层的实际数量与实际占比，以及用符号串表示的层布局。

分配算法与覆盖模式

底层实现在 mamba_hybrid_layer_allocation.py。Symbols定义了四种层的符号：M（Mamba）、*（Attention）、-（MLP）、E（MoE）。

默认的_allocate_auto算法分两步：

1. 先均匀分配注意力层（起始与结尾均为 Mamba 层，注意力层等间距插入）；
2. 再把 MLP 层等间距地替换掉部分 Mamba 层（右对齐分布，且不会替换注意力层）。

最终由allocate_layers(total_layers_count, target_attention_ratio, target_mlp_ratio, override_pattern)返回一个符号列表，MambaStack依据它逐层构建（见 mamba_block.py）。

若需要手工指定与算法生成不同的层布局，用--hybrid-override-pattern传入一个由Symbols字符组成的字符串。_allocate_override会校验两件事：

• 字符串长度必须等于总层数，否则抛ValueError（"The hybrid override pattern is the wrong length"）；
• 每个字符必须是Symbols.VALID中的一员，否则报错。

此外，当任一比例参数大于 0.0 时，覆盖模式的各类层数量必须与算法默认分配的数量一致，否则allocate_layers会抛出ValueError（"The number of each type of layer in the override pattern must match..."）。日志中会打印原模式 A、覆盖模式 B 及差异警告（覆盖后不再重新做层计数校验）。文件底部自带一组test_cases演示了各种比例组合与覆盖模式（例如 48 层、0.25/0.25比例下用MM*-MM*-...模式覆盖），可直接python mamba_hybrid_layer_allocation.py运行查看输出。

与混合比例配套的还有三个 Mamba 层维度参数（同样定义于 megatron/training/arguments.py）：--mamba-state-dim（SSM 状态维度，默认 128）、--mamba-head-dim（每个头维度，默认 64）、--mamba-num-groups（分组数，默认 8）；--mamba-num-heads未设置时按hidden-size * expand // mamba-head-dim推算。这些参数会被 mamba_mixer.py 的MambaMixer读取（d_state、headdim、ngroups，以及默认d_conv=4、expand=2、chunk_size=128），并约束融合核的维度对齐（如(2*d_inner + 2*ngroups*d_state + nheads) % 16 == 0）。修改这些参数时必须保证与检查点一致。

Mamba vs Mamba-2：版本支持说明

当前main分支代码只支持 Mamba-2，不支持最初的 Mamba（Mamba-1）。MambaMixer内部直接构建 Mamba-2 风格的多头 SSM 状态（ssm_states_shape = (nheads_local_tp, headdim, d_state)），并依赖mamba-ssm包中的 Mamba-2 融合算子；混合模型的"注意力 + Mamba"组合也与 Mamba-2 的 state 管理机制（见 megatron/core/ssm/mamba_mixer.py 与megatron/core/inference/contexts中的 Mamba 推理状态）深度耦合。

如果需要运行最初的 Mamba 版本，应使用技术报告对应的固定代码快照（该快照的examples/mamba目录可配置为运行 Mamba-1）。这同时意味着：以本文配置训练得到的权重，应视为 Mamba-2 时代的产物；若要与Mamba2-*系列旧权重对接，需回到快照分支加载。

小结

本目录提供了一条从零训练到推理部署的完整 Mamba-2 流水线：Dockerfile解决依赖环境，train.sh覆盖 800M/8B 两档规模的混合架构预训练，两个生成脚本分别对应混合模型与 Transformer 基线的推理，hybrid_conversion.py打通不同并行规模间的检查点迁移，而--hybrid-attention-ratio/--hybrid-mlp-ratio/--hybrid-override-pattern三件套则提供了从纯 Mamba 到纯 Transformer 之间任意混合形态的建模能力。理解这些脚本背后的MambaModel→MambaStack→allocate_layers调用链，可以帮助你在不修改核心代码的情况下，为自定义的混合层比例、模型规模与并行配置快速搭建训练与推理环境。

• 人工智能
• 大模型
• 强化学习
• AI Agent
• 微调