functorch.compile.min_cut_rematerialization_partition¶

functorch.compile.min_cut_rematerialization_partition(joint_module, _joint_inputs, compiler='nvfuser', recomputable_ops=None, *, num_fwd_outputs)[source]¶

对联合图进行分区，使得向后重新计算向前。重新计算有助于在内存带宽和计算之间进行权衡。

为了创建前向和后向图，我们复制联合图，手动将输出设置为仅原始前向或后向输出。然后我们通过死代码消除来运行生成的图。

警告

此API是实验性的，可能会发生变化。

Parameters

joint_module (fx.GraphModule) – 联合的前向和后向图。这是AOT Autograd追踪的结果。
_joint_inputs – 联合图的输入。此参数未使用。
compiler – 此选项决定了可重新计算操作符的默认集合。目前，有两个选项：nvfuser 和 inductor。
recomputable_ops – 这是一个可选的重新计算操作集。如果这不是None，那么将使用这组操作而不是默认的操作集。
num_fwd_outputs – 前向图的输出数量。

Returns

返回生成的前向和后向Fx图模块。