【大模型开发】使用工具（TensorRT、TorchScript、ONNX 和 FSDP 等）对大模型进行优化与部署

将为你提供一套完整的大模型优化与部署方案，涵盖 TensorRT、TorchScript、ONNX 和 FSDP 的使用方法，并结合硬件特性进行编译和并行优化。方案将包括这些工具的介绍、它们之间的关系、优化步骤，以及实际代码示例。。

1. 工具介绍与关系梳理

TensorRT：NVIDIA GPU上的高性能推理优化

TensorRT 是 NVIDIA 推出的高性能深度学习推理优化器和运行时，它能够在 GPU 上以低精度（如 FP16 和 INT8）执行深度学习模型推理，从而显著提高推理速度 (英伟达PyTorch优化神器TensorRT重磅更新！10亿参数大模型实时运行，GPT推理加速21倍_腾讯新闻)。通过对训练后的神经网络模型进行层融合、内存优化和算子高效实现等手段，TensorRT 可以大幅降低模型推理的延迟并提高吞吐量。例如，与未优化的原始模型相比，使用 TensorRT 可将大型模型的推理延迟降低数倍 (英伟达PyTorch优化神器TensorRT重磅更新！10亿参数大模型实时运行，GPT推理加速21倍_腾讯新闻)。TensorRT 提供了丰富的接口来支持模型精度削减（如半精度 FP16 和8比特整精度 INT8）以及针对具体硬件架构的内核优化，使模型在不明显损失精度的情况下获得更快的推理性能

TensorRT 通常与其他工具结合使用：一方面可直接解析 ONNX 格式模型并将其编译为高度优化的推理引擎；另一方面，通过 Torch-TensorRT 等与 PyTorch 的集成，TensorRT 可以无缝对接 TorchScript 模型，在保留 PyTorch 易用性的同时获得更高的推理性能（后文详述）。总的来说，TensorRT 专注于模型推理阶段的极致性能优化，特别适用于在 NVIDIA GPU 上部署对延迟和吞吐要求极高的深度学习服务。

TorchScript：PyTorch静态图模式优化与部署

TorchScript 是 PyTorch 提供的将动态图模型转换为静态计算图的工具和中间表示形式。通过 TorchScript，我们可以将用 Python 编写的模型“脚本化”（Script）或“跟踪”（Trace）为一个可序列化的静态计算图，从而脱离Python解释器在其他环境（如C++应用、移动设备）中高效运行 (PyTorch JIT和TorchScript，一个API提升推理性能50%-CSDN博客)。TorchScript 的静态图表示允许编译器优化（例如算子融合、常量折叠）在编译阶段就对模型进行优化，而无需每次推理时再构建计算图 (PyTorch JIT和TorchScript，一个API提升推理性能50%-CSDN博客)。这降低了 Python 层的开销，在某些情况下可以提高推理性能；有测试表明，在GPU上使用TorchScript能将PyTorch模型推理性能提升约50% (PyTorch JIT和TorchScript，一个API提升推理性能50%-CSDN博客)（提升幅度依模型和硬件而定）。

TorchScript 提供两种生成方式：torch.jit.trace 基于范例输入追踪模型的执行路径，适用于纯前向计算的模型；torch.jit.script 则直接分析模型代码，可处理动态控制流。生成的 TorchScript 模型可保存为 .pt 文件并加载到无需Python的运行时中执行，实现跨平台部署 (PyTorch JIT和TorchScript，一个API提升推理性能50%-CSDN博客)。此外，TorchScript 与 PyTorch JIT 编译器结合，使我们能够在需要时将模型从动态图切换到静态图，提高部署效率 (PyTorch JIT和TorchScript，一个API提升推理性能50%-CSDN博客)。总之，TorchScript 聚焦于模型表示和执行方式的优化，使PyTorch模型在保持灵活开发的同时，也能获得接近静态图框架的性能和可移植性。

ONNX：开放的跨框架模型格式与优化

ONNX（Open Neural Network Exchange）是一个开放的深度学习模型交换格式，旨在促进不同深度学习框架之间的互操作性 (pytorch模型转ONNX、并进行比较推理_onnx和pt效果对比-CSDN博客)。通过将训练得到的模型导出为ONNX格式，我们可以在不重新实现模型的情况下，将其从一个框架（如 PyTorch、TensorFlow 等）迁移到另一个框架或运行时环境中执行 (pytorch模型转ONNX、并进行比较推理_onnx和pt效果对比-CSDN博客)。ONNX 定义了一套通用的算子集合和 IR（中间表示）结构，支持包括卷积神经网络、循环神经网络、Transformer 在内的各种模型类型 (pytorch模型转ONNX、并进行比较推理_onnx和pt效果对比-CSDN博客)。自2017年由微软和Facebook发布以来，ONNX 已获得众多公司和开源社区的支持，使其成为事实上的跨平台模型表示标准。

(ONNX Runtime and models - Azure Machine Learning | Microsoft Learn)ONNX 作为开放格式充当不同框架与不同部署平台之间的桥梁：左侧是支持导出为 ONNX 的训练框架（如 PyTorch、TensorFlow、Keras 等），右侧是部署平台（云服务、Windows/Linux 设备、移动端等）。通过 ONNX 模型，训练与部署解耦，实现“一次训练，随处推理” (ONNX Runtime and models - Azure Machine Learning | Microsoft Learn)。

使用 ONNX 时通常包含两个步骤：首先在原生框架中导出模型为 ONNX 格式文件；然后在目标环境中导入模型并使用 ONNX Runtime 或其他兼容引擎进行推理 (pytorch模型转ONNX、并进行比较推理_onnx和pt效果对比-CSDN博客)。ONNX Runtime 是微软主导的高性能推理引擎，针对 ONNX 模型做了大量图优化，并可以对接不同硬件加速器（如 TensorRT、OpenVINO、CUDA、TensorRT 等）以提升性能 (ONNX Runtime and models - Azure Machine Learning | Microsoft Learn) (ONNX Runtime and models - Azure Machine Learning | Microsoft Learn)。值得注意的是，在实际转换中需要确保模型的算子获得 ONNX 支持，某些框架特有的算子可能需要自定义导出或转换；但总体而言，ONNX 极大地简化了模型在不同框架和硬件间部署的流程，是大模型优化部署流程中常用的中间表示与交换格式。

FSDP：分布式训练与推理拆分

FSDP（Fully Sharded Data Parallel，全分片数据并行）是 PyTorch 官方提供的大模型分布式训练和推理策略，用于将模型的参数和状态拆分到多块 GPU 上 (pytorch FSDP特性-CSDN博客)。传统的数据并行（Data Parallel，如 PyTorch DDP）在每张卡上保留一份完整的模型副本，而 FSDP 则将模型的参数、梯度和优化器状态在参与并行的所有 GPU 间按块划分，每张卡只保留其中一部分 (pytorch FSDP特性-CSDN博客)。在前向和反向计算每一层时，FSDP 会动态地将该层的完整权重从各卡碎片 All-Gather 到一起进行计算，用完后再释放；反向传播完成后，对梯度进行 Reduce-Scatter 来更新碎片 (pytorch FSDP特性-CSDN博客) (pytorch FSDP特性-CSDN博客)。这种设计使得每张卡的静态内存开销约为原来的 1/N（N 为卡数） (pytorch FSDP特性-CSDN博客)，从而大幅降低单卡显存占用，可以训练和推理更大的模型。

除了内存拆分之外，FSDP 还支持混合精度训练和 CPU Offload 等特性，将暂时不用的参数页交换到 CPU 内存，以进一步减少显存占用 (pytorch FSDP特性-CSDN博客) (pytorch FSDP特性-CSDN博客)。通过合理设置重叠通信和计算，FSDP 在扩展到多卡、多机时也能获得接近线性的加速比。需要注意 FSDP 的通信开销：由于每层都需要在 GPU 间同步权重，网络带宽会成为瓶颈之一，但在大模型场景下这一开销通常是可以接受的权衡。总的来说，FSDP 专注于大模型的分布式并行，允许我们在现有硬件资源下高效地训练超出单卡显存限制的模型，并支持推理时按需调入权重，使超大模型的推理成为可能。

这些工具如何协同优化大模型

上述工具各有侧重，但可以协同作用于大模型的训练和部署流程，形成端到端的优化方案：

• 训练阶段：我们可以使用 FSDP 在多GPU/多机环境下高效训练大模型。FSDP负责拆分模型、降低单卡显存需求，并利用混合精度等提升训练速度。训练完成后，往往会得到分布式存储的权重，需要汇总（gather）回单机得到完整模型权重用于后续部署。

• 模型导出与表示：训练好的模型可以通过 TorchScript 或 ONNX 导出为静态格式。TorchScript 让模型以 PyTorch JIT 支持的形式存在，方便直接在PyTorch生态中部署或者结合 Torch-TensorRT 优化。ONNX 则作为独立于框架的开放格式，便于后续借助各类推理引擎优化。我们可以视具体需求选择一种或同时使用两种方式：例如先用 TorchScript 优化模型，然后再导出 ONNX 以比较不同后端性能。

• 推理优化阶段：针对部署硬件选择合适的引擎。如果是 NVIDIA GPU 环境，优先考虑 TensorRT 对模型进一步优化，加上 FP16/INT8 精度削减以取得最大性能提升。可以将上一阶段得到的 TorchScript 或 ONNX 模型交给 TensorRT 编译引擎。借助 Torch-TensorRT，PyTorch 用户甚至无需手动导出 ONNX，只需一行代码即可将 TorchScript 模型编译为 TensorRT 引擎 (英伟达PyTorch优化神器TensorRT重磅更新！10亿参数大模型实时运行，GPT推理加速21倍_腾讯新闻) (Accelerating Inference Up to 6x Faster in PyTorch with Torch-TensorRT | NVIDIA Technical Blog)。对于非 NVIDIA 硬件（如CPU或其他厂商AI加速器），则可使用 ONNX Runtime 等根据目标设备选择最优的执行后端。

• 分布式推理：如果模型体积过大（例如数十亿参数级别）以至于单卡显存难以容纳，即使在推理阶段也可以采取类似 FSDP 的策略进行模型拆分推理。也就是说，在多卡上部署模型的不同分片，各卡各自运行模型的一部分计算，通过通信汇总输出。这可以与上述优化结合使用：例如将每个分片编译成 TensorRT 引擎在对应卡上运行。同时，也可以在部署层面通过像 Ray Serve 这样的框架将不同机器上的分片协同起来完成一次完整推理。不过，需要权衡通信开销和实现复杂度，因此通常在模型无法单卡放下时才采用分布式推理，并尽可能结合硬件加速以减轻分布成本。

综上，大模型的优化部署往往是多管齐下的：FSDP解决训练和推理内存瓶颈，TorchScript/ONNX提供优化和跨环境的模型表示，TensorRT等引擎榨取推理性能极限。它们相辅相成，共同服务于降低大模型从训练到部署整个生命周期的成本。下面将详细介绍不同优化方案的具体步骤。

2. 优化与部署方案

使用 TorchScript 进行 PyTorch 模型优化

在完成对模型的训练和调试后，首先可以利用 TorchScript 将 PyTorch 的动态图模型转换为静态图，以优化其推理性能和部署便携性。基本步骤如下：

1. 脚本化模型（Script or Trace）：选择合适的方法将模型转换为 TorchScript。对于不含复杂控制流的模型，可以使用 torch.jit.trace(model, example_inputs) 跟踪模型的计算图；如果模型内部有条件判断、循环等动态行为，则应使用 torch.jit.script(model) 直接将模型脚本化。两种方式都会返回一个 ScriptModule（即 TorchScript 模型）。

2. 图优化：PyTorch 在生成 TorchScript 时会应用内部的图优化，如常量折叠、简化计算节点等。此外，我们可以调用 torch.jit.optimize_for_inference(script_module) 对模型做进一步优化（比如融合 BatchNorm 到卷积）。这些优化在保证输出一致的前提下删去多余的运算，提高执行效率。

3. 保存与加载：将优化后的 TorchScript 模型通过 script_module.save("model.ts") 保存为文件。在部署端，通过 torch.jit.load("model.ts") 加载模型。由于 TorchScript 模型不依赖 Python 解释器，因此可以在纯C++环境下运行（例如使用LibTorch加载），这对跨语言部署非常有利。

4. 推理执行：使用 TorchScript 模型进行推理时的接口与原始 PyTorch 模型相同，但此时模型的执行是在C++后端完成的。由于省去了Python层开销，并利用了静态图优化，推理速度往往更快。特别是在CPU环境，TorchScript 可以避免 GIL 锁影响，实现多线程并发推理；在GPU环境，TorchScript 则保证了更高的算子执行效率和调度性能。

部署建议：当目标部署环境不便于使用Python时（如移动端或仅提供C++环境的服务器），TorchScript 是首选方案。同时在对实时性要求较高的场景，即使仍在Python环境中运行，提前将模型TorchScript化也能减少每次推理的开销。如果模型需要进一步结合硬件优化（如 TensorRT），TorchScript 也可以作为中间步骤，先将模型静态化和优化，再交由下游工具处理。

ONNX 作为中间表示的优化路径

使用 ONNX 作为中间表示进行优化是一种通用且灵活的方案，适用于跨框架部署和利用多种推理加速库。其典型流程如下：

1. 模型导出为 ONNX：在训练完成的框架中（这里假设是 PyTorch），调用诸如 torch.onnx.export(model, example_inputs, "model.onnx", export_params=True, opset_version=14, do_constant_folding=True) 将模型权重和计算图导出为ONNX格式文件。需要提供一个代表模型输入形状的 example_inputs 张量。设置 do_constant_folding=True 可以在导出时执行常量折叠等简单优化。opset_version 则选择ONNX算子集的版本，通常使用最新稳定版以获得更丰富的算子支持。

2. ONNX 模型检查与优化：导出后可以使用 ONNX 官方工具检查模型的正确性，例如 onnx.checker.check_model(onnx_model). 另外，可以借助 onnx.optimizer 模块或 onnx-simplifier 等第三方工具对 ONNX 图进行优化简化，如删除多余的Transpose、合并冗余节点等。这一步可以精简模型结构，为后续推理提速。

3. 选择推理引擎：有了 ONNX 模型，可以选择合适的 ONNX 推理运行时。ONNX Runtime (ORT) 是常用选择，支持在 CPU、GPU 上高效执行 ONNX 模型，并根据后端自动进行图划分和加速 (ONNX Runtime and models - Azure Machine Learning | Microsoft Learn)。例如，在CPU上 ORT 会使用线程池和向量化优化，在GPU上则可调用CUDA内核或 TensorRT (通过 TensorRT Execution Provider) 来加速计算。如果部署在Intel硬件上，还可以使用 ORT 的 OpenVINO 执行提供器，或直接使用 OpenVINO 工具链加载ONNX。总之，ONNX 模型具有“一次导出，多种执行”的优势。

4. 推理部署：将 ONNX 模型和相应的运行时集成到应用中。在Python中，可直接用 onnxruntime：

   import onnxruntime as ort  
   sess = ort.InferenceSession("model.onnx", providers=['CPUExecutionProvider'])  
   preds = sess.run(None, {input_name: input_data})
   

   在C++中，可使用 onnxruntime C++ API 或其他支持ONNX的推理库。在云端部署时，还可以考虑 NVIDIA Triton Inference Server 等，它能原生支持加载 ONNX 模型并提供HTTP/gRPC接口进行服务。

部署建议：如果你的应用需要跨平台兼容（例如训练用PyTorch但推理需要在TensorFlow或其他环境执行），ONNX 是理想选择。在硬件利用方面，ONNX Runtime 已针对主流CPU/GPU做了大量优化，并提供了灵活的后端切换，可以根据部署环境动态选择最佳加速方案。因此，当面临多样化的硬件平台，或需要与现有系统集成而不便直接依赖PyTorch时，ONNX 路线能提供最大的灵活性。尽管将模型导出再导入略有繁琐，但对大模型部署来说，这一步带来的性能和兼容性收益往往是值得的。

使用 TensorRT 进行 GPU 推理优化

当确定部署环境为 NVIDIA GPU 时，TensorRT 通常是争取极致推理性能的必经步骤。利用 TensorRT 对模型进行优化和部署的流程如下：

1. 模型获取：准备需要优化的模型，可以是通过 PyTorch 等训练得到的模型。此时有两种路径：一是将模型转换为 ONNX，然后使用 TensorRT 解析ONNX；二是直接使用 Torch-TensorRT 将 PyTorch 模型编译为 TensorRT 引擎。两种方式最终都会生成 TensorRT 可以执行的引擎(plan文件)。

2. 构建 TensorRT 引擎：如果采用 ONNX 路径，使用 TensorRT 的 Python/C++ API 读取 ONNX 并构建引擎。例如，在Python中：

   import tensorrt as trt
   logger = trt.Logger(trt.Logger.WARNING)
   builder = trt.Builder(logger)
   network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
   parser = trt.OnnxParser(network, logger)
   with open("model.onnx", "rb") as f:
       parser.parse(f.read())
   config = builder.create_builder_config()
   config.max_workspace_size = 1 << 30  # 分配工作空间
   config.flags = trt.BuilderFlag.FP16  # 启用FP16精度
   engine = builder.build_engine(network, config)
   

   上述流程解析了ONNX模型，并开启FP16模式来构建优化引擎。max_workspace_size 用于指定优化时可用的GPU内存，用于层融合、张量格式转换等优化。对于INT8量化，还需要设置 config.flags = trt.BuilderFlag.INT8 并提供校准器（config.int8_calibrator = calibrator），以便TensorRT评估量化缩放因子。如果使用 Torch-TensorRT，则更加简洁，只需：

   import torch_tensorrt
   trt_model = torch_tensorrt.compile(model, inputs=[torch_tensorrt.Input(example_tensor.shape)], enabled_precisions={torch.float, torch.half})
   

   这样直接得到的 trt_model 就是一个 TensorRT 加速的 TorchScript 模块。其中通过设置 enabled_precisions 实现FP16支持（同时包含 torch.float 以允许fallback），如果要INT8则加入 torch.int8并提供校准数据。

3. 执行推理：无论通过何种方式获得 TensorRT 引擎，最终都可在C++或Python中使用该引擎执行推理。在 TensorRT API 中，需要创建执行上下文（context = engine.create_execution_context()），并为输入输出分配CUDA内存，然后调用 context.execute_v2(bindings) 来运行。Torch-TensorRT 编译的模型则可直接像PyTorch模型一样调用，例如 outputs = trt_model(inputs)，其内部实际会运行TensorRT引擎。由于TensorRT已经对模型进行了高度优化，推理过程将比原始框架更高效，常常能充分利用GPU的计算能力和带宽。

4. 性能调优：针对TensorRT引擎，我们可以进一步调优。例如调整批大小（TensorRT引擎通常对固定max batch优化）、使用异步流并行执行推理、多线程管理多个执行上下文以提高GPU利用率等。另外，通过TensorRT的 profiler 可以分析瓶颈层并尝试调整网络结构或TensorRT配置。需要注意引擎的兼容性：TensorRT 引擎 plan 文件在不同GPU架构之间可能不通用，一般在目标设备上生成引擎以确保最佳兼容和性能。

部署建议：在拥有NVIDIA GPU的环境下，优先考虑 TensorRT 来部署大模型以获取最高推理性能。特别是对于大批量和低延迟场景（如在线服务），使用 FP16 可充分利用张量核心，加速矩阵运算；有充裕校准数据且对精度影响可控时，引入 INT8 可以进一步降低计算和内存带宽需求，实现更高吞吐。TensorRT 适合对性能要求极致的部署，但需要一定的工程工作来集成。如果希望简化流程，可以使用 Torch-TensorRT 或 ONNX Runtime 的 TensorRT 执行器等方案快速上手。在部署过程中，还应结合显存情况调整批处理大小和并发策略，以发挥 GPU 最大效能。

使用 FSDP 进行分布式训练和推理加速

当模型规模超出单机显存或需要加速训练时，FSDP 分布式并行是一个强有力的工具。其使用方式和优化效果如下：

1. 环境初始化：在使用 FSDP 前，需要先初始化分布式训练环境。对于单机多卡，可以使用 torch.distributed.init_process_group(backend='nccl') 等方法建立进程通信（多机则需指定 init_method 或使用 torchrun 等工具）。每个进程代表一块 GPU，并通过设置 LOCAL_RANK 等环境变量来区分。

2. 包装模型：构建模型后，将其移动到当前进程对应的设备上（如 model.to(device)），然后用 FSDP 对模型进行包装：

   from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
   model = FSDP(model, cpu_offload=CPUOffload(offload_params=True), auto_wrap_policy=size_based_auto_wrap_policy(1e6))
   

   这里我们示例性地开启了参数CPU Offload（将不活动参数移至CPU内存）来进一步节省显存，并使用 size_based_auto_wrap_policy(1e6) 自动包装策略，将大于1e6参数的子模块单独拆分 shard（避免一次all-gather过多数据）。FSDP 会将模型权重按策略切分到各GPU，并在内部管理通信。

3. 训练循环：使用 FSDP 后的模型与常规单机模型的训练流程基本一致：前向计算 outputs = model(data)，计算损失并反向传播 loss.backward()，再调用优化器更新参数。不过需注意：优化器应当只维护 FSDP 模型的 .parameters()（它实际上是本rank的碎片），且通常使用 optim.AdamW(model.parameters(), ...) 即可，FSDP 会在后台处理参数同步和梯度聚合。在反向传播过程中，FSDP 将逐层执行 all-gather 和 reduce-scatter (pytorch FSDP特性-CSDN博客) (pytorch FSDP特性-CSDN博客)，开发者只需按正常顺序编写训练代码。通过FSDP，显存高峰将被压制在单卡能够容纳的范围内，从而成功训练原本无法单卡训练的大模型。

4. 推理或微调：训练完毕后，可以选择在分布式环境中直接进行推理，此时保持模型为 FSDP 包装状态，输入数据也需要在多个进程上分别提供（比如对一个batch手动拆分到不同rank）。FSDP 会按需要all-gather权重计算并得到各自的输出，最终需聚合结果。如果想回到单机使用模型，则需要使用 FSDP.full_state_dict(model) 收集完整权重，再加载到单机模型中。对于超大模型（如部署大型语言模型），也可以选择在推理阶段继续使用 FSDP，将模型碎片化到多台机器的多块GPU上共同完成推理任务，这种情况下需要一个额外的协调层来汇总输出。

部署建议：FSDP 在模型训练阶段意义重大，它最大化利用集群算力训练大模型且开销相对低于模型并行。对于部署，如果模型大小尚在单机GPU可承受范围内，尽量还是在单机上做推理（结合前述 TorchScript/TensorRT 等优化）以避免跨设备通信延迟；只有当模型实在太大单机无法容纳时，再考虑分布式推理。那种情况下，可以借助 FSDP 或类似的 ZeRO 技术在推理时按需加载权重。总之，优先在训练中使用 FSDP 加速和省显存，推理时则追求简化部署、充分优化单机性能，二者配合可以覆盖大模型从训练到部署的各种场景需求。

结合多个工具的优化流程

针对超大模型的实际项目，我们常常需要综合运用上述工具来达到最佳效果。下面给出一个将 TorchScript、ONNX、TensorRT 和 FSDP 结合的优化与部署流程示例：

1. 分布式训练大模型：利用 FSDP 在多卡上训练一个巨量参数模型：

   import torch.distributed as dist
   from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
   dist.init_process_group(backend='nccl', world_size=4, rank=local_rank)
   model = LargeModel().to(local_rank)
   model = FSDP(model)
   # 正常训练循环...
   for data, target in train_loader:
       output = model(data.to(local_rank))
       loss = loss_fn(output, target.to(local_rank))
       loss.backward()
       optimizer.step()
       optimizer.zero_grad()
   

   训练结束后，各进程通过协作得到模型的完整权重快照：

   full_state = FSDP.full_state_dict(model)
   if local_rank == 0:
       torch.save(full_state, "model_full.pth")
   

2. 模型汇总与静态图转换：在单机环境下（例如只使用 rank0），加载完整权重并转换模型为 TorchScript：

   full_model = LargeModel()
   full_model.load_state_dict(torch.load("model_full.pth"))
   full_model.eval()
   # 将模型脚本化
   scripted_model = torch.jit.script(full_model)
   scripted_model.save("model_scripted.pt")
   

   通过脚本化，我们获得了模型的静态图表示，为后续导出ONNX和TensorRT打下基础。

3. 导出 ONNX 并优化：加载脚本化模型，导出为 ONNX：

   import torch.onnx
   import onnx
   model = torch.jit.load("model_scripted.pt")
   dummy_input = torch.randn(1, 3, 224, 224)  # 示例输入
   torch.onnx.export(model, dummy_input, "model.onnx", opset_version=14, do_constant_folding=True, input_names=["input"], output_names=["output"])
   onnx_model = onnx.load("model.onnx")
   onnx.checker.check_model(onnx_model)
   # 可选: 应用图优化（省略具体代码调用）
   

   现在得到了经过初步优化的 ONNX 模型。

4. 构建 TensorRT 引擎并量化：使用 TensorRT 将 ONNX 编译为高度优化的引擎：

   import tensorrt as trt
   logger = trt.Logger(trt.Logger.INFO)
   builder = trt.Builder(logger)
   network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
   parser = trt.OnnxParser(network, logger)
   with open("model.onnx", "rb") as f:
       parser.parse(f.read())
   config = builder.create_builder_config()
   config.max_workspace_size = 1 << 31  # 分配足够工作内存
   config.set_flag(trt.BuilderFlag.FP16)          # 开启FP16精度
   config.set_flag(trt.BuilderFlag.INT8)           # 开启INT8精度
   config.int8_calibrator = MyCalibrator(data_loader)  # 设置INT8校准器 (假设已实现)
   engine = builder.build_engine(network, config)
   

   在这个过程中，我们启用了 FP16 和 INT8 模式，并假设提供了校准数据来生成量化信息。生成的 engine 就包含了针对目标GPU的最优内核和低精度变体。此时可以将引擎序列化保存，供推理服务直接加载。

5. 部署推理服务：在推理服务器上反序列化加载 TensorRT 引擎或使用 Torch-TensorRT 编译模型，然后执行推理：

   # 方法1: TensorRT 原生执行
   import pycuda.driver as cuda
   import numpy as np
   # 分配GPU内存并拷贝输入
   d_input = cuda.mem_alloc(input_nbytes)
   d_output = cuda.mem_alloc(output_nbytes)
   stream = cuda.Stream()
   context = engine.create_execution_context()
   cuda.memcpy_htod_async(d_input, input_np.astype(np.float16), stream)  # FP16输入
   context.execute_async_v2(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle)
   cuda.memcpy_dtoh_async(output_np, d_output, stream)
   stream.synchronize()
   # output_np 即为推理结果
   
   # 方法2: Torch-TensorRT 集成执行
   import torch_tensorrt
   trt_ts_module = torch_tensorrt.compile(full_model, inputs=[torch_tensorrt.Input(dummy_input.shape)], enabled_precisions={torch.float, torch.half, torch.int8}, calibrator=MyCalibrator(data_loader))
   result = trt_ts_module(torch.from_numpy(input_np))
   

   无论采用哪种方式，最终都利用 TensorRT 实现了高性能推理。方法1展示了独立使用TensorRT API执行推理的流程，方法2则直接在PyTorch中获得了集成TensorRT的模块，使用起来更为简洁。

在这个完整流程中，FSDP确保了我们可以成功训练超大模型，TorchScript和ONNX让模型转换为可优化的静态表示，TensorRT将模型推理性能推向极致。实际应用中可以根据需要裁剪流程：例如模型不特别大则无需FSDP，或者部署环境单一也可跳过ONNX。然而这个范例展示了将多种工具串联以应对从大模型训练到部署的最严苛需求：既要训练得了，又要跑得快。

4. 优化方向与未来建议

算子融合、自动图优化与编译优化

未来的大模型优化将更多地依赖编译器技术和自动优化。一方面，深度学习框架会更智能地进行算子融合，将多个算子合并为一个内核以减少中间内存访问和调度开销。比如 TensorRT 和 XLA 已经广泛应用算子融合，但未来在动态图框架中（如 PyTorch 2.0 引入的 TorchDynamo）也会通过图捕获来实现融合优化。另一方面，自动图优化将更加普遍：框架能够在不干预模型准确性的前提下重排计算顺序、移除冗余计算、自动混合精度等。Just-In-Time 编译和Ahead-Of-Time 编译的结合将被用来针对模型和硬件生成高度优化的代码。例如 PyTorch 2.x 系列引入的 torch.compile（基于 TorchDynamo 和 NVFuser 等）正是这种趋势的体现——开发者无需手动调整模型，实现即可由系统在后台自动优化为高效的低级代码。

硬件加速与特性利用（CUDA、Tensor Cores）

针对硬件特性的优化仍然是提升性能的重要方向。以 NVIDIA GPU 为例，未来优化应当充分利用 Tensor Cores 等专用单元：这意味着在可能的情况下使用 FP16/BF16 甚至 INT8 张量运算，并配合稀疏矩阵乘法（Ampere架构起支持2:4结构化稀疏）等特性来进一步提速。此外，新硬件不断涌现，如 GPU 引入更多寄存器和共享内存、CPU 引入更宽的向量指令和AI加速单元，这些都需要优化工具链紧跟适配。开发者在部署时应根据实际硬件配置调整策略：例如在Ampere/Ada架构GPU上开启 TF32、FP16 混合精度以获得几乎无损的加速，在Intel CPU上利用 OpenVINO 或 oneDNN 的 INT8 优化。在未来，异构计算也可能更多应用于大模型：利用 GPU + DSP/NPU 协同，甚至采用定制芯片。相应地，模型需要切分和调度到最合适的计算单元，这对编译器和运行时提出了更高要求。总之，紧贴硬件发展并充分发掘硬件特性（如CUDA中的流并行、张量核心、高速内存等）将使大模型部署受益匪浅。

多机多卡训练与推理优化

随着模型和数据规模爆炸式增长，多机多卡将成为常态，无论训练还是推理都需要在分布式环境下优化。未来在软件层面，我们会看到更加自动化的并行策略选择和调度。例如，根据模型结构自动选择数据并行、张量模型并行或流水线并行的组合（混合并行）来达到最佳效率，这方面 Megatron-LM 等已经进行了探索。像 DeepSpeed 提供的 ZeRO 技术将继续发展，可能进一步减少通信开销和内存占用。针对推理，大规模模型的分布式推理也会得到更多支持，例如跨节点的参数高效加载、分布式批处理推理的动态负载均衡等。在网络通信方面，新兴的互连技术（如NVLink、InfiniBand、以太网RDMA等）和拓扑感知的调度将降低多机通信成本，使得接近线性的加速更加可期。开发者在展望未来时，应当考虑构建弹性可扩展的训练和推理架构，使得当硬件资源扩充时，能够通过水平扩展继续提高大模型的训练速度和服务吞吐。

未来发展方向（如 Triton、DeepSpeed 与 Ray Serve）

除了上述优化手段，行业中还出现了一些针对大模型的新型工具和框架，值得关注并在未来加以利用：

• NVIDIA Triton Inference Server：Triton是NVIDIA提供的开源推理服务框架，它支持部署多种格式的模型（TorchScript、ONNX、TensorRT、TensorFlow 等），并提供批量请求合并、动态路由、模型版本管理等功能。对于需要部署多个大模型并发服务的场景，Triton 可显著简化架构并提升利用率。例如，它可以自动将多个请求打包成更大的批次送入GPU，以提高吞吐。未来 Triton 预计会更紧密地集成 TensorRT 和新硬件特性，成为大模型线上部署的标准方案之一。

• DeepSpeed：来自微软的 DeepSpeed 框架在大模型优化领域扮演重要角色。其提供的 ZeRO并行优化系列已经被集成到PyTorch（即FSDP），但DeepSpeed本身还在演进，包含如 DeepSpeed-Inference（专门优化Transformer类模型的推理，支持张量并行和高效的KV缓存管理）、量化加速等前沿特性。未来我们可能看到DeepSpeed与PyTorch更加融合，比如更简单的接口来启用混合并行、自动并行切分等。同时DeepSpeed团队也在探索高效内存管理、调度方面，这些都直接影响大模型训练/推理效率。对于开发者来说，保持对DeepSpeed更新的关注能够获得最新的大模型优化技巧，例如通过几行配置开启MoE模型高效训练、8-bit精度推理加速等。

• Ray Serve：Ray Serve 是基于 Ray 的分布式服务框架，擅长于将模型部署为大规模分布式在线服务。对于大模型，Ray Serve 提供了方便的批量调度、异步请求处理和多节点扩展能力。它可以与上述优化工具协同：例如我们可以在每个节点利用 TensorRT 优化模型，在多个节点间用 Ray Serve 统一路由请求，实现既快又扩展良好的推理服务。未来，随着模型变得更大以及应用对延迟敏感度提高，像 Ray Serve 这样的高级调度层可能引入更多智能策略（如基于模型负载的弹性伸缩、热点模型自动缓存等）。对于希望将大模型以服务形式对外提供的团队，关注和学习这类部署框架将有助于构建稳定高效的在线服务。

技术发展趋势：综合来看，大模型优化部署正在向“自动化”和“一体化”方向发展——自动化指减少人工调整，更多依靠编译器和调度系统智能地为模型匹配最佳优化；一体化则体现为训练和推理流程的融合同调，比如一些框架可以直接从训练生成最优推理格式（编译好的模型）甚至直接部署到服务。未来我们很可能能够直接在高层声明性能/资源需求，底层系统就帮我们挑选 TorchScript、TensorRT、分布式并行等组合，使得大模型的性能优化像使用高阶API一样简单。作为开发者，应及时跟进行业新工具的出现，并积极尝试将其纳入自己的大模型开发部署流程中，以保持技术优势。

【哈佛博后带小白玩转机器学习】 哔哩哔哩_bilibili

总课时超400+，时长75+小时