…phs (deepspeedai#7609)

# Motivation

PyTorch provides `min_cut_rematerialization_partition()` to partition a
joint graph while respecting recomputation annotation. That algorithm
forms a data-flow-like graph from the joint graph, adds to edges weights
from some recomputation-cost-related heuristics and applies the min-cut
algorithm to determine which nodes to recompute. Users can force
recomputation of a node by annotating its `node.meta["recompute"]` to
MUST_RECOMPUTE or PREFER_RECOMPUTE, as is implemented in [1].

While originally designed for activation checkpointing,
min_cut_rematerialization can also be used to recompute param aliases.
When partitioning a joint graph, we don't want to save for backward the
gathered parameters and values computed from them via aliasing ops, as
that essentially means the gathered parameter will be saved. Instead of
customizing the partitioner or patching `choose_saved_values_set`, we
can achieve that by annotating such nodes to be MUST_RECOMPUTE.

Both eager and inductor backends can use min_cut_rematerialization
easily. The eager backend can use min-cut by customizing the
partition_fn for `aot_module_simplified`, and is already using that for
graphs with activation checkpointing enabled. The inductor backend uses
that algorithm since torch 2.0.0 [2] and is still the default after the
inductor partitioner is made configurable a few weeks ago [3].

That approach also helps DeepCompile + torch autocast nicely. When
autocast is enabled, downcasted parameters are preferred to be
recomputed. It suffices to mark such casting nodes as must-recompute.

[1]
https://github.com/pytorch/pytorch/blob/main/torch/_functorch/partitioners.py#L1813
[2]
https://github.com/pytorch/pytorch/blob/v2.0.0/torch/_inductor/compile_fx.py#L459
[3] pytorch/pytorch#157580

# Proposal

Motivated by the flexibility and the requirement for optimizing
DeepCompile + autocast, I propose to switch to the min-cut-based
partitioner for both backends. This PR implements that switch, cleans up
dead code and also recomputes downcasted parameters in the backward.

# Preliminary Evaluation

Here's a summary of the tests using
https://gist.github.com/eternalNight/3c2cf8c703f1e9e7742d3b7f9e1edae3 on
a 8x RTX 5090 node.

| Configuration | Base Time (ms) | Base Mem (GB) | Time with this PR
(ms) | Mem with this PR (GB) |

|---------------------|----------------|---------------|------------------------|-----------------------|
| eager + autocast | 551.92 | 12.07 | 571.24 | 9.96 |
| eager + bf16 | 419.87 | 9.47 | 445.76 | 7.30 |
| inductor + autocast | 546.97 | 12.84 | 570.09 | 13.04 |
| inductor + bf16 | 444.03 | 10.01 | 444.70 | 10.19 |

## Reduced memory with eager backend

The initial goal of this PR is to reduce peak memory usage when torch
autocast is enabled. That is achieved according to the first row of the
table, but in two different ways simultaneously.

1. Downcasted parameters during forward are throwed away and recomputed
(by the fused cast + allgather) in the backward pass.
2. Without this PR, `fast_free_schedule` will arange most allgather at
the beginning of the graph. That leads to a even higher peak during
forward, but is no longer seen with PR.
3. By diffing the graphs passed to `add_z3_gather_release`, I noticed
that recomputations selected by min-cut is slightly different (that test
script has activation checkpointing enabled for the LLM module). That
can also impact computation time and memory usage.

Here's the shape of memory usage before this PR with eager backend +
torch autocast. eager + BF16 shows similar shapes. Numbers reported in
the table are peak during forward. The peak memory usage during backend
reduces ~0.7GB in both cases.

<img width="1482" height="629" alt="image"
src="https://github.com/user-attachments/assets/7e7ec859-9a04-4ddd-ba37-c2d475a81058"
/>

After this PR:

<img width="1482" height="453" alt="image"
src="https://github.com/user-attachments/assets/f15c71b8-f823-4aa5-801a-a36188c5e866"
/>

## Similar memory with inductor backend

Unlike eager backend, the inductor backend uses similar memory with or
without this PR. The memory usage pattern is as follows, which requires
further analysis.

Before this PR:

<img width="1070" height="613" alt="image"
src="https://github.com/user-attachments/assets/317b9a58-d4ef-459f-ac7b-67ef2318a9de"
/>

After this PR:

<img width="911" height="536" alt="image"
src="https://github.com/user-attachments/assets/7e737a81-cf27-402c-aeea-dfe661043fc1"
/>

Signed-off-by: Junjie Mao <[email protected]>

…phs (deepspeedai#7609)

# Motivation

PyTorch provides `min_cut_rematerialization_partition()` to partition a
joint graph while respecting recomputation annotation. That algorithm
forms a data-flow-like graph from the joint graph, adds to edges weights
from some recomputation-cost-related heuristics and applies the min-cut
algorithm to determine which nodes to recompute. Users can force
recomputation of a node by annotating its `node.meta["recompute"]` to
MUST_RECOMPUTE or PREFER_RECOMPUTE, as is implemented in [1].

While originally designed for activation checkpointing,
min_cut_rematerialization can also be used to recompute param aliases.
When partitioning a joint graph, we don't want to save for backward the
gathered parameters and values computed from them via aliasing ops, as
that essentially means the gathered parameter will be saved. Instead of
customizing the partitioner or patching `choose_saved_values_set`, we
can achieve that by annotating such nodes to be MUST_RECOMPUTE.

Both eager and inductor backends can use min_cut_rematerialization
easily. The eager backend can use min-cut by customizing the
partition_fn for `aot_module_simplified`, and is already using that for
graphs with activation checkpointing enabled. The inductor backend uses
that algorithm since torch 2.0.0 [2] and is still the default after the
inductor partitioner is made configurable a few weeks ago [3].

That approach also helps DeepCompile + torch autocast nicely. When
autocast is enabled, downcasted parameters are preferred to be
recomputed. It suffices to mark such casting nodes as must-recompute.

[1]
https://github.com/pytorch/pytorch/blob/main/torch/_functorch/partitioners.py#L1813
[2]
https://github.com/pytorch/pytorch/blob/v2.0.0/torch/_inductor/compile_fx.py#L459
[3] pytorch/pytorch#157580

# Proposal

Motivated by the flexibility and the requirement for optimizing
DeepCompile + autocast, I propose to switch to the min-cut-based
partitioner for both backends. This PR implements that switch, cleans up
dead code and also recomputes downcasted parameters in the backward.

# Preliminary Evaluation

Here's a summary of the tests using
https://gist.github.com/eternalNight/3c2cf8c703f1e9e7742d3b7f9e1edae3 on
a 8x RTX 5090 node.

| Configuration | Base Time (ms) | Base Mem (GB) | Time with this PR
(ms) | Mem with this PR (GB) |

|---------------------|----------------|---------------|------------------------|-----------------------|
| eager + autocast | 551.92 | 12.07 | 571.24 | 9.96 |
| eager + bf16 | 419.87 | 9.47 | 445.76 | 7.30 |
| inductor + autocast | 546.97 | 12.84 | 570.09 | 13.04 |
| inductor + bf16 | 444.03 | 10.01 | 444.70 | 10.19 |

## Reduced memory with eager backend

The initial goal of this PR is to reduce peak memory usage when torch
autocast is enabled. That is achieved according to the first row of the
table, but in two different ways simultaneously.

1. Downcasted parameters during forward are throwed away and recomputed
(by the fused cast + allgather) in the backward pass.
2. Without this PR, `fast_free_schedule` will arange most allgather at
the beginning of the graph. That leads to a even higher peak during
forward, but is no longer seen with PR.
3. By diffing the graphs passed to `add_z3_gather_release`, I noticed
that recomputations selected by min-cut is slightly different (that test
script has activation checkpointing enabled for the LLM module). That
can also impact computation time and memory usage.

Here's the shape of memory usage before this PR with eager backend +
torch autocast. eager + BF16 shows similar shapes. Numbers reported in
the table are peak during forward. The peak memory usage during backend
reduces ~0.7GB in both cases.

<img width="1482" height="629" alt="image"
src="https://github.com/user-attachments/assets/7e7ec859-9a04-4ddd-ba37-c2d475a81058"
/>

After this PR:

<img width="1482" height="453" alt="image"
src="https://github.com/user-attachments/assets/f15c71b8-f823-4aa5-801a-a36188c5e866"
/>

## Similar memory with inductor backend

Unlike eager backend, the inductor backend uses similar memory with or
without this PR. The memory usage pattern is as follows, which requires
further analysis.

Before this PR:

<img width="1070" height="613" alt="image"
src="https://github.com/user-attachments/assets/317b9a58-d4ef-459f-ac7b-67ef2318a9de"
/>

After this PR:

<img width="911" height="536" alt="image"
src="https://github.com/user-attachments/assets/7e737a81-cf27-402c-aeea-dfe661043fc1"
/>

Signed-off-by: Junjie Mao <[email protected]>
Signed-off-by: Ma, Liangliang <[email protected]>