FasterTransformer GPT-NeoX 推理指南：20B 模型部署、权重转换与多卡并行实践-平芜编程栈

推理引擎
算子库
大模型

【免费下载链接】FasterTransformer

Transformer related optimization, including BERT, GPT

项目地址：https://gitcode.com/gh_mirrors/fa/FasterTransformer

导读

本文围绕 FasterTransformer 仓库中的 docs/gptneox_guide.md 展开，系统讲解如何在 FasterTransformer 上运行 EleutherAI 的 GPT-NeoX 模型（目前仅验证过 20B 参数版本）。你将掌握：从 EleutherAI / HuggingFace 两种来源获取并转换检查点、配置 tokenizer、生成 GEMM 配置、以 C++ 与 PyTorch 两种方式执行推理，以及通过tensor_para_size/pipeline_para_size实现单节点内多卡并行（TP/PP）的完整操作链路，并结合仓库源码理解其注意力优化与模型结构设计。

背景：为什么在 FasterTransformer 上运行 GPT-NeoX

GPT-NeoX 是 EleutherAI 开发的 GPT 类自回归语言模型，其 20B 版本是目前仓库中验证过的唯一 GPT-NeoX 规格。FasterTransformer 对 GPT-NeoX 的优化与 GPT 系列基本一致，可参见 docs/gpt_guide.md 中的优化说明；更多通用介绍可参考 docs/gptj_guide.md。

从源码结构看，FasterTransformer 在 src/fastertransformer/models/gptneox/ 下实现了完整的 GPT-NeoX 推理模型：

GptNeoX.h：模型主入口，组合上下文解码器GptNeoXContextDecoder、自回归解码器GptNeoXDecoder与动态解码层DynamicDecodeLayer；
GptNeoXContextDecoder.cc：负责 context（提示词）部分的并行前向；
GptNeoXDecoder.cc：负责逐 token 自回归解码；
GptNeoXWeight.h 与 GptNeoXWeight.cc：定义权重加载与划分逻辑。

在 GptNeoX.h 中可以看到几个关键结构事实：

neox_rotary_style_ = true：GPT-NeoX 采用 Neox 风格的旋转位置编码（rotary embedding）；
use_gptj_residual_ = true：残差结构与 GPT-J 保持一致；
layernorm_eps_ = 1e-5f：LayerNorm 的 epsilon 固定为 1e-5；
通过CONTEXT_ATTENTION_BMM1_HALF_ACCUM环境变量控制 context 阶段 QK GEMM 的累加精度（默认 fp32）。

推理调优环境变量

针对具体使用场景，FasterTransformer 提供以下环境变量（与 GPT-J 指南中的设置一致）：

名称	描述	默认值	可选值
`FMHA_ENABLE`	启用 fused multi-head attention 核函数（fp16 累加）	禁用	`ON`表示启用，其余值均表示禁用
`CONTEXT_ATTENTION_BMM1_HALF_ACCUM`	对 QK GEMM 使用 fp16 累加；仅对非融合 multi-head attention 核函数生效	fp32 累加	`ON`表示 fp32 累加，其余值表示 fp16 累加

从 GptNeoX.h 的源码可见，CONTEXT_ATTENTION_BMM1_HALF_ACCUM是通过std::getenv在构造阶段读取的：当环境变量为nullptr或取值不等于"ON"时，is_context_qk_buf_float_为true（即默认采用 fp32 累加缓冲区）。

支持的功能清单

根据指南，当前支持能力如下：

检查点转换器：EleutherAI 格式、HuggingFace 格式；
数据类型：FP32、FP16；
推理特性：多 GPU 多节点推理、动态随机种子（dynamic random seed）、停止 token（stop tokens）、坏词列表（bad words list）、同时支持 beam search 与采样（sampling）。

环境准备与构建

通用依赖要求与 GPT-J 一致，详见 docs/gptj_guide.md 的 Requirements 一节，核心包括：

CMake >= 3.13（PyTorch 场景）；
CUDA 11.0 及以上；
NCCL 2.10 及以上（多卡推理必需）；
Python 3（仅在 Python 3 上验证）；
PyTorch：在 1.8.0 上验证，>= 1.5.0 应可运行。

官方推荐使用 NGC 镜像，例如nvcr.io/nvidia/pytorch:22.09-py3。构建时按 GPU 计算能力设置-DSM（如 V100 为 70、T4 为 75、A100 为 80、A10 为 86），并以-DBUILD_MULTI_GPU=ON打开多卡支持：

git clone https://github.com/NVIDIA/FasterTransformer.git mkdir -p FasterTransformer/build cd FasterTransformer/build git submodule init && git submodule update pip3 install fire jax jaxlib cmake -DSM=xx -DCMAKE_BUILD_TYPE=Release -DBUILD_MULTI_GPU=ON .. make -j12

从 EleutherAI 检查点部署

下载模型权重

使用wget递归下载 EleutherAI 公开的 GPT-NeoX-20B slim_weights：

wget --cut-dirs=5 -nH -r --no-parent --reject "index.html*" https://mystic.the-eye.eu/public/AI/models/GPT-NeoX-20B/slim_weights/ -P 20B_checkpoints

下载完成后，20B_checkpoints目录内是分片保存的layer_XX-model_YY-model_states.pt文件。从 examples/pytorch/gptneox/utils/eleutherai_gpt_neox_convert.py 的文档注释可以看到其权重组织方式：

layer_00-model_00-model_states.pt：word_embeddings.weight（embedding 表，按张量并行切分）；
layer_02-model_00-model_states.pt ~ layer_45-model_01-model_states.pt：每层包含input_layernorm.*、attention.query_key_value.*、attention.rotary_emb.inv_freq、attention.dense.*、post_attention_layernorm.*、mlp.dense_h_to_4h.*、mlp.dense_4h_to_h.*；
layer_47-model_00-model_states.pt：最终 LayerNorm（model_00与model_01权重相同，取其一即可）；
layer_48-model_00-model_states.pt：final_linear.weight（logit GEMM 权重）；
mp_rank_xx_model_states.pt：训练状态，推理无需使用。

转换脚本还维护了一份weights_skip_tensor_split列表（如input_layernorm.bias/weight、attention.dense.bias、mlp.dense_4h_to_h.bias、post_attention_layernorm.bias/weight），这些张量不做张量并行切分。

转换为 FT 权重

python ../examples/pytorch/gptneox/utils/eleutherai_gpt_neox_convert.py 20B_checkpoints ../models/gptneox -t 2

参数含义：

第一个位置参数20B_checkpoints：下载的 EleutherAI 检查点目录；
第二个位置参数../models/gptneox：FT 权重输出目录；
-t 2：tensor parallelism 大小（此处为 2 卡张量并行）。后续 C++ 推理时mpirun -n的进程数必须与之匹配。

转换过程会执行权重重排（如将attention.query_key_value.weight重排为 FT 需要的[hidden_size, 3, num_heads, head_hidden]布局），并按-t指定的 GPU 数把权重切分为%d-gpu/子目录。同时，脚本会生成 FT 推理所需的 config.ini 描述文件，其内容与该模型的核心超参数一一对应（见下文“模型配置”小节）。

Tokenizer

单独下载 tokenizer 配置：

wget https://mystic.the-eye.eu/public/AI/models/GPT-NeoX-20B/slim_weights/20B_tokenizer.json

对文本文件进行 tokenize / detokenize 时，使用仓库脚本 examples/pytorch/gptneox/utils/hftokenizer.py，并通过--tokenizer传入 tokenizer 配置路径：

python ../examples/pytorch/gptneox/utils/hftokenizer.py out --tokenizer 20B_tokenizer.json

从 HuggingFace 检查点部署

先克隆 HuggingFace 上的 GPT-NeoX 模型仓库（需安装 git-lfs）：

git lfs clone https://huggingface.co/<MODEL_GROUP>/<MODEL_NAME>

再使用转换脚本 examples/pytorch/gptneox/utils/huggingface_gptneox_convert.py 转为 FT 原始权重：

python ../examples/pytorch/gptneox/utils/huggingface_gptneox_convert.py -i ../path/to/your/model -o ../../path/to/fastertransformer/model -i_g 1 -m_n gptneox

各参数含义：

-i：HuggingFace 模型所在目录（GPTNeoXForCausalLM.from_pretrained读取）；
-o：FT 权重输出目录；
-i_g：inference GPU 数，即张量并行大小，权重会按该数值切分并存入<saved_dir>/<infer_gpu_num>-gpu/；
-m_n：模型名称，此处为gptneox。

从 huggingface_gptneox_convert.py 的源码可以看到，转换时对 QKV 权重做了关键布局变换：HuggingFace 中attention.query_key_value.weight的存储形状为[hidden_size, num_heads, 3, head_hidden]，转换脚本先 reshape 再转置为 FT 需要的[hidden_size, 3, num_heads, head_hidden]；同时按-i_g对dense/mlp权重分别沿行、列方向切分，保证每张卡只持有本地分片。转换同时支持--weight_data_type指定fp32/fp16（get_weight_data_type）。

运行 GPT-NeoX

第一步：生成 gemm_config.in

在编译产物build/bin目录下，先用gpt_gemm为当前硬件与模型规格生成 GEMM 调优配置：

./bin/gpt_gemm <batch_size> <beam_width> <max_input_len> <head_number> <size_per_head> <inter_size> <vocab_size> <data_type> <tensor_para_size> E.g., ./bin/gpt_gemm 8 1 32 64 96 24576 50432 1 2

其中data_type：0 表示 FP32，1 表示 FP16，2 表示 BF16。示例命令对应 GPT-NeoX-20B 的规格（head_number=64、size_per_head=96、inter_size=24576、vocab_size=50432），tensor_para_size=2，与后续 C++ 示例的mpirun -n 2保持一致。

第二步：C++ 推理

配置文件 examples/cpp/gptneox/gptneox_config.ini 控制模型路径、模型规模、张量并行大小及采样超参数，运行：

mpirun -n 2 --allow-run-as-root ./bin/gptneox_example

将gptneox_config.ini中[ft_instance_hyperparameter]段的data_type设置为fp16，即可在 FP16 下运行。

gptneox_example的编译入口位于 examples/cpp/gptneox/gptneox_example.cc，其CMakeLists.txt定义于 examples/cpp/gptneox/CMakeLists.txt；同目录下还提供了 gptneox_triton_example.cc 用于以 Triton 模型方式启动。

推理完成后，out文件中的 token id 序列可用 hftokenizer 还原为文本：

wget https://mystic.the-eye.eu/public/AI/models/GPT-NeoX-20B/slim_weights/20B_tokenizer.json python ../examples/pytorch/gptneox/utils/hftokenizer.py out --tokenizer 20B_tokenizer.json

第三步：PyTorch 推理

examples/pytorch/gptneox/gptneox_example.py 演示了如何声明模型（examples/pytorch/gptneox/utils/gptneox.py中的GptNeoX封装）、加载检查点、前向 context 输入并得到生成结果。运行python .../gptneox_example.py -h可查看全部参数。

常用命令行参数（摘自脚本的 argparse 定义，见 gptneox_example.py）：

参数	默认值	说明
`--output_len`	32	生成序列长度
`--beam_width`	1	beam search 的 beam 宽度；为 1 时走采样
`--top_k`	1	top-k 采样候选数
`--top_p`	0.0	top-p 采样概率阈值
`--temperature`	1.0	采样温度
`--len_penalty`	0.0	长度惩罚（仅 beam search 生效）
`--beam_search_diversity_rate`	0.0	beam search diversity rate
`--tensor_para_size`	1	张量并行大小
`--pipeline_para_size`	1	流水线并行大小
`--ckpt_path`	`../models/gptneox/c-model/NeoX-1.3B/1-gpu`	转换后的 FT 权重目录（内含 config.ini）
`--tokenizer_path`	`../models/gptneox/model/NeoX-1.3B`	tokenizer 所在目录
`--lib_path`	`./lib/libth_transformer.so`	`libth_transformer.so`动态库路径
`--sample_input_file`	无	上下文输入文件路径（每行一条）
`--max_batch_size`	8	最大 batch 大小
`--repetition_penalty`	1.0	重复惩罚
`--max_seq_len`	1024	位置编码表最大序列长度
`--inference_data_type`(`--data_type`)	fp16	可选`fp32`/`fp16`
`--time`	关闭	是否测量推理耗时（输出 tokens/sec 吞吐）
`--enable_random_seed`	关闭	是否启用动态随机种子

关于输入与模型配置，脚本从ckpt_path下的config.ini读取head_num、size_per_head、vocab_size、num_layer、rotary_embedding、start_id、end_id、use_gptj_residual、weight_data_type等超参数（见 gptneox_example.py）。

带上下文输入（context）的生成

要基于上下文生成输出，需创建一个每行一条 prompt 的文本文件，并通过--sample_input_file指向它。脚本按max_batch_size截取行数，将每行经 tokenizer encode 后作为start_ids（gptneox_example.py）。若未指定该参数，脚本默认以batch_size = max_batch_size的<|endoftext|>（即end_id）作为无条件生成输入。

模型配置（config.ini）

转换脚本生成的config.ini采用如下结构（参见 eleutherai_gpt_neox_convert.py），与 C++ 端 gptneox_config.ini 中[gptneox_20B]段一一对应：

[gptneox] model_name=gptneox_20B head_num=64 size_per_head=96 vocab_size=50432 num_layer=44 rotary_embedding=24 start_id=0 end_id=2 inter_size=24576 use_gptj_residual=1 weight_data_type=fp32

各字段含义：

head_num=64、size_per_head=96：64 头、每头 96 维，隐藏层维度为 64×96=6144；
vocab_size=50432：词表大小；
num_layer=44：Decoder 层数；
rotary_embedding=24：旋转位置编码维度（GPT-NeoX 对 Query 和 Key 都施加旋转，Neox 风格）；
start_id=0/end_id=2：起始 / 结束 token id，必须与 tokenizer 保持一致；
inter_size=24576：FFN 中间层维度；
use_gptj_residual=1：采用 GPT-J 风格残差结构；
weight_data_type=fp32：权重存储精度（可在运行时以--data_type fp16做 FP16 推理）。

多卡 TP / PP 运行

PyTorch 版支持单节点内张量并行（TP）与流水线并行（PP）的组合。进程数必须严格等于tensor_para_size * pipeline_para_size，且转换权重时-i_g（或 EleutherAI 转换的-t）要与目标tensor_para_size对应：

# 无并行（tensor_para_size=1, pipeline_para_size=1） python ../examples/pytorch/gptneox/gptneox_example.py # TP（tensor_para_size=2, pipeline_para_size=1） mpirun -n 2 --allow-run-as-root python ../examples/pytorch/gptneox/gptneox_example.py --tensor_para_size=2 --pipeline_para_size=1 --ckpt_path="/path/to/your/model/2-gpu" # PP（tensor_para_size=1, pipeline_para_size=2） mpirun -n 2 --allow-run-as-root python ../examples/pytorch/gptneox/gptneox_example.py --tensor_para_size=1 --pipeline_para_size=2 --ckpt_path="/path/to/your/model/1-gpu" # TP + PP（tensor_para_size=2, pipeline_para_size=2） mpirun -n 4 --allow-run-as-root python ../examples/pytorch/gptneox/gptneox_example.py --tensor_para_size=2 --pipeline_para_size=2 --ckpt_path="/path/to/your/model/2-gpu"

注意：PP 场景下各 rank 使用相同的1-gpu权重（模型按层数均匀切分给流水线各阶段）；TP 场景下每 rank 使用切分后的2-gpu权重。当tensor_para_size * pipeline_para_size > 1时，脚本会调用dist.init_process_group(backend=dist.Backend.MPI)初始化 MPI 进程组（gptneox_example.py）。

在 C++ 端，tensor_para_size与pipeline_para_size同样在 gptneox_config.ini 的[ft_instance_hyperparameter]段中控制，且进程数必须等于两者之积。

采样与解码参数

gptneox_config.ini 的[request]段集中了推理请求级的解码参数：

参数	示例值	说明
`beam_width`	1	beam search 宽度；为 1 时使用采样
`top_k`	1	top-k 采样 k 值
`top_p`	0.0	top-p 采样阈值
`temperature`	1.0	采样温度
`repetition_penalty`	1.0	重复惩罚
`presence_penalty`	0.0	presence 惩罚；与 repetition_penalty 二者只能选一
`len_penalty`	0.0	长度惩罚
`beam_search_diversity_rate`	0.0	beam search diversity rate
`request_batch_size`	8	请求 batch 大小
`request_output_len`	32	请求输出长度

从 GPT-J 指南（docs/gptj_guide.md）可知该系列模型解码的通用行为：当beam_width > 1时执行 beam search；beam_width == 1时按top_k/top_p走采样；若 beam width 为 1 且 top_k 为 0、top_p 为 0.0，则自动退化为 greedy search。repetition_penalty与presence_penalty为互斥关系，同时设置时只有其一被允许。

附加能力与扩展说明

Prompt Learning（前缀提示词）

原指南中以注释形式保留了prefix_prompt的实验性配置说明（适用于定制检查点，通过examples/pytorch/gptneox/utils/huggingface_jp_gptneox_convert.py转换权重）。其思路是在[gptneox_20B]段中设置num_tasks与prompt_learning_type，并为每个任务声明[gptneox_20B_task_N]段（task_name、prompt_length），prompt_learning_type取值：

0：no prompt；
1：soft_prompt；
2：prefix_prompt；
3：p/prompt_tuning。

从 GptNeoX.h 源码可以看到，模型结构体中确实内置了PromptLearningType prompt_learning_type_、prompt_learning_start_id_、has_prefix_prompt_、has_prefix_soft_prompt_等成员，且 GptNeoX.h 引入了utils/prompt_learning.h，说明 prompt learning 机制在该模型实现中是可扩展的。需要提醒的是，原指南将该部分标注为仅适用于定制检查点的未正式发布流程，实际使用前请自行核对 tokenizer 的start_id/end_id一致性。

停止 token 与坏词列表

examples/cpp/gptneox/ 目录下随示例提供了 stop_words.csv、bad_words.csv 与 start_ids.csv，分别用于指定停止 token、禁止生成的坏词与起始 token 序列。坏词列表的构造可参考 GPT 系列通用的word_list.py思路，注意 tokenizer 对word与<Space>word通常会映射为两个不同 token，构造列表时需要区分。

推理结果解码

无论 C++ 还是 PyTorch 路径，最终输出的 token id 均需通过 tokenizer 还原为自然语言：

python ../examples/pytorch/gptneox/utils/hftokenizer.py out --tokenizer 20B_tokenizer.json

小结与检查清单

在 FasterTransformer 上跑通 GPT-NeoX-20B 的完整路径可归纳为四步：

获取权重：从 EleutherAI（slim_weights）或 HuggingFace（git lfs clone）下载检查点；
转换格式：分别使用 eleutherai_gpt_neox_convert.py 或 huggingface_gptneox_convert.py 转出含config.ini的 FT 权重目录，并按目标tensor_para_size切分；
生成 GEMM 配置：用./bin/gpt_gemm按模型规格（64 头 / 96 每头 / 24576 inter / 50432 vocab）生成gemm_config.in；
运行推理：C++ 端通过 gptneox_config.ini +mpirun运行gptneox_example；PyTorch 端通过 gptneox_example.py 运行，并按需组合 TP / PP。

常见踩坑点：mpirun -n进程数必须等于tensor_para_size * pipeline_para_size；多卡时权重切分数必须与tensor_para_size一致；start_id/end_id必须与所选 tokenizer 一致；repetition_penalty与presence_penalty不可同时启用。上述任何一步的参数校验，都可以回溯到 GptNeoX.h 与两个转换脚本的源码中得到印证。