DeepSpeed

2a769889 - DeepCompile: Use min_cut_rematerialization for partitioning joint graphs (#7609)

DeepCompile: Use min_cut_rematerialization for partitioning joint graphs (#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] https://github.com/pytorch/pytorch/pull/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 <junjie.mao@linux.alibaba.com>