尧图网站建设 
Megatron Core 导出到 TensorRT-LLM 实战指南从权重转换到引擎构建【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM本指南基于当前开源仓库Megatron-LM中的examples/export/trtllm_export/README.md及其配套源码完整讲解如何将 Megatron Core 训练出的模型以 GPT 为例转换为 TensorRT-LLMTRTLLM格式并构建推理引擎。读完本文你将掌握单设备CPU/GPU导出与分布式 GPU 导出的完整调用流程、TRTLLMHelper两大核心 API 的用法以及底层权重重命名、张量并行切分、MoE 专家切分等转换原理可直接在自有环境下跑通「Megatron Core 模型 → TRTLLM Engine」的整条链路。1. 导出功能全景Megatron Core 的 TRTLLM 导出模块Megatron Core 在megatron/core/export/目录下提供了面向 TensorRT-LLM 的模型导出能力其目录结构与职责如下megatron/core/export/ ├── model_type.py # ModelType 枚举支持的模型类型 ├── data_type.py # DataType 枚举导出精度 ├── export_config.py # ExportConfig导出时的推理 TP/PP 配置 └── trtllm/ ├── trtllm_helper.py # TRTLLMHelper对外两大 API 的入口 ├── trtllm_layers.py # TRTLLMLayersmcore 层名与 TRTLLM 层名的映射表 ├── trt_model_config.py # TRT_MODEL_CONFIGTRTLLM PretrainedConfig 构造 ├── trt_model_type.py # TRT_MODEL_TYPE_STRING架构字符串映射 ├── model_to_trllm_mapping/ │ └── default_conversion_dict.py # 默认层名转换字典 ├── trtllm_weights_converter/ # 权重转换器单设备 / 分布式 │ ├── single_device_trtllm_model_weights_converter.py │ ├── distributed_trtllm_model_weights_converter.py │ └── utils.py └── engine_builder/ └── trtllm_engine_builder.py # TRTLLMEngineBuilder引擎构建从 trtllm_helper.py 的源码注释可以看到TRTLLMHelper对外暴露两个公共 API它们是整个导出流程的骨架get_trtllm_pretrained_config_and_model_weights(...)把 Megatron Core 的model_state_dict转换成 TRTLLM 格式的模型权重与 PretrainedConfigbuild_and_save_engine(...)基于转换后的权重与配置调用 TensorRT-LLM 构建引擎并保存到指定目录。完整的调用链为Megatron Core 模型 →TRTLLMHelper.get_trtllm_pretrained_config_and_model_weights→TRTLLMHelper.build_and_save_engine→ TRTLLM Engine落盘。TRTLLMHelper在初始化时会检查环境中是否安装了tensorrt_llm未安装会直接抛出ImportError并提示pip install tensorrt-llm见 trtllm_helper.py。2. Quick Start单设备 CPU 导出 GPT 模型官方示例脚本位于 examples/export/trtllm_export/single_device_export/gpt_single_device_cpu_export.py。该脚本在**单设备单进程**上构建一个 2 层小 GPT 模型将其权重转换到 TRTLLM 格式并基于inference_tp_size2切分构建引擎。下面按文档给出的 5 个步骤逐一拆解。2.1 STEP 1初始化模型并行与默认参数导出要求先初始化分布式环境。文档强调初始化 TP1、PP1以便在 CPU 上拿到完整的模型 state dictinitialize_distributed(tensor_model_parallel_size1, pipeline_model_parallel_size1)脚本中的initialize_distributed实际做了三件事见 gpt_single_device_cpu_export.py调用parallel_state.destroy_model_parallel()清理旧的并行状态用torch.distributed.init_process_group初始化 PyTorch 分布式进程组world_size取torch.cuda.device_count()rank取环境变量LOCAL_RANK并用torch.cuda.set_device绑定当前卡调用parallel_state.initialize_model_parallel(tensor_model_parallel_size, pipeline_model_parallel_size)初始化 Megatron Core 的模型并行状态。2.2 STEP 2用 model_provider 加载模型示例通过TransformerConfigGPTModel构建一个极小的 GPT 模型源码见 gpt_single_device_cpu_export.pytransformer_config TransformerConfig( num_layers2, hidden_size64, # Needs to be atleast 32 times num_attn_heads num_attention_heads2, use_cpu_initializationTrue, pipeline_dtypetorch.float32, ) gpt_model GPTModel( configtransformer_config, transformer_layer_specget_gpt_layer_local_spec(), vocab_size100, max_sequence_length_SEQUENCE_LENGTH, )参数要点hidden_size64且代码注释明确hidden_size 至少要是num_attention_heads的 32 倍本示例 64 2 × 32use_cpu_initializationTrue表示权重在 CPU 上初始化便于单设备收集完整 state dictpipeline_dtypetorch.float32指定管线计算精度_SEQUENCE_LENGTH在脚本顶部定义为64见 gpt_single_device_cpu_export.py。文档还给出了从已有分布式检查点加载真实模型的可选路径脚本中封装为load_distributed_checkpoint见 gpt_single_device_cpu_export.py# Optionally you can also load a model using this code # sharded_state_dictgpt_model.sharded_state_dict(prefix) # checkpoint dist_checkpointing.load(sharded_state_dictsharded_state_dict, checkpoint_dircheckpoint_path) # gpt_model.load_state_dict(checkpoint)即先用gpt_model.sharded_state_dict()获取分片 state dict 的 schema再用dist_checkpointing.load从checkpoint_dir加载最后load_state_dict回填到模型。脚本主流程中这行代码被注释可通过传入ckpt_path启用。2.3 STEP 3实例化 TRTLLM Helper接着从模型属性中提取导出所需信息实例化TRTLLMHelper完整源码见 gpt_single_device_cpu_export.pyseq_len_interpolation_factor None if hasattr(gpt_model, rotary_pos_emb): seq_len_interpolation_factor gpt_model.rotary_pos_emb.seq_len_interpolation_factor trtllm_helper TRTLLMHelper( transformer_configgpt_model.config, model_typeModelType.gpt, position_embedding_type gpt_model.position_embedding_type, max_position_embeddings gpt_model.max_position_embeddings, rotary_percentage gpt_model.rotary_percent, rotary_base gpt_model.rotary_base, moe_tp_mode 2, multi_query_mode False, activation gelu, seq_len_interpolation_factor seq_len_interpolation_factor, share_embeddings_and_output_weightsgpt_model.share_embeddings_and_output_weights )TRTLLMHelper.__init__的完整参数签名及默认值见 trtllm_helper.py参数默认值说明transformer_config必填Megatron Core 的TransformerConfigmodel_type必填模型类型枚举见下文ModelTypetrtllm_conversion_dict{}自定义层名转换字典会与默认字典合并见 3.2 节position_embedding_typelearned_absolute仅支持learned_absolute或rope传入其他值会触发断言失败max_position_embeddingsNone最大位置编码长度rotary_percentage1.0RoPE 旋转百分比rotary_base10000RoPE 的 theta 基数rope_scaling_factor8.0仅 Nemotron-NAS 等模型使用moe_tp_mode2TRTLLM 的 MoE 张量并行模式multi_query_modeFalse多查询注意力模式activationgelu激活函数名seq_len_interpolation_factorNoneRoPE 序列长度插值因子moe_renorm_modeNoneMoE 专家权重归一化模式share_embeddings_and_output_weightsFalse输入输出层是否共享权重tied embedding其中seq_len_interpolation_factor从gpt_model.rotary_pos_emb上动态读取这正是 RoPE 位置编码的「序列长度插值」Sequence Length Interpolation因子导出时会写入 TRTLLM 配置的rotary_scaling字段见 trtllm_helper.py。2.4 STEP 4转换 TRTLLM 权重与配置这一步调用核心 APIget_trtllm_pretrained_config_and_model_weights将模型 state dict 转换成 TRTLLM 权重与配置。文档代码与脚本 gpt_single_device_cpu_export.py 一致model_state_dict{} for key , val in gpt_model.state_dict().items(): # val is non for _extra_state layers . We filter it out if val is not None: model_state_dict[key] val export_config ExportConfig(inference_tp_size 2) weight_list, config_list trtllm_helper.get_trtllm_pretrained_config_and_model_weights( model_state_dict model_state_dict, dtype DataType.bfloat16, export_configexport_config )要点手动过滤掉值为None的_extra_state层实际上TRTLLMHelper内部也会过滤所有含extra_state的键见 trtllm_helper.pyExportConfig(inference_tp_size2)声明推理时的张量并行规模为 2即最终的 TRTLLM 引擎会按 TP2 切分权重dtypeDataType.bfloat16指定导出精度文档特别提示如果整个模型能塞进 GPU 显存先把 state dict 搬到 GPU 再调用该函数性能更快。ExportConfig的全部字段见 export_config.py字段默认值说明inference_tp_size1推理张量并行大小inference_pp_size1推理流水线并行大小use_parallel_embeddingFalse是否使用并行 embeddinguse_embedding_sharingNone已废弃会抛出DeprecationWarning改用TRTLLMHelper的share_embeddings_and_output_weights单设备转换模式下返回值是两个列表weight_list和config_list列表长度等于world_size inference_tp_size × inference_pp_size即每个 GPU rank 一份权重与一份配置见 trtllm_helper.py。2.5 STEP 5构建 TRTLLM Engine最后遍历(weight, config)对调用build_and_save_engine为每个 rank 构建引擎for trtllm_model_weights, trtllm_model_config in zip(weight_list, config_list): trtllm_helper.build_and_save_engine( max_input_len256, max_output_len256, max_batch_size8, engine_dir/opt/megatron-lm/engine, trtllm_model_weightstrtllm_model_weights, trtllm_model_configtrtllm_model_config, lora_ckpt_listNone, use_lora_pluginNone, max_lora_rank64, lora_target_modulesNone, max_prompt_embedding_table_size0, paged_kv_cacheTrue, remove_input_paddingTrue, paged_context_fmhaFalse, use_refitFalse, max_num_tokensNone, max_seq_len512, opt_num_tokensNone, max_beam_width1, tokens_per_block128, multiple_profilesFalse, gpt_attention_pluginauto, gemm_pluginauto, )build_and_save_engine的完整参数与源码默认值见 trtllm_helper.py参数默认值示例值说明engine_dir必填/opt/megatron-lm/engine引擎保存目录trtllm_model_weights必填—转换后的 TRTLLM 权重字典trtllm_model_config必填—转换后的 TRTLLM 配置max_input_len1024256最大输入长度max_output_len1024256最大输出长度max_batch_size48最大批大小lora_ckpt_listNoneNoneLoRA 检查点列表use_lora_pluginNoneNone是否启用 LoRA 插件max_lora_rank6464最大 LoRA ranklora_target_modulesNoneNoneLoRA 目标模块max_prompt_embedding_table_size00prompt embedding 表最大尺寸paged_kv_cacheTrueTrue启用 Paged KV Cacheremove_input_paddingTrueTrue移除输入 paddingpaged_context_fmhaFalseFalse启用 paged context FMHAuse_refitFalseFalse启用 Refitmax_num_tokensNoneNone最大 token 数由 TRTLLM 推导max_seq_lenNone512最大序列长度为None时取max_input_len max_output_lenopt_num_tokensNoneNone最优 token 数max_beam_width11最大 beam 宽度tokens_per_block128128Paged KV Cache 每块 token 数multiple_profilesFalseFalse是否生成多个 profilegpt_attention_pluginautoautoGPT Attention 插件gemm_pluginauto |autoGEMM 插件在引擎构建的底层实现中见 trtllm_engine_builder.py根据paged_kv_cache调用plugin_config.enable_paged_kv_cache(tokens_per_block...)当max_seq_len为None时自动计算max_input_len max_output_len通过tensorrt_llm._common.check_max_num_tokens校验max_num_tokens/opt_num_tokens与各长度参数的一致性用model_cls.from_config(trtllm_model_config)实例化 TRTLLM 模型optimize_model处理并行 embedding 与共享 embedding 表preprocess_weights预处理权重后model.load(...)载入最后build_trtllm(model, build_config)构建并engine.save(engine_dir)。3. 运行导出脚本3.1 Docker 容器内运行官方示例假设在一个预装了 TRTLLM 的容器内运行文档原文命令# In a workstation MLM_PATH/path/to/megatron-lm CONTAINER_IMAGEgitlab-master.nvidia.com:5005/dl/joc/nemo-ci/trtllm_0.12/train:pipe.17669124-x86 docker run -it --gpusall --ipchost -v $MLM_PATH/:/opt/megatron-lm $CONTAINER_IMAGE bash # Inside the container run the following. cd /opt/megatron-lm/ CUDA_VISIBLE_DEVICES0 torchrun --nproc-per-node 1 examples/export/trtllm_export/single_device_export/gpt_single_device_cpu_export.py说明示例中的CONTAINER_IMAGE是 NVIDIA 内部 CI 镜像版本 TRTLLM 0.12实际使用时应替换为你自己的、已安装tensorrt-llm的镜像通过-v $MLM_PATH/:/opt/megatron-lm将仓库挂载到容器内/opt/megatron-lm因此示例脚本中engine_dir/opt/megatron-lm/engine指向的就是宿主机仓库下的engine目录单设备导出使用--nproc-per-node 1即单进程。3.2 关于运行前提运行导出脚本需要满足环境中已安装tensorrt-llmTRTLLMHelper与TRTLLMEngineBuilder都会在导入阶段检测缺失时抛ImportError已安装torch、megatron-core依赖含tqdm单设备转换器明确要求pip install tqdm见 single_device_trtllm_model_weights_converter.py使用torchrun启动以满足initialize_distributed中对LOCAL_RANK环境变量的读取。4. GPU 分布式导出on-device 转换文档指出单设备版本是把全部权重收集到 CPU/GPU 上再转换而 examples/export/trtllm_export/distributed_export/gpt_distributed_gpu_export.py 提供更优化的 on-device 分布式导出每个 GPU 只持有自己的那部分 state dict在设备本地完成转换最后把引擎写盘避免整模型权重集中到单一设备上的显存/内存压力。运行命令文档原文CUDA_VISIBLE_DEVICES0,1 torchrun --nproc-per-node 2 examples/export/trtllm_export/distributed_export/gpt_distributed_gpu_export.py与单设备版本的关键差异对照 gpt_distributed_gpu_export.py 与 trtllm_helper.py并行初始化不同分布式脚本用initialize_distributed(tensor_model_parallel_size2, pipeline_model_parallel_size1)即训练侧的模型就按 TP2 切分加载每个 rank 只持有自己的分片模型搬到 GPUgpt_model.to(device)将模型移动到 CUDA转换 API 参数不同调用时传on_device_distributed_conversionTrue并额外传入vocab_size与gpus_per_node不再传ExportConfigtrtllm_model_weights, trtllm_model_config trtllm_helper.get_trtllm_pretrained_config_and_model_weights( model_state_dict gpt_model.state_dict(), dtype DataType.bfloat16, on_device_distributed_conversionTrue, vocab_size_VOCAB_SIZE, gpus_per_node2, )返回值不同on-device 模式返回的是单元素列表[trtllm_model_weights_on_device]与[trtllm_model_config]构建引擎时直接取[0]见 gpt_distributed_gpu_export.py。get_trtllm_pretrained_config_and_model_weights在 on-device 模式下有几条强约束源码断言见 trtllm_helper.pyvocab_size必须显式传入不能为None因为单设备模式是从输入层推导词表大小的仅支持ModelType.gpt、gptnext、llama、nemotron_nas四种模型类型export_config必须为None——推理 TP 规模直接由模型加载时的并行状态推断例如要以 TP2 导出就用 TP2 加载模型而不是通过ExportConfig指定gpus_per_node必须显式传入用于 TRTLLM 的Mapping与配置生成。on-device 模式内部使用DistributedTRTLLMModelWeightsConverter见 distributed_trtllm_model_weights_converter.py并从转换器上反推inference_pp_size、inference_tp_size强制use_parallel_embeddingTrue然后为每个 rank 构造带tensorrt_llm.Mapping含world_size、rank、tp_size、pp_size的模型配置见 trtllm_helper.py。5. 源码级原理权重如何被转换与切分5.1 支持的模型类型与架构字符串ModelType枚举见 model_type.py支持gpt、gptnext、llama、falcon、starcoder、mixtral、gemma、nemotron_nas。它们对应 TRTLLM 的架构字符串见 trt_model_type.pyModelTypeTRTLLM 架构gpt/gptnext/starcoderGPTForCausalLMllama/mixtralLlamaForCausalLMgemmaGemmaForCausalLMfalconFalconForCausalLMnemotron_nasDeciLMForCausalLMDataType枚举见 data_type.py支持三种导出精度bfloat16、float16、float32。5.2 层名重命名从 mcore 命名到 TRTLLM 命名转换的第一步是把 mcore 的层名如decoder.layers.2.self_attention.linear_qkv.weight重命名为 TRTLLM 的层名如transformer.layers.2.attention.qkv.weight。其机制是trtllm_layers.py 中的TRTLLMLayers枚举定义了 TRTLLM 侧的标准层名如attention_qkv_weight transformer.layers.attention.qkv.weight、mlp_fc_weight transformer.layers.mlp.fc.weight等并覆盖 embedding、final layernorm、attention、MLP、MoE router 与 expert 等全部层default_conversion_dict.py 中的DEFAULT_CONVERSION_DICT定义了 mcore → TRTLLM 的默认映射例如decoder.layers.self_attention.linear_qkv.weight → attention_qkv_weight、decoder.layers.mlp.linear_fc1.weight → mlp_fc_weight、decoder.layers.mlp.experts.experts.linear_fc1.weight → mlp_fc_weight_mixture_of_experts等同时覆盖了 Transformer Engine 风格的layer_norm_weight命名TRTLLMLayers.rename_input_layer_names_to_trtllm_layer_names用正则(?layers\.)\d(?\.)先抽取层号映射完成后再把层号插回见 trtllm_layers.py遇到映射表里不存在的键会抛ValueError提示在初始化TRTLLMHelper时通过trtllm_conversion_dict补充映射_extra_state与adapter_layer相关键会被直接丢弃。TRTLLMHelper默认使用DEFAULT_CONVERSION_DICT当model_type ModelType.nemotron_nas时额外叠加NEMOTRON_NAS_CONVERSION_DICTDeci 的线性注意力与线性 FFN 命名用户传入的trtllm_conversion_dict最后合并覆盖见 trtllm_helper.py。注意映射字典的键必须去掉层号。5.3 权重转换与 TP/专家切分SingleDeviceTRTLLMModelWeightsConverter.convert完成实际转换见 single_device_trtllm_model_weights_converter.py核心逻辑包括非 transformer 层embedding、lm_head、final layernorm直接转成目标 dtype 并存入trtllm_model_weightsembedding 层在use_parallel_embedding时按inference_tp_size补齐pad_vocab_size权重转置二维权重统一val val.T对齐 TRTLLM 的权重布局dense / proj 类权重沿 dim 0 用torch.chunk按inference_tp_size切分保存为layer_name.{rank}.binMLP fc / gate 类权重若激活函数为 gated 形式如 SwiGLU 等通过is_gated_activation判断先把fc拆成fc与gate两半再各自沿最后一维切分QKV 权重按[hidden_dim, num_kv_heads, q_num 2, size_per_head]重构拆出 Q/K/V当num_kv_heads inference_tp_size时对 K/V 做_duplicate_kv_head复制以适配 GQA/MQA 的 TP 切分并校验 TP 大小必须能整除或复制 query groups否则抛异常见 single_device_trtllm_model_weights_converter.pyMoE expert 权重fc拆成w1/w3两半后各自 TP 切分再拼接proj沿最后一维切分保存为layer_name.{expert_idx}.bin的 expert_split 格式layernorm_zero_centered_gamma兼容当配置开启layernorm_zero_centered_gamma且归一化为LayerNorm时对layernorm.weight做1.0还原与 NeMo 的 layernorm1p 语义一致。转换完成后get_local_model_weights_per_gpu(mapping, config)根据tensorrt_llm.Mapping的tp_rank、pp_rank与pp_layers(num_layers)把已切分的.bin权重分发给每个 GPU rank只保留tp_rank对应的分片、只保留落在本 PP 段内的层并把层号重排为从 0 开始embedding 与 lm_head 只在首/末 PP rank 上分发见 single_device_trtllm_model_weights_converter.py。5.4 TRTLLM 配置的生成_get_trtllm_config从TransformerConfig映射出 TRTLLM 的PretrainedConfig见 trtllm_helper.py关键映射包括num_attention_heads、num_key_value_heads取num_query_groups为空时退化为num_attention_heads、head_size取kv_channels、hidden_size、intermediate_size取ffn_hidden_size、norm_epsilon取layernorm_epsilonposition_embedding_typerope映射为rope_gpt_neox否则learned_absoluterotary_pct/rotary_base来自 helper 参数存在seq_len_interpolation_factor时写入rotary_scaling {type: linear, factor: ...}MoE 相关moe_num_expertsmoe_router_topk 0时置 0、moe_top_k、moe_normalization_mode默认RENORMALIZE、moe_tp_modeworld_size、tp_size、pp_size、gpus_per_node写入配置随后为每个 rank 挂上tensorrt_llm.MappingFP8 导出时配置quant_algo与kv_cache_quant_algo详见下文falcon模型额外设置new_decoder_architecture与parallel_attentionnemotron_nas从heterogeneous_layers_config_encoded_json解析block_configs并设置 llama3 风格rotary_scaling。5.5 FP8 量化导出get_trtllm_pretrained_config_and_model_weights还支持 FP8 量化导出传入fp8_quantizedTrue/fp8_kvcacheTrue时_load_scaling_factors会从 state dict 中的_extra_state读取 Megatron 的scale_inv_fwd/scale_fwd缩放因子重命名为 TRTLLM 的activation_scaling_factor/weights_scaling_factorgated 激活还会为gate层复制一份再通过_add_scales_to_converter合并进权重字典见 trtllm_helper.py。转换器在遇到带缩放因子的权重层时会先乘上weight_multiplier再转成torch.float8_e4m3fn见 single_device_trtllm_model_weights_converter.py。6. 当前限制与未来工作原文档明确列出了当前实现的能力边界与规划不做任何夸大流水线并行导出导出流程的流水线并行支持仍在开发中Work in progress更多模型的 GPU 导出on-device 分布式转换目前仅支持gpt、gptnext、llama、nemotron_nas四类见 4 节断言其余模型类型正在推进Refit 功能build_and_save_engine已预留use_refit参数但完整 Refit 工作流仍在规划中VLLM 支持计划中尚未实现。此外从源码可以确认两点隐含前提单设备转换的权重切分规模由ExportConfig.inference_tp_size / inference_pp_size决定而 on-device 转换则由加载模型时的并行规模决定两者都以「导出引擎的推理并行布局」为准与训练时的并行布局解耦。7. 总结与进一步探索本文以官方 README 为骨架完整复现了 Megatron Core → TRTLLM 的单设备与分布式导出两条路径并深入到TRTLLMHelper、TRTLLMLayers、默认转换字典、单设备/分布式权重转换器与引擎构建器展示了层名重命名、TP 切分、GQA 复制、MoE 专家切分与 FP8 缩放因子注入等底层细节。如需继续深入推荐阅读仓库中的以下文件单设备导出示例examples/export/trtllm_export/single_device_export/gpt_single_device_cpu_export.py分布式导出示例examples/export/trtllm_export/distributed_export/gpt_distributed_gpu_export.py导出入口 Helpermegatron/core/export/trtllm/trtllm_helper.py层名与映射表megatron/core/export/trtllm/trtllm_layers.py、megatron/core/export/trtllm/model_to_trllm_mapping/default_conversion_dict.py权重转换器single_device_trtllm_model_weights_converter.py、distributed_trtllm_model_weights_converter.py引擎构建器megatron/core/export/trtllm/engine_builder/trtllm_engine_builder.py导出配置与枚举megatron/core/export/export_config.py、megatron/core/export/model_type.py、megatron/core/export/data_type.py【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 Megatron Core 导出到 TensorRT-LLM 实战指南：从权重转换到引擎构建
Megatron Core 导出到 TensorRT-LLM 实战指南从权重转换到引擎构建【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM本指南基于当前开源仓库Megatron-LM中的examples/export/trtllm_export/README.md及其配套源码完整讲解如何将 Megatron Core 训练出的模型以 GPT 为例转换为 TensorRT-LLMTRTLLM格式并构建推理引擎。读完本文你将掌握单设备CPU/GPU导出与分布式 GPU 导出的完整调用流程、TRTLLMHelper两大核心 API 的用法以及底层权重重命名、张量并行切分、MoE 专家切分等转换原理可直接在自有环境下跑通「Megatron Core 模型 → TRTLLM Engine」的整条链路。1. 导出功能全景Megatron Core 的 TRTLLM 导出模块Megatron Core 在megatron/core/export/目录下提供了面向 TensorRT-LLM 的模型导出能力其目录结构与职责如下megatron/core/export/ ├── model_type.py # ModelType 枚举支持的模型类型 ├── data_type.py # DataType 枚举导出精度 ├── export_config.py # ExportConfig导出时的推理 TP/PP 配置 └── trtllm/ ├── trtllm_helper.py # TRTLLMHelper对外两大 API 的入口 ├── trtllm_layers.py # TRTLLMLayersmcore 层名与 TRTLLM 层名的映射表 ├── trt_model_config.py # TRT_MODEL_CONFIGTRTLLM PretrainedConfig 构造 ├── trt_model_type.py # TRT_MODEL_TYPE_STRING架构字符串映射 ├── model_to_trllm_mapping/ │ └── default_conversion_dict.py # 默认层名转换字典 ├── trtllm_weights_converter/ # 权重转换器单设备 / 分布式 │ ├── single_device_trtllm_model_weights_converter.py │ ├── distributed_trtllm_model_weights_converter.py │ └── utils.py └── engine_builder/ └── trtllm_engine_builder.py # TRTLLMEngineBuilder引擎构建从 trtllm_helper.py 的源码注释可以看到TRTLLMHelper对外暴露两个公共 API它们是整个导出流程的骨架get_trtllm_pretrained_config_and_model_weights(...)把 Megatron Core 的model_state_dict转换成 TRTLLM 格式的模型权重与 PretrainedConfigbuild_and_save_engine(...)基于转换后的权重与配置调用 TensorRT-LLM 构建引擎并保存到指定目录。完整的调用链为Megatron Core 模型 →TRTLLMHelper.get_trtllm_pretrained_config_and_model_weights→TRTLLMHelper.build_and_save_engine→ TRTLLM Engine落盘。TRTLLMHelper在初始化时会检查环境中是否安装了tensorrt_llm未安装会直接抛出ImportError并提示pip install tensorrt-llm见 trtllm_helper.py。2. Quick Start单设备 CPU 导出 GPT 模型官方示例脚本位于 examples/export/trtllm_export/single_device_export/gpt_single_device_cpu_export.py。该脚本在**单设备单进程**上构建一个 2 层小 GPT 模型将其权重转换到 TRTLLM 格式并基于inference_tp_size2切分构建引擎。下面按文档给出的 5 个步骤逐一拆解。2.1 STEP 1初始化模型并行与默认参数导出要求先初始化分布式环境。文档强调初始化 TP1、PP1以便在 CPU 上拿到完整的模型 state dictinitialize_distributed(tensor_model_parallel_size1, pipeline_model_parallel_size1)脚本中的initialize_distributed实际做了三件事见 gpt_single_device_cpu_export.py调用parallel_state.destroy_model_parallel()清理旧的并行状态用torch.distributed.init_process_group初始化 PyTorch 分布式进程组world_size取torch.cuda.device_count()rank取环境变量LOCAL_RANK并用torch.cuda.set_device绑定当前卡调用parallel_state.initialize_model_parallel(tensor_model_parallel_size, pipeline_model_parallel_size)初始化 Megatron Core 的模型并行状态。2.2 STEP 2用 model_provider 加载模型示例通过TransformerConfigGPTModel构建一个极小的 GPT 模型源码见 gpt_single_device_cpu_export.pytransformer_config TransformerConfig( num_layers2, hidden_size64, # Needs to be atleast 32 times num_attn_heads num_attention_heads2, use_cpu_initializationTrue, pipeline_dtypetorch.float32, ) gpt_model GPTModel( configtransformer_config, transformer_layer_specget_gpt_layer_local_spec(), vocab_size100, max_sequence_length_SEQUENCE_LENGTH, )参数要点hidden_size64且代码注释明确hidden_size 至少要是num_attention_heads的 32 倍本示例 64 2 × 32use_cpu_initializationTrue表示权重在 CPU 上初始化便于单设备收集完整 state dictpipeline_dtypetorch.float32指定管线计算精度_SEQUENCE_LENGTH在脚本顶部定义为64见 gpt_single_device_cpu_export.py。文档还给出了从已有分布式检查点加载真实模型的可选路径脚本中封装为load_distributed_checkpoint见 gpt_single_device_cpu_export.py# Optionally you can also load a model using this code # sharded_state_dictgpt_model.sharded_state_dict(prefix) # checkpoint dist_checkpointing.load(sharded_state_dictsharded_state_dict, checkpoint_dircheckpoint_path) # gpt_model.load_state_dict(checkpoint)即先用gpt_model.sharded_state_dict()获取分片 state dict 的 schema再用dist_checkpointing.load从checkpoint_dir加载最后load_state_dict回填到模型。脚本主流程中这行代码被注释可通过传入ckpt_path启用。2.3 STEP 3实例化 TRTLLM Helper接着从模型属性中提取导出所需信息实例化TRTLLMHelper完整源码见 gpt_single_device_cpu_export.pyseq_len_interpolation_factor None if hasattr(gpt_model, rotary_pos_emb): seq_len_interpolation_factor gpt_model.rotary_pos_emb.seq_len_interpolation_factor trtllm_helper TRTLLMHelper( transformer_configgpt_model.config, model_typeModelType.gpt, position_embedding_type gpt_model.position_embedding_type, max_position_embeddings gpt_model.max_position_embeddings, rotary_percentage gpt_model.rotary_percent, rotary_base gpt_model.rotary_base, moe_tp_mode 2, multi_query_mode False, activation gelu, seq_len_interpolation_factor seq_len_interpolation_factor, share_embeddings_and_output_weightsgpt_model.share_embeddings_and_output_weights )TRTLLMHelper.__init__的完整参数签名及默认值见 trtllm_helper.py参数默认值说明transformer_config必填Megatron Core 的TransformerConfigmodel_type必填模型类型枚举见下文ModelTypetrtllm_conversion_dict{}自定义层名转换字典会与默认字典合并见 3.2 节position_embedding_typelearned_absolute仅支持learned_absolute或rope传入其他值会触发断言失败max_position_embeddingsNone最大位置编码长度rotary_percentage1.0RoPE 旋转百分比rotary_base10000RoPE 的 theta 基数rope_scaling_factor8.0仅 Nemotron-NAS 等模型使用moe_tp_mode2TRTLLM 的 MoE 张量并行模式multi_query_modeFalse多查询注意力模式activationgelu激活函数名seq_len_interpolation_factorNoneRoPE 序列长度插值因子moe_renorm_modeNoneMoE 专家权重归一化模式share_embeddings_and_output_weightsFalse输入输出层是否共享权重tied embedding其中seq_len_interpolation_factor从gpt_model.rotary_pos_emb上动态读取这正是 RoPE 位置编码的「序列长度插值」Sequence Length Interpolation因子导出时会写入 TRTLLM 配置的rotary_scaling字段见 trtllm_helper.py。2.4 STEP 4转换 TRTLLM 权重与配置这一步调用核心 APIget_trtllm_pretrained_config_and_model_weights将模型 state dict 转换成 TRTLLM 权重与配置。文档代码与脚本 gpt_single_device_cpu_export.py 一致model_state_dict{} for key , val in gpt_model.state_dict().items(): # val is non for _extra_state layers . We filter it out if val is not None: model_state_dict[key] val export_config ExportConfig(inference_tp_size 2) weight_list, config_list trtllm_helper.get_trtllm_pretrained_config_and_model_weights( model_state_dict model_state_dict, dtype DataType.bfloat16, export_configexport_config )要点手动过滤掉值为None的_extra_state层实际上TRTLLMHelper内部也会过滤所有含extra_state的键见 trtllm_helper.pyExportConfig(inference_tp_size2)声明推理时的张量并行规模为 2即最终的 TRTLLM 引擎会按 TP2 切分权重dtypeDataType.bfloat16指定导出精度文档特别提示如果整个模型能塞进 GPU 显存先把 state dict 搬到 GPU 再调用该函数性能更快。ExportConfig的全部字段见 export_config.py字段默认值说明inference_tp_size1推理张量并行大小inference_pp_size1推理流水线并行大小use_parallel_embeddingFalse是否使用并行 embeddinguse_embedding_sharingNone已废弃会抛出DeprecationWarning改用TRTLLMHelper的share_embeddings_and_output_weights单设备转换模式下返回值是两个列表weight_list和config_list列表长度等于world_size inference_tp_size × inference_pp_size即每个 GPU rank 一份权重与一份配置见 trtllm_helper.py。2.5 STEP 5构建 TRTLLM Engine最后遍历(weight, config)对调用build_and_save_engine为每个 rank 构建引擎for trtllm_model_weights, trtllm_model_config in zip(weight_list, config_list): trtllm_helper.build_and_save_engine( max_input_len256, max_output_len256, max_batch_size8, engine_dir/opt/megatron-lm/engine, trtllm_model_weightstrtllm_model_weights, trtllm_model_configtrtllm_model_config, lora_ckpt_listNone, use_lora_pluginNone, max_lora_rank64, lora_target_modulesNone, max_prompt_embedding_table_size0, paged_kv_cacheTrue, remove_input_paddingTrue, paged_context_fmhaFalse, use_refitFalse, max_num_tokensNone, max_seq_len512, opt_num_tokensNone, max_beam_width1, tokens_per_block128, multiple_profilesFalse, gpt_attention_pluginauto, gemm_pluginauto, )build_and_save_engine的完整参数与源码默认值见 trtllm_helper.py参数默认值示例值说明engine_dir必填/opt/megatron-lm/engine引擎保存目录trtllm_model_weights必填—转换后的 TRTLLM 权重字典trtllm_model_config必填—转换后的 TRTLLM 配置max_input_len1024256最大输入长度max_output_len1024256最大输出长度max_batch_size48最大批大小lora_ckpt_listNoneNoneLoRA 检查点列表use_lora_pluginNoneNone是否启用 LoRA 插件max_lora_rank6464最大 LoRA ranklora_target_modulesNoneNoneLoRA 目标模块max_prompt_embedding_table_size00prompt embedding 表最大尺寸paged_kv_cacheTrueTrue启用 Paged KV Cacheremove_input_paddingTrueTrue移除输入 paddingpaged_context_fmhaFalseFalse启用 paged context FMHAuse_refitFalseFalse启用 Refitmax_num_tokensNoneNone最大 token 数由 TRTLLM 推导max_seq_lenNone512最大序列长度为None时取max_input_len max_output_lenopt_num_tokensNoneNone最优 token 数max_beam_width11最大 beam 宽度tokens_per_block128128Paged KV Cache 每块 token 数multiple_profilesFalseFalse是否生成多个 profilegpt_attention_pluginautoautoGPT Attention 插件gemm_pluginauto |autoGEMM 插件在引擎构建的底层实现中见 trtllm_engine_builder.py根据paged_kv_cache调用plugin_config.enable_paged_kv_cache(tokens_per_block...)当max_seq_len为None时自动计算max_input_len max_output_len通过tensorrt_llm._common.check_max_num_tokens校验max_num_tokens/opt_num_tokens与各长度参数的一致性用model_cls.from_config(trtllm_model_config)实例化 TRTLLM 模型optimize_model处理并行 embedding 与共享 embedding 表preprocess_weights预处理权重后model.load(...)载入最后build_trtllm(model, build_config)构建并engine.save(engine_dir)。3. 运行导出脚本3.1 Docker 容器内运行官方示例假设在一个预装了 TRTLLM 的容器内运行文档原文命令# In a workstation MLM_PATH/path/to/megatron-lm CONTAINER_IMAGEgitlab-master.nvidia.com:5005/dl/joc/nemo-ci/trtllm_0.12/train:pipe.17669124-x86 docker run -it --gpusall --ipchost -v $MLM_PATH/:/opt/megatron-lm $CONTAINER_IMAGE bash # Inside the container run the following. cd /opt/megatron-lm/ CUDA_VISIBLE_DEVICES0 torchrun --nproc-per-node 1 examples/export/trtllm_export/single_device_export/gpt_single_device_cpu_export.py说明示例中的CONTAINER_IMAGE是 NVIDIA 内部 CI 镜像版本 TRTLLM 0.12实际使用时应替换为你自己的、已安装tensorrt-llm的镜像通过-v $MLM_PATH/:/opt/megatron-lm将仓库挂载到容器内/opt/megatron-lm因此示例脚本中engine_dir/opt/megatron-lm/engine指向的就是宿主机仓库下的engine目录单设备导出使用--nproc-per-node 1即单进程。3.2 关于运行前提运行导出脚本需要满足环境中已安装tensorrt-llmTRTLLMHelper与TRTLLMEngineBuilder都会在导入阶段检测缺失时抛ImportError已安装torch、megatron-core依赖含tqdm单设备转换器明确要求pip install tqdm见 single_device_trtllm_model_weights_converter.py使用torchrun启动以满足initialize_distributed中对LOCAL_RANK环境变量的读取。4. GPU 分布式导出on-device 转换文档指出单设备版本是把全部权重收集到 CPU/GPU 上再转换而 examples/export/trtllm_export/distributed_export/gpt_distributed_gpu_export.py 提供更优化的 on-device 分布式导出每个 GPU 只持有自己的那部分 state dict在设备本地完成转换最后把引擎写盘避免整模型权重集中到单一设备上的显存/内存压力。运行命令文档原文CUDA_VISIBLE_DEVICES0,1 torchrun --nproc-per-node 2 examples/export/trtllm_export/distributed_export/gpt_distributed_gpu_export.py与单设备版本的关键差异对照 gpt_distributed_gpu_export.py 与 trtllm_helper.py并行初始化不同分布式脚本用initialize_distributed(tensor_model_parallel_size2, pipeline_model_parallel_size1)即训练侧的模型就按 TP2 切分加载每个 rank 只持有自己的分片模型搬到 GPUgpt_model.to(device)将模型移动到 CUDA转换 API 参数不同调用时传on_device_distributed_conversionTrue并额外传入vocab_size与gpus_per_node不再传ExportConfigtrtllm_model_weights, trtllm_model_config trtllm_helper.get_trtllm_pretrained_config_and_model_weights( model_state_dict gpt_model.state_dict(), dtype DataType.bfloat16, on_device_distributed_conversionTrue, vocab_size_VOCAB_SIZE, gpus_per_node2, )返回值不同on-device 模式返回的是单元素列表[trtllm_model_weights_on_device]与[trtllm_model_config]构建引擎时直接取[0]见 gpt_distributed_gpu_export.py。get_trtllm_pretrained_config_and_model_weights在 on-device 模式下有几条强约束源码断言见 trtllm_helper.pyvocab_size必须显式传入不能为None因为单设备模式是从输入层推导词表大小的仅支持ModelType.gpt、gptnext、llama、nemotron_nas四种模型类型export_config必须为None——推理 TP 规模直接由模型加载时的并行状态推断例如要以 TP2 导出就用 TP2 加载模型而不是通过ExportConfig指定gpus_per_node必须显式传入用于 TRTLLM 的Mapping与配置生成。on-device 模式内部使用DistributedTRTLLMModelWeightsConverter见 distributed_trtllm_model_weights_converter.py并从转换器上反推inference_pp_size、inference_tp_size强制use_parallel_embeddingTrue然后为每个 rank 构造带tensorrt_llm.Mapping含world_size、rank、tp_size、pp_size的模型配置见 trtllm_helper.py。5. 源码级原理权重如何被转换与切分5.1 支持的模型类型与架构字符串ModelType枚举见 model_type.py支持gpt、gptnext、llama、falcon、starcoder、mixtral、gemma、nemotron_nas。它们对应 TRTLLM 的架构字符串见 trt_model_type.pyModelTypeTRTLLM 架构gpt/gptnext/starcoderGPTForCausalLMllama/mixtralLlamaForCausalLMgemmaGemmaForCausalLMfalconFalconForCausalLMnemotron_nasDeciLMForCausalLMDataType枚举见 data_type.py支持三种导出精度bfloat16、float16、float32。5.2 层名重命名从 mcore 命名到 TRTLLM 命名转换的第一步是把 mcore 的层名如decoder.layers.2.self_attention.linear_qkv.weight重命名为 TRTLLM 的层名如transformer.layers.2.attention.qkv.weight。其机制是trtllm_layers.py 中的TRTLLMLayers枚举定义了 TRTLLM 侧的标准层名如attention_qkv_weight transformer.layers.attention.qkv.weight、mlp_fc_weight transformer.layers.mlp.fc.weight等并覆盖 embedding、final layernorm、attention、MLP、MoE router 与 expert 等全部层default_conversion_dict.py 中的DEFAULT_CONVERSION_DICT定义了 mcore → TRTLLM 的默认映射例如decoder.layers.self_attention.linear_qkv.weight → attention_qkv_weight、decoder.layers.mlp.linear_fc1.weight → mlp_fc_weight、decoder.layers.mlp.experts.experts.linear_fc1.weight → mlp_fc_weight_mixture_of_experts等同时覆盖了 Transformer Engine 风格的layer_norm_weight命名TRTLLMLayers.rename_input_layer_names_to_trtllm_layer_names用正则(?layers\.)\d(?\.)先抽取层号映射完成后再把层号插回见 trtllm_layers.py遇到映射表里不存在的键会抛ValueError提示在初始化TRTLLMHelper时通过trtllm_conversion_dict补充映射_extra_state与adapter_layer相关键会被直接丢弃。TRTLLMHelper默认使用DEFAULT_CONVERSION_DICT当model_type ModelType.nemotron_nas时额外叠加NEMOTRON_NAS_CONVERSION_DICTDeci 的线性注意力与线性 FFN 命名用户传入的trtllm_conversion_dict最后合并覆盖见 trtllm_helper.py。注意映射字典的键必须去掉层号。5.3 权重转换与 TP/专家切分SingleDeviceTRTLLMModelWeightsConverter.convert完成实际转换见 single_device_trtllm_model_weights_converter.py核心逻辑包括非 transformer 层embedding、lm_head、final layernorm直接转成目标 dtype 并存入trtllm_model_weightsembedding 层在use_parallel_embedding时按inference_tp_size补齐pad_vocab_size权重转置二维权重统一val val.T对齐 TRTLLM 的权重布局dense / proj 类权重沿 dim 0 用torch.chunk按inference_tp_size切分保存为layer_name.{rank}.binMLP fc / gate 类权重若激活函数为 gated 形式如 SwiGLU 等通过is_gated_activation判断先把fc拆成fc与gate两半再各自沿最后一维切分QKV 权重按[hidden_dim, num_kv_heads, q_num 2, size_per_head]重构拆出 Q/K/V当num_kv_heads inference_tp_size时对 K/V 做_duplicate_kv_head复制以适配 GQA/MQA 的 TP 切分并校验 TP 大小必须能整除或复制 query groups否则抛异常见 single_device_trtllm_model_weights_converter.pyMoE expert 权重fc拆成w1/w3两半后各自 TP 切分再拼接proj沿最后一维切分保存为layer_name.{expert_idx}.bin的 expert_split 格式layernorm_zero_centered_gamma兼容当配置开启layernorm_zero_centered_gamma且归一化为LayerNorm时对layernorm.weight做1.0还原与 NeMo 的 layernorm1p 语义一致。转换完成后get_local_model_weights_per_gpu(mapping, config)根据tensorrt_llm.Mapping的tp_rank、pp_rank与pp_layers(num_layers)把已切分的.bin权重分发给每个 GPU rank只保留tp_rank对应的分片、只保留落在本 PP 段内的层并把层号重排为从 0 开始embedding 与 lm_head 只在首/末 PP rank 上分发见 single_device_trtllm_model_weights_converter.py。5.4 TRTLLM 配置的生成_get_trtllm_config从TransformerConfig映射出 TRTLLM 的PretrainedConfig见 trtllm_helper.py关键映射包括num_attention_heads、num_key_value_heads取num_query_groups为空时退化为num_attention_heads、head_size取kv_channels、hidden_size、intermediate_size取ffn_hidden_size、norm_epsilon取layernorm_epsilonposition_embedding_typerope映射为rope_gpt_neox否则learned_absoluterotary_pct/rotary_base来自 helper 参数存在seq_len_interpolation_factor时写入rotary_scaling {type: linear, factor: ...}MoE 相关moe_num_expertsmoe_router_topk 0时置 0、moe_top_k、moe_normalization_mode默认RENORMALIZE、moe_tp_modeworld_size、tp_size、pp_size、gpus_per_node写入配置随后为每个 rank 挂上tensorrt_llm.MappingFP8 导出时配置quant_algo与kv_cache_quant_algo详见下文falcon模型额外设置new_decoder_architecture与parallel_attentionnemotron_nas从heterogeneous_layers_config_encoded_json解析block_configs并设置 llama3 风格rotary_scaling。5.5 FP8 量化导出get_trtllm_pretrained_config_and_model_weights还支持 FP8 量化导出传入fp8_quantizedTrue/fp8_kvcacheTrue时_load_scaling_factors会从 state dict 中的_extra_state读取 Megatron 的scale_inv_fwd/scale_fwd缩放因子重命名为 TRTLLM 的activation_scaling_factor/weights_scaling_factorgated 激活还会为gate层复制一份再通过_add_scales_to_converter合并进权重字典见 trtllm_helper.py。转换器在遇到带缩放因子的权重层时会先乘上weight_multiplier再转成torch.float8_e4m3fn见 single_device_trtllm_model_weights_converter.py。6. 当前限制与未来工作原文档明确列出了当前实现的能力边界与规划不做任何夸大流水线并行导出导出流程的流水线并行支持仍在开发中Work in progress更多模型的 GPU 导出on-device 分布式转换目前仅支持gpt、gptnext、llama、nemotron_nas四类见 4 节断言其余模型类型正在推进Refit 功能build_and_save_engine已预留use_refit参数但完整 Refit 工作流仍在规划中VLLM 支持计划中尚未实现。此外从源码可以确认两点隐含前提单设备转换的权重切分规模由ExportConfig.inference_tp_size / inference_pp_size决定而 on-device 转换则由加载模型时的并行规模决定两者都以「导出引擎的推理并行布局」为准与训练时的并行布局解耦。7. 总结与进一步探索本文以官方 README 为骨架完整复现了 Megatron Core → TRTLLM 的单设备与分布式导出两条路径并深入到TRTLLMHelper、TRTLLMLayers、默认转换字典、单设备/分布式权重转换器与引擎构建器展示了层名重命名、TP 切分、GQA 复制、MoE 专家切分与 FP8 缩放因子注入等底层细节。如需继续深入推荐阅读仓库中的以下文件单设备导出示例examples/export/trtllm_export/single_device_export/gpt_single_device_cpu_export.py分布式导出示例examples/export/trtllm_export/distributed_export/gpt_distributed_gpu_export.py导出入口 Helpermegatron/core/export/trtllm/trtllm_helper.py层名与映射表megatron/core/export/trtllm/trtllm_layers.py、megatron/core/export/trtllm/model_to_trllm_mapping/default_conversion_dict.py权重转换器single_device_trtllm_model_weights_converter.py、distributed_trtllm_model_weights_converter.py引擎构建器megatron/core/export/trtllm/engine_builder/trtllm_engine_builder.py导出配置与枚举megatron/core/export/export_config.py、megatron/core/export/model_type.py、megatron/core/export/data_type.py【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考