PyTorch Joint with Descriptors导出带语义描述的 AOTAutograd 联合图【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch导读Joint with descriptors是 PyTorch 提供的一套实验性导出 API它把torch.compile通过aot_module_simplified支持的所有特性多可微输出、输入原地修改、张量子类、符号形状等以联合前向反向图joint graph的形式完整导出并用**描述符descriptors**为图中每个输入/输出标注其语义含义。读完本文你将掌握aot_export_joint_with_descriptors/aot_compile_joint_with_descriptors两个 API 的用法、全部描述符类型的含义与数据流规则以及如何在联合图上做自定义优化后再将其恢复为可正常执行、可微调用的函数——这是实现 autoparallel 这类分布式 SPMD 重分片系统的关键前置技术。什么是 Joint with Descriptorsaot_export_joint_with_descriptors导出的联合图与aot_export_joint_simple或aot_export_module(trace_jointTrue)有一个根本区别生成的联合图没有固定的位置化调用约定。例如你不能假设 traced 联合图的第二个参数一定对应被 trace 模块的第二个参数。取而代之的是图的输入与输出通过描述符来结构化描述描述符被标注在 placeholder 节点和 return 节点的meta[desc]上详见 aot_autograd.py 的 docstring。这套导出 API 相比传统aot_export_joint_simple的主要优势是与torch.compile支持的所有情况达到特性对齐feature parity包括多个可微输出multiple differentiable outputs必须在图外处理的输入原地修改input mutations张量子类tensor subclasses视图/别名view/alias处理副作用效应 tokeneffect tokens在导出并完成自定义优化之后还可以把联合图重新转换回可微调用的函数像普通函数一样执行甚至再次交给torch.compile。文档明确说明autoparallel——一个接收模型、对输入和参数重新分片以使其成为分布式 SPMD 程序的系统——正是基于这套 API 实现的。为什么普通 torch.compile 联合图 pass 做不到文档给出了两个核心原因解释了为何 autoparallel 不能作为传统的 torch.compile joint graph pass 实现参数分片决策的时机太早参数的分片sharding必须在参数初始化 / checkpoint 加载之前决定远早于 torch.compile 通常运行的时间点允许改变输入语义需要替换参数的含义例如把复制参数替换为分片版本改变其输入尺寸。而 torch.compile 通常是语义保持的semantics preserving不允许改变输入的含义。核心 APIaot_export_joint_with_descriptors该函数位于 aot_autograd.py签名如下def aot_export_joint_with_descriptors( stack: contextlib.ExitStack, mod: nn.Module, args: tuple[Any, ...], kwargs: dict[str, Any] | None None, *, decompositions: dict[OpOverload, Callable[..., Any]] | None None, keep_inference_input_mutations: bool False, ignore_shape_env: bool False, disable_functionalization: bool False, _record_nn_module_stack: bool False, _disable_torch_fn_metadata_mode: bool False, ) - JointWithDescriptors各参数的核心作用结合源码 aot_autograd.py参数说明stack必须由调用方创建并进入的contextlib.ExitStack会传入函数内部。该 context manager 在调用编译函数完成编译前必须保持激活源码注释注明这是一个 TODO未来可能放宽此要求mod被导出的nn.Moduleargs/kwargs模块的示例输入fake 化后用于 tracedecompositions算子分解表core aten 分解源码注释强调这与编译阶段不同在此阶段就需要keep_inference_input_mutations是否保留推理期输入修改ignore_shape_env是否忽略 shape envdisable_functionalization是否禁用 functionalization_record_nn_module_stack/_disable_torch_fn_metadata_mode内部选项关于参数与 buffer 的说明文档与源码给出一个重要提醒如果传入的nn.Module自身带有参数parameters和缓冲区buffers导出时会生成额外的隐式参数/buffer 参数并为其分配ParamAOTInput和BufferAOTInput描述符。但如果输入模块来自 Dynamo 之类的机制则不会得到这些描述符——因为 Dynamo 已经负责把参数/buffer 提升为参数了。此时需要分析输入节点的 Sources 来判断输入是否为参数及其 FQN详见 aot_autograd.py。返回的 JointWithDescriptors函数返回JointWithDescriptors定义于 schemas.py其关键成员包括_aot_state/_aot_graph_capture内部 AOT 状态与图捕获结果params_spec/buffers_spec最终编译函数中参数与 buffer 的期望传入顺序参数在前、buffer 在后in_spec/out_spec输入输出的 pytree TreeSpecgraph_module属性可读写的torch.fx.GraphModule这是拿到联合图后做自定义优化如重新分片的入口fake_mode属性导出时的 FakeTensorModecache_hash()返回适合作为缓存键的哈希字符串基于str(self)的 SHA-256。此外源码明确注明这些 API 不命中缓存——只缓存最终编译结果不缓存中间的导出结果。Descriptors 描述符体系描述符定义在 descriptors.py 中。模块 docstring 开宗明义AOTAutograd descriptors 是一种类似路径path-like的数据结构类似 pytree paths 与 sources用于描述 FX 图输入/输出的语义含义。之所以需要它是因为存在大量改变调用约定的图捕获包装器graph capture wrappers再加上 tangent、梯度等额外参数/输出很难直接看出 FX 图上的参数与原始函数参数的对应关系。所有描述符都是dataclasses.dataclass(frozenTrue)的不可变数据类。形式化语义文档给出严格定义假设def wrapped_graph(*args): ret graph(*in_transform(args)) return out_transform(ret)那么图输入input[i]的描述符描述了一个函数fin_i满足fin_i(args) in_transform(args)[i]——输入描述符告诉你如何从外层输入得到内层输入图输出output[j]的描述符描述了一个函数fout_j满足fout_j(out_transform(ret)) ret[j]——输出描述符告诉你如何从外层输出得到内层输出逆向数据流。直观示例文档给出几组输入描述符示例来自 descriptors.pyPlainAOTInput(idx0)原始 callable 的第一个输入原样保留ParamAOTInput(targetmod.weight)FQN 为mod.weight的参数TangentAOTInput(outputPlainAOTOutput(idx1))对应前向图第二个输出的梯度tangent输入ViewBaseAOTInput(base_ofPlainAOTInput(idx0))第一个输入实际上是某个与其他输入张量别名alias的张量的可微视图AOTAutograd 用一个代表所有这些输入 base 的输入替换了原输入此时会为该输入而非原 PlainAOTInput生成GradAOTOutput。如果联合图中出现不希望存在的视图 base可以在编译区域外克隆这些视图前提是不修改该张量来消除它SubclassGetAttrAOTInput(baseAOTInput(idx0), attrinner)该张量对应第一个索引处张量子类的inner张量。AOTAutograd 产生的联合图从不直接接收张量子类作为输入总是把它们解包unpack成组成它的普通张量用描述符识别相关部分嵌套子类可以递归嵌套。输出描述符示例PlainAOTOutput(idx0)原始前向函数的第一个输出原样保留GradAOTOutput(grad_ofPlainAOTInput(idx1))图第二个输入的已计算梯度是反向图的输出InputMutationAOTOutput(mutated_inputPlainAOTInput(idx0))当第一个输入被修改时需要拷回图第一个输入的新值。有时这些输出可以被省略、copy_直接在图中完成由keep_input_mutations控制但当输入修改必须被求导穿透时总会生成这样的输出IntermediateBaseAOTOutput(base_ofPlainAOTOutput(idx0))当多个输出相互别名时用一个代表所有别名 base 的输出张量替换它们该输出表示它是其中一个原始输出的 base。若不想这样在返回图之前克隆所有输出即可SubclassGetAttrAOTOutput(basePlainAOTOutput(idx0), attrinner)该张量对应第一个原始输出是张量子类的 inner 张量将与其他子类组件一起被重新打包repack为张量子类。五大构建阶段pipeline 顺序AOTAutograd 对函数做三大变换按此顺序执行视图/修改处理 → Autograd → 子类处理。因此描述符的构建顺序为详见 descriptors.pyPlainAOTInput/PlainAOTOutput起点描述原始展平用户函数的精确输入/输出假设已展平如需处理未展平函数可在描述符之上从用户侧链式拼接 pytree KeyPathSyntheticBaseAOTInput/ViewBaseAOTInput/MetadataMutationAOTOutput/InputMutationAOTOutput/IntermediateBaseAOTOutput通过移除重复的 PlainAOTInput 并引入新的人工输入/输出来处理修改与别名。这些输入与原始用户输入没有直接对应关系但如果实现一个不关心输入精确语义的 pass应当把它们与普通输入统一处理TangentAOTInput/GradAOTOutput为每个可微 AOTOutput 引入一个 tangent 输入、为每个可微 AOTInput 引入一个梯度输出来处理 autograd包括步骤 1-2 新引入的。这些描述符的参数只能是步骤 1-2 已有的描述符——因为 AOTAutograd 当前不支持双重反向double backwards永远不会出现 grad 的 tangent 或反之SubclassGetAttrAOTInput/SubclassGetAttrAOTOutput等为每个是子类的 AOTInput/AOTOutput 引入展平的输入/输出可能包含符号 size/stride。参数只能是步骤 1-3 已有的描述符支持递归子类因此步骤 4 的描述符之间可以互相嵌套ForwardTokenAOTInput/ForwardTokenAOTOutput/BackwardTokenAOTInput/BackwardTokenAOTOutput额外添加的合成 token 输入/输出仅用于防止 DCE死代码消除与重排。关键约束描述符只能自上而下top-to-bottom构建。例如SubclassGetAttrAOTInput(TangentAOTInput(PlainAOTOutput(...))) # OK是合法的PlainAOTOutput → TangentAOTInput → SubclassGetAttrAOTInput 符合 pipeline 顺序而以下写法非法TangentAOTInput(SubclassGetAttrAOTOutput(PlainAOTOutput(...))) # BADautograd 必须先于子类处理 GradAOTOutput(SubclassGetAttrAOTInput(PlainAOTInput(...))) # BAD子类在 create joint 之后才处理后者应写为SubclassGetAttrAOTOutput(GradAOTOutput(PlainAOTInput(...)))这直观地体现了我们总是直接对子类做 autograd而不是先把子类解构成 inner 张量再做 autograd。描述符索引按处理优先级分类模块 docstring 提供了按你多大概率需要处理它分类的完整索引AOTInput 类重要必须处理PlainAOTInputprimalsParamAOTInputTangentAOTInputSubclassGetAttrAOTInput等使用子类时视图相关可通过克隆图输入来消除不消除则务必处理它们与 GradAOTOutput 的配对ViewBaseAOTInputSyntheticBaseAOTInput非张量大多可直接忽略DummyAOTInputPhiloxForwardSeedAOTInput、PhiloxForwardBaseOffsetAOTInputPhiloxBackwardSeedAOTInput、PhiloxBackwardBaseOffsetAOTInputForwardTokenAOTInput、BackwardTokenAOTInputAOTOutput 类重要PlainAOTOutputGradAOTOutputSubclassGetAttrAOTOutput等使用子类时较冷门不消除则务必处理它们与 TangentAOTInput 的配对InputMutationAOTOutput修改不可微时可消除IntermediateBaseAOTOutput可通过克隆图输出来消除MetadataMutationAOTOutput避免修改元数据即可消除非张量大多可直接忽略PhiloxUpdatedForwardOffsetAOTOutput、PhiloxUpdatedBackwardOffsetAOTOutputForwardTokenAOTOutput、BackwardTokenAOTOutputDummyAOTOutput文档还提示DifferentiableAOTInput/DifferentiableAOTOutput是便捷分类基类用于判断哪些输入/输出可以被GradAOTOutput/TangentAOTInput包裹——本质上是除子类描述符之外的所有张量型 AOTInput/AOTOutput。全部描述符参考下表完整列出 descriptors.py 中定义的所有描述符及其expr()表达式expr()是基类AOTInput/AOTOutput声明的抽象方法子类必须实现输入描述符AOTInput描述符语义expr() 示例PlainAOTInput(idx)展平调用约定下的第 idx 个普通输入args[0]ParamAOTInput(target)FQN 为 target 的参数self.get_parameter(mod.weight)BufferAOTInput(target)FQN 为 target 的 bufferself.get_buffer(buf)TangentAOTInput(output)某个可微输出的 tangent 输入__output_tangent(output[0])ViewBaseAOTInput(base_of)多个可微输入是同一输入视图时的 base 输入args[0]._baseSyntheticBaseAOTInput(base_of)视图都不可微时为 autograd 构造的合成 base__make_synthetic_base(args[0])SubclassGetAttrAOTInput(base, attr)子类解包后的某属性张量可嵌套args[0].innerSubclassSizeAOTInput(base, idx)子类在第 idx 维的外层 size SymInt 输入args[0].size(0)SubclassStrideAOTInput(base, idx)子类在第 idx 维的外层 stride SymInt 输入args[0].stride(0)PhiloxForwardSeedAOTInput前向图函数化 Philox RNG 调用的 seed__philox_forward_seedPhiloxForwardBaseOffsetAOTInput前向图函数化 Philox RNG 调用的 offset__philox_forward_base_offsetPhiloxBackwardSeedAOTInput反向图函数化 Philox RNG 调用的 seed__philox_backward_seedPhiloxBackwardBaseOffsetAOTInput反向图函数化 Philox RNG 调用的 offset__philox_backward_base_offsetForwardTokenAOTInput(idx)穿线经过副作用操作的 world token__forward_token0BackwardTokenAOTInput(idx)反向的 world token__backward_token0DummyAOTInput(idx)占位描述符不应在正常情况使用__dummy0输出描述符AOTOutput描述符语义expr() 示例PlainAOTOutput(idx)输出元组第 idx 个普通张量输出output[0]GradAOTOutput(grad_of)可微输入对应的已计算梯度__grad(args[0])InputMutationAOTOutput(mutated_input)输入张量的修改后值用于正确传播 autograd__input_mutation(args[0])IntermediateBaseAOTOutput(base_of)多个互相别名输出的中间 base只报告其中一个输出__intermediate_base(output[0])MetadataMutationAOTOutput(idx)元数据被修改的别名参数__aliased_arg_with_metadata_mutation0SubclassGetAttrAOTOutput(base, attr)该输出将被打包进此位置的子类output[0].innerSubclassSizeAOTOutput(base, idx)该输出 size 将被打包进此位置的子类output[0].size(0)SubclassStrideAOTOutput(base, idx)该输出 stride 将被打包进此位置的子类output[0].stride(0)PhiloxUpdatedForwardOffsetAOTOutput函数化 RNG 调用的最终 offset仅前向__philox_updated_forward_offsetPhiloxUpdatedBackwardOffsetAOTOutput函数化 RNG 调用的最终 offset仅反向__philox_updated_backward_offsetForwardTokenAOTOutput(idx)副作用调用的 world token 输出防止被 DCE仅前向__forward_token0BackwardTokenAOTOutput(idx)副作用调用的 world token 输出仅反向__backward_token0SavedForBackwardsAOTOutput(idx)为反向保存的激活输出__saved_for_backwards_0SavedForBackwardsNoVcCheckAOTOutput(idx)反向时不检查版本计数器的激活输出__saved_for_backwards_no_vc_check_0DummyAOTOutput(idx)占位描述符不应在正常情况使用__dummy0其中SavedForBackwardsNoVcCheckAOTOutput对应一个有意思的 eager 语义细节见源码中 Note [Activations with no version counter checks in eager]在 eager 模式下若用户通过ctx.foo foo而非ctx.save_for_backward(foo)存张量autograd 引擎不会做版本计数器检查torch.compile 为了与之对齐需要知道哪些激活享受了无 VC 检查待遇从而在运行时避免对它们调用ctx.save_for_backward——这个描述符就负责传递这一信息。基类提供的判别方法AOTInputexpr()、is_param()是否为参数或派生自参数、is_buffer()是否为 buffer 或派生自 buffer、is_tangent()是否为 tangent 或派生自 tangentAOTOutputexpr()、is_grad()是否为 grad 或派生自 gradDifferentiableAOTInput/DifferentiableAOTOutput分别标记可被 GradAOTOutput 包裹的 AOTInput与可被 TangentAOTInput 包裹的 AOTOutput源码注释提示目前可微与不可微的类型纪律还不够好可依赖运行时测试。注意TangentAOTInput.__post_init__与GradAOTOutput.__post_init__都会校验参数类型不合法时抛出AssertionError。实现细节描述符如何生成实际生成过程比上面的风格化视图更复杂。AOTAutograd 被组织为一系列包裹在原始用户函数外的 wrapper它们组合成最终被 trace 的函数。因此描述符的构建是先在前向构建 wrappers、修改展平参数以适配新输入签名时逐步建立 AOTInput然后在 trace 过程中逆向建立 AOTOutput。一个主要的例外是TangentAOTInput的创建它需要引用尚未创建的 AOTOutput。两种解决方式详见 descriptors.py在 precompile 步骤dedupe 与 synthetic base 处理之后做一次前向元数据收集的初始 pass产生初始 PlainAOTOutput 集合用于创建 tangent 输入偶尔直接违反因果律在构建某个 AOTInput 时预测某个 AOTOutput 稍后会被创建。下表源码原样提供截至 2025 年 7 月穷举了输入/输出在各 wrapper 阶段的流转与各阶段可能引入的描述符Build wrappers (FLOWS DOWN) Run trace (FLOWS UP) ------------------------------------------------------------------------------------------------- Begin PlainAOTInput (n/a) ParamAOTInput Precompile dedupe (remove dupes) (nothing) Precompile synthetic base SyntheticBaseAOTInput MetadataMutationAOTOutput ViewBaseAOTInput Forward metadata trace PlainAOTOutput (n/a) MetadataMutationAOTOutput Prepare for autograd (nothing) InputMutationAOTOutput IntermediateBaseAOTOutput Create joint TangentAOTInput GradAOTOutput w/ InputMutationAOTOutput w/ IntermediateBaseAOTOutput Precompile subclass SubclassGetAttrAOTInput et al. SubclassGetAttrAOTOutput et al. Effect tokens ForwardTokenAOTInput ForwardTokenAOTOutput BackwardTokenAOTInput BackwardTokenAOTOutput End (n/a) PlainAOTOutput输入与输出流分开来看输入描述符传播构建 wrapper 时发生[IN] Begin原始调用约定PlainAOTInput、ParamAOTInput→[IN]Precompile dedupe移除重复 AOTInput→[IN]Precompile synthetic baseSyntheticBaseAOTInput、ViewBaseAOTInput→ Forward metadata trace微型输出描述符传播[OUT]原始输出约定 PlainAOTOutput[OUT]Precompile synthetic base 产生 MetadataMutationAOTOutput→[IN]Prepare for autograd无新增→[IN]Create jointTangentAOTInput可能带 IntermediateBaseAOTOutput、InputMutationAOTOutput→[IN]Precompile subclassSubclassGetAttrAOTInput 等→[IN]Effect tokensForwardTokenAOTInput、BackwardTokenAOTInput注意 BackwardTokenAOTInput 严格说不是 wrapper 生成的而是token_discovery在 FX trace 过程中隐式附加参数完成的触发 trace用修改后的输入对 wrapper 做一次 trace输出描述符传播从用户函数调用处 unwind 时发生[OUT] BeginPlainAOTOutput →[OUT]Effect tokensForward/BackwardTokenAOTOutput→[OUT]Precompile subclassSubclassGetAttrAOTOutput 等→[OUT]Create jointGradAOTOutput→[OUT]Prepare for autogradInputMutationAOTOutput、IntermediateBaseAOTOutput→[OUT]Precompile synthetic baseMetadataMutationAOTOutput→[OUT]Precompile dedupe无新增。编译回可微调用aot_compile_joint_with_descriptors作为aot_export_joint_with_descriptors的配套函数它把联合图编译成遵循标准调用约定的 callable定义于 aot_autograd.pydef aot_compile_joint_with_descriptors( jd: JointWithDescriptors, *, partition_fn: Callable[..., Any] default_partition, fw_compiler: AOTDispatchCompiler boxed_nop_preserve_node_meta, bw_compiler: AOTDispatchCompiler | None boxed_nop_preserve_node_meta, serializable: bool False, ) - Callable[..., Any]:参数说明jd上一步导出的JointWithDescriptorspartition_fn默认default_partition负责把联合图切分为前向/反向图fw_compiler/bw_compiler前向/反向图编译器默认boxed_nop_preserve_node_meta用torch.fx.Interpreter的boxed_run直接运行图并保留节点元数据见 aot_autograd.py。源码特别说明不强制做完整编译可以让 fw/bw compiler 留空分区后的 FX 图直接运行整体 autograd Function 可以允许在图内allowed in graph以便后续在更大编译区域中重新处理serializable为 True 时配置编译以产出可序列化 callable借助 AOTAutogradCache 机制设置 cache_info 并 patch 相关 config此时函数总是返回BundledAOTAutogradSerializableCallable。函数内部会用aot_stage2_compile完成第二阶段编译最后用jd.in_spec/jd.out_spec把展平输入/输出做 pytree 展开与重组unflatten并包裹torch._dynamo.nonstrict_trace以允许递归编译见 aot_autograd.py。注意编译函数不实例化模块这给了你子类化并自定义行为的灵活性而不必担心 FQN 重新绑定。FX 工具函数文档的最后一节是torch._functorch._aot_autograd.fx_utils模块见 fx_utils.py。该模块专门提供处理带描述符联合图的工具函数源码注释明确它们不适用于普通 FX 图。get_all_input_and_grad_nodes(g)给定带描述符的联合图返回每个输入及其对应梯度输出节点的映射——字典以描述该输入的AOTInput描述符为键值为(input_node, grad_node_or_None)。注意所有前向张量输入都会被返回包括不可微输入其 grad 为 None因此可安全地用它处理所有输入SymInt、token、RNG 状态等非张量输入不会被遍历。若联合图含子类输入/输出则抛RuntimeError因为子类场景下输入与梯度之间不一定有一一对应关系例如 primal 是普通张量而 tangent 是解包成多个普通张量的子类get_all_output_and_tangent_nodes(g)与前者对称返回每个可微输出节点及其 tangent 输入节点若存在的映射用于前向模式自动微分。模块中还有一个内部辅助_raise_autograd_subclass_not_implemented(n, desc)在遇到解包后的子类输入时抛出带详细解释的RuntimeError并鼓励有具体用例的开发者提 issue 以设计合适的 API。实战验证从导出到编译再到二次 compile仓库测试文件 test_aot_joint_with_descriptors.py 提供了可直接运行的完整用法这里给出两个最有代表性的模式。基础导出线性模块class SimpleLinear(nn.Module): def __init__(self): super().__init__() self.linear nn.Linear(3, 2) def forward(self, x): return self.linear(x) model SimpleLinear() inputs (torch.randn(4, 3),) with ExitStack() as stack: joint_with_descriptors aot_export_joint_with_descriptors( stack, model, inputs, decompositionsdecomposition_table ) graph_code joint_with_descriptors.graph_module.print_readable( print_outputFalse, expanded_defTrue )从测试的期望输出可以看到导出图的参数都被标注了描述符注释例如# ParamAOTInput(targetlinear1.weight)、# ParamAOTInput(targetlinear1.bias)等——这印证了 placeholder 节点meta[desc]中的描述符信息可以直接用来解读每个 primals 参数的语义。导出 编译 二次 torch.compile测试test_export_and_compile展示了完整的端到端流程with ExitStack() as stack: joint_with_descriptors aot_export_joint_with_descriptors( stack, model, inputs ) model_fn aot_compile_joint_with_descriptors(joint_with_descriptors) compiled_fn torch.compile(fullgraphTrue)(model_fn) compiled_fn(*dict(model.named_parameters()).values(), inputs).sum().backward() assert model.linear.weight.grad is not None关键点ExitStack的存活范围覆盖导出和编译两个阶段aot_compile_joint_with_descriptors返回的函数按params_spec/buffers_spec顺序接收展平参数参数在前、buffer 在后因此调用时用dict(model.named_parameters()).values()依次传入参数编译结果可以再次被torch.compile(fullgraphTrue)包裹并且能正常做反向传播、正确累积参数梯度——这正是文档所说的处理后可转换回可正常执行的可微调用。多输出模块测试test_multiple_outputs_module表明多可微输出场景torch.compile完整特性之一同样受支持MultiOutputModule.forward返回out1, out2两个输出导出图的 primals 参数同样带有各自的ParamAOTInput描述符说明这套 API 可以处理aot_export_joint_simple无法覆盖的更复杂情况。使用注意事项与边界综合文档与源码使用该 API 时有几点需要特别留意必须使用 ExitStack 上下文stack必须由调用方创建并保持激活直至编译完成这是当前版本的硬性要求不命中缓存导出结果不会被缓存只有最终编译结果可走缓存描述符可能很奇特文档建议仔细思考对不理解的描述符是否有安全的 fallback——例如应该有能力处理某个输入在最终 FX 图输入中无法原样找到的情况子类输入的约束fx_utils工具函数暂不支持子类输入/输出的联合图无双重反向当前不会生成 grad 的 tangent 或 tangent 的 gradDynamo 生成模块的差异若模块来自 Dynamo参数/buffer 已被提升为参数不会得到ParamAOTInput/BufferAOTInput描述符需自行分析 Sources 判断参数与 FQN非张量描述符RNG 状态Philox、token、Dummy 等非张量描述符在大多数自定义 pass 中可以直接忽略但不要误把它们当作可微输入处理。总而言之Joint with descriptors为在完整 torch.compile 语义下拿到可解释、可重写、可还原的联合图提供了统一而灵活的通道是构建 autoparallel 等需要先改语义、再编译的高级系统的正确起点。建议进一步阅读 descriptors.py 的模块级 docstring其中包含截至 2025 年 7 月的完整 wrapper 流转表与 test_aot_joint_with_descriptors.py 中的全部测试以获得对描述符体系最准确、最完整的认知。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考