PyTorch Static Runtime 静态运行时面向 CPU 推理的 TorchScript 优化执行引擎【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch导读Static Runtime 是 PyTorch 提供的一个面向 CPU 推理场景的优化运行时可作为 TorchScript JIT 解释器的即插即用替代方案drop-in replacement在 C 与 Python 中均可使用。它的核心思路是通过**内存规划Memory Planning**把中间张量的分配合并到单块连续内存slab中从而大幅降低推理过程中的动态分配开销。本文以官方文档 torch/csrc/jit/runtime/static/README.md 为骨架结合仓库源码impl.h、impl.cpp、ops.h、init.cpp与测试用例test/test_static_runtime.py深入讲解其适用场景、线程模型、内存规划原理、算子注册机制与核心数据结构。读完本文你将掌握如何用 Static Runtime 替换 JIT 解释器、如何配置其选项以获得最优推理性能并理解其底层实现原理。这是一个实验性功能experimental feature。在使用前请确认你的模型与运行环境满足下文所述的前提条件。一、为什么需要 Static Runtime适用场景与前提假设Static Runtime 专门为 CPU 推理优化其价值体现在以下三个条件下最为明显对应官方文档的适用场景模型几乎不含控制流very little control flow。Static Runtime 在加载时对整个图做静态分析与规划控制流prim::If、prim::Loop、prim::fork虽受支持但会引入额外的运行时决策与子块管理开销。PyTorch 的运行时开销如张量创建占模型运行时间不可忽视的比例。尤其是当张量分配消耗大量时间时Static Runtime 的帮助最大——中间张量的内存被合并进单块 slab推理过程中绝大多数动态分配被消除。推理性能极其重要inference performance is extremely important。使用前提假设Assumptions官方文档明确了当前使用该功能的一组假设仅推理执行Inference only且仅 CPU静态输入 dtypeStatic input dtypes静态输入 shape运行时虽然支持动态 shape但过度的动态 shape 会显著降低性能原因见下文标准尺寸调整一节。此外从源码实现看还有一条硬性约束impl.cpp 中的canEnableStaticRuntime会检查图是否可被 Static Runtime 支持其注释明确指出Graph must be frozen——在构造StaticModule时若传入的模块尚未冻结构造函数会自动调用freeze_module完成冻结见impl.cpp第 337-340 行以保证图中没有可变的 Python 状态干扰静态分析。二、线程模型与两种执行模式Static Runtime 支持两种执行模式对应文档 Threading model 章节两者的差别在于并发模型不同。模式 1单线程执行仅算子内部并行单线程执行除算子内并行intra-op parallelism外不做任何线程级并行。此时有两种等价写法// m 是 TorchScript 模块 auto runtime StaticRuntime(m, opts); auto output runtime.run(args, kwargs);或先调用PrepareForStaticRuntime对图做预处理再构造运行时auto mod PrepareForStaticRuntime(m); auto runtime StaticRuntime(mod, opts); auto output runtime.run(args, kwargs);模式 2数据并行式多线程每线程一个运行时实例类似于数据并行同一模型在不同线程上同时处理不同输入。此时必须先用PrepareForStaticRuntime准备图。一个模型对应一个StaticModule或官方术语 InferenceModule实例每个运行线程对应一个StaticRuntime实例。为避免频繁构造StaticRuntime建议用同步栈如boost::lockfree::stack缓存所有运行时实例// 初始化阶段每个模型只做一次 auto mod PrepareForStaticRuntime(m); // 128 对大多数场景足够可按需调整 boost::lockfree::stackstd::shared_ptrStaticRuntime, boost::lockfree::fixed_sizedtrue pool(128); // 推理阶段 std::shared_ptrStaticRuntime runtime nullptr; pool.pop(runtime); if (!runtime) { runtime std::make_sharedStaticRuntime(mod, opts); } auto output runtime-run(args, kwargs); pool.push(runtime);在 impl.h 的类注释中可以看到与之对应的StaticModule/StaticRuntime用法StaticRuntime(*module)构造并自行管理内存。两种模式共同的硬性约束为官方文档加粗强调的In both modes,StaticRuntimemay not be used after its associatedStaticModuleis destructed!两种模式下StaticRuntime都不能在其关联的StaticModule析构之后继续使用原因在于StaticRuntime内部以非拥有unowned引用的方式指向StaticModule持有的共享数据结构见下文结构图与生命周期。三、内存规划Memory Planning性能优化的核心Static Runtime 的加速本质来自其内存规划器Memory Planner做的两件事对应文档 Memory Planning 章节合并内部张量存储的分配coalesce internal allocations for tensor storage通过静态分析确定如何高效复用内存static analysis to figure out how to efficiently reuse memory。3.1 标准尺寸调整Standard ResizingStatic Runtime 会在第一次推理迭代时记录它见到的每个受管中间张量managed intermediate tensor所需的空间。一个中间张量成为受管managed需满足两个条件产生它的算子具有 out variant。Out variant 是对算子的包装概念上把算子签名从Tensor some_op(const Tensor some_arg)变换为void some_op(Tensor output, const Tensor some_arg)。Out variant 通过REGISTER_OPERATOR_FUNCTOR宏注册详见下文注册算子。该张量不对图输出产生别名does not alias a graph output。输出张量由内存规划器单独处理见受管输出张量。在上述算法下静态分析把张量归入若干StorageGroup存储组。同一存储组内的张量共享内存而两个张量能够进入同一存储组的前提是它们的生命周期lifetime不重叠。在后续迭代中Static Runtime 在每次 run 开始时分配张量缓冲区分配的内存大小为sum([max(tensor.size()) for tensor in storage_groups])即各存储组内最大张量尺寸之和。如果后续运行中某张量需要超过已分配空间就会发生一次动态分配——这正是动态 shape 会降低性能的原因。在标准尺寸调整策略下Static Runtime 会在迭代结束时记录每个存储组内新的最大张量尺寸并在下一次迭代分配可能更大的缓冲区。3.2 受管输出张量Managed Output TensorsStaticRuntime可以通过StaticModuleOptions中的manage_output_tensors选项选择性地管理输出张量。开启该选项后输出张量的分配会被合并在一起。注意包含输出张量的缓冲区与包含中间张量的缓冲区是分离的——前者需要存活到推理 run 结束之后因为调用方还要使用输出而后者需要在 run 结束时被释放。底层实现上每个返回的Tensor中都保存了一个指向输出 arena 的引用计数指针refcounted pointerarena 被显式销毁。3.3StaticModuleOptions内存相关配置项StaticModuleOptions定义于 impl.h其中与内存规划直接相关的字段如下全部带默认值选项默认值作用enable_out_varianttrue启用 out variant是 Static Runtime 进行内存规划的前提optimize_memorytrue为生命周期不重叠的张量复用存储以降低内存占用要求enable_out_varianttruemanage_output_tensorsfalse为图输出张量批量分配存储存储的释放发生在 Static Runtime 之外要求enable_out_varianttrueuse_copy_variantstrue门控 ReplaceWithCopy pass把有时会别名输出的算子替换为总是拷贝的 out variant使输出可参与内存规划。注意由于替换在 TensorExpr 融合之后进行结果图不再符合融合器假设因此即使该开关打开只要启用了 TensorExpr 融合该 pass 也不会执行use_maybe_copy_variantstrue门控 ReplaceWithMaybeCopy pass把有时会别名输出的算子替换为包含 out variant 的子图。与use_copy_variants同理启用 TensorExpr 融合时该 pass 同样不执行enable_tensorexpr_fusionfalse模型加载时对算子做 TensorExpr 融合需要提供 sample inputs见下文在 impl.cpp 的OptimizeGraph中可以看到这些选项如何被消费若enable_tensorexpr_fusion为真且提供了sample_inputs则先执行performTensorExprFusion随后依次执行Inline、ConstantPropagation、Canonicalize、RemoveTensorMutation等 pass 完成图优化。四、注册算子Registering Ops三种执行模式Static Runtime 对图中的每个torch::jit::Node选择执行实现共有三种模式对应文档 Registering Ops 章节Out variants返回张量且输出可被内存规划器管理的算子。通过 ops.h 中的REGISTER_OPERATOR_FUNCTOR宏注册宏的使用方式如下REGISTER_OPERATOR_FUNCTOR( aten::op_name, aten_op_name, // 该宏会生成一个 struct此字段为其命名 [](torch::jit::Node* n) - SROperator { // 该机制让我们可以支持 schema 的子集 if (n-matches(some_schema)) { return some_overload; } else if (n-matches(another_schema)) { return another_overload; } return nullptr; })其中SROperator是std::functionvoid(ProcessedNode*)的类型别名定义见 impl.hProcessedNode的细节见下文实现细节。Native functions与 out variant 类似区别在于输出无法被管理——因为算子的返回类型不是张量或者它是视图算子view op返回张量别名而非新张量。通过REGISTER_NATIVE_OPERATOR_FUNCTOR宏注册ops.h用法与REGISTER_OPERATOR_FUNCTOR相同。JIT fallbackStatic Runtime 对该算子没有实现因此回退选用 JIT 解释器使用的实现。加载模型时的算子选择流程对图中的每个torch::jit::Node按以下顺序决策若注册了 out variant把节点传给生成SROperator的函数结果非nullptr则使用该算子若注册了 native function把节点传给生成SROperator的函数结果非nullptr则使用该算子使用 JIT 实现若 JIT 实现也不存在Static Runtime 会抛出异常。从源码看这一选择逻辑落在ProcessedFunction上impl.h它用枚举Kind { kOutVariant, kNativeFunction, kInterpreterFallback }标记每个节点最终选定的执行方式ProcessedFunction在StaticModule构造时初始化。五、实现细节Implementation Details5.1 核心数据结构与生命周期官方文档给出了核心数据结构的引用关系图箭头A→B表示A存储了对B的引用若是非拥有引用A不能活得比B或B间接引用的任何对象更久若是拥有引用则A与B生命周期相同IValue array◄────────────────┐─────────────────────────────────────────┐ ▲ │ Owns │ Owns │ │ ┌───────────────────────────────►ProcessedNode───────►BlockRunner │Owns │ │ │ │ │ Owns │ │ Owns │ │ StaticModule◄───────────StaticRuntime───────────►BlockRunner────────►MemoryPlanner │ ▼ │ │ │ │ │ ... Owns│ │ │ │ │ ▼ │ │ │ │ BlockInfo◄├───────────────────────────────────────────┘──────────────────┘ │ │ │ Owns│ │ ▼ │ ProcessedFunction ◄─────────────────────────────────────────────────────────────────────────────┘5.2StaticModule与StaticRuntimeStaticModule由torch::jit::Module构造用于生成StaticRuntime实例。每个StaticModule恰好缓存一个StaticRuntime实例通过runtime()访问时惰性初始化。StaticModule::operator()可直接用于推理底层只是转发给缓存 runtime 的StaticRuntime::operator()。由此带来的一个后果是StaticModule::operator()不是线程安全的。多线程场景下的正确用法是给每个线程各自一个StaticRuntime实例。新实例可以直接创建StaticRuntime(static_module)也可以从已有实例clone()。StaticModule接收一组选项StaticModuleOptions控制其派生的所有 runtime 实例的行为。内部地StaticRuntime拥有一个IValue数组被所有BlockRunner和ProcessedNode引用运行时产生的所有值都存储在该数组中。5.3BlockRunnerBlockRunner代表图中的单个子块sub-block。每个图至少有一个对应顶层块的BlockRunnerStaticRuntime通过调用(*top_level_block)(args, kwargs)开始推理。每个BlockRunner拥有自己的MemoryPlanner和一组ProcessedNode。带子块的特殊节点如prim::If可能拥有BlockRunner由算子实现负责调用对应子块的BlockRunner。5.4MemoryPlannerMemoryPlanner是抽象基类每个子类实现一种内存规划算法详见内存规划一节。除了张量内存规划外它还封装了另两类优化受管输出张量Managed output tensors见 3.2 节借用IValueBorrowed IValues只做输入解包的算子如dict_unpack可能产生弱引用以避免引用计数 bumpMemoryPlanner需要负责适当地销毁这些借用。5.5ProcessedNode与ProcessedFunctionProcessedNode是单个算子的抽象。每个ProcessedNode存储对StaticRuntime的IValue数组的非拥有引用并知道如何把输入/输出索引映射到数组下标processed_node-output(i)返回ivalue_array[some_set_of_indices[i]]的引用。每个ProcessedNode持有一个ProcessedFunction表示实际要执行的算子。ProcessedFunction在StaticModule构造时按上文算子选择流程初始化。注意所有ProcessedFunction在全部 runtime 实例之间共享因此所有ProcessedFunction必须是线程安全的。5.6ProcessedNodeMetadataProcessedNodeMetadata为ProcessedNode保存各种额外字段通常未被使用但少数算子需要额外机制prim::If根据条件判断持有 true/false 两个子块的BlockRunnerprim::Loop持有循环子块执行的BlockRunnerprim::fork持有torch::jit::TaskLauncher即std::functionvoid(std::functionvoid())负责分叉图的执行。5.7 异步执行Asynchronous ExecutionStaticRuntime::runAsync()API 允许把异步操作放到传入的TaskLauncher上执行父图parent graph在调用线程上内联执行inline execution像prim::fork这样的异步操作在传入的 launcher 上执行若未提供 launcher则通过at::launch执行即运行在inter-op 线程池上。六、Python 端使用与测试佐证Static Runtime 的 Python 绑定通过 init.cpp 暴露核心入口是torch::jit::StaticModule类提供__call__直接推理内部把 py::args/kwargs 转成IValue后调用self(arg_ivalues, kwarg_ivalues)benchmark(args, kwargs, warmup_runs, main_runs)整体基准测试benchmark_individual_ops(args, kwargs, warmup_runs, main_runs)逐算子基准测试返回IndividualMetrics包含setup_time、memory_alloc_time、memory_dealloc_time、output_dealloc_time、first_iter_time、total_time、out_nodes_count、total_nodes_count、time_per_node、time_per_node_type、percent_per_node_type、instances_per_node_type、out_nodes等统计字段runAsync(args, kwargs)异步执行入口。Python 侧的用户包装类位于 torch/testing/_internal/static_module.py其StaticModule类对 C 绑定做了轻量封装__call__、benchmark、runAsync、benchmark_individual_ops。仓库测试 test/test_static_runtime.py 是理解实际用法的绝佳参考覆盖了fork/wait 系列test_fork_wait_1至test_fork_wait_4及对应_async变体验证StaticModule(torch_graph)后调用static_runtime_module.runAsync((inp1, inp2), {})的异步执行路径包括嵌套 fork/wait、多 fork/wait、异常传播真实模型test_multihead_attention_layer直接把attention模块传入StaticModule进行推理test_multihead_attention_layer_benchmark展示benchmark([src, src, src, src_mask], {}, 2, 2)与benchmark_individual_ops的基准用法test_mlp复用了./bench/dlrm_s_benchmark.sh的参数构造多层感知机验证。七、使用建议与限制总结基于官方文档与源码使用 Static Runtime 时请留意以下要点确认模型适合控制流极少、张量分配耗时占比高、推理性能要求苛刻的 CPU 模型收益最大输入 shape 尽量静态化避免触发动态分配回退。先冻结再构造虽然构造StaticModule时会自动freeze_module若传入的模块未冻结见 impl.cpp但建议在模型侧预先eval()与冻结使图优化更可控。多线程推理遵循一模型一StaticModule、一线程一StaticRuntime并用 lockfree 栈缓存 runtime 实例切勿在StaticModule析构后继续使用其StaticRuntime。按需开启选项默认enable_out_varianttrue、optimize_memorytrue若模型输出需在 run 之后长期存活可考虑manage_output_tensorstrue如需 TensorExpr 算子融合开启enable_tensorexpr_fusion并提供sample_inputs但需知悉它会抑制use_copy_variants/use_maybe_copy_variants对应的替换 pass。善用基准工具通过 Python 绑定的benchmark与benchmark_individual_ops配合IndividualMetrics可以量化 setup、内存分配/释放、逐节点耗时快速定位收益点与瓶颈。Static Runtime 作为实验性功能其适用范围与前提假设仅推理、仅 CPU、静态 dtype/shape 优先直接决定了收益上限。理解其内存规划与算子注册机制后你可以在自己的推理服务中精准判断它是否适用并针对性地配置StaticModuleOptions以获得最佳的 CPU 推理性能。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考