Stabilize DeepCompile ZeRO-3 memory scheduling - #8169
Conversation
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
tohtana
changed the title
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 55779855e3
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
| m.def("register_graph_z3", | ||
| &dc::register_graph_z3, | ||
| "Register graph with a list of ds parameter ids"); | ||
| m.def("set_z3_gather_buffer_pool_budget_for_test", |
There was a problem hiding this comment.
wondering if these new items are just test purpose only?
|
|
||
|
|
||
| def deepcompile_z3_inductor_config_patch(enabled: bool): | ||
| """Disable reduction heuristics that create oversized kernels for DeepCompile ZeRO-3 graphs.""" |
There was a problem hiding this comment.
How does inductor-generated mix-order or persistent reductions bother ZeRO-3? Is that related to GPU memory use?
| def all_reduce(tensor, op, process_group=None): | ||
| if process_group is None: | ||
| return dist.all_reduce(tensor, op) | ||
| return dist.all_reduce(tensor, op, group=process_group) |
There was a problem hiding this comment.
That wrapper is unnecessary. dist.all_reduce already takes group as optional and default to None (which means the default process group will be used). Same to get_rank and barrier.
| if (!enabled_ || pressure_recovery_in_progress_) { return at::Tensor(); } | ||
|
|
||
| Entry* best = nullptr; | ||
| for (auto& entry : entries_) { |
There was a problem hiding this comment.
is entries_ ordered? if not, is the full loop expensive vs an max heap like solution?
|
general question, do we have experiments results showcasing the changes' benefits? |
3 participants
Depends on #8159.
DeepCompile's ZeRO-3 scheduler does not consistently account for memory pressure across gather/release ordering, prefetch, selective gathering, and native gather-buffer reuse. This can retain excess gathered storage or trigger a full-parameter gather while Dynamo evaluates guards.
This PR adds rank-consistent scheduler budgeting and diagnostics, coordinates prefetch and selective gathering with graph profiling, reuses native gather storage within a bounded pressure-aware lifecycle, and avoids guard-time full-parameter gathers while the DeepCompile eager fallback is active.