DeepSpeed
615e6d5e - Scope DeepCompile compiler state to graph and engine lifecycles (#8159)

Commit
3 days ago
Scope DeepCompile compiler state to graph and engine lifecycles (#8159) DeepCompile uses several process-global PyTorch compiler mechanisms while building graph-specific state. Scheduled recompilation also replaces previously compiled graphs. These lifetimes did not have explicit ownership boundaries. - Forward inputs could be consumed by a different graph - Dynamo frame IDs could collide across engines - Cleanup could either leave graph-specific patches installed or release shared compiler state that another engine still needed This PR gives those existing mechanisms explicit graph and engine lifecycle ownership. Fixes: - Keep forward inputs in a graph-local one-shot queue and `InputStorage`. - Qualify compiled-backward frames with an owner token and track them per engine. - Keep the `torch.autograd.Function` patch active until the last owned backward frame is released, then clear captured backward inputs. - Reference-count the ZeRO-3 Dynamo configuration overrides across engines and restore the original values after the last owner exits. - Limit the graph-specific `AotAutograd.__init__` patch to the Inductor compilation call and restore it in `finally`. - Release engine-owned compiler state during scheduled recompilation, deactivation, and destruction. --------- Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Author
Parents
Loading