functorch.compile.min_cut_rematerialization_partition

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

将联合图进行划分，以便反向计算重新计算正向。重新计算有助于在内存带宽和计算之间进行权衡。

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

警告

此 API 处于实验阶段，可能会发生变化。

参数

• joint_module (fx.GraphModule) – 联合正向和反向图。这是 AOT Autograd 追踪的结果。

• _joint_inputs – 联合图的输入。这未使用。

• compiler – 此选项决定默认的可重新计算操作集。目前，有两种选项：nvfuser 和 inductor。

• recomputable_ops – 这是一个可选的可重新计算操作集。如果此值不为 None，则将使用此操作集，而不是默认操作集。

• num_fwd_outputs – 正向图的输出数量。

返回值

返回生成的正向和反向 Fx 图模块。