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 圖形模組。